diff --git a/src/hnsw.h b/src/hnsw.h index e6a5be5..878e02d 100644 --- a/src/hnsw.h +++ b/src/hnsw.h @@ -177,12 +177,13 @@ struct HnswNeighborArray HnswCandidate items[FLEXIBLE_ARRAY_MEMBER]; }; -typedef struct HnswPairingHeapNode +typedef struct HnswSearchCandidate { - HnswCandidate *inner; pairingheap_node c_node; pairingheap_node w_node; -} HnswPairingHeapNode; + HnswElementPtr element; + float distance; +} HnswSearchCandidate; /* HNSW index options */ typedef struct HnswOptions @@ -414,7 +415,7 @@ void *HnswAlloc(HnswAllocator * allocator, Size size); HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc); HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno); void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing); -HnswCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec); +HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec); void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building); void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); diff --git a/src/hnswscan.c b/src/hnswscan.c index 87337f3..95125a2 100644 --- a/src/hnswscan.c +++ b/src/hnswscan.c @@ -188,14 +188,14 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir) so->first = false; #if defined(HNSW_MEMORY) - elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(so->tmpCtx, false) / (1024 * 1024)); + elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024); #endif } for (;;) { char *base = NULL; - HnswCandidate *hc; + HnswSearchCandidate *hc; HnswElement element; ItemPointer heaptid; diff --git a/src/hnswutils.c b/src/hnswutils.c index 89c2eaa..da7d8bc 100644 --- a/src/hnswutils.c +++ b/src/hnswutils.c @@ -606,10 +606,10 @@ GetElementDistance(char *base, HnswElement element, Datum q, FmgrInfo *procinfo, /* * Create a candidate for the entry point */ -HnswCandidate * +HnswSearchCandidate * HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec) { - HnswCandidate *hc = palloc(sizeof(HnswCandidate)); + HnswSearchCandidate *hc = palloc(sizeof(HnswSearchCandidate)); HnswPtrStore(base, hc->element, entryPoint); if (index == NULL) @@ -619,8 +619,8 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, return hc; } -#define HnswGetPairingHeapCandidate(membername, ptr) (pairingheap_container(HnswPairingHeapNode, membername, ptr)->inner) -#define HnswGetPairingHeapCandidateConst(membername, ptr) (pairingheap_const_container(HnswPairingHeapNode, membername, ptr)->inner) +#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr) +#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr) /* * Compare candidate distances @@ -628,10 +628,10 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, static int CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg) { - if (HnswGetPairingHeapCandidateConst(c_node, a)->distance < HnswGetPairingHeapCandidateConst(c_node, b)->distance) + if (HnswGetSearchCandidateConst(c_node, a)->distance < HnswGetSearchCandidateConst(c_node, b)->distance) return 1; - if (HnswGetPairingHeapCandidateConst(c_node, a)->distance > HnswGetPairingHeapCandidateConst(c_node, b)->distance) + if (HnswGetSearchCandidateConst(c_node, a)->distance > HnswGetSearchCandidateConst(c_node, b)->distance) return -1; return 0; @@ -643,27 +643,15 @@ CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, v static int CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg) { - if (HnswGetPairingHeapCandidateConst(w_node, a)->distance < HnswGetPairingHeapCandidateConst(w_node, b)->distance) + if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance) return -1; - if (HnswGetPairingHeapCandidateConst(w_node, a)->distance > HnswGetPairingHeapCandidateConst(w_node, b)->distance) + if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance) return 1; return 0; } -/* - * Create a pairing heap node for a candidate - */ -static HnswPairingHeapNode * -CreatePairingHeapNode(HnswCandidate * c) -{ - HnswPairingHeapNode *node = palloc(sizeof(HnswPairingHeapNode)); - - node->inner = c; - return node; -} - /* * Init visited */ @@ -840,16 +828,14 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F /* Add entry points to v, C, and W */ foreach(lc2, ep) { - HnswCandidate *hc = (HnswCandidate *) lfirst(lc2); + HnswSearchCandidate *hc = (HnswSearchCandidate *) lfirst(lc2); bool found; - HnswPairingHeapNode *node; if (initVisited) AddToVisited(base, v, hc->element, index, &found); - node = CreatePairingHeapNode(hc); - pairingheap_add(C, &node->c_node); - pairingheap_add(W, &node->w_node); + pairingheap_add(C, &hc->c_node); + pairingheap_add(W, &hc->w_node); /* * Do not count elements being deleted towards ef when vacuuming. It @@ -862,8 +848,8 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F while (!pairingheap_is_empty(C)) { - HnswCandidate *c = HnswGetPairingHeapCandidate(c_node, pairingheap_remove_first(C)); - HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W)); + HnswSearchCandidate *c = HnswGetSearchCandidate(c_node, pairingheap_remove_first(C)); + HnswSearchCandidate *f = HnswGetSearchCandidate(w_node, pairingheap_first(W)); HnswElement cElement; if (c->distance > f->distance) @@ -879,12 +865,11 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F for (int i = 0; i < unvisitedLength; i++) { HnswElement eElement; - HnswCandidate *e; - HnswPairingHeapNode *node; + HnswSearchCandidate *e; float eDistance; bool alwaysAdd = wlen < ef; - f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W)); + f = HnswGetSearchCandidate(w_node, pairingheap_first(W)); if (index == NULL) { @@ -907,7 +892,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F if (discarded != NULL) { /* Create a new candidate */ - e = palloc(sizeof(HnswCandidate)); + e = palloc(sizeof(HnswSearchCandidate)); HnswPtrStore(base, e->element, eElement); e->distance = eDistance; @@ -922,13 +907,11 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F continue; /* Create a new candidate */ - e = palloc(sizeof(HnswCandidate)); + e = palloc(sizeof(HnswSearchCandidate)); HnswPtrStore(base, e->element, eElement); e->distance = eDistance; - - node = CreatePairingHeapNode(e); - pairingheap_add(C, &node->c_node); - pairingheap_add(W, &node->w_node); + pairingheap_add(C, &e->c_node); + pairingheap_add(W, &e->w_node); /* * Do not count elements being deleted towards ef when vacuuming. @@ -942,14 +925,10 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F /* No need to decrement wlen */ if (wlen > ef) { - HnswPairingHeapNode *furthestNode = pairingheap_container(HnswPairingHeapNode, w_node, pairingheap_remove_first(W)); - HnswCandidate *hc = furthestNode->inner; + HnswSearchCandidate *d = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W)); if (discarded != NULL) - *discarded = lappend(*discarded, hc); - - /* TODO */ - /* pfree(furthestNode); */ + *discarded = lappend(*discarded, d); } } } @@ -958,12 +937,9 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F /* Add each element of W to w */ while (!pairingheap_is_empty(W)) { - HnswPairingHeapNode *node = pairingheap_container(HnswPairingHeapNode, w_node, pairingheap_remove_first(W)); - HnswCandidate *hc = node->inner; + HnswSearchCandidate *hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W)); w = lappend(w, hc); - - pfree(node); } return w; @@ -1342,16 +1318,27 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint { int lm = HnswGetLayerM(m, lc); List *neighbors; - List *lw; + List *lw = NIL; + ListCell *lc2; w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true); + /* Convert search candidates to candidates */ + foreach(lc2, w) + { + HnswSearchCandidate *sc = lfirst(lc2); + HnswCandidate *hc = palloc(sizeof(HnswCandidate)); + + hc->element = sc->element; + hc->distance = sc->distance; + + lw = lappend(lw, hc); + } + /* Elements being deleted or skipped can help with search */ /* but should be removed before selecting neighbors */ if (index != NULL) - lw = RemoveElements(base, w, skipElement); - else - lw = w; + lw = RemoveElements(base, lw, skipElement); /* * Candidates are sorted, but not deterministically. Could set