mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-22 20:15:46 +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;
|
void *state;
|
||||||
} HnswAllocator;
|
} HnswAllocator;
|
||||||
|
|
||||||
typedef struct HnswSupport
|
|
||||||
{
|
|
||||||
FmgrInfo *procinfo;
|
|
||||||
Oid collation;
|
|
||||||
} HnswSupport;
|
|
||||||
|
|
||||||
typedef struct HnswBuildState
|
typedef struct HnswBuildState
|
||||||
{
|
{
|
||||||
/* Info */
|
/* Info */
|
||||||
@@ -258,7 +252,7 @@ typedef struct HnswBuildState
|
|||||||
double reltuples;
|
double reltuples;
|
||||||
|
|
||||||
/* Support functions */
|
/* Support functions */
|
||||||
HnswSupport support;
|
FmgrInfo *procinfo;
|
||||||
FmgrInfo *normprocinfo;
|
FmgrInfo *normprocinfo;
|
||||||
Oid collation;
|
Oid collation;
|
||||||
|
|
||||||
@@ -335,7 +329,7 @@ typedef struct HnswScanOpaqueData
|
|||||||
MemoryContext tmpCtx;
|
MemoryContext tmpCtx;
|
||||||
|
|
||||||
/* Support functions */
|
/* Support functions */
|
||||||
HnswSupport support;
|
FmgrInfo *procinfo;
|
||||||
FmgrInfo *normprocinfo;
|
FmgrInfo *normprocinfo;
|
||||||
Oid collation;
|
Oid collation;
|
||||||
} HnswScanOpaqueData;
|
} HnswScanOpaqueData;
|
||||||
@@ -355,7 +349,8 @@ typedef struct HnswVacuumState
|
|||||||
int efConstruction;
|
int efConstruction;
|
||||||
|
|
||||||
/* Support functions */
|
/* Support functions */
|
||||||
HnswSupport support;
|
FmgrInfo *procinfo;
|
||||||
|
Oid collation;
|
||||||
|
|
||||||
/* Variables */
|
/* Variables */
|
||||||
struct tidhash_hash *deleted;
|
struct tidhash_hash *deleted;
|
||||||
@@ -371,29 +366,28 @@ 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, 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);
|
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, HnswSupport * support, 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, HnswSupport * support, bool loadVec);
|
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 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, 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 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 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);
|
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);
|
||||||
|
|
||||||
|
|||||||
@@ -373,7 +373,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
|
|||||||
* Update neighbors
|
* Update neighbors
|
||||||
*/
|
*/
|
||||||
static void
|
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--)
|
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 */
|
/* 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, support);
|
HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, procinfo, collation);
|
||||||
LWLockRelease(&neighborElement->lock);
|
LWLockRelease(&neighborElement->lock);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -400,7 +400,7 @@ UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
|
|||||||
* Update graph in memory
|
* Update graph in memory
|
||||||
*/
|
*/
|
||||||
static void
|
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;
|
HnswGraph *graph = buildstate->graph;
|
||||||
char *base = buildstate->hnswarea;
|
char *base = buildstate->hnswarea;
|
||||||
@@ -413,7 +413,7 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
|
|||||||
AddElementInMemory(base, graph, element);
|
AddElementInMemory(base, graph, element);
|
||||||
|
|
||||||
/* Update neighbors */
|
/* Update neighbors */
|
||||||
UpdateNeighborsInMemory(base, support, element, m);
|
UpdateNeighborsInMemory(base, procinfo, collation, 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,6 +426,8 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
|
|||||||
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;
|
||||||
@@ -451,10 +453,10 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
|
|||||||
}
|
}
|
||||||
|
|
||||||
/* Find neighbors for 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 */
|
/* Update graph in memory */
|
||||||
UpdateGraphInMemory(&buildstate->support, element, m, efConstruction, entryPoint, buildstate);
|
UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate);
|
||||||
|
|
||||||
/* Release entry lock */
|
/* Release entry lock */
|
||||||
LWLockRelease(entryLock);
|
LWLockRelease(entryLock);
|
||||||
@@ -667,10 +669,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
|
|||||||
buildstate->efConstruction = HnswGetEfConstruction(index);
|
buildstate->efConstruction = HnswGetEfConstruction(index);
|
||||||
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
|
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)
|
if (buildstate->dimensions > HNSW_MAX_DIM)
|
||||||
elog(ERROR, "column cannot have more than %d dimensions for hnsw index", 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;
|
buildstate->indtuples = 0;
|
||||||
|
|
||||||
/* Get support functions */
|
/* Get support functions */
|
||||||
HnswInitSupport(&buildstate->support, index);
|
buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||||
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];
|
||||||
|
|
||||||
|
|||||||
@@ -340,7 +340,7 @@ ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
|
|||||||
* Update neighbors
|
* Update neighbors
|
||||||
*/
|
*/
|
||||||
void
|
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;
|
char *base = NULL;
|
||||||
|
|
||||||
@@ -373,7 +373,7 @@ HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e,
|
|||||||
*/
|
*/
|
||||||
|
|
||||||
/* Select neighbors */
|
/* 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 */
|
/* 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, 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;
|
BlockNumber newInsertPage = InvalidBlockNumber;
|
||||||
|
|
||||||
@@ -543,7 +543,7 @@ UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, in
|
|||||||
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
|
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
|
||||||
|
|
||||||
/* Update neighbors */
|
/* Update neighbors */
|
||||||
HnswUpdateNeighborsOnDisk(index, support, element, m, false, building);
|
HnswUpdateNeighborsOnDisk(index, procinfo, collation, 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,11 +560,10 @@ 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
|
||||||
@@ -595,10 +594,10 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
|
|||||||
}
|
}
|
||||||
|
|
||||||
/* Find neighbors for element */
|
/* 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 */
|
/* Update graph on disk */
|
||||||
UpdateGraphOnDisk(index, &support, element, m, efConstruction, entryPoint, building);
|
UpdateGraphOnDisk(index, procinfo, collation, element, m, efConstruction, entryPoint, building);
|
||||||
|
|
||||||
/* Release lock */
|
/* Release lock */
|
||||||
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
|
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
|
||||||
|
|||||||
@@ -15,6 +15,8 @@ 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;
|
||||||
@@ -27,15 +29,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
|
|||||||
if (entryPoint == NULL)
|
if (entryPoint == NULL)
|
||||||
return NIL;
|
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--)
|
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;
|
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);
|
ALLOCSET_DEFAULT_SIZES);
|
||||||
|
|
||||||
/* Set support functions */
|
/* Set support functions */
|
||||||
HnswInitSupport(&so->support, index);
|
so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||||
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];
|
||||||
|
|
||||||
|
|||||||
@@ -149,16 +149,6 @@ 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
|
||||||
*
|
*
|
||||||
@@ -565,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, HnswSupport * support, bool loadVec)
|
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
|
||||||
{
|
{
|
||||||
Buffer buf;
|
Buffer buf;
|
||||||
Page page;
|
Page page;
|
||||||
@@ -585,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(support->procinfo, support->collation, *q, PointerGetDatum(&etup->data)));
|
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->data)));
|
||||||
|
|
||||||
UnlockReleaseBuffer(buf);
|
UnlockReleaseBuffer(buf);
|
||||||
}
|
}
|
||||||
@@ -594,27 +584,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, HnswSupport * support)
|
GetCandidateDistance(char *base, HnswCandidate * hc, Datum 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(support->procinfo, support->collation, q, value));
|
return DatumGetFloat8(FunctionCall2Coll(procinfo, 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, HnswSupport * support, bool loadVec)
|
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, 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, support);
|
hc->distance = GetCandidateDistance(base, hc, q, procinfo, collation);
|
||||||
else
|
else
|
||||||
HnswLoadElement(entryPoint, &hc->distance, &q, index, support, loadVec);
|
HnswLoadElement(entryPoint, &hc->distance, &q, index, procinfo, collation, loadVec);
|
||||||
return hc;
|
return hc;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -732,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, 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;
|
List *w = NIL;
|
||||||
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
|
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;
|
f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
|
||||||
|
|
||||||
if (index == NULL)
|
if (index == NULL)
|
||||||
eDistance = GetCandidateDistance(base, e, q, support);
|
eDistance = GetCandidateDistance(base, e, q, procinfo, collation);
|
||||||
else
|
else
|
||||||
HnswLoadElement(eElement, &eDistance, &q, index, support, inserting);
|
HnswLoadElement(eElement, &eDistance, &q, index, procinfo, collation, inserting);
|
||||||
|
|
||||||
Assert(!eElement->deleted);
|
Assert(!eElement->deleted);
|
||||||
|
|
||||||
@@ -924,19 +914,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, HnswSupport * support)
|
HnswGetDistance(char *base, HnswElement a, HnswElement b, FmgrInfo *procinfo, Oid collation)
|
||||||
{
|
{
|
||||||
Datum aValue = HnswGetValue(base, a);
|
Datum aValue = HnswGetValue(base, a);
|
||||||
Datum bValue = HnswGetValue(base, b);
|
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
|
* Check if an element is closer to q than any element from R
|
||||||
*/
|
*/
|
||||||
static bool
|
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);
|
HnswElement eElement = HnswPtrAccess(base, e->element);
|
||||||
ListCell *lc2;
|
ListCell *lc2;
|
||||||
@@ -945,7 +935,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support
|
|||||||
{
|
{
|
||||||
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, support);
|
float distance = HnswGetDistance(base, eElement, riElement, procinfo, collation);
|
||||||
|
|
||||||
if (distance <= e->distance)
|
if (distance <= e->distance)
|
||||||
return false;
|
return false;
|
||||||
@@ -958,7 +948,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support
|
|||||||
* Algorithm 4 from paper
|
* Algorithm 4 from paper
|
||||||
*/
|
*/
|
||||||
static List *
|
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 *r = NIL;
|
||||||
List *w = list_copy(c);
|
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 */
|
/* Use previous state of r and wd to skip work when possible */
|
||||||
if (mustCalculate)
|
if (mustCalculate)
|
||||||
e->closer = CheckElementCloser(base, e, r, support);
|
e->closer = CheckElementCloser(base, e, r, procinfo, collation);
|
||||||
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 */
|
||||||
@@ -1004,7 +994,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, HnswSupport * support, Hnsw
|
|||||||
*/
|
*/
|
||||||
if (e->closer)
|
if (e->closer)
|
||||||
{
|
{
|
||||||
e->closer = CheckElementCloser(base, e, added, support);
|
e->closer = CheckElementCloser(base, e, added, procinfo, collation);
|
||||||
|
|
||||||
if (!e->closer)
|
if (!e->closer)
|
||||||
removedAny = true;
|
removedAny = true;
|
||||||
@@ -1017,7 +1007,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, HnswSupport * support, Hnsw
|
|||||||
*/
|
*/
|
||||||
if (removedAny)
|
if (removedAny)
|
||||||
{
|
{
|
||||||
e->closer = CheckElementCloser(base, e, r, support);
|
e->closer = CheckElementCloser(base, e, r, procinfo, collation);
|
||||||
if (e->closer)
|
if (e->closer)
|
||||||
added = lappend(added, e);
|
added = lappend(added, e);
|
||||||
}
|
}
|
||||||
@@ -1025,7 +1015,7 @@ SelectNeighbors(char *base, List *c, int lm, int lc, HnswSupport * support, Hnsw
|
|||||||
}
|
}
|
||||||
else if (e == newCandidate)
|
else if (e == newCandidate)
|
||||||
{
|
{
|
||||||
e->closer = CheckElementCloser(base, e, r, support);
|
e->closer = CheckElementCloser(base, e, r, procinfo, collation);
|
||||||
if (e->closer)
|
if (e->closer)
|
||||||
added = lappend(added, e);
|
added = lappend(added, e);
|
||||||
}
|
}
|
||||||
@@ -1076,7 +1066,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, 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);
|
HnswElement hce = HnswPtrAccess(base, hc->element);
|
||||||
HnswNeighborArray *currentNeighbors = HnswGetNeighbors(base, hce, lc);
|
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);
|
HnswElement hc3Element = HnswPtrAccess(base, hc3->element);
|
||||||
|
|
||||||
if (HnswPtrIsNull(base, hc3Element->value))
|
if (HnswPtrIsNull(base, hc3Element->value))
|
||||||
HnswLoadElement(hc3Element, &hc3->distance, &q, index, support, true);
|
HnswLoadElement(hc3Element, &hc3->distance, &q, index, procinfo, collation, true);
|
||||||
else
|
else
|
||||||
hc3->distance = GetCandidateDistance(base, hc3, q, support);
|
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)
|
||||||
@@ -1131,7 +1121,7 @@ HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm
|
|||||||
c = lappend(c, ¤tNeighbors->items[i]);
|
c = lappend(c, ¤tNeighbors->items[i]);
|
||||||
c = lappend(c, &hc2);
|
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 */
|
/* Should not happen */
|
||||||
if (pruned == NULL)
|
if (pruned == NULL)
|
||||||
@@ -1205,7 +1195,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, 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 *ep;
|
||||||
List *w;
|
List *w;
|
||||||
@@ -1225,13 +1215,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, support, 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, support, m, true, skipElement);
|
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement);
|
||||||
ep = w;
|
ep = w;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1249,7 +1239,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
|||||||
List *neighbors;
|
List *neighbors;
|
||||||
List *lw;
|
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 */
|
/* Elements being deleted or skipped can help with search */
|
||||||
/* but should be removed before selecting neighbors */
|
/* 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
|
* 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, support, element, NULL, NULL, false);
|
neighbors = SelectNeighbors(base, lw, lm, lc, procinfo, collation, element, NULL, NULL, false);
|
||||||
|
|
||||||
AddConnections(base, element, neighbors, lc);
|
AddConnections(base, element, neighbors, lc);
|
||||||
|
|
||||||
|
|||||||
@@ -189,6 +189,8 @@ 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);
|
||||||
@@ -203,7 +205,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, &vacuumstate->support, m, efConstruction, true);
|
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, 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);
|
||||||
@@ -227,7 +229,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
|
|||||||
UnlockReleaseBuffer(buf);
|
UnlockReleaseBuffer(buf);
|
||||||
|
|
||||||
/* Update neighbors */
|
/* 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);
|
LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
|
||||||
|
|
||||||
/* Load element */
|
/* Load element */
|
||||||
HnswLoadElement(highestPoint, NULL, NULL, index, &vacuumstate->support, true);
|
HnswLoadElement(highestPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
|
||||||
|
|
||||||
/* Repair if needed */
|
/* Repair if needed */
|
||||||
if (NeedsUpdated(vacuumstate, highestPoint))
|
if (NeedsUpdated(vacuumstate, highestPoint))
|
||||||
@@ -292,7 +294,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->support, true);
|
HnswLoadElement(entryPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
|
||||||
|
|
||||||
if (NeedsUpdated(vacuumstate, entryPoint))
|
if (NeedsUpdated(vacuumstate, entryPoint))
|
||||||
{
|
{
|
||||||
@@ -571,8 +573,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->support.procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
vacuumstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||||
vacuumstate->support.collation = index->rd_indcollation[0];
|
vacuumstate->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",
|
||||||
|
|||||||
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