liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-14 00:26:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added hnsw.search_mem_multiplier option

Browse Source
This commit is contained in:
Andrew Kane
2024-10-24 18:02:20 -07:00
parent 67eff41c44
commit ac6576e53a
4 changed files with 15 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/hnsw.c
View File

@@ -28,6 +28,7 @@ static const struct config_enum_entry hnsw_iterative_search_options[] = {
int hnsw_ef_search; int hnsw_ef_search;
int hnsw_iterative_search; int hnsw_iterative_search;
int hnsw_max_search_tuples; int hnsw_max_search_tuples;
double hnsw_search_mem_multiplier;
int hnsw_lock_tranche_id; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
@@ -87,6 +88,11 @@ HnswInit(void)
NULL, &hnsw_max_search_tuples, NULL, &hnsw_max_search_tuples,
20000, 1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL); 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"); MarkGUCPrefixReserved("hnsw");
} }

2
src/hnsw.h
View File

@@ -111,6 +111,7 @@
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_iterative_search; extern int hnsw_iterative_search;
extern int hnsw_max_search_tuples; extern int hnsw_max_search_tuples;
extern double hnsw_search_mem_multiplier;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef enum HnswIterativeSearchMode typedef enum HnswIterativeSearchMode
@@ -372,6 +373,7 @@ typedef struct HnswScanOpaqueData
int m; int m;
int64 tuples; int64 tuples;
double previousDistance; double previousDistance;
Size maxMemory;
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Support functions */ /* Support functions */

8
src/hnswscan.c
View File

@@ -138,6 +138,8 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
/* Set support functions */ /* Set support functions */
HnswInitSupport(&so->support, index); HnswInitSupport(&so->support, index);
so->maxMemory = (Size) (work_mem * 1024.0 * hnsw_search_mem_multiplier);
scan->opaque = so; scan->opaque = so;
return scan; return scan;
@@ -244,13 +246,13 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded))); so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
} }
/* Prevent scans from consuming too much memory */ /* 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)) if (pairingheap_is_empty(so->discarded))
{ {
ereport(DEBUG1, ereport(DEBUG1,
(errmsg("hnsw index scan exceeded work_mem after " INT64_FORMAT " tuples", so->tuples), (errmsg("hnsw index scan reached memory limit after " INT64_FORMAT " tuples", so->tuples),
errhint("Increase work_mem to scan more tuples."))); errhint("Increase hnsw.search_mem_multiplier to scan more tuples.")));
break; break;
} }

5
test/t/043_hnsw_iterative_search.pl
View File

@@ -28,7 +28,6 @@ my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.iterative_search = relaxed_order; SET hnsw.iterative_search = relaxed_order;
SET hnsw.max_search_tuples = 100000; 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; 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); is($count, 10);
@@ -45,7 +44,6 @@ foreach ((30000, 50000, 70000))
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.iterative_search = relaxed_order; SET hnsw.iterative_search = relaxed_order;
SET hnsw.max_search_tuples = $max_tuples; 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; 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; $sum += $count;
@@ -61,8 +59,9 @@ my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET hnsw.iterative_search = relaxed_order; SET hnsw.iterative_search = relaxed_order;
SET client_min_messages = debug1; SET client_min_messages = debug1;
SET work_mem = '1MB'; 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; 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(); done_testing();