liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-17 01:34:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use inMemory for conditionals

Browse Source
This commit is contained in:
Andrew Kane
2024-10-09 21:49:32 -07:00
parent a98534e5ab
commit 17266ed409
1 changed files with 16 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/hnswutils.c
View File

@@ -598,9 +598,10 @@ HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec) HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{ {
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate)); HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
bool inMemory = index == NULL;
HnswPtrStore(base, sc->element, entryPoint); HnswPtrStore(base, sc->element, entryPoint);
if (index == NULL) if (inMemory)
sc->distance = GetElementDistance(base, entryPoint, q, procinfo, collation); sc->distance = GetElementDistance(base, entryPoint, q, procinfo, collation);
else else
HnswLoadElement(entryPoint, &sc->distance, &q, index, procinfo, collation, loadVec, NULL); HnswLoadElement(entryPoint, &sc->distance, &q, index, procinfo, collation, loadVec, NULL);
@@ -644,9 +645,9 @@ CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b,
* Init visited * Init visited
*/ */
static inline void static inline void
InitVisited(char *base, visited_hash * v, Relation index, int ef, int m) InitVisited(char *base, visited_hash * v, bool inMemory, int ef, int m)
{ {
if (index != NULL) if (!inMemory)
v->tids = tidhash_create(CurrentMemoryContext, ef * m * 2, NULL); v->tids = tidhash_create(CurrentMemoryContext, ef * m * 2, NULL);
else if (base != NULL) else if (base != NULL)
v->offsets = offsethash_create(CurrentMemoryContext, ef * m * 2, NULL); v->offsets = offsethash_create(CurrentMemoryContext, ef * m * 2, NULL);
@@ -658,9 +659,9 @@ InitVisited(char *base, visited_hash * v, Relation index, int ef, int m)
* Add to visited * Add to visited
*/ */
static inline void static inline void
AddToVisited(char *base, visited_hash * v, HnswElementPtr elementPtr, Relation index, bool *found) AddToVisited(char *base, visited_hash * v, HnswElementPtr elementPtr, bool inMemory, bool *found)
{ {
if (index != NULL) if (!inMemory)
{ {
HnswElement element = HnswPtrAccess(base, elementPtr); HnswElement element = HnswPtrAccess(base, elementPtr);
ItemPointerData indextid; ItemPointerData indextid;
@@ -721,7 +722,7 @@ HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswUnvisited * unv
HnswCandidate *hc = &localNeighborhood->items[i]; HnswCandidate *hc = &localNeighborhood->items[i];
bool found; bool found;
AddToVisited(base, v, hc->element, NULL, &found); AddToVisited(base, v, hc->element, true, &found);
if (!found) if (!found)
unvisited[(*unvisitedLength)++].element = HnswPtrAccess(base, hc->element); unvisited[(*unvisitedLength)++].element = HnswPtrAccess(base, hc->element);
@@ -805,11 +806,12 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited)); HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
int unvisitedLength; int unvisitedLength;
bool inMemory = index == NULL;
InitVisited(base, &v, index, ef, m); InitVisited(base, &v, inMemory, ef, m);
/* Create local memory for neighborhood if needed */ /* Create local memory for neighborhood if needed */
if (index == NULL) if (inMemory)
{ {
neighborhoodSize = HNSW_NEIGHBOR_ARRAY_SIZE(lm); neighborhoodSize = HNSW_NEIGHBOR_ARRAY_SIZE(lm);
localNeighborhood = palloc(neighborhoodSize); localNeighborhood = palloc(neighborhoodSize);
@@ -821,7 +823,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
HnswSearchCandidate *sc = (HnswSearchCandidate *) lfirst(lc2); HnswSearchCandidate *sc = (HnswSearchCandidate *) lfirst(lc2);
bool found; bool found;
AddToVisited(base, &v, sc->element, index, &found); AddToVisited(base, &v, sc->element, inMemory, &found);
pairingheap_add(C, &sc->c_node); pairingheap_add(C, &sc->c_node);
pairingheap_add(W, &sc->w_node); pairingheap_add(W, &sc->w_node);
@@ -846,7 +848,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
cElement = HnswPtrAccess(base, c->element); cElement = HnswPtrAccess(base, c->element);
if (index == NULL) if (inMemory)
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, &v, lc, localNeighborhood, neighborhoodSize); HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, &v, lc, localNeighborhood, neighborhoodSize);
else else
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, &v, index, m, lm, lc); HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, &v, index, m, lm, lc);
@@ -860,7 +862,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
f = HnswGetSearchCandidate(w_node, pairingheap_first(W)); f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
if (index == NULL) if (inMemory)
{ {
eElement = unvisited[i].element; eElement = unvisited[i].element;
eDistance = GetElementDistance(base, eElement, q, procinfo, collation); eDistance = GetElementDistance(base, eElement, q, procinfo, collation);
@@ -1222,9 +1224,10 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
int entryLevel; int entryLevel;
Datum q = HnswGetValue(base, element); Datum q = HnswGetValue(base, element);
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
bool inMemory = index == NULL;
/* Precompute hash */ /* Precompute hash */
if (index == NULL) if (inMemory)
PrecomputeHash(base, element); PrecomputeHash(base, element);
/* No neighbors if no entry point */ /* No neighbors if no entry point */
@@ -1273,7 +1276,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
/* 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 */
if (index != NULL) if (!inMemory)
lw = RemoveElements(base, lw, skipElement); lw = RemoveElements(base, lw, skipElement);
/* /*