diff --git a/src/hnsw.h b/src/hnsw.h index 2f45039..9fb650a 100644 --- a/src/hnsw.h +++ b/src/hnsw.h @@ -155,12 +155,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 @@ -381,7 +382,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 0463a89..30815af 100644 --- a/src/hnswscan.c +++ b/src/hnswscan.c @@ -161,14 +161,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 } while (list_length(so->w) > 0) { char *base = NULL; - HnswCandidate *hc = llast(so->w); + HnswSearchCandidate *hc = llast(so->w); HnswElement element = HnswPtrAccess(base, hc->element); ItemPointer heaptid; diff --git a/src/hnswutils.c b/src/hnswutils.c index 6e01cf4..ac1e7de 100644 --- a/src/hnswutils.c +++ b/src/hnswutils.c @@ -608,10 +608,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) @@ -621,8 +621,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 @@ -630,10 +630,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; @@ -645,27 +645,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 */ @@ -825,15 +813,13 @@ 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; 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 @@ -846,8 +832,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) @@ -863,12 +849,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) { @@ -899,13 +884,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. @@ -926,7 +909,7 @@ 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)) { - HnswCandidate *hc = HnswGetPairingHeapCandidate(w_node, pairingheap_remove_first(W)); + HnswSearchCandidate *hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W)); w = lappend(w, hc); } @@ -1307,16 +1290,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); + /* 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