Added hnsw.search_mem_multiplier option

ankane · ankane · commit ac6576e53a56 · 2024-10-24T18:02:20.000-07:00
diff --git 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
@@ -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
@@ -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
@@ -28,7 +28,6 @@
 	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 @@
 			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 @@
 	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();

Original file line number	Diff line number	Diff line change
`@@ -138,6 +138,8 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)`
`138`	`138`	`/* Set support functions */`
`139`	`139`	`HnswInitSupport(&so->support, index);`
`140`	`140`
	`141`	`+ so->maxMemory = (Size) (work_mem * 1024.0 * hnsw_search_mem_multiplier);`
	`142`	`+`
`141`	`143`	`scan->opaque = so;`
`142`	`144`
`143`	`145`	`return scan;`
`@@ -244,13 +246,13 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)`
`244`	`246`	`so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));`
`245`	`247`	`}`
`246`	`248`	`/* Prevent scans from consuming too much memory */`
`247`		`- else if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)`
	`249`	`+ else if (MemoryContextMemAllocated(so->tmpCtx, false) >= so->maxMemory)`
`248`	`250`	`{`
`249`	`251`	`if (pairingheap_is_empty(so->discarded))`
`250`	`252`	`{`
`251`	`253`	`ereport(DEBUG1,`
`252`		`- (errmsg("hnsw index scan exceeded work_mem after " INT64_FORMAT " tuples", so->tuples),`
`253`		`- errhint("Increase work_mem to scan more tuples.")));`
	`254`	`+ (errmsg("hnsw index scan reached memory limit after " INT64_FORMAT " tuples", so->tuples),`
	`255`	`+ errhint("Increase hnsw.search_mem_multiplier to scan more tuples.")));`
`254`	`256`
`255`	`257`	`break;`
`256`	`258`	`}`