Skip to content

Commit 495041e

Browse files
ankaneankane
committed
Added option to limit tuples [skip ci]
1 parent 52c385c commit 495041e

File tree

4 files changed

+34
-11
lines changed

4 files changed

+34
-11
lines changed

‎src/hnsw.c

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#endif
1919

2020
int hnsw_ef_search;
21+
int hnsw_ef_stream;
2122
bool hnsw_streaming;
2223
int hnsw_lock_tranche_id;
2324
static relopt_kind hnsw_relopt_kind;
@@ -74,7 +75,10 @@ HnswInit(void)
7475
NULL, &hnsw_streaming,
7576
HNSW_DEFAULT_STREAMING, PGC_USERSET, 0, NULL, NULL, NULL);
7677

77-
/* TODO Add option for limiting iterative search */
78+
/* TODO Figure out name */
79+
DefineCustomIntVariable("hnsw.ef_stream", "Sets the max number of additional candidates to visit for streaming search",
80+
"-1 means all", &hnsw_ef_stream,
81+
HNSW_DEFAULT_EF_STREAM, HNSW_MIN_EF_STREAM, HNSW_MAX_EF_STREAM, PGC_USERSET, 0, NULL, NULL, NULL);
7882

7983
MarkGUCPrefixReserved("hnsw");
8084
}

‎src/hnsw.h

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,9 @@
4747
#define HNSW_MIN_EF_SEARCH 1
4848
#define HNSW_MAX_EF_SEARCH 1000
4949
#define HNSW_DEFAULT_STREAMING false
50+
#define HNSW_DEFAULT_EF_STREAM -1
51+
#define HNSW_MIN_EF_STREAM -1
52+
#define HNSW_MAX_EF_STREAM INT_MAX
5053

5154
/* Tuple types */
5255
#define HNSW_ELEMENT_TUPLE_TYPE 1
@@ -126,6 +129,7 @@
126129

127130
/* Variables */
128131
extern int hnsw_ef_search;
132+
extern int hnsw_ef_stream;
129133
extern bool hnsw_streaming;
130134
extern int hnsw_lock_tranche_id;
131135

@@ -412,7 +416,7 @@ bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
412416
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
413417
void HnswInitPage(Buffer buf, Page page);
414418
void HnswInit(void);
415-
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited);
419+
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
416420
HnswElement HnswGetEntryPoint(Relation index);
417421
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
418422
void *HnswAlloc(HnswAllocator * allocator, Size size);

‎src/hnswscan.c

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,11 +36,11 @@ GetScanItems(IndexScanDesc scan, Datum q)
3636

3737
for (int lc = entryPoint->level; lc >= 1; lc--)
3838
{
39-
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true);
39+
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true, NULL);
4040
ep = w;
4141
}
4242

43-
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, hnsw_streaming ? &so->discarded : NULL, true);
43+
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);
4444
}
4545

4646
/*
@@ -73,7 +73,7 @@ ResumeScanItems(IndexScanDesc scan)
7373
ep = lappend(ep, hc);
7474
}
7575

76-
return HnswSearchLayer(base, so->q, ep, batch_size, 0, index, procinfo, collation, so->m, false, NULL, &so->v, &so->discarded, false);
76+
return HnswSearchLayer(base, so->q, ep, batch_size, 0, index, procinfo, collation, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples);
7777
}
7878

7979
/*
@@ -219,8 +219,17 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
219219
if (!hnsw_streaming)
220220
break;
221221

222+
/* Reached max number of additional tuples */
223+
if (hnsw_ef_stream != -1 && so->tuples >= hnsw_ef_search + hnsw_ef_stream)
224+
{
225+
if (pairingheap_is_empty(so->discarded))
226+
break;
227+
228+
/* Return remaining tuples */
229+
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
230+
}
222231
/* Prevent scans from consuming too much memory */
223-
if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)
232+
else if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)
224233
{
225234
if (pairingheap_is_empty(so->discarded))
226235
{
@@ -278,8 +287,6 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
278287
continue;
279288
}
280289

281-
so->tuples++;
282-
283290
heaptid = &element->heaptids[--element->heaptidsLength];
284291

285292
MemoryContextSwitchTo(oldCtx);

‎src/hnswutils.c

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -807,7 +807,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
807807
* Algorithm 2 from paper
808808
*/
809809
List *
810-
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited)
810+
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples)
811811
{
812812
List *w = NIL;
813813
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -849,8 +849,13 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
849849
bool found;
850850

851851
if (initVisited)
852+
{
852853
AddToVisited(base, v, hc->element, index, &found);
853854

855+
if (tuples != NULL)
856+
(*tuples)++;
857+
}
858+
854859
pairingheap_add(C, &hc->c_node);
855860
pairingheap_add(W, &hc->w_node);
856861

@@ -879,6 +884,9 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
879884
else
880885
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, v, index, m, lm, lc);
881886

887+
if (tuples != NULL)
888+
(*tuples) += unvisitedLength;
889+
882890
for (int i = 0; i < unvisitedLength; i++)
883891
{
884892
HnswElement eElement;
@@ -1318,7 +1326,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
13181326
/* 1st phase: greedy search to insert level */
13191327
for (int lc = entryLevel; lc >= level + 1; lc--)
13201328
{
1321-
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true);
1329+
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true, NULL);
13221330
ep = w;
13231331
}
13241332

@@ -1337,7 +1345,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
13371345
List *lw = NIL;
13381346
ListCell *lc2;
13391347

1340-
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true);
1348+
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true, NULL);
13411349

13421350
/* Convert search candidates to candidates */
13431351
foreach(lc2, w)

0 commit comments

Comments
 (0)