Compare commits

..

3 Commits

Author SHA1 Message Date
Andrew Kane
dad5dffe8e Fixed issue 2024-09-20 22:29:23 -07:00
Andrew Kane
06b641407f Added time budget for HNSW index scans 2024-09-20 22:06:09 -07:00
Andrew Kane
6d2af6d3f9 Improved code [skip ci] 2024-09-20 15:21:57 -07:00
6 changed files with 41 additions and 44 deletions

View File

@@ -18,6 +18,7 @@
#endif #endif
int hnsw_ef_search; int hnsw_ef_search;
int hnsw_time_budget;
int hnsw_lock_tranche_id; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
@@ -68,6 +69,10 @@ HnswInit(void)
"Valid range is 1..1000.", &hnsw_ef_search, "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); 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"); MarkGUCPrefixReserved("hnsw");
} }

View File

@@ -106,6 +106,7 @@
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_time_budget;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef struct HnswElementData HnswElementData; typedef struct HnswElementData HnswElementData;
@@ -129,7 +130,6 @@ struct HnswElementData
uint8 heaptidsLength; uint8 heaptidsLength;
uint8 level; uint8 level;
uint8 deleted; uint8 deleted;
uint8 version;
uint32 hash; uint32 hash;
HnswNeighborsPtr neighbors; HnswNeighborsPtr neighbors;
BlockNumber blkno; BlockNumber blkno;
@@ -306,10 +306,10 @@ typedef struct HnswElementTupleData
uint8 type; uint8 type;
uint8 level; uint8 level;
uint8 deleted; uint8 deleted;
uint8 version; uint8 unused;
ItemPointerData heaptids[HNSW_HEAPTIDS]; ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid; ItemPointerData neighbortid;
uint16 unused; uint16 unused2;
Vector data; Vector data;
} HnswElementTupleData; } HnswElementTupleData;
@@ -318,7 +318,7 @@ typedef HnswElementTupleData * HnswElementTuple;
typedef struct HnswNeighborTupleData typedef struct HnswNeighborTupleData
{ {
uint8 type; uint8 type;
uint8 version; uint8 unused;
uint16 count; uint16 count;
ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER]; ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborTupleData; } HnswNeighborTupleData;
@@ -331,6 +331,7 @@ typedef struct HnswScanOpaqueData
bool first; bool first;
List *w; List *w;
MemoryContext tmpCtx; MemoryContext tmpCtx;
instr_time start;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; FmgrInfo *procinfo;
@@ -375,7 +376,7 @@ bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page); void HnswInitPage(Buffer buf, Page page);
void HnswInit(void); 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); HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint); void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
void *HnswAlloc(HnswAllocator * allocator, Size size); void *HnswAlloc(HnswAllocator * allocator, Size size);

View File

@@ -36,7 +36,7 @@ GetInsertPage(Relation index)
* Check for a free offset * Check for a free offset
*/ */
static bool static bool
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage, uint8 *tupleVersion) HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage)
{ {
OffsetNumber offno; OffsetNumber offno;
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page); OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
@@ -98,7 +98,6 @@ HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size
{ {
*freeOffno = offno; *freeOffno = offno;
*freeNeighborOffno = neighborOffno; *freeNeighborOffno = neighborOffno;
*tupleVersion = etup->version;
return true; return true;
} }
else if (*nbuf != buf) else if (*nbuf != buf)
@@ -154,7 +153,6 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
OffsetNumber freeOffno = InvalidOffsetNumber; OffsetNumber freeOffno = InvalidOffsetNumber;
OffsetNumber freeNeighborOffno = InvalidOffsetNumber; OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
uint8 tupleVersion;
char *base = NULL; char *base = NULL;
/* Calculate sizes */ /* Calculate sizes */
@@ -204,7 +202,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
} }
/* Next, try space from a deleted element */ /* Next, try space from a deleted element */
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage, &tupleVersion)) if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage))
{ {
if (nbuf != buf) if (nbuf != buf)
{ {
@@ -214,10 +212,6 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
npage = GenericXLogRegisterBuffer(state, nbuf, 0); npage = GenericXLogRegisterBuffer(state, nbuf, 0);
} }
/* Set tuple version */
etup->version = tupleVersion;
ntup->version = tupleVersion;
break; break;
} }

View File

