Added HnswInitSearchCandidate function

This commit is contained in:
Andrew Kane
2024-10-10 19:30:47 -07:00
parent 35b252a3e3
commit 124018b8dd

View File

@@ -581,21 +581,34 @@ GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport *
return HnswGetDistance(q->value, value, support); return HnswGetDistance(q->value, value, support);
} }
/*
* Allocate a search candidate
*/
static HnswSearchCandidate *
HnswInitSearchCandidate(char *base, HnswElement element, double distance)
{
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, sc->element, element);
sc->distance = distance;
return sc;
}
/* /*
* Create a candidate for the entry point * Create a candidate for the entry point
*/ */
HnswSearchCandidate * HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * 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; bool inMemory = index == NULL;
double distance;
HnswPtrStore(base, sc->element, entryPoint);
if (inMemory) if (inMemory)
sc->distance = GetElementDistance(base, entryPoint, q, support); distance = GetElementDistance(base, entryPoint, q, support);
else else
HnswLoadElement(entryPoint, &sc->distance, q, index, support, loadVec, NULL); HnswLoadElement(entryPoint, &distance, q, index, support, loadVec, NULL);
return sc;
return HnswInitSearchCandidate(base, entryPoint, distance);
} }
/* /*
@@ -912,9 +925,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
if (discarded != NULL) if (discarded != NULL)
{ {
/* Create a new candidate */ /* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate)); e = HnswInitSearchCandidate(base, eElement, eDistance);
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
pairingheap_add(*discarded, &e->w_node); pairingheap_add(*discarded, &e->w_node);
} }
@@ -926,9 +937,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
continue; continue;
/* Create a new candidate */ /* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate)); e = HnswInitSearchCandidate(base, eElement, eDistance);
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
pairingheap_add(C, &e->c_node); pairingheap_add(C, &e->c_node);
pairingheap_add(W, &e->w_node); pairingheap_add(W, &e->w_node);