Skip to content

Commit 5c9429a

Browse files
ankaneankane
committed
Reduced memory usage for HNSW index scans
1 parent 4b44d6e commit 5c9429a

File tree

1 file changed

+43
-12
lines changed

1 file changed

+43
-12
lines changed

‎src/hnswutils.c‎

Lines changed: 43 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,12 @@ typedef union
112112
tidhash_hash *tids;
113113
} visited_hash;
114114

115+
typedef union
116+
{
117+
HnswElement element;
118+
ItemPointerData indextid;
119+
} HnswUnvisited;
120+
115121
/*
116122
* Get the max number of connections in an upper layer for each element in the index
117123
*/
@@ -547,19 +553,19 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
547553
/*
548554
* Load an element and optionally get its distance from q
549555
*/
550-
void
551-
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, float *maxDistance)
556+
static void
557+
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, float *maxDistance, HnswElement * element)
552558
{
553559
Buffer buf;
554560
Page page;
555561
HnswElementTuple etup;
556562

557563
/* Read vector */
558-
buf = ReadBuffer(index, element->blkno);
564+
buf = ReadBuffer(index, blkno);
559565
LockBuffer(buf, BUFFER_LOCK_SHARE);
560566
page = BufferGetPage(buf);
561567

562-
etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, element->offno));
568+
etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
563569

564570
Assert(HnswIsElementTuple(etup));
565571

@@ -574,11 +580,25 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
574580

575581
/* Load element */
576582
if (distance == NULL || maxDistance == NULL || *distance < *maxDistance)
577-
HnswLoadElementFromTuple(element, etup, true, loadVec);
583+
{
584+
if (*element == NULL)
585+
*element = HnswInitElementFromBlock(blkno, offno);
586+
587+
HnswLoadElementFromTuple(*element, etup, true, loadVec);
588+
}
578589

579590
UnlockReleaseBuffer(buf);
580591
}
581592

593+
/*
594+
* Load an element and optionally get its distance from q
595+
*/
596+
void
597+
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, float *maxDistance)
598+
{
599+
HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, procinfo, collation, loadVec, maxDistance, &element);
600+
}
601+
582602
/*
583603
* Get the distance for an element
584604
*/
@@ -720,7 +740,7 @@ CountElement(char *base, HnswElement skipElement, HnswElement e)
720740
* Load unvisited neighbors from memory
721741
*/
722742
static void
723-
HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswElement * unvisited, int *unvisitedLength, visited_hash * v, int lc, HnswNeighborArray * neighborhoodData, Size neighborhoodSize)
743+
HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswUnvisited * unvisited, int *unvisitedLength, visited_hash * v, int lc, HnswNeighborArray * neighborhoodData, Size neighborhoodSize)
724744
{
725745
/* Get the neighborhood at layer lc */
726746
HnswNeighborArray *neighborhood = HnswGetNeighbors(base, element, lc);
@@ -741,15 +761,15 @@ HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswElement * unvis
741761
AddToVisited(base, v, hc->element, NULL, &found);
742762

743763
if (!found)
744-
unvisited[(*unvisitedLength)++] = HnswPtrAccess(base, hc->element);
764+
unvisited[(*unvisitedLength)++].element = HnswPtrAccess(base, hc->element);
745765
}
746766
}
747767

748768
/*
749769
* Load unvisited neighbors from disk
750770
*/
751771
static void
752-
HnswLoadUnvisitedFromDisk(HnswElement element, HnswElement * unvisited, int *unvisitedLength, visited_hash * v, Relation index, int m, int lm, int lc)
772+
HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *unvisitedLength, visited_hash * v, Relation index, int m, int lm, int lc)
753773
{
754774
Buffer buf;
755775
Page page;
@@ -782,7 +802,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswElement * unvisited, int *unv
782802
tidhash_insert(v->tids, *indextid, &found);
783803

784804
if (!found)
785-
unvisited[(*unvisitedLength)++] = HnswInitElementFromBlock(ItemPointerGetBlockNumber(indextid), ItemPointerGetOffsetNumber(indextid));
805+
unvisited[(*unvisitedLength)++].indextid = *indextid;
786806
}
787807
}
788808

@@ -801,7 +821,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
801821
HnswNeighborArray *neighborhoodData = NULL;
802822
Size neighborhoodSize = 0;
803823
int lm = HnswGetLayerM(m, lc);
804-
HnswElement *unvisited = palloc(lm * sizeof(HnswElement));
824+
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
805825
int unvisitedLength;
806826

807827
InitVisited(base, &v, index, ef, m);
@@ -853,16 +873,27 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
853873

854874
for (int i = 0; i < unvisitedLength; i++)
855875
{
856-
HnswElement eElement = unvisited[i];
876+
HnswElement eElement;
857877
float eDistance;
858878
bool alwaysAdd = wlen < ef;
859879

860880
f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
861881

862882
if (index == NULL)
883+
{
884+
eElement = unvisited[i].element;
863885
eDistance = GetElementDistance(base, eElement, q, procinfo, collation);
886+
}
864887
else
865-
HnswLoadElement(eElement, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd ? NULL : &f->distance);
888+
{
889+
ItemPointer indextid = &unvisited[i].indextid;
890+
BlockNumber blkno = ItemPointerGetBlockNumber(indextid);
891+
OffsetNumber offno = ItemPointerGetOffsetNumber(indextid);
892+
893+
/* Avoid any allocations if not adding */
894+
eElement = NULL;
895+
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
896+
}
866897

867898
if (eDistance < f->distance || alwaysAdd)
868899
{

0 commit comments

Comments
 (0)