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

Speed up HNSW index build

Browse Source
This commit is contained in:
Andrew Kane
2023-10-04 19:43:25 -07:00
parent 8d7abb6590
commit 97ac01773d
2 changed files with 21 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/hnsw.h
View File

@@ -110,12 +110,14 @@ typedef struct HnswCandidate
{ {
HnswElement element; HnswElement element;
float distance; float distance;
bool closer;
} HnswCandidate; } HnswCandidate;
typedef struct HnswNeighborArray typedef struct HnswNeighborArray
{ {
int length; int length;
HnswCandidate *items; HnswCandidate *items;
bool closerSet;
} HnswNeighborArray; } HnswNeighborArray;
typedef struct HnswPairingHeapNode typedef struct HnswPairingHeapNode

25
src/hnswutils.c
View File

@@ -139,6 +139,7 @@ HnswInitNeighbors(HnswElement element, int m)
a = &element->neighbors[lc]; a = &element->neighbors[lc];
a->length = 0; a->length = 0;
a->items = palloc(sizeof(HnswCandidate) * lm); a->items = palloc(sizeof(HnswCandidate) * lm);
a->closerSet = false;
} }
} }
@@ -748,11 +749,12 @@ CheckElementCloser(HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid c
* Algorithm 4 from paper * Algorithm 4 from paper
*/ */
static List * static List *
SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswCandidate * *pruned) SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned)
{ {
List *r = NIL; List *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
pairingheap *wd; pairingheap *wd;
bool mustCalculate = !e2->neighbors[lc].closerSet;
if (list_length(w) <= m) if (list_length(w) <= m)
return w; return w;
@@ -763,18 +765,29 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswC
{ {
/* Assumes w is already ordered desc */ /* Assumes w is already ordered desc */
HnswCandidate *e = llast(w); HnswCandidate *e = llast(w);
bool closer;
w = list_delete_last(w); w = list_delete_last(w);
closer = CheckElementCloser(e, r, lc, procinfo, collation); /*
* r and wd will be the same as previous calls until the new
* candidate, so can skip distance calculations
*/
if (mustCalculate)
e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
else if (e == newCandidate)
{
e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
mustCalculate = true;
}
if (closer) if (e->closer)
r = lappend(r, e); r = lappend(r, e);
else else
pairingheap_add(wd, &(CreatePairingHeapNode(e)->ph_node)); pairingheap_add(wd, &(CreatePairingHeapNode(e)->ph_node));
} }
e2->neighbors[lc].closerSet = true;
/* Keep pruned connections */ /* Keep pruned connections */
while (!pairingheap_is_empty(wd) && list_length(r) < m) while (!pairingheap_is_empty(wd) && list_length(r) < m)
r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner); r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner);
@@ -909,7 +922,7 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
c = lappend(c, &hc2); c = lappend(c, &hc2);
list_sort(c, CompareCandidateDistances); list_sort(c, CompareCandidateDistances);
SelectNeighbors(c, m, lc, procinfo, collation, &pruned); SelectNeighbors(c, m, lc, procinfo, collation, hc->element, &hc2, &pruned);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -1008,7 +1021,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
else else
lw = w; lw = w;
neighbors = SelectNeighbors(lw, lm, lc, procinfo, collation, NULL); neighbors = SelectNeighbors(lw, lm, lc, procinfo, collation, element, NULL, NULL);
AddConnections(element, neighbors, lm, lc); AddConnections(element, neighbors, lm, lc);