liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-06-06 05:51:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Moved insert-specific code to hnswinsert.c

Browse Source
This commit is contained in:
Andrew Kane
2024-09-29 19:44:11 -07:00
parent 4c72f91206
commit 8eb8cdf0f3
2 changed files with 46 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
src/hnswinsert.c
View File

@@ -364,6 +364,32 @@ HnswLoadNeighbors(HnswElement element, Relation index, int m, int lm, int lc)
return neighbors;
}
/*
* Load elements for insert
*/
static void
LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, int *idx, Relation index, FmgrInfo *procinfo, Oid collation)
{
char *base = NULL;
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *hc = &neighbors->items[i];
HnswElement element = HnswPtrAccess(base, hc->element);
double distance;
HnswLoadElement(element, &distance, &q, index, procinfo, collation, true, NULL);
hc->distance = distance;
/* Prune element if being deleted */
if (element->heaptidsLength == 0)
{
*idx = i;
break;
}
}
}
/*
* Check if connection already exists
*/
@@ -426,7 +452,17 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, Hns
*/
/* Select neighbors */
HnswUpdateConnection(NULL, e, hc, neighborNeighbors, lm, &idx, index, procinfo, collation);
if (neighbors->length < lm)
idx = -2;
else
{
Datum q = HnswGetValue(base, neighborElement);
LoadElementsForInsert(neighborNeighbors, q, &idx, index, procinfo, collation);
if (idx == -1)
HnswUpdateConnection(base, e, hc, neighborNeighbors, lm, &idx, index, procinfo, collation);
}
/* New element was not selected as a neighbor */
if (idx == -1)

53
src/hnswutils.c
View File

@@ -1074,39 +1074,12 @@ AddConnections(char *base, HnswElement element, List *neighbors, int lc)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
}
/*
* Load elements for insert
*/
static void
LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, HnswCandidate * *pruned, Relation index, FmgrInfo *procinfo, Oid collation)
{
char *base = NULL;
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *hc = &neighbors->items[i];
HnswElement element = HnswPtrAccess(base, hc->element);
double distance;
HnswLoadElement(element, &distance, &q, index, procinfo, collation, true, NULL);
hc->distance = distance;
/* Prune element if being deleted */
if (element->heaptidsLength == 0)
{
*pruned = &neighbors->items[i];
break;
}
}
}
/*
* Update connections
*/
void
HnswUpdateConnection(char *base, HnswElement newElement, HnswCandidate * hc, HnswNeighborArray * neighbors, int lm, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
{
HnswElement hce = HnswPtrAccess(base, hc->element);
HnswCandidate newHc;
HnswPtrStore(base, newHc.element, newElement);
@@ -1123,27 +1096,19 @@ HnswUpdateConnection(char *base, HnswElement newElement, HnswCandidate * hc, Hns
else
{
/* Shrink connections */
List *c = NIL;
HnswCandidate *pruned = NULL;
/* Load elements on insert */
if (index != NULL)
LoadElementsForInsert(neighbors, HnswGetValue(base, hce), &pruned, index, procinfo, collation);
/* Add candidates */
for (int i = 0; i < neighbors->length; i++)
c = lappend(c, &neighbors->items[i]);
c = lappend(c, &newHc);
SelectNeighbors(base, c, lm, procinfo, collation, neighbors, &newHc, &pruned, true);
/* Should not happen */
if (pruned == NULL)
{
List *c = NIL;
/* Add candidates */
for (int i = 0; i < neighbors->length; i++)
c = lappend(c, &neighbors->items[i]);
c = lappend(c, &newHc);
SelectNeighbors(base, c, lm, procinfo, collation, neighbors, &newHc, &pruned, true);
/* Should not happen */
if (pruned == NULL)
return;
}
return;
/* Find and replace the pruned element */
for (int i = 0; i < neighbors->length; i++)