From 06b641407f786cce3bf47f22fcdb1fdfd9a957da Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 20 Sep 2024 22:06:09 -0700 Subject: [PATCH] Added time budget for HNSW index scans --- src/hnsw.c | 5 +++++ src/hnsw.h | 4 +++- src/hnswscan.c | 8 ++++++-- src/hnswutils.c | 18 +++++++++++++++--- 4 files changed, 29 insertions(+), 6 deletions(-) diff --git a/src/hnsw.c b/src/hnsw.c index a7b1e5f..f3fa00c 100644 --- a/src/hnsw.c +++ b/src/hnsw.c @@ -18,6 +18,7 @@ #endif int hnsw_ef_search; +int hnsw_time_budget; int hnsw_lock_tranche_id; static relopt_kind hnsw_relopt_kind; @@ -68,6 +69,10 @@ HnswInit(void) "Valid range is 1..1000.", &hnsw_ef_search, HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL); + DefineCustomIntVariable("hnsw.time_budget", "Sets the time budget for search", + NULL, &hnsw_time_budget, + -1, -1, INT_MAX, PGC_USERSET, GUC_UNIT_MS, NULL, NULL, NULL); + MarkGUCPrefixReserved("hnsw"); } diff --git a/src/hnsw.h b/src/hnsw.h index 2f45039..d3a5e6f 100644 --- a/src/hnsw.h +++ b/src/hnsw.h @@ -106,6 +106,7 @@ /* Variables */ extern int hnsw_ef_search; +extern int hnsw_time_budget; extern int hnsw_lock_tranche_id; typedef struct HnswElementData HnswElementData; @@ -330,6 +331,7 @@ typedef struct HnswScanOpaqueData bool first; List *w; MemoryContext tmpCtx; + instr_time start; /* Support functions */ FmgrInfo *procinfo; @@ -374,7 +376,7 @@ bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); void HnswInitPage(Buffer buf, Page page); void HnswInit(void); -List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement); +List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, instr_time *start); HnswElement HnswGetEntryPoint(Relation index); void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint); void *HnswAlloc(HnswAllocator * allocator, Size size); diff --git a/src/hnswscan.c b/src/hnswscan.c index 0463a89..4ba29cf 100644 --- a/src/hnswscan.c +++ b/src/hnswscan.c @@ -3,6 +3,7 @@ #include "access/relscan.h" #include "hnsw.h" #include "pgstat.h" +#include "portability/instr_time.h" #include "storage/bufmgr.h" #include "storage/lmgr.h" #include "utils/memutils.h" @@ -33,11 +34,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); + w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, &so->start); ep = w; } - return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL); + return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->start); } /* @@ -90,6 +91,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys) so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); so->collation = index->rd_indcollation[0]; + if (hnsw_time_budget != -1) + INSTR_TIME_SET_CURRENT(so->start); + scan->opaque = so; return scan; diff --git a/src/hnswutils.c b/src/hnswutils.c index 6e01cf4..cb46446 100644 --- a/src/hnswutils.c +++ b/src/hnswutils.c @@ -9,6 +9,7 @@ #include "fmgr.h" #include "hnsw.h" #include "lib/pairingheap.h" +#include "portability/instr_time.h" #include "sparsevec.h" #include "storage/bufmgr.h" #include "utils/datum.h" @@ -799,7 +800,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u * Algorithm 2 from paper */ List * -HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement) +HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, instr_time *start) { List *w = NIL; pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); @@ -850,6 +851,17 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W)); HnswElement cElement; + /* Check time budget */ + if (hnsw_time_budget != -1) + { + instr_time duration; + + INSTR_TIME_SET_CURRENT(duration); + INSTR_TIME_SUBTRACT(duration, *start); + if (INSTR_TIME_GET_MILLISEC(duration) >= hnsw_time_budget) + break; + } + if (c->distance > f->distance) break; @@ -1291,7 +1303,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint /* 1st phase: greedy search to insert level */ for (int lc = entryLevel; lc >= level + 1; lc--) { - w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement); + w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL); ep = w; } @@ -1309,7 +1321,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint List *neighbors; List *lw; - w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement); + w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL); /* Elements being deleted or skipped can help with search */ /* but should be removed before selecting neighbors */