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

Added option to limit tuples [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-09-22 18:10:19 -07:00
parent 52c385c03a
commit 495041e43b
4 changed files with 34 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
src/hnswscan.c
View File

@@ -36,11 +36,11 @@ GetScanItems(IndexScanDesc scan, Datum q)
for (int lc = entryPoint->level; lc >= 1; lc--)
{
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true);
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true, NULL);
ep = w;
}
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, hnsw_streaming ? &so->discarded : NULL, true);
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, hnsw_streaming ? &so->discarded : NULL, true, &so->tuples);
}
/*
@@ -73,7 +73,7 @@ ResumeScanItems(IndexScanDesc scan)
ep = lappend(ep, hc);
}
return HnswSearchLayer(base, so->q, ep, batch_size, 0, index, procinfo, collation, so->m, false, NULL, &so->v, &so->discarded, false);
return HnswSearchLayer(base, so->q, ep, batch_size, 0, index, procinfo, collation, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples);
}
/*
@@ -219,8 +219,17 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
if (!hnsw_streaming)
break;
/* Reached max number of additional tuples */
if (hnsw_ef_stream != -1 && so->tuples >= hnsw_ef_search + hnsw_ef_stream)
{
if (pairingheap_is_empty(so->discarded))
break;
/* Return remaining tuples */
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
}
/* Prevent scans from consuming too much memory */
if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)
else if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)
{
if (pairingheap_is_empty(so->discarded))
{
@@ -278,8 +287,6 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
continue;
}
so->tuples++;
heaptid = &element->heaptids[--element->heaptidsLength];
MemoryContextSwitchTo(oldCtx);