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

Reduced allocations for pairing heap

Browse Source
This commit is contained in:
Andrew Kane
2024-09-19 01:59:46 -07:00
parent a1b80faa67
commit d74d3065bc
2 changed files with 21 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/hnsw.h
View File

@@ -162,8 +162,9 @@ struct HnswNeighborArray
typedef struct HnswPairingHeapNode typedef struct HnswPairingHeapNode
{ {
pairingheap_node ph_node;
HnswCandidate *inner; HnswCandidate *inner;
pairingheap_node c_node;
pairingheap_node w_node;
} HnswPairingHeapNode; } HnswPairingHeapNode;
/* HNSW index options */ /* HNSW index options */

31
src/hnswutils.c
View File

@@ -607,16 +607,19 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
return hc; return hc;
} }
#define HnswGetPairingHeapCandidate(membername, ptr) (pairingheap_container(HnswPairingHeapNode, membername, ptr)->inner)
#define HnswGetPairingHeapCandidateConst(membername, ptr) (pairingheap_const_container(HnswPairingHeapNode, membername, ptr)->inner)
/* /*
* Compare candidate distances * Compare candidate distances
*/ */
static int static int
CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg) CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{ {
if (((const HnswPairingHeapNode *) a)->inner->distance < ((const HnswPairingHeapNode *) b)->inner->distance) if (HnswGetPairingHeapCandidateConst(c_node, a)->distance < HnswGetPairingHeapCandidateConst(c_node, b)->distance)
return 1; return 1;
if (((const HnswPairingHeapNode *) a)->inner->distance > ((const HnswPairingHeapNode *) b)->inner->distance) if (HnswGetPairingHeapCandidateConst(c_node, a)->distance > HnswGetPairingHeapCandidateConst(c_node, b)->distance)
return -1; return -1;
return 0; return 0;
@@ -628,10 +631,10 @@ CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, v
static int static int
CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg) CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{ {
if (((const HnswPairingHeapNode *) a)->inner->distance < ((const HnswPairingHeapNode *) b)->inner->distance) if (HnswGetPairingHeapCandidateConst(w_node, a)->distance < HnswGetPairingHeapCandidateConst(w_node, b)->distance)
return -1; return -1;
if (((const HnswPairingHeapNode *) a)->inner->distance > ((const HnswPairingHeapNode *) b)->inner->distance) if (HnswGetPairingHeapCandidateConst(w_node, a)->distance > HnswGetPairingHeapCandidateConst(w_node, b)->distance)
return 1; return 1;
return 0; return 0;
@@ -746,11 +749,13 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
{ {
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2); HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
bool found; bool found;
HnswPairingHeapNode *node;
AddToVisited(base, &v, hc, index, &found); AddToVisited(base, &v, hc, index, &found);
pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node)); node = CreatePairingHeapNode(hc);
pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node)); pairingheap_add(C, &node->c_node);
pairingheap_add(W, &node->w_node);
/* /*
* Do not count elements being deleted towards ef when vacuuming. It * Do not count elements being deleted towards ef when vacuuming. It
@@ -764,8 +769,8 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
while (!pairingheap_is_empty(C)) while (!pairingheap_is_empty(C))
{ {
HnswNeighborArray *neighborhood; HnswNeighborArray *neighborhood;
HnswCandidate *c = ((HnswPairingHeapNode *) pairingheap_remove_first(C))->inner; HnswCandidate *c = HnswGetPairingHeapCandidate(c_node, pairingheap_remove_first(C));
HnswCandidate *f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner; HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
HnswElement cElement; HnswElement cElement;
if (c->distance > f->distance) if (c->distance > f->distance)
@@ -801,7 +806,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
HnswElement eElement = HnswPtrAccess(base, e->element); HnswElement eElement = HnswPtrAccess(base, e->element);
bool alwaysAdd = wlen < ef; bool alwaysAdd = wlen < ef;
f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner; f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
if (index == NULL) if (index == NULL)
eDistance = GetCandidateDistance(base, e, q, procinfo, collation); eDistance = GetCandidateDistance(base, e, q, procinfo, collation);
@@ -811,6 +816,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
if (eDistance < f->distance || alwaysAdd) if (eDistance < f->distance || alwaysAdd)
{ {
HnswCandidate *ec; HnswCandidate *ec;
HnswPairingHeapNode *node;
Assert(!eElement->deleted); Assert(!eElement->deleted);
@@ -823,8 +829,9 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
HnswPtrStore(base, ec->element, eElement); HnswPtrStore(base, ec->element, eElement);
ec->distance = eDistance; ec->distance = eDistance;
pairingheap_add(C, &(CreatePairingHeapNode(ec)->ph_node)); node = CreatePairingHeapNode(ec);
pairingheap_add(W, &(CreatePairingHeapNode(ec)->ph_node)); pairingheap_add(C, &node->c_node);
pairingheap_add(W, &node->w_node);
/* /*
* Do not count elements being deleted towards ef when * Do not count elements being deleted towards ef when
@@ -847,7 +854,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
/* Add each element of W to w */ /* Add each element of W to w */
while (!pairingheap_is_empty(W)) while (!pairingheap_is_empty(W))
{ {
HnswCandidate *hc = ((HnswPairingHeapNode *) pairingheap_remove_first(W))->inner; HnswCandidate *hc = HnswGetPairingHeapCandidate(w_node, pairingheap_remove_first(W));
w = lappend(w, hc); w = lappend(w, hc);
} }