mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-22 03:57:34 +08:00
Compare commits
1 Commits
hnsw-suppo
...
hnsw-datum
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c8e4bf1c3c |
26
src/hnsw.h
26
src/hnsw.h
@@ -234,12 +234,6 @@ typedef struct HnswAllocator
|
||||
void *state;
|
||||
} HnswAllocator;
|
||||
|
||||
typedef struct HnswSupport
|
||||
{
|
||||
FmgrInfo *procinfo;
|
||||
Oid collation;
|
||||
} HnswSupport;
|
||||
|
||||
typedef struct HnswBuildState
|
||||
{
|
||||
/* Info */
|
||||
@@ -258,7 +252,7 @@ typedef struct HnswBuildState
|
||||
double reltuples;
|
||||
|
||||
/* Support functions */
|
||||
HnswSupport support;
|
||||
FmgrInfo *procinfo;
|
||||
FmgrInfo *normprocinfo;
|
||||
Oid collation;
|
||||
|
||||
@@ -335,7 +329,7 @@ typedef struct HnswScanOpaqueData
|
||||
MemoryContext tmpCtx;
|
||||
|
||||
/* Support functions */
|
||||
HnswSupport support;
|
||||
FmgrInfo *procinfo;
|
||||
FmgrInfo *normprocinfo;
|
||||
Oid collation;
|
||||
} HnswScanOpaqueData;
|
||||
@@ -355,7 +349,8 @@ typedef struct HnswVacuumState
|
||||
int efConstruction;
|
||||
|
||||
/* Support functions */
|
||||
HnswSupport support;
|
||||
FmgrInfo *procinfo;
|
||||
Oid collation;
|
||||
|
||||
/* Variables */
|
||||
struct tidhash_hash *deleted;
|
||||
@@ -371,29 +366,28 @@ typedef struct HnswVacuumState
|
||||
int HnswGetM(Relation index);
|
||||
int HnswGetEfConstruction(Relation index);
|
||||
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum);
|
||||
void HnswInitSupport(HnswSupport * support, Relation index);
|
||||
bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
|
||||
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, HnswSupport * support, 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);
|
||||
HnswElement HnswGetEntryPoint(Relation index);
|
||||
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
|
||||
void *HnswAlloc(HnswAllocator * allocator, Size size);
|
||||
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
|
||||
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
|
||||
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, HnswSupport * support, bool loadVec);
|
||||
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);
|
||||
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 HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
|
||||
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);
|
||||
void HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, 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 HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, HnswSupport * support, bool loadVec);
|
||||
void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
|
||||
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
|
||||
void HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, HnswSupport * support);
|
||||
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);
|
||||
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
|
||||
|
||||
|
||||
@@ -373,7 +373,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
|
||||
* Update neighbors
|
||||
*/
|
||||
static void
|
||||
UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
|
||||
UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswElement e, int m)
|
||||
{
|
||||
for (int lc = e->level; lc >= 0; lc--)
|
||||
{
|
||||
@@ -390,7 +390,7 @@ UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
|
||||
|
||||
/* Use element for lock instead of hc since hc can be replaced */
|
||||
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
|
||||
HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, support);
|
||||
HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, procinfo, collation);
|
||||
LWLockRelease(&neighborElement->lock);
|
||||
}
|
||||
}
|
||||
@@ -400,7 +400,7 @@ UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
|
||||
* Update graph in memory
|
||||
*/
|
||||
static void
|
||||
UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
|
||||
UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
|
||||
{
|
||||
HnswGraph *graph = buildstate->graph;
|
||||
char *base = buildstate->hnswarea;
|
||||
@@ -413,7 +413,7 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
|
||||
AddElementInMemory(base, graph, element);
|
||||
|
||||
/* Update neighbors */
|
||||
UpdateNeighborsInMemory(base, support, element, m);
|
||||
UpdateNeighborsInMemory(base, procinfo, collation, element, m);
|
||||
|
||||
/* Update entry point if needed (already have lock) */
|
||||
if (entryPoint == NULL || element->level > entryPoint->level)
|
||||
@@ -426,6 +426,8 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
|
||||
static void
|
||||
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
|
||||
{
|
||||
FmgrInfo *procinfo = buildstate->procinfo;
|
||||
Oid collation = buildstate->collation;
|
||||
HnswGraph *graph = buildstate->graph;
|
||||
HnswElement entryPoint;
|
||||
LWLock *entryLock = &graph->entryLock;
|
||||
@@ -451,10 +453,10 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
|
||||
}
|
||||
|
||||
/* Find neighbors for element */
|
||||
HnswFindElementNeighbors(base, element, entryPoint, NULL, &buildstate->support, m, efConstruction, false);
|
||||
HnswFindElementNeighbors(base, element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
|
||||
|
||||
/* Update graph in memory */
|
||||
UpdateGraphInMemory(&buildstate->support, element, m, efConstruction, entryPoint, buildstate);
|
||||
UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate);
|
||||
|
||||
/* Release entry lock */
|
||||
LWLockRelease(entryLock);
|
||||
@@ -667,10 +669,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
|
||||
buildstate->efConstruction = HnswGetEfConstruction(index);
|
||||
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
|
||||
|
||||
/* Require column to have dimensions to be indexed */
|
||||
if (buildstate->dimensions < 0)
|
||||
elog(ERROR, "column does not have dimensions");
|
||||
|
||||
if (buildstate->dimensions > HNSW_MAX_DIM)
|
||||
elog(ERROR, "column cannot have more than %d dimensions for hnsw index", HNSW_MAX_DIM);
|
||||
|
||||
@@ -681,7 +679,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
|
||||
buildstate->indtuples = 0;
|
||||
|
||||
/* Get support functions */
|
||||
HnswInitSupport(&buildstate->support, index);
|
||||
buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
||||
buildstate->collation = index->rd_indcollation[0];
|
||||
|
||||
|
||||
@@ -340,7 +340,7 @@ ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
|
||||
* Update neighbors
|
||||
*/
|
||||
void
|
||||
HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building)
|
||||
HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building)
|
||||
{
|
||||
char *base = NULL;
|
||||
|
||||
@@ -373,7 +373,7 @@ HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e,
|
||||
*/
|
||||
|
||||
/* Select neighbors */
|
||||
HnswUpdateConnection(NULL, e, hc, lm, lc, &idx, index, support);
|
||||
HnswUpdateConnection(NULL, e, hc, lm, lc, &idx, index, procinfo, collation);
|
||||
|
||||
/* New element was not selected as a neighbor */
|
||||
if (idx == -1)
|
||||
@@ -527,7 +527,7 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
|
||||
* Update graph on disk
|
||||
*/
|
||||
static void
|
||||
UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
|
||||
UpdateGraphOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
|
||||
{
|
||||
BlockNumber newInsertPage = InvalidBlockNumber;
|
||||
|
||||
@@ -543,7 +543,7 @@ UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, in
|
||||
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
|
||||
|
||||
/* Update neighbors */
|
||||
HnswUpdateNeighborsOnDisk(index, support, element, m, false, building);
|
||||
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, false, building);
|
||||
|
||||
/* Update entry point if needed */
|
||||
if (entryPoint == NULL || element->level > entryPoint->level)
|
||||
@@ -560,11 +560,10 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
|
||||
HnswElement element;
|
||||
int m;
|
||||
int efConstruction = HnswGetEfConstruction(index);
|
||||
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
Oid collation = index->rd_indcollation[0];
|
||||
LOCKMODE lockmode = ShareLock;
|
||||
char *base = NULL;
|
||||
HnswSupport support;
|
||||
|
||||
HnswInitSupport(&support, index);
|
||||
|
||||
/*
|
||||
* Get a shared lock. This allows vacuum to ensure no in-flight inserts
|
||||
@@ -595,10 +594,10 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
|
||||
}
|
||||
|
||||
/* Find neighbors for element */
|
||||
HnswFindElementNeighbors(base, element, entryPoint, index, &support, m, efConstruction, false);
|
||||
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false);
|
||||
|
||||
/* Update graph on disk */
|
||||
UpdateGraphOnDisk(index, &support, element, m, efConstruction, entryPoint, building);
|
||||
UpdateGraphOnDisk(index, procinfo, collation, element, m, efConstruction, entryPoint, building);
|
||||
|
||||
/* Release lock */
|
||||
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
|
||||
|
||||
@@ -15,6 +15,8 @@ GetScanItems(IndexScanDesc scan, Datum q)
|
||||
{
|
||||
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
|
||||
Relation index = scan->indexRelation;
|
||||
FmgrInfo *procinfo = so->procinfo;
|
||||
Oid collation = so->collation;
|
||||
List *ep;
|
||||
List *w;
|
||||
int m;
|
||||
@@ -27,15 +29,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
|
||||
if (entryPoint == NULL)
|
||||
return NIL;
|
||||
|
||||
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, &so->support, false));
|
||||
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, false));
|
||||
|
||||
for (int lc = entryPoint->level; lc >= 1; lc--)
|
||||
{
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, &so->support, m, false, NULL);
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL);
|
||||
ep = w;
|
||||
}
|
||||
|
||||
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, &so->support, m, false, NULL);
|
||||
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -106,7 +108,7 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
|
||||
ALLOCSET_DEFAULT_SIZES);
|
||||
|
||||
/* Set support functions */
|
||||
HnswInitSupport(&so->support, index);
|
||||
so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
||||
so->collation = index->rd_indcollation[0];
|
||||
|
||||
|
||||
@@ -149,16 +149,6 @@ HnswOptionalProcInfo(Relation index, uint16 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
|
||||
*
|
||||
@@ -565,7 +555,7 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
|
||||
* Load an element and optionally get its distance from q
|
||||
*/
|
||||
void
|
||||
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, HnswSupport * support, bool loadVec)
|
||||
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
|
||||
{
|
||||
Buffer buf;
|
||||
Page page;
|
||||
@@ -585,7 +575,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
|
||||
|
||||
/* Calculate distance */
|
||||
if (distance != NULL)
|
||||
*distance = (float) DatumGetFloat8(FunctionCall2Coll(support->procinfo, support->collation, *q, PointerGetDatum(&etup->data)));
|
||||
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->data)));
|
||||
|
||||
UnlockReleaseBuffer(buf);
|
||||
}
|
||||
@@ -594,27 +584,27 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
|
||||
* Get the distance for a candidate
|
||||
*/
|
||||
static float
|
||||
GetCandidateDistance(char *base, HnswCandidate * hc, Datum q, HnswSupport * support)
|
||||
GetCandidateDistance(char *base, HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
|
||||
{
|
||||
HnswElement hce = HnswPtrAccess(base, hc->element);
|
||||
Datum value = HnswGetValue(base, hce);
|
||||
|
||||
return DatumGetFloat8(FunctionCall2Coll(support->procinfo, support->collation, q, value));
|
||||
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, value));
|
||||
}
|
||||
|
||||
/*
|
||||
* Create a candidate for the entry point
|
||||
*/
|
||||
HnswCandidate *
|
||||
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, HnswSupport * support, bool loadVec)
|
||||
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
|
||||
{
|
||||
HnswCandidate *hc = palloc(sizeof(HnswCandidate));
|
||||
|
||||
HnswPtrStore(base, hc->element, entryPoint);
|
||||
if (index == NULL)
|
||||
hc->distance = GetCandidateDistance(base, hc, q, support);
|
||||
hc->distance = GetCandidateDistance(base, hc, q, procinfo, collation);
|
||||
else
|
||||
HnswLoadElement(entryPoint, &hc->distance, &q, index, support, loadVec);
|
||||
HnswLoadElement(entryPoint, &hc->distance, &q, index, procinfo, collation, loadVec);
|
||||
return hc;
|
||||
}
|
||||
|
||||
@@ -732,7 +722,7 @@ CountElement(char *base, HnswElement skipElement, HnswCandidate * hc)
|
||||
* Algorithm 2 from paper
|
||||
*/
|
||||
List *
|
||||
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, HnswSupport * support, 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)
|
||||
{
|
||||
List *w = NIL;
|
||||
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
|
||||
@@ -814,9 +804,9 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, H
|
||||
f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
|
||||
|
||||
if (index == NULL)
|
||||
eDistance = GetCandidateDistance(base, e, q, support);
|
||||
eDistance = GetCandidateDistance(base, e, q, procinfo, collation);
|
||||
else
|
||||
HnswLoadElement(eElement, &eDistance, &q, index, support, inserting);
|
||||
HnswLoadElement(eElement, &eDistance, &q, index, procinfo, collation, inserting);
|
||||
|
||||
Assert(!eElement->deleted);
|
||||
|
||||
@@ -924,19 +914,19 @@ CompareCandidateDistancesOffset(const void *a, const void *b)
|
||||
* Calculate the distance between elements
|
||||
*/
|
||||
static float
|
||||
HnswGetDistance(char *base, HnswElement a, HnswElement b, HnswSupport * support)
|
||||
HnswGetDistance(char *base, HnswElement a, HnswElement b, FmgrInfo *procinfo, Oid collation)
|
||||
{
|
||||
Datum aValue = HnswGetValue(base, a);
|
||||
Datum bValue = HnswGetValue(base, b);
|
||||
|
||||
return DatumGetFloat8(FunctionCall2Coll(support->procinfo, support->collation, aValue, bValue));
|
||||
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, aValue, bValue));
|
||||
}
|
||||
|
||||
/*
|
||||
* Check if an element is closer to q than any element from R
|
||||
*/
|
||||
static bool
|
||||
CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support)
|
||||
CheckElementCloser(char *base, HnswCandidate * e, List *r, FmgrInfo *procinfo, Oid collation)
|
||||
{
|
||||
HnswElement eElement = HnswPtrAccess(base, e->element);
|
||||
ListCell *lc2;
|
||||
@@ -945,7 +935,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support
|
||||
{
|
||||
HnswCandidate *ri = lfirst(lc2);
|
||||
HnswElement riElement = HnswPtrAccess(base, ri->element);
|
||||
float distance = HnswGetDistance(base, eElement, riElement, support);
|
||||
float distance = HnswGetDistance(base, eElement, riElement, procinfo, collation);
|
||||
|
||||
if (distance <= e->distance)
|
||||
return false;
|
||||
@@ -958,7 +948,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support
|
||||
* Algorithm 4 from paper
|
||||
*/
|
||||
static List *
|
||||
SelectNeighbors(char *base, List *c, int lm, int lc, HnswSupport * support, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
|
||||
SelectNeighbors(char *base, List *c, int lm, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
|
||||
{
|
||||
List *r = NIL;
|
||||
List *w = list_copy(c);
|
||||
@@ -991,7 +981,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, HnswSupport * support, Hnsw
|
||||
|
||||
/* Use previous state of r and wd to skip work when possible */
|
||||
if (mustCalculate)
|
||||
e->closer = CheckElementCloser(base, e, r, support);
|
||||
e->closer = CheckElementCloser(base, e, r, procinfo, collation);
|
||||
else if (list_length(added) > 0)
|
||||
{
|
||||
/* Keep Valgrind happy for in-memory, parallel builds */
|
||||
@@ -1004,7 +994,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, HnswSupport * support, Hnsw
|
||||
*/
|
||||
if (e->closer)
|
||||
{
|
||||
e->closer = CheckElementCloser(base, e, added, support);
|
||||
e->closer = CheckElementCloser(base, e, added, procinfo, collation);
|
||||
|
||||
if (!e->closer)
|
||||
removedAny = true;
|
||||
@@ -1017,7 +1007,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, HnswSupport * support, Hnsw
|
||||
*/
|
||||
if (removedAny)
|
||||
{
|
||||
e->closer = CheckElementCloser(base, e, r, support);
|
||||
e->closer = CheckElementCloser(base, e, r, procinfo, collation);
|
||||
if (e->closer)
|
||||
added = lappend(added, e);
|
||||
}
|
||||
@@ -1025,7 +1015,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, HnswSupport * support, Hnsw
|
||||
}
|
||||
else if (e == newCandidate)
|
||||
{
|
||||
e->closer = CheckElementCloser(base, e, r, support);
|
||||
e->closer = CheckElementCloser(base, e, r, procinfo, collation);
|
||||
if (e->closer)
|
||||
added = lappend(added, e);
|
||||
}
|
||||
@@ -1076,7 +1066,7 @@ AddConnections(char *base, HnswElement element, List *neighbors, int lc)
|
||||
* Update connections
|
||||
*/
|
||||
void
|
||||
HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, HnswSupport * support)
|
||||
HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
|
||||
{
|
||||
HnswElement hce = HnswPtrAccess(base, hc->element);
|
||||
HnswNeighborArray *currentNeighbors = HnswGetNeighbors(base, hce, lc);
|
||||
@@ -1109,9 +1099,9 @@ HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm
|
||||
HnswElement hc3Element = HnswPtrAccess(base, hc3->element);
|
||||
|
||||
if (HnswPtrIsNull(base, hc3Element->value))
|
||||
HnswLoadElement(hc3Element, &hc3->distance, &q, index, support, true);
|
||||
HnswLoadElement(hc3Element, &hc3->distance, &q, index, procinfo, collation, true);
|
||||
else
|
||||
hc3->distance = GetCandidateDistance(base, hc3, q, support);
|
||||
hc3->distance = GetCandidateDistance(base, hc3, q, procinfo, collation);
|
||||
|
||||
/* Prune element if being deleted */
|
||||
if (hc3Element->heaptidsLength == 0)
|
||||
@@ -1131,7 +1121,7 @@ HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm
|
||||
c = lappend(c, ¤tNeighbors->items[i]);
|
||||
c = lappend(c, &hc2);
|
||||
|
||||
SelectNeighbors(base, c, lm, lc, support, hce, &hc2, &pruned, true);
|
||||
SelectNeighbors(base, c, lm, lc, procinfo, collation, hce, &hc2, &pruned, true);
|
||||
|
||||
/* Should not happen */
|
||||
if (pruned == NULL)
|
||||
@@ -1205,7 +1195,7 @@ PrecomputeHash(char *base, HnswElement element)
|
||||
* Algorithm 1 from paper
|
||||
*/
|
||||
void
|
||||
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing)
|
||||
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing)
|
||||
{
|
||||
List *ep;
|
||||
List *w;
|
||||
@@ -1225,13 +1215,13 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
||||
return;
|
||||
|
||||
/* Get entry point and level */
|
||||
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, true));
|
||||
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, true));
|
||||
entryLevel = entryPoint->level;
|
||||
|
||||
/* 1st phase: greedy search to insert level */
|
||||
for (int lc = entryLevel; lc >= level + 1; lc--)
|
||||
{
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, true, skipElement);
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement);
|
||||
ep = w;
|
||||
}
|
||||
|
||||
@@ -1249,7 +1239,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
||||
List *neighbors;
|
||||
List *lw;
|
||||
|
||||
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, support, 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 */
|
||||
/* but should be removed before selecting neighbors */
|
||||
@@ -1263,7 +1253,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
||||
* sortCandidates to true for in-memory builds to enable closer
|
||||
* caching, but there does not seem to be a difference in performance.
|
||||
*/
|
||||
neighbors = SelectNeighbors(base, lw, lm, lc, support, element, NULL, NULL, false);
|
||||
neighbors = SelectNeighbors(base, lw, lm, lc, procinfo, collation, element, NULL, NULL, false);
|
||||
|
||||
AddConnections(base, element, neighbors, lc);
|
||||
|
||||
|
||||
@@ -189,6 +189,8 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
|
||||
GenericXLogState *state;
|
||||
int m = vacuumstate->m;
|
||||
int efConstruction = vacuumstate->efConstruction;
|
||||
FmgrInfo *procinfo = vacuumstate->procinfo;
|
||||
Oid collation = vacuumstate->collation;
|
||||
BufferAccessStrategy bas = vacuumstate->bas;
|
||||
HnswNeighborTuple ntup = vacuumstate->ntup;
|
||||
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
|
||||
@@ -203,7 +205,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
|
||||
element->heaptidsLength = 0;
|
||||
|
||||
/* Find neighbors for element, skipping itself */
|
||||
HnswFindElementNeighbors(base, element, entryPoint, index, &vacuumstate->support, m, efConstruction, true);
|
||||
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, true);
|
||||
|
||||
/* Zero memory for each element */
|
||||
MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);
|
||||
@@ -227,7 +229,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
|
||||
UnlockReleaseBuffer(buf);
|
||||
|
||||
/* Update neighbors */
|
||||
HnswUpdateNeighborsOnDisk(index, &vacuumstate->support, element, m, true, false);
|
||||
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, true, false);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -254,7 +256,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
|
||||
LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
|
||||
|
||||
/* Load element */
|
||||
HnswLoadElement(highestPoint, NULL, NULL, index, &vacuumstate->support, true);
|
||||
HnswLoadElement(highestPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
|
||||
|
||||
/* Repair if needed */
|
||||
if (NeedsUpdated(vacuumstate, highestPoint))
|
||||
@@ -292,7 +294,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
|
||||
* is outdated, this can remove connections at higher levels in
|
||||
* the graph until they are repaired, but this should be fine.
|
||||
*/
|
||||
HnswLoadElement(entryPoint, NULL, NULL, index, &vacuumstate->support, true);
|
||||
HnswLoadElement(entryPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
|
||||
|
||||
if (NeedsUpdated(vacuumstate, entryPoint))
|
||||
{
|
||||
@@ -571,8 +573,8 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
|
||||
vacuumstate->callback_state = callback_state;
|
||||
vacuumstate->efConstruction = HnswGetEfConstruction(index);
|
||||
vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD);
|
||||
vacuumstate->support.procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
vacuumstate->support.collation = index->rd_indcollation[0];
|
||||
vacuumstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
vacuumstate->collation = index->rd_indcollation[0];
|
||||
vacuumstate->ntup = palloc0(HNSW_TUPLE_ALLOC_SIZE);
|
||||
vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
|
||||
"Hnsw vacuum temporary context",
|
||||
|
||||
113
test/t/020_hnsw_array.pl
Normal file
113
test/t/020_hnsw_array.pl
Normal file
@@ -0,0 +1,113 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
my $node;
|
||||
my @queries = ();
|
||||
my @expected;
|
||||
my $limit = 20;
|
||||
|
||||
sub test_recall
|
||||
{
|
||||
my ($min, $operator) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v $operator '$queries[0]' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan/);
|
||||
|
||||
for my $i (0 .. $#queries)
|
||||
{
|
||||
my $actual = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
|
||||
));
|
||||
my @actual_ids = split("\n", $actual);
|
||||
my %actual_set = map { $_ => 1 } @actual_ids;
|
||||
|
||||
my @expected_ids = split("\n", $expected[$i]);
|
||||
|
||||
foreach (@expected_ids)
|
||||
{
|
||||
if (exists($actual_set{$_}))
|
||||
{
|
||||
$correct++;
|
||||
}
|
||||
$total++;
|
||||
}
|
||||
}
|
||||
|
||||
cmp_ok($correct / $total, ">=", $min, $operator);
|
||||
}
|
||||
|
||||
# Initialize node
|
||||
$node = get_new_node('node');
|
||||
$node->init;
|
||||
$node->start;
|
||||
|
||||
# Create table
|
||||
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
|
||||
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v float4[3]);");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 10000) i;"
|
||||
);
|
||||
|
||||
$node->safe_psql("postgres", qq(
|
||||
CREATE FUNCTION float4_l2_distance(float4[], float4[]) RETURNS float8
|
||||
AS 'BEGIN RETURN l2_distance(\$1::vector, \$2::vector); END;'
|
||||
LANGUAGE plpgsql IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION float4_l2_squared_distance(float4[], float4[]) RETURNS float8
|
||||
AS 'BEGIN RETURN vector_l2_squared_distance(\$1::vector, \$2::vector); END;'
|
||||
LANGUAGE plpgsql IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE OPERATOR <-> (
|
||||
LEFTARG = float4[], RIGHTARG = float4[], PROCEDURE = float4_l2_distance,
|
||||
COMMUTATOR = '<->'
|
||||
);
|
||||
|
||||
CREATE OPERATOR CLASS float4_l2_ops
|
||||
FOR TYPE float4[] USING hnsw AS
|
||||
OPERATOR 1 <-> (float4[], float4[]) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 float4_l2_squared_distance(float4[], float4[]);
|
||||
));
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my $r1 = rand();
|
||||
my $r2 = rand();
|
||||
my $r3 = rand();
|
||||
push(@queries, "{$r1,$r2,$r3}");
|
||||
}
|
||||
|
||||
# Check each index type
|
||||
my @operators = ("<->");
|
||||
my @opclasses = ("float4_l2_ops");
|
||||
|
||||
for my $i (0 .. $#operators)
|
||||
{
|
||||
my $operator = $operators[$i];
|
||||
my $opclass = $opclasses[$i];
|
||||
|
||||
# Get exact results
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
# Add index
|
||||
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v $opclass);");
|
||||
|
||||
my $min = $operator eq "<#>" ? 0.80 : 0.99;
|
||||
test_recall($min, $operator);
|
||||
}
|
||||
|
||||
done_testing();
|
||||
Reference in New Issue
Block a user