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

DRY HNSW distance calculations

Browse Source
This commit is contained in:
Andrew Kane
2024-10-09 17:01:49 -07:00
parent 77688b4309
commit f4b67b078f
1 changed files with 14 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/hnswutils.c
View File

@@ -482,6 +482,15 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
} }
} }
/*
* Calculate the distance between values
*/
static inline float
HnswGetDistance(Datum a, Datum b, FmgrInfo *procinfo, Oid collation)
{
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, a, b));
}
/* /*
* Load an element and optionally get its distance from q * Load an element and optionally get its distance from q
*/ */
@@ -507,7 +516,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Dat
if (DatumGetPointer(*q) == NULL) if (DatumGetPointer(*q) == NULL)
*distance = 0; *distance = 0;
else else
*distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->data))); *distance = HnswGetDistance(*q, PointerGetDatum(&etup->data), procinfo, collation);
} }
/* Load element */ /* Load element */
@@ -539,7 +548,7 @@ GetElementDistance(char *base, HnswElement element, Datum q, FmgrInfo *procinfo,
{ {
Datum value = HnswGetValue(base, element); Datum value = HnswGetValue(base, element);
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, value)); return HnswGetDistance(q, value, procinfo, collation);
} }
/* /*
@@ -921,18 +930,6 @@ CompareCandidateDistancesOffset(const ListCell *a, const ListCell *b)
return 0; return 0;
} }
/*
* Calculate the distance between elements
*/
static float
HnswGetDistance(char *base, HnswElement a, HnswElement b, FmgrInfo *procinfo, Oid collation)
{
Datum aValue = HnswGetValue(base, a);
Datum bValue = HnswGetValue(base, b);
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, aValue, bValue));
}
/* /*
* Check if an element is closer to q than any element from R * Check if an element is closer to q than any element from R
*/ */
@@ -940,13 +937,15 @@ static bool
CheckElementCloser(char *base, HnswCandidate * e, List *r, FmgrInfo *procinfo, Oid collation) CheckElementCloser(char *base, HnswCandidate * e, List *r, FmgrInfo *procinfo, Oid collation)
{ {
HnswElement eElement = HnswPtrAccess(base, e->element); HnswElement eElement = HnswPtrAccess(base, e->element);
Datum eValue = HnswGetValue(base, eElement);
ListCell *lc2; ListCell *lc2;
foreach(lc2, r) foreach(lc2, r)
{ {
HnswCandidate *ri = lfirst(lc2); HnswCandidate *ri = lfirst(lc2);
HnswElement riElement = HnswPtrAccess(base, ri->element); HnswElement riElement = HnswPtrAccess(base, ri->element);
float distance = HnswGetDistance(base, eElement, riElement, procinfo, collation); Datum riValue = HnswGetValue(base, riElement);
float distance = HnswGetDistance(eValue, riValue, procinfo, collation);
if (distance <= e->distance) if (distance <= e->distance)
return false; return false;