@@ -3,6 +3,7 @@
#include "access/relscan.h" #include "access/relscan.h"
#include "hnsw.h" #include "hnsw.h"
#include "pgstat.h" #include "pgstat.h"
#include "portability/instr_time.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "storage/lmgr.h" #include "storage/lmgr.h"
#include "utils/memutils.h" #include "utils/memutils.h"
@@ -22,6 +23,7 @@ GetScanItems(IndexScanDesc scan, Datum q)
int m; int m;
HnswElement entryPoint; HnswElement entryPoint;
char *base = NULL; char *base = NULL;
instr_time *start = hnsw_time_budget == -1 ? NULL : &so->start;
/* Get m and entry point */ /* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
@@ -33,11 +35,11 @@ GetScanItems(IndexScanDesc scan, Datum q)
for (int lc = entryPoint->level; lc >= 1; lc--) 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, start);
ep = w; 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, start);
} }
/* /*
@@ -90,6 +92,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
so->collation = index->rd_indcollation[0]; so->collation = index->rd_indcollation[0];
if (hnsw_time_budget != -1)
INSTR_TIME_SET_CURRENT(so->start);
scan->opaque = so; scan->opaque = so;
return scan; return scan;

View File

@@ -9,6 +9,7 @@
#include "fmgr.h" #include "fmgr.h"
#include "hnsw.h" #include "hnsw.h"
#include "lib/pairingheap.h" #include "lib/pairingheap.h"
#include "portability/instr_time.h"
#include "sparsevec.h" #include "sparsevec.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h" #include "utils/datum.h"
@@ -253,8 +254,6 @@ HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel,
element->level = level; element->level = level;
element->deleted = 0; element->deleted = 0;
/* Start at one to make it easier to find issues */
element->version = 1;
HnswInitNeighbors(base, element, m, allocator); HnswInitNeighbors(base, element, m, allocator);
@@ -407,7 +406,6 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
etup->type = HNSW_ELEMENT_TUPLE_TYPE; etup->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level; etup->level = element->level;
etup->deleted = 0; etup->deleted = 0;
etup->version = element->version;
for (int i = 0; i < HNSW_HEAPTIDS; i++) for (int i = 0; i < HNSW_HEAPTIDS; i++)
{ {
if (i < element->heaptidsLength) if (i < element->heaptidsLength)
@@ -450,7 +448,6 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
} }
ntup->count = idx; ntup->count = idx;
ntup->version = e->version;
} }
/* /*
@@ -524,7 +521,6 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
{ {
element->level = etup->level; element->level = etup->level;
element->deleted = etup->deleted; element->deleted = etup->deleted;
element->version = etup->version;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid); element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid); element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->heaptidsLength = 0; element->heaptidsLength = 0;
@@ -771,30 +767,20 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
int start; int start;
ItemPointerData indextids[HNSW_MAX_M * 2]; ItemPointerData indextids[HNSW_MAX_M * 2];
*unvisitedLength = 0;
buf = ReadBuffer(index, element->neighborPage); buf = ReadBuffer(index, element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_SHARE); LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf); page = BufferGetPage(buf);
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno)); ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
start = (element->level - lc) * m;
/*
* Ensure the neighbor tuple has not been deleted or replaced between
* index scan iterations
*/
if (ntup->version != element->version)
{
UnlockReleaseBuffer(buf);
return;
}
/* Copy to minimize lock time */ /* Copy to minimize lock time */
start = (element->level - lc) * m;
memcpy(&indextids, ntup->indextids + start, lm * sizeof(ItemPointerData)); memcpy(&indextids, ntup->indextids + start, lm * sizeof(ItemPointerData));
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
*unvisitedLength = 0;
for (int i = 0; i < lm; i++) for (int i = 0; i < lm; i++)
{ {
ItemPointer indextid = &indextids[i]; ItemPointer indextid = &indextids[i];
@@ -814,7 +800,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
* Algorithm 2 from paper * Algorithm 2 from paper
*/ */
List * 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; List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -865,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)); HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
HnswElement cElement; HnswElement cElement;
/* Check time budget */
if (start != NULL)
{
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) if (c->distance > f->distance)
break; break;
@@ -889,9 +886,6 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
{ {
eElement = unvisited[i].element; eElement = unvisited[i].element;
eDistance = GetElementDistance(base, eElement, q, procinfo, collation); eDistance = GetElementDistance(base, eElement, q, procinfo, collation);
if (!(eDistance < f->distance || alwaysAdd))
continue;
} }
else else
{ {
@@ -907,6 +901,9 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
continue; continue;
} }
if (!(eDistance < f->distance || alwaysAdd))
continue;
Assert(!eElement->deleted); Assert(!eElement->deleted);
/* Make robust to issues */ /* Make robust to issues */
@@ -1306,7 +1303,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
/* 1st phase: greedy search to insert level */ /* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--) 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; ep = w;
} }
@@ -1324,7 +1321,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
List *neighbors; List *neighbors;
List *lw; 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 */ /* Elements being deleted or skipped can help with search */
/* but should be removed before selecting neighbors */ /* but should be removed before selecting neighbors */

View File

@@ -527,11 +527,6 @@ MarkDeleted(HnswVacuumState * vacuumstate)
for (int i = 0; i < ntup->count; i++) for (int i = 0; i < ntup->count; i++)
ItemPointerSetInvalid(&ntup->indextids[i]); ItemPointerSetInvalid(&ntup->indextids[i]);
/* Increment version */
/* This is used to avoid incorrect reads for iterative scans */
etup->version++;
ntup->version = etup->version;
/* /*
* We modified the tuples in place, no need to call * We modified the tuples in place, no need to call
* PageIndexTupleOverwrite * PageIndexTupleOverwrite