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

Moved sorting logic into SelectNeighbors

Browse Source
This commit is contained in:
Andrew Kane
2023-10-06 12:56:15 -07:00
parent cae162ffc6
commit 8085d3e538
1 changed files with 36 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
src/hnswutils.c
View File

@@ -693,6 +693,34 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
return w; return w;
} }
/*
* Compare candidate distances
*/
static int
#if PG_VERSION_NUM >= 130000
CompareCandidateDistances(const ListCell *a, const ListCell *b)
#else
CompareCandidateDistances(const void *a, const void *b)
#endif
{
HnswCandidate *hca = lfirst((ListCell *) a);
HnswCandidate *hcb = lfirst((ListCell *) b);
if (hca->distance < hcb->distance)
return 1;
if (hca->distance > hcb->distance)
return -1;
if (hca->element < hcb->element)
return 1;
if (hca->element > hcb->element)
return -1;
return 0;
}
/* /*
* Calculate the distance between elements * Calculate the distance between elements
*/ */
@@ -749,7 +777,7 @@ 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, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned) SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool needsSorting)
{ {
List *r = NIL; List *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
@@ -763,6 +791,10 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
wd = pairingheap_allocate(CompareNearestCandidates, NULL); wd = pairingheap_allocate(CompareNearestCandidates, NULL);
/* Ensure order is deterministic for closer caching */
if (needsSorting)
list_sort(w, CompareCandidateDistances);
while (list_length(w) > 0 && list_length(r) < m) while (list_length(w) > 0 && list_length(r) < m)
{ {
/* Assumes w is already ordered desc */ /* Assumes w is already ordered desc */
@@ -868,34 +900,6 @@ AddConnections(HnswElement element, List *neighbors, int m, int lc)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2)); a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
} }
/*
* Compare candidate distances
*/
static int
#if PG_VERSION_NUM >= 130000
CompareCandidateDistances(const ListCell *a, const ListCell *b)
#else
CompareCandidateDistances(const void *a, const void *b)
#endif
{
HnswCandidate *hca = lfirst((ListCell *) a);
HnswCandidate *hcb = lfirst((ListCell *) b);
if (hca->distance < hcb->distance)
return 1;
if (hca->distance > hcb->distance)
return -1;
if (hca->element < hcb->element)
return 1;
if (hca->element > hcb->element)
return -1;
return 0;
}
/* /*
* Update connections * Update connections
*/ */
@@ -949,13 +953,12 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
{ {
List *c = NIL; List *c = NIL;
/* Add and sort candidates */ /* Add candidates */
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
c = lappend(c, &currentNeighbors->items[i]); c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2); c = lappend(c, &hc2);
list_sort(c, CompareCandidateDistances);
SelectNeighbors(c, m, lc, procinfo, collation, hc->element, &hc2, &pruned); SelectNeighbors(c, m, lc, procinfo, collation, hc->element, &hc2, &pruned, true);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -1054,11 +1057,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
else else
lw = w; lw = w;
/* Ensure order is deterministic for SelectNeighbors closer caching */ neighbors = SelectNeighbors(lw, lm, lc, procinfo, collation, element, NULL, NULL, index == NULL);
if (index == NULL)
list_sort(lw, CompareCandidateDistances);
neighbors = SelectNeighbors(lw, lm, lc, procinfo, collation, element, NULL, NULL);
AddConnections(element, neighbors, lm, lc); AddConnections(element, neighbors, lm, lc);