liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-12 15:46:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

Reduced memory usage for HNSW index scans

Browse Source
This commit is contained in:
Andrew Kane
2024-09-19 03:27:35 -07:00
parent 4b44d6e745
commit 5c9429a0f8
1 changed files with 43 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
src/hnswutils.c
View File

@@ -112,6 +112,12 @@ typedef union
tidhash_hash *tids; tidhash_hash *tids;
} visited_hash; } visited_hash;
typedef union
{
HnswElement element;
ItemPointerData indextid;
} HnswUnvisited;
/* /*
* Get the max number of connections in an upper layer for each element in the index * Get the max number of connections in an upper layer for each element in the index
*/ */
@@ -547,19 +553,19 @@ 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 static void
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, float *maxDistance) HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, float *maxDistance, HnswElement * element)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
HnswElementTuple etup; HnswElementTuple etup;
/* Read vector */ /* Read vector */
buf = ReadBuffer(index, element->blkno); buf = ReadBuffer(index, blkno);
LockBuffer(buf, BUFFER_LOCK_SHARE); LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf); page = BufferGetPage(buf);
etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, element->offno)); etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
Assert(HnswIsElementTuple(etup)); Assert(HnswIsElementTuple(etup));
@@ -574,11 +580,25 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
/* Load element */ /* Load element */
if (distance == NULL || maxDistance == NULL || *distance < *maxDistance) if (distance == NULL || maxDistance == NULL || *distance < *maxDistance)
HnswLoadElementFromTuple(element, etup, true, loadVec); {
if (*element == NULL)
*element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(*element, etup, true, loadVec);
}
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
/*
* Load an element and optionally get its distance from q
*/
void
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, float *maxDistance)
{
HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, procinfo, collation, loadVec, maxDistance, &element);
}
/* /*
* Get the distance for an element * Get the distance for an element
*/ */
@@ -720,7 +740,7 @@ CountElement(char *base, HnswElement skipElement, HnswElement e)
* Load unvisited neighbors from memory * Load unvisited neighbors from memory
*/ */
static void static void
HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswElement * unvisited, int *unvisitedLength, visited_hash * v, int lc, HnswNeighborArray * neighborhoodData, Size neighborhoodSize) HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswUnvisited * unvisited, int *unvisitedLength, visited_hash * v, int lc, HnswNeighborArray * neighborhoodData, Size neighborhoodSize)
{ {
/* Get the neighborhood at layer lc */ /* Get the neighborhood at layer lc */
HnswNeighborArray *neighborhood = HnswGetNeighbors(base, element, lc); HnswNeighborArray *neighborhood = HnswGetNeighbors(base, element, lc);
@@ -741,7 +761,7 @@ HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswElement * unvis
AddToVisited(base, v, hc->element, NULL, &found); AddToVisited(base, v, hc->element, NULL, &found);
if (!found) if (!found)
unvisited[(*unvisitedLength)++] = HnswPtrAccess(base, hc->element); unvisited[(*unvisitedLength)++].element = HnswPtrAccess(base, hc->element);
} }
} }
@@ -749,7 +769,7 @@ HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswElement * unvis
* Load unvisited neighbors from disk * Load unvisited neighbors from disk
*/ */
static void static void
HnswLoadUnvisitedFromDisk(HnswElement element, HnswElement * unvisited, int *unvisitedLength, visited_hash * v, Relation index, int m, int lm, int lc) HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *unvisitedLength, visited_hash * v, Relation index, int m, int lm, int lc)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -782,7 +802,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswElement * unvisited, int *unv
tidhash_insert(v->tids, *indextid, &found); tidhash_insert(v->tids, *indextid, &found);
if (!found) if (!found)
unvisited[(*unvisitedLength)++] = HnswInitElementFromBlock(ItemPointerGetBlockNumber(indextid), ItemPointerGetOffsetNumber(indextid)); unvisited[(*unvisitedLength)++].indextid = *indextid;
} }
} }
@@ -801,7 +821,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
HnswNeighborArray *neighborhoodData = NULL; HnswNeighborArray *neighborhoodData = NULL;
Size neighborhoodSize = 0; Size neighborhoodSize = 0;
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
HnswElement *unvisited = palloc(lm * sizeof(HnswElement)); HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
int unvisitedLength; int unvisitedLength;
InitVisited(base, &v, index, ef, m); InitVisited(base, &v, index, ef, m);
@@ -853,16 +873,27 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
for (int i = 0; i < unvisitedLength; i++) for (int i = 0; i < unvisitedLength; i++)
{ {
HnswElement eElement = unvisited[i]; HnswElement eElement;
float eDistance; float eDistance;
bool alwaysAdd = wlen < ef; bool alwaysAdd = wlen < ef;
f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W)); f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
if (index == NULL) if (index == NULL)
{
eElement = unvisited[i].element;
eDistance = GetElementDistance(base, eElement, q, procinfo, collation); eDistance = GetElementDistance(base, eElement, q, procinfo, collation);
}
else else
HnswLoadElement(eElement, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd ? NULL : &f->distance); {
ItemPointer indextid = &unvisited[i].indextid;
BlockNumber blkno = ItemPointerGetBlockNumber(indextid);
OffsetNumber offno = ItemPointerGetOffsetNumber(indextid);
/* Avoid any allocations if not adding */
eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
}
if (eDistance < f->distance || alwaysAdd) if (eDistance < f->distance || alwaysAdd)
{ {