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

Added HnswSupport struct [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-02-04 23:28:32 -08:00
parent 2716a223a6
commit eb9fa5cf8c
6 changed files with 80 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
src/hnsw.h
View File

@@ -234,6 +234,12 @@ typedef struct HnswAllocator
void *state; void *state;
} HnswAllocator; } HnswAllocator;
typedef struct HnswSupport
{
FmgrInfo *procinfo;
Oid collation;
} HnswSupport;
typedef struct HnswBuildState typedef struct HnswBuildState
{ {
/* Info */ /* Info */
@@ -252,7 +258,7 @@ typedef struct HnswBuildState
double reltuples; double reltuples;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; HnswSupport support;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
Oid collation; Oid collation;
@@ -329,7 +335,7 @@ typedef struct HnswScanOpaqueData
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; HnswSupport support;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
Oid collation; Oid collation;
} HnswScanOpaqueData; } HnswScanOpaqueData;
@@ -349,8 +355,7 @@ typedef struct HnswVacuumState
int efConstruction; int efConstruction;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; HnswSupport support;
Oid collation;
/* Variables */ /* Variables */
struct tidhash_hash *deleted; struct tidhash_hash *deleted;
@@ -366,28 +371,29 @@ typedef struct HnswVacuumState
int HnswGetM(Relation index); int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index); int HnswGetEfConstruction(Relation index);
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum); FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum);
void HnswInitSupport(HnswSupport * support, Relation index);
bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result); bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
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, HnswSupport * support, 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, HnswSupport * support, 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, Datum q, Relation rel, HnswSupport * support, 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);
void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc); void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
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, HnswSupport * support, 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, Datum *q, Relation index, HnswSupport * support, 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, HnswSupport * support);
void HnswLoadNeighbors(HnswElement element, Relation index, int m); void HnswLoadNeighbors(HnswElement element, Relation index, int m);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc); PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);

16
src/hnswbuild.c
View File

@@ -373,7 +373,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
* Update neighbors * Update neighbors
*/ */
static void static void
UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswElement e, int m) UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
{ {
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
@@ -390,7 +390,7 @@ UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswEleme
/* Use element for lock instead of hc since hc can be replaced */ /* Use element for lock instead of hc since hc can be replaced */
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE); LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, procinfo, collation); HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, support);
LWLockRelease(&neighborElement->lock); LWLockRelease(&neighborElement->lock);
} }
} }
@@ -400,7 +400,7 @@ UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswEleme
* Update graph in memory * Update graph in memory
*/ */
static void static void
UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate) UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
{ {
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
@@ -413,7 +413,7 @@ UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int
AddElementInMemory(base, graph, element); AddElementInMemory(base, graph, element);
/* Update neighbors */ /* Update neighbors */
UpdateNeighborsInMemory(base, procinfo, collation, element, m); UpdateNeighborsInMemory(base, support, element, m);
/* Update entry point if needed (already have lock) */ /* Update entry point if needed (already have lock) */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -426,8 +426,6 @@ UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int
static void static void
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element) InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
{ {
FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation;
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
HnswElement entryPoint; HnswElement entryPoint;
LWLock *entryLock = &graph->entryLock; LWLock *entryLock = &graph->entryLock;
@@ -453,10 +451,10 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
} }
/* Find neighbors for element */ /* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, NULL, procinfo, collation, m, efConstruction, false); HnswFindElementNeighbors(base, element, entryPoint, NULL, &buildstate->support, m, efConstruction, false);
/* Update graph in memory */ /* Update graph in memory */
UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate); UpdateGraphInMemory(&buildstate->support, element, m, efConstruction, entryPoint, buildstate);
/* Release entry lock */ /* Release entry lock */
LWLockRelease(entryLock); LWLockRelease(entryLock);
@@ -683,7 +681,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->indtuples = 0; buildstate->indtuples = 0;
/* Get support functions */ /* Get support functions */
buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); HnswInitSupport(&buildstate->support, index);
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->collation = index->rd_indcollation[0]; buildstate->collation = index->rd_indcollation[0];

17
src/hnswinsert.c
View File

