diff --git a/src/hnsw.c b/src/hnsw.c index be90f0b..11bd507 100644 --- a/src/hnsw.c +++ b/src/hnsw.c @@ -28,6 +28,7 @@ static const struct config_enum_entry hnsw_iterative_search_options[] = { int hnsw_ef_search; int hnsw_iterative_search; int hnsw_max_search_tuples; +double hnsw_search_mem_multiplier; int hnsw_lock_tranche_id; static relopt_kind hnsw_relopt_kind; @@ -87,6 +88,11 @@ HnswInit(void) NULL, &hnsw_max_search_tuples, 20000, 1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL); + /* Same range and default as hash_mem_multiplier */ + DefineCustomRealVariable("hnsw.search_mem_multiplier", "Sets the multiple of work_mem to use for iterative search", + NULL, &hnsw_search_mem_multiplier, + 2, 1, 1000, PGC_USERSET, 0, NULL, NULL, NULL); + MarkGUCPrefixReserved("hnsw"); } diff --git a/src/hnsw.h b/src/hnsw.h index 6b184ec..a224a6e 100644 --- a/src/hnsw.h +++ b/src/hnsw.h @@ -111,6 +111,7 @@ extern int hnsw_ef_search; extern int hnsw_iterative_search; extern int hnsw_max_search_tuples; +extern double hnsw_search_mem_multiplier; extern int hnsw_lock_tranche_id; typedef enum HnswIterativeSearchMode @@ -372,6 +373,7 @@ typedef struct HnswScanOpaqueData int m; int64 tuples; double previousDistance; + Size maxMemory; MemoryContext tmpCtx; /* Support functions */ diff --git a/src/hnswscan.c b/src/hnswscan.c index 3006a05..6ed8af5 100644 --- a/src/hnswscan.c +++ b/src/hnswscan.c @@ -138,6 +138,8 @@ hnswbeginscan(Relation index, int nkeys, int norderbys) /* Set support functions */ HnswInitSupport(&so->support, index); + so->maxMemory = (Size) (work_mem * 1024.0 * hnsw_search_mem_multiplier); + scan->opaque = so; return scan; @@ -244,13 +246,13 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir) so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded))); } /* Prevent scans from consuming too much memory */ - else if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L) + else if (MemoryContextMemAllocated(so->tmpCtx, false) >= so->maxMemory) { if (pairingheap_is_empty(so->discarded)) { ereport(DEBUG1, - (errmsg("hnsw index scan exceeded work_mem after " INT64_FORMAT " tuples", so->tuples), - errhint("Increase work_mem to scan more tuples."))); + (errmsg("hnsw index scan reached memory limit after " INT64_FORMAT " tuples", so->tuples), + errhint("Increase hnsw.search_mem_multiplier to scan more tuples."))); break; } diff --git a/test/t/043_hnsw_iterative_search.pl b/test/t/043_hnsw_iterative_search.pl index ea48010..ad0e681 100644 --- a/test/t/043_hnsw_iterative_search.pl +++ b/test/t/043_hnsw_iterative_search.pl @@ -28,7 +28,6 @@ my $count = $node->safe_psql("postgres", qq( SET enable_seqscan = off; SET hnsw.iterative_search = relaxed_order; SET hnsw.max_search_tuples = 100000; - SET work_mem = '8MB'; SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t; )); is($count, 10); @@ -45,7 +44,6 @@ foreach ((30000, 50000, 70000)) SET enable_seqscan = off; SET hnsw.iterative_search = relaxed_order; SET hnsw.max_search_tuples = $max_tuples; - SET work_mem = '8MB'; SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t; )); $sum += $count; @@ -61,8 +59,9 @@ my ($ret, $stdout, $stderr) = $node->psql("postgres", qq( SET hnsw.iterative_search = relaxed_order; SET client_min_messages = debug1; SET work_mem = '1MB'; + SET hnsw.search_mem_multiplier = 1; SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t; )); -like($stderr, qr/hnsw index scan exceeded work_mem after \d+ tuples/); +like($stderr, qr/hnsw index scan reached memory limit after \d+ tuples/); done_testing();