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

Added HnswQuery struct for query data

Browse Source
This commit is contained in:
Andrew Kane
2024-10-09 23:45:47 -07:00
parent 32ab27d72a
commit fa6782985a
4 changed files with 37 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
src/hnswutils.c
View File

@@ -533,7 +533,7 @@ HnswGetDistance(Datum a, Datum b, HnswSupport * support)
* Load an element and optionally get its distance from q
*/
static void
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Datum *q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
{
Buffer buf;
Page page;
@@ -551,10 +551,10 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Dat
/* Calculate distance */
if (distance != NULL)
{
if (DatumGetPointer(*q) == NULL)
if (DatumGetPointer(q->value) == NULL)
*distance = 0;
else
*distance = HnswGetDistance(*q, PointerGetDatum(&etup->data), support);
*distance = HnswGetDistance(q->value, PointerGetDatum(&etup->data), support);
}
/* Load element */
@@ -573,7 +573,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Dat
* Load an element and optionally get its distance from q
*/
void
HnswLoadElement(HnswElement element, double *distance, Datum *q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
{
HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, support, loadVec, maxDistance, &element);
}
@@ -582,18 +582,18 @@ HnswLoadElement(HnswElement element, double *distance, Datum *q, Relation index,
* Get the distance for an element
*/
static double
GetElementDistance(char *base, HnswElement element, Datum q, HnswSupport * support)
GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport * support)
{
Datum value = HnswGetValue(base, element);
return HnswGetDistance(q, value, support);
return HnswGetDistance(q->value, value, support);
}
/*
* Create a candidate for the entry point
*/
HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, HnswSupport * support, bool loadVec)
HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec)
{
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
bool inMemory = index == NULL;
@@ -602,7 +602,7 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
if (inMemory)
sc->distance = GetElementDistance(base, entryPoint, q, support);
else
HnswLoadElement(entryPoint, &sc->distance, &q, index, support, loadVec, NULL);
HnswLoadElement(entryPoint, &sc->distance, q, index, support, loadVec, NULL);
return sc;
}
@@ -791,7 +791,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
* 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, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement)
{
List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -873,7 +873,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, H
/* Avoid any allocations if not adding */
eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, support, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
if (eElement == NULL)
continue;
@@ -1220,10 +1220,12 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
List *w;
int level = element->level;
int entryLevel;
Datum q = HnswGetValue(base, element);
HnswQuery q;
HnswElement skipElement = existing ? element : NULL;
bool inMemory = index == NULL;
q.value = HnswGetValue(base, element);
/* Precompute hash */
if (inMemory)
PrecomputeHash(base, element);
@@ -1233,13 +1235,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, support, 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, support, m, true, skipElement);
ep = w;
}
@@ -1258,7 +1260,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
List *lw = NIL;
ListCell *lc2;
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, support, m, true, skipElement);
w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement);
/* Convert search candidates to candidates */
foreach(lc2, w)