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

Added time budget for HNSW index scans

Browse Source
This commit is contained in:
Andrew Kane
2024-09-20 22:06:09 -07:00
parent 6d2af6d3f9
commit 06b641407f
4 changed files with 29 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/hnsw.c
View File

@@ -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");
}

4
src/hnsw.h
View File

@@ -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);

8
src/hnswscan.c
View File

@@ -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;

18
src/hnswutils.c
View File

@@ -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 */