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

Merge branch 'master' into ivfflat-streaming

Browse Source
This commit is contained in:
Andrew Kane
2024-09-21 18:37:31 -07:00
parent e2ba6cf38f be4e9a9df2
commit 88f56dc234
4 changed files with 53 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/hnsw.h
View File

@@ -155,12 +155,13 @@ struct HnswNeighborArray
HnswCandidate items[FLEXIBLE_ARRAY_MEMBER]; HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
}; };
typedef struct HnswPairingHeapNode typedef struct HnswSearchCandidate
{ {
HnswCandidate *inner;
pairingheap_node c_node; pairingheap_node c_node;
pairingheap_node w_node; pairingheap_node w_node;
} HnswPairingHeapNode; HnswElementPtr element;
float distance;
} HnswSearchCandidate;
/* HNSW index options */ /* HNSW index options */
typedef struct HnswOptions 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 HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno); 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); 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 HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m); void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);

4
src/hnswscan.c
View File

@@ -161,14 +161,14 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->first = false; so->first = false;
#if defined(HNSW_MEMORY) #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 #endif
} }
while (list_length(so->w) > 0) while (list_length(so->w) > 0)
{ {
char *base = NULL; char *base = NULL;
HnswCandidate *hc = llast(so->w); HnswSearchCandidate *hc = llast(so->w);
HnswElement element = HnswPtrAccess(base, hc->element); HnswElement element = HnswPtrAccess(base, hc->element);
ItemPointer heaptid; ItemPointer heaptid;

80
src/hnswutils.c
View File

@@ -608,10 +608,10 @@ GetElementDistance(char *base, HnswElement element, Datum q, FmgrInfo *procinfo,
/* /*
* Create a candidate for the entry point * Create a candidate for the entry point
*/ */
HnswCandidate * HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec) 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); HnswPtrStore(base, hc->element, entryPoint);
if (index == NULL) if (index == NULL)
@@ -621,8 +621,8 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
return hc; return hc;
} }
#define HnswGetPairingHeapCandidate(membername, ptr) (pairingheap_container(HnswPairingHeapNode, membername, ptr)->inner) #define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
#define HnswGetPairingHeapCandidateConst(membername, ptr) (pairingheap_const_container(HnswPairingHeapNode, membername, ptr)->inner) #define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
/* /*
* Compare candidate distances * Compare candidate distances
@@ -630,10 +630,10 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
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 (HnswGetPairingHeapCandidateConst(c_node, a)->distance < HnswGetPairingHeapCandidateConst(c_node, b)->distance) if (HnswGetSearchCandidateConst(c_node, a)->distance < HnswGetSearchCandidateConst(c_node, b)->distance)
return 1; 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 -1;
return 0; return 0;
@@ -645,27 +645,15 @@ 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 (HnswGetPairingHeapCandidateConst(w_node, a)->distance < HnswGetPairingHeapCandidateConst(w_node, b)->distance) if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
return -1; 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 1;
return 0; 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 * 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 */ /* Add entry points to v, C, and W */
foreach(lc2, ep) foreach(lc2, ep)
{ {
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2); HnswSearchCandidate *hc = (HnswSearchCandidate *) lfirst(lc2);
bool found; bool found;
HnswPairingHeapNode *node;
AddToVisited(base, &v, hc->element, index, &found); AddToVisited(base, &v, hc->element, index, &found);
node = CreatePairingHeapNode(hc); pairingheap_add(C, &hc->c_node);
pairingheap_add(C, &node->c_node); pairingheap_add(W, &hc->w_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
@@ -846,8 +832,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))
{ {
HnswCandidate *c = HnswGetPairingHeapCandidate(c_node, pairingheap_remove_first(C)); HnswSearchCandidate *c = HnswGetSearchCandidate(c_node, pairingheap_remove_first(C));
HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W)); HnswSearchCandidate *f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
HnswElement cElement; HnswElement cElement;
if (c->distance > f->distance) if (c->distance > f->distance)
@@ -863,20 +849,16 @@ 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; HnswElement eElement;
HnswCandidate *e; HnswSearchCandidate *e;
HnswPairingHeapNode *node;
float eDistance; float eDistance;
bool alwaysAdd = wlen < ef; bool alwaysAdd = wlen < ef;
f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W)); f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
if (index == NULL) if (index == NULL)
{ {
eElement = unvisited[i].element; eElement = unvisited[i].element;
eDistance = GetElementDistance(base, eElement, q, procinfo, collation); eDistance = GetElementDistance(base, eElement, q, procinfo, collation);
if (!(eDistance < f->distance || alwaysAdd))
continue;
} }
else else
{ {
@@ -892,6 +874,9 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
continue; continue;
} }
if (!(eDistance < f->distance || alwaysAdd))
continue;
Assert(!eElement->deleted); Assert(!eElement->deleted);
/* Make robust to issues */ /* Make robust to issues */
@@ -899,13 +884,11 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
continue; continue;
/* Create a new candidate */ /* Create a new candidate */
e = palloc(sizeof(HnswCandidate)); e = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, e->element, eElement); HnswPtrStore(base, e->element, eElement);
e->distance = eDistance; e->distance = eDistance;
pairingheap_add(C, &e->c_node);
node = CreatePairingHeapNode(e); pairingheap_add(W, &e->w_node);
pairingheap_add(C, &node->c_node);
pairingheap_add(W, &node->w_node);
/* /*
* Do not count elements being deleted towards ef when vacuuming. * 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 */ /* Add each element of W to w */
while (!pairingheap_is_empty(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); w = lappend(w, hc);
} }
@@ -1307,16 +1290,27 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
{ {
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
List *neighbors; List *neighbors;
List *lw; List *lw = NIL;
ListCell *lc2;
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement); 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 */ /* Elements being deleted or skipped can help with search */
/* but should be removed before selecting neighbors */ /* but should be removed before selecting neighbors */
if (index != NULL) if (index != NULL)
lw = RemoveElements(base, w, skipElement); lw = RemoveElements(base, lw, skipElement);
else
lw = w;
/* /*
* Candidates are sorted, but not deterministically. Could set * Candidates are sorted, but not deterministically. Could set

15
src/ivfscan.c
View File

@@ -15,16 +15,19 @@
#include "utils/memutils.h" #include "utils/memutils.h"
#endif #endif
#define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr)
#define GetScanListConst(ptr) pairingheap_const_container(IvfflatScanList, ph_node, ptr)
/* /*
* Compare list distances * Compare list distances
*/ */
static int static int
CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg) CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{ {
if (((const IvfflatScanList *) a)->distance > ((const IvfflatScanList *) b)->distance) if (GetScanListConst(a)->distance > GetScanListConst(b)->distance)
return 1; return 1;
if (((const IvfflatScanList *) a)->distance < ((const IvfflatScanList *) b)->distance) if (GetScanListConst(a)->distance < GetScanListConst(b)->distance)
return -1; return -1;
return 0; return 0;
@@ -91,14 +94,14 @@ GetScanLists(IndexScanDesc scan, Datum value)
/* Calculate max distance */ /* Calculate max distance */
if (listCount == so->probes) if (listCount == so->probes)
maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance; maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
} }
else if (distance < maxDistance) else if (distance < maxDistance)
{ {
IvfflatScanList *scanlist; IvfflatScanList *scanlist;
/* Remove */ /* Remove */
scanlist = (IvfflatScanList *) pairingheap_remove_first(so->listQueue); scanlist = GetScanList(pairingheap_remove_first(so->listQueue));
/* Reuse */ /* Reuse */
scanlist->startPage = list->startPage; scanlist->startPage = list->startPage;
@@ -106,7 +109,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
pairingheap_add(so->listQueue, &scanlist->ph_node); pairingheap_add(so->listQueue, &scanlist->ph_node);
/* Update max distance */ /* Update max distance */
maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance; maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
} }
} }
@@ -140,7 +143,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
/* Search closest probes lists */ /* Search closest probes lists */
while (!pairingheap_is_empty(so->listQueue) && (++batchProbes) <= so->probes) while (!pairingheap_is_empty(so->listQueue) && (++batchProbes) <= so->probes)
{ {
BlockNumber searchPage = ((IvfflatScanList *) pairingheap_remove_first(so->listQueue))->startPage; BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
/* Search all entry pages for list */ /* Search all entry pages for list */
while (BlockNumberIsValid(searchPage)) while (BlockNumberIsValid(searchPage))