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

Added HnswQuery struct [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-02-04 22:30:59 -08:00
parent 2716a223a6
commit 0d8cc82908
3 changed files with 33 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/hnsw.h
View File

@@ -362,6 +362,11 @@ typedef struct HnswVacuumState
MemoryContext tmpCtx; MemoryContext tmpCtx;
} HnswVacuumState; } HnswVacuumState;
typedef struct HnswQuery
{
Datum value;
} HnswQuery;
/* Methods */ /* Methods */
int HnswGetM(Relation index); int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index); int HnswGetEfConstruction(Relation index);
@@ -370,14 +375,14 @@ bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * re
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, HnswQuery * q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement);
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);
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc); HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno); HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing); void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec); HnswCandidate *HnswEntryCandidate(char *base, HnswElement em, HnswQuery * q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building); void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m); void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
@@ -385,7 +390,7 @@ void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator *
bool HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, bool building); bool HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, bool building);
void HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building); void HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec); void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec); void HnswLoadElement(HnswElement element, float *distance, HnswQuery * q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element); void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation); void HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index, int m); void HnswLoadNeighbors(HnswElement element, Relation index, int m);

11
src/hnswscan.c
View File

@@ -11,7 +11,7 @@
* Algorithm 5 from paper * Algorithm 5 from paper
*/ */
static List * static List *
GetScanItems(IndexScanDesc scan, Datum q) GetScanItems(IndexScanDesc scan, Datum value)
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation; Relation index = scan->indexRelation;
@@ -22,6 +22,9 @@ GetScanItems(IndexScanDesc scan, Datum q)
int m; int m;
HnswElement entryPoint; HnswElement entryPoint;
char *base = NULL; char *base = NULL;
HnswQuery q;
q.value = value;
/* Get m and entry point */ /* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
@@ -29,15 +32,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
if (entryPoint == NULL) if (entryPoint == NULL)
return NIL; return NIL;
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, false)); ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, procinfo, collation, false));
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);
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);
} }
/* /*

32
src/hnswutils.c
View File

@@ -555,7 +555,7 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
* Load an element and optionally get its distance from q * Load an element and optionally get its distance from q
*/ */
void void
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec) HnswLoadElement(HnswElement element, float *distance, HnswQuery * q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -575,7 +575,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
/* Calculate distance */ /* Calculate distance */
if (distance != NULL) if (distance != NULL)
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->data))); *distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q->value, PointerGetDatum(&etup->data)));
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -584,19 +584,19 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
* Get the distance for a candidate * Get the distance for a candidate
*/ */
static float static float
GetCandidateDistance(char *base, HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation) GetCandidateDistance(char *base, HnswCandidate * hc, HnswQuery * q, FmgrInfo *procinfo, Oid collation)
{ {
HnswElement hce = HnswPtrAccess(base, hc->element); HnswElement hce = HnswPtrAccess(base, hc->element);
Datum value = HnswGetValue(base, hce); Datum value = HnswGetValue(base, hce);
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, value)); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q->value, value));
} }
/* /*
* Create a candidate for the entry point * Create a candidate for the entry point
*/ */
HnswCandidate * HnswCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec) HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{ {
HnswCandidate *hc = palloc(sizeof(HnswCandidate)); HnswCandidate *hc = palloc(sizeof(HnswCandidate));
@@ -604,7 +604,7 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
if (index == NULL) if (index == NULL)
hc->distance = GetCandidateDistance(base, hc, q, procinfo, collation); hc->distance = GetCandidateDistance(base, hc, q, procinfo, collation);
else else
HnswLoadElement(entryPoint, &hc->distance, &q, index, procinfo, collation, loadVec); HnswLoadElement(entryPoint, &hc->distance, q, index, procinfo, collation, loadVec);
return hc; return hc;
} }
@@ -722,7 +722,7 @@ CountElement(char *base, HnswElement skipElement, HnswCandidate * hc)
* 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, HnswQuery * q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
{ {
List *w = NIL; List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -806,7 +806,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
if (index == NULL) if (index == NULL)
eDistance = GetCandidateDistance(base, e, q, procinfo, collation); eDistance = GetCandidateDistance(base, e, q, procinfo, collation);
else else
HnswLoadElement(eElement, &eDistance, &q, index, procinfo, collation, inserting); HnswLoadElement(eElement, &eDistance, q, index, procinfo, collation, inserting);
Assert(!eElement->deleted); Assert(!eElement->deleted);
@@ -1091,7 +1091,9 @@ HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm
/* Load elements on insert */ /* Load elements on insert */
if (index != NULL) if (index != NULL)
{ {
Datum q = HnswGetValue(base, hce); HnswQuery q;
q.value = HnswGetValue(base, hce);
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
{ {
@@ -1101,7 +1103,7 @@ HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm
if (HnswPtrIsNull(base, hc3Element->value)) if (HnswPtrIsNull(base, hc3Element->value))
HnswLoadElement(hc3Element, &hc3->distance, &q, index, procinfo, collation, true); HnswLoadElement(hc3Element, &hc3->distance, &q, index, procinfo, collation, true);
else else
hc3->distance = GetCandidateDistance(base, hc3, q, procinfo, collation); hc3->distance = GetCandidateDistance(base, hc3, &q, procinfo, collation);
/* Prune element if being deleted */ /* Prune element if being deleted */
if (hc3Element->heaptidsLength == 0) if (hc3Element->heaptidsLength == 0)
@@ -1201,9 +1203,11 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
Datum q = HnswGetValue(base, element); HnswQuery q;
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
q.value = HnswGetValue(base, element);
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
/* Precompute hash */ /* Precompute hash */
if (index == NULL) if (index == NULL)
@@ -1215,13 +1219,13 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
return; return;
/* Get entry point and level */ /* Get entry point and level */
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, true)); ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, procinfo, collation, true));
entryLevel = entryPoint->level; entryLevel = entryPoint->level;
/* 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);
ep = w; ep = w;
} }
@@ -1239,7 +1243,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);
/* 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 */