@@ -340,7 +340,7 @@ ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
* Update neighbors * Update neighbors
*/ */
void void
HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building) HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building)
{ {
char *base = NULL; char *base = NULL;
@@ -373,7 +373,7 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, Hns
*/ */
/* Select neighbors */ /* Select neighbors */
HnswUpdateConnection(NULL, e, hc, lm, lc, &idx, index, procinfo, collation); HnswUpdateConnection(NULL, e, hc, lm, lc, &idx, index, support);
/* New element was not selected as a neighbor */ /* New element was not selected as a neighbor */
if (idx == -1) if (idx == -1)
@@ -527,7 +527,7 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
* Update graph on disk * Update graph on disk
*/ */
static void static void
UpdateGraphOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building) UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
{ {
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
@@ -543,7 +543,7 @@ UpdateGraphOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building); HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, false, building); HnswUpdateNeighborsOnDisk(index, support, element, m, false, building);
/* Update entry point if needed */ /* Update entry point if needed */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -560,10 +560,11 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
HnswElement element; HnswElement element;
int m; int m;
int efConstruction = HnswGetEfConstruction(index); int efConstruction = HnswGetEfConstruction(index);
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid collation = index->rd_indcollation[0];
LOCKMODE lockmode = ShareLock; LOCKMODE lockmode = ShareLock;
char *base = NULL; char *base = NULL;
HnswSupport support;
HnswInitSupport(&support, index);
/* /*
* Get a shared lock. This allows vacuum to ensure no in-flight inserts * Get a shared lock. This allows vacuum to ensure no in-flight inserts
@@ -594,10 +595,10 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
} }
/* Find neighbors for element */ /* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false); HnswFindElementNeighbors(base, element, entryPoint, index, &support, m, efConstruction, false);
/* Update graph on disk */ /* Update graph on disk */
UpdateGraphOnDisk(index, procinfo, collation, element, m, efConstruction, entryPoint, building); UpdateGraphOnDisk(index, &support, element, m, efConstruction, entryPoint, building);
/* Release lock */ /* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode); UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);

10
src/hnswscan.c
View File

@@ -15,8 +15,6 @@ GetScanItems(IndexScanDesc scan, Datum q)
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation; Relation index = scan->indexRelation;
FmgrInfo *procinfo = so->procinfo;
Oid collation = so->collation;
List *ep; List *ep;
List *w; List *w;
int m; int m;
@@ -29,15 +27,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, &so->support, 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, &so->support, 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, &so->support, m, false, NULL);
} }
/* /*
@@ -108,7 +106,7 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
/* Set support functions */ /* Set support functions */
so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); HnswInitSupport(&so->support, index);
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];

66
src/hnswutils.c
View File

@@ -149,6 +149,16 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
return index_getprocinfo(index, 1, procnum); return index_getprocinfo(index, 1, procnum);
} }
/*
* Init support function
*/
void
HnswInitSupport(HnswSupport * support, Relation index)
{
support->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
support->collation = index->rd_indcollation[0];
}
/* /*
* Divide by the norm * Divide by the norm
* *
@@ -555,7 +565,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, Datum *q, Relation index, HnswSupport * support, bool loadVec)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -575,7 +585,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(support->procinfo, support->collation, *q, PointerGetDatum(&etup->data)));
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -584,27 +594,27 @@ 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, Datum q, HnswSupport * support)
{ {
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(support->procinfo, support->collation, q, 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, Datum q, Relation index, HnswSupport * support, bool loadVec)
{ {
HnswCandidate *hc = palloc(sizeof(HnswCandidate)); HnswCandidate *hc = palloc(sizeof(HnswCandidate));
HnswPtrStore(base, hc->element, entryPoint); HnswPtrStore(base, hc->element, entryPoint);
if (index == NULL) if (index == NULL)
hc->distance = GetCandidateDistance(base, hc, q, procinfo, collation); hc->distance = GetCandidateDistance(base, hc, q, support);
else else
HnswLoadElement(entryPoint, &hc->distance, &q, index, procinfo, collation, loadVec); HnswLoadElement(entryPoint, &hc->distance, &q, index, support, loadVec);
return hc; return hc;
} }
@@ -722,7 +732,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, Datum q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement)
{ {
List *w = NIL; List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -804,9 +814,9 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner; f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
if (index == NULL) if (index == NULL)
eDistance = GetCandidateDistance(base, e, q, procinfo, collation); eDistance = GetCandidateDistance(base, e, q, support);
else else
HnswLoadElement(eElement, &eDistance, &q, index, procinfo, collation, inserting); HnswLoadElement(eElement, &eDistance, &q, index, support, inserting);
Assert(!eElement->deleted); Assert(!eElement->deleted);
@@ -914,19 +924,19 @@ CompareCandidateDistancesOffset(const void *a, const void *b)
* Calculate the distance between elements * Calculate the distance between elements
*/ */
static float static float
HnswGetDistance(char *base, HnswElement a, HnswElement b, FmgrInfo *procinfo, Oid collation) HnswGetDistance(char *base, HnswElement a, HnswElement b, HnswSupport * support)
{ {
Datum aValue = HnswGetValue(base, a); Datum aValue = HnswGetValue(base, a);
Datum bValue = HnswGetValue(base, b); Datum bValue = HnswGetValue(base, b);
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, aValue, bValue)); return DatumGetFloat8(FunctionCall2Coll(support->procinfo, support->collation, aValue, bValue));
} }
/* /*
* Check if an element is closer to q than any element from R * Check if an element is closer to q than any element from R
*/ */
static bool static bool
CheckElementCloser(char *base, HnswCandidate * e, List *r, FmgrInfo *procinfo, Oid collation) CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support)
{ {
HnswElement eElement = HnswPtrAccess(base, e->element); HnswElement eElement = HnswPtrAccess(base, e->element);
ListCell *lc2; ListCell *lc2;
@@ -935,7 +945,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, FmgrInfo *procinfo, O
{ {
HnswCandidate *ri = lfirst(lc2); HnswCandidate *ri = lfirst(lc2);
HnswElement riElement = HnswPtrAccess(base, ri->element); HnswElement riElement = HnswPtrAccess(base, ri->element);
float distance = HnswGetDistance(base, eElement, riElement, procinfo, collation); float distance = HnswGetDistance(base, eElement, riElement, support);
if (distance <= e->distance) if (distance <= e->distance)
return false; return false;
@@ -948,7 +958,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, FmgrInfo *procinfo, O
* Algorithm 4 from paper * Algorithm 4 from paper
*/ */
static List * static List *
SelectNeighbors(char *base, List *c, int lm, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates) SelectNeighbors(char *base, List *c, int lm, int lc, HnswSupport * support, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
{ {
List *r = NIL; List *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
@@ -981,7 +991,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, FmgrInfo *procinfo, Oid col
/* Use previous state of r and wd to skip work when possible */ /* Use previous state of r and wd to skip work when possible */
if (mustCalculate) if (mustCalculate)
e->closer = CheckElementCloser(base, e, r, procinfo, collation); e->closer = CheckElementCloser(base, e, r, support);
else if (list_length(added) > 0) else if (list_length(added) > 0)
{ {
/* Keep Valgrind happy for in-memory, parallel builds */ /* Keep Valgrind happy for in-memory, parallel builds */
@@ -994,7 +1004,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, FmgrInfo *procinfo, Oid col
*/ */
if (e->closer) if (e->closer)
{ {
e->closer = CheckElementCloser(base, e, added, procinfo, collation); e->closer = CheckElementCloser(base, e, added, support);
if (!e->closer) if (!e->closer)
removedAny = true; removedAny = true;
@@ -1007,7 +1017,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, FmgrInfo *procinfo, Oid col
*/ */
if (removedAny) if (removedAny)
{ {
e->closer = CheckElementCloser(base, e, r, procinfo, collation); e->closer = CheckElementCloser(base, e, r, support);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -1015,7 +1025,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, FmgrInfo *procinfo, Oid col
} }
else if (e == newCandidate) else if (e == newCandidate)
{ {
e->closer = CheckElementCloser(base, e, r, procinfo, collation); e->closer = CheckElementCloser(base, e, r, support);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -1066,7 +1076,7 @@ AddConnections(char *base, HnswElement element, List *neighbors, int lc)
* Update connections * Update connections
*/ */
void void
HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation) HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, HnswSupport * support)
{ {
HnswElement hce = HnswPtrAccess(base, hc->element); HnswElement hce = HnswPtrAccess(base, hc->element);
HnswNeighborArray *currentNeighbors = HnswGetNeighbors(base, hce, lc); HnswNeighborArray *currentNeighbors = HnswGetNeighbors(base, hce, lc);
@@ -1099,9 +1109,9 @@ HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm
HnswElement hc3Element = HnswPtrAccess(base, hc3->element); HnswElement hc3Element = HnswPtrAccess(base, hc3->element);
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, support, true);
else else
hc3->distance = GetCandidateDistance(base, hc3, q, procinfo, collation); hc3->distance = GetCandidateDistance(base, hc3, q, support);
/* Prune element if being deleted */ /* Prune element if being deleted */
if (hc3Element->heaptidsLength == 0) if (hc3Element->heaptidsLength == 0)
@@ -1121,7 +1131,7 @@ HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm
c = lappend(c, &currentNeighbors->items[i]); c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2); c = lappend(c, &hc2);
SelectNeighbors(base, c, lm, lc, procinfo, collation, hce, &hc2, &pruned, true); SelectNeighbors(base, c, lm, lc, support, hce, &hc2, &pruned, true);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -1195,7 +1205,7 @@ PrecomputeHash(char *base, HnswElement element)
* Algorithm 1 from paper * Algorithm 1 from paper
*/ */
void void
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing) HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing)
{ {
List *ep; List *ep;
List *w; List *w;
@@ -1215,13 +1225,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, support, 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, support, m, true, skipElement);
ep = w; ep = w;
} }
@@ -1239,7 +1249,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, support, 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 */
@@ -1253,7 +1263,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
* sortCandidates to true for in-memory builds to enable closer * sortCandidates to true for in-memory builds to enable closer
* caching, but there does not seem to be a difference in performance. * caching, but there does not seem to be a difference in performance.
*/ */
neighbors = SelectNeighbors(base, lw, lm, lc, procinfo, collation, element, NULL, NULL, false); neighbors = SelectNeighbors(base, lw, lm, lc, support, element, NULL, NULL, false);
AddConnections(base, element, neighbors, lc); AddConnections(base, element, neighbors, lc);

14
src/hnswvacuum.c
View File

@@ -189,8 +189,6 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
GenericXLogState *state; GenericXLogState *state;
int m = vacuumstate->m; int m = vacuumstate->m;
int efConstruction = vacuumstate->efConstruction; int efConstruction = vacuumstate->efConstruction;
FmgrInfo *procinfo = vacuumstate->procinfo;
Oid collation = vacuumstate->collation;
BufferAccessStrategy bas = vacuumstate->bas; BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup; HnswNeighborTuple ntup = vacuumstate->ntup;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m); Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
@@ -205,7 +203,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
element->heaptidsLength = 0; element->heaptidsLength = 0;
/* Find neighbors for element, skipping itself */ /* Find neighbors for element, skipping itself */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, true); HnswFindElementNeighbors(base, element, entryPoint, index, &vacuumstate->support, m, efConstruction, true);
/* Zero memory for each element */ /* Zero memory for each element */
MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE); MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);
@@ -229,7 +227,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, true, false); HnswUpdateNeighborsOnDisk(index, &vacuumstate->support, element, m, true, false);
} }
/* /*
@@ -256,7 +254,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
LockPage(index, HNSW_UPDATE_LOCK, ShareLock); LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
/* Load element */ /* Load element */
HnswLoadElement(highestPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true); HnswLoadElement(highestPoint, NULL, NULL, index, &vacuumstate->support, true);
/* Repair if needed */ /* Repair if needed */
if (NeedsUpdated(vacuumstate, highestPoint)) if (NeedsUpdated(vacuumstate, highestPoint))
@@ -294,7 +292,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
* is outdated, this can remove connections at higher levels in * is outdated, this can remove connections at higher levels in
* the graph until they are repaired, but this should be fine. * the graph until they are repaired, but this should be fine.
*/ */
HnswLoadElement(entryPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true); HnswLoadElement(entryPoint, NULL, NULL, index, &vacuumstate->support, true);
if (NeedsUpdated(vacuumstate, entryPoint)) if (NeedsUpdated(vacuumstate, entryPoint))
{ {
@@ -573,8 +571,8 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
vacuumstate->callback_state = callback_state; vacuumstate->callback_state = callback_state;
vacuumstate->efConstruction = HnswGetEfConstruction(index); vacuumstate->efConstruction = HnswGetEfConstruction(index);
vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD); vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD);
vacuumstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); vacuumstate->support.procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
vacuumstate->collation = index->rd_indcollation[0]; vacuumstate->support.collation = index->rd_indcollation[0];
vacuumstate->ntup = palloc0(HNSW_TUPLE_ALLOC_SIZE); vacuumstate->ntup = palloc0(HNSW_TUPLE_ALLOC_SIZE);
vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw vacuum temporary context", "Hnsw vacuum temporary context",