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

Limit work_mem [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-09-21 22:37:03 -07:00
parent 96ff1b992f
commit e07d17fe38
3 changed files with 46 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/hnsw.h
View File

@@ -363,6 +363,7 @@ typedef struct HnswScanOpaqueData
List *discarded; List *discarded;
Datum q; Datum q;
int m; int m;
int64 tuples;
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Support functions */ /* Support functions */

36
src/hnswscan.c
View File

@@ -131,6 +131,7 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
tidhash_reset(so->v.tids); tidhash_reset(so->v.tids);
so->first = true; so->first = true;
so->discarded = NIL; so->discarded = NIL;
so->tuples = 0;
MemoryContextReset(so->tmpCtx); MemoryContextReset(so->tmpCtx);
if (keys && scan->numberOfKeys > 0) if (keys && scan->numberOfKeys > 0)
@@ -204,15 +205,34 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
if (!hnsw_streaming) if (!hnsw_streaming)
break; break;
/* if (MemoryContextMemAllocated(so->tmpCtx, false) > (work_mem * 1024L))
* Ensure vacuum does not mark tuples as deleted during an {
* iteration if (list_length(so->discarded) == 0)
*/ {
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock); ereport(NOTICE,
(errmsg("iterative search exceeded work_mem after " INT64_FORMAT " tuples", so->tuples),
errhint("Increase work_mem to scan more tuples.")));
HnswBench("scan iteration", so->w = ResumeScanItems(scan)); break;
}
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock); /* Return remaining tuples and exit */
/* TODO sort */
so->w = so->discarded;
so->discarded = NIL;
}
else
{
/*
* Ensure vacuum does not mark tuples as deleted during an
* iteration
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
HnswBench("scan iteration", so->w = ResumeScanItems(scan));
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
}
#if defined(HNSW_MEMORY) #if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024); elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
@@ -240,6 +260,8 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
continue; continue;
} }
so->tuples++;
heaptid = &element->heaptids[--element->heaptidsLength]; heaptid = &element->heaptids[--element->heaptidsLength];
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);

19
test/t/039_hnsw_streaming.pl
View File

@@ -16,15 +16,28 @@ $node->start;
$node->safe_psql("postgres", "CREATE EXTENSION vector;"); $node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));"); $node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres", $node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 10000) i;" "INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
); );
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);"); $node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
SET max_parallel_maintenance_workers = 2;
CREATE INDEX ON tst USING hnsw (v vector_l2_ops)
));
my $count = $node->safe_psql("postgres", qq( my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.streaming = on; SET hnsw.streaming = on;
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 1000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t; 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); is($count, 10);
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.streaming = on;
SET work_mem = '2MB';
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/iterative search exceeded work_mem after \d+ tuples/);
done_testing(); done_testing();