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

Improved concurrent inserts

Browse Source
This commit is contained in:
Andrew Kane
2023-08-09 23:51:35 -07:00
parent d63d430af8
commit 7c0d94c99c
3 changed files with 151 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/hnsw.h
View File

@@ -106,13 +106,6 @@ typedef struct HnswNeighborArray
HnswCandidate *items; HnswCandidate *items;
} HnswNeighborArray; } HnswNeighborArray;
typedef struct HnswUpdate
{
HnswCandidate hc;
int level;
int index;
} HnswUpdate;
typedef struct HnswPairingHeapNode typedef struct HnswPairingHeapNode
{ {
pairingheap_node ph_node; pairingheap_node ph_node;
@@ -265,7 +258,7 @@ List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, Fmgr
HnswElement HnswGetEntryPoint(Relation index); HnswElement HnswGetEntryPoint(Relation index);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel); HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
void HnswFreeElement(HnswElement element); void HnswFreeElement(HnswElement element);
HnswElement HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, List **updates, bool vacuuming); HnswElement HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, List ***updateNeighbors, bool vacuuming);
HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadvec); HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadvec);
void HnswUpdateMetaPage(Relation index, bool updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum); void HnswUpdateMetaPage(Relation index, bool updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum);
void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m); void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m);
@@ -274,6 +267,8 @@ void HnswInitNeighbors(HnswElement element, int m);
bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel); bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel);
void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec); void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswSetElementTuple(HnswElementTuple etup, HnswElement element); void HnswSetElementTuple(HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index);
/* Index access methods */ /* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo); IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);

52
src/hnswinsert.c
View File

@@ -269,39 +269,40 @@ WriteNewElementPages(Relation index, HnswElement e, int m)
HnswUpdateMetaPage(index, false, NULL, insertPage, MAIN_FORKNUM); HnswUpdateMetaPage(index, false, NULL, insertPage, MAIN_FORKNUM);
} }
/*
* Calculate index for update
*/
static int
HnswGetIndex(HnswUpdate * update, int m)
{
return (update->hc.element->level - update->level) * m + update->index;
}
/* /*
* Update neighbors * Update neighbors
*/ */
static void static void
UpdateNeighborPages(Relation index, HnswElement e, int m, List *updates) UpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, List **neighbors)
{ {
ListCell *lc; for (int lc = e->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
List *levelNeighbors = neighbors[lc];
ListCell *lc2;
/* Could update multiple at once for same element */ foreach(lc2, levelNeighbors)
/* but should only happen a low percent of time, so keep simple for now */
foreach(lc, updates)
{ {
HnswCandidate *hc = lfirst(lc2);
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
HnswUpdate *update = lfirst(lc);
ItemId itemid; ItemId itemid;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
Size ntupSize; Size ntupSize;
int idx; int idx = -1;
OffsetNumber offno = update->hc.element->neighborOffno; OffsetNumber offno = hc->element->neighborOffno;
/* Get latest neighbors */
HnswLoadNeighbors(hc->element, index);
HnswUpdateConnection(e, hc, lm, lc, &idx, index, procinfo, collation);
if (idx == -1)
continue;
/* Register page */ /* Register page */
buf = ReadBuffer(index, update->hc.element->neighborPage); buf = ReadBuffer(index, hc->element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index); state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0); page = GenericXLogRegisterBuffer(state, buf, 0);
@@ -311,8 +312,8 @@ UpdateNeighborPages(Relation index, HnswElement e, int m, List *updates)
ntup = (HnswNeighborTuple) PageGetItem(page, itemid); ntup = (HnswNeighborTuple) PageGetItem(page, itemid);
ntupSize = ItemIdGetLength(itemid); ntupSize = ItemIdGetLength(itemid);
/* Calculate index */ /* Calculate index for update */
idx = HnswGetIndex(update, m); idx += (hc->element->level - lc) * m;
/* Make robust to issues */ /* Make robust to issues */
if (idx < ntup->count) if (idx < ntup->count)
@@ -336,6 +337,7 @@ UpdateNeighborPages(Relation index, HnswElement e, int m, List *updates)
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
} }
}
/* /*
* Add a heap TID to an existing element * Add a heap TID to an existing element
@@ -391,7 +393,7 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
* Write changes to disk * Write changes to disk
*/ */
static void static void
WriteElement(Relation index, HnswElement element, int m, List *updates, HnswElement dup, HnswElement entryPoint) WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, List **neighbors, HnswElement dup, HnswElement entryPoint)
{ {
/* Try to add to existing page */ /* Try to add to existing page */
if (dup != NULL) if (dup != NULL)
@@ -402,7 +404,7 @@ WriteElement(Relation index, HnswElement element, int m, List *updates, HnswElem
/* If fails, take this path */ /* If fails, take this path */
WriteNewElementPages(index, element, m); WriteNewElementPages(index, element, m);
UpdateNeighborPages(index, element, m, updates); UpdateNeighborPages(index, procinfo, collation, element, m, neighbors);
/* Update metapage if needed */ /* Update metapage if needed */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -424,7 +426,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
double ml = HnswGetMl(m); double ml = HnswGetMl(m);
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
List *updates = NIL; List **neighbors;
HnswElement dup; HnswElement dup;
/* Detoast once for all calls */ /* Detoast once for all calls */
@@ -446,10 +448,10 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
entryPoint = HnswGetEntryPoint(index); entryPoint = HnswGetEntryPoint(index);
/* Insert element in graph */ /* Insert element in graph */
dup = HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, &updates, false); dup = HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, &neighbors, false);
/* Write to disk */ /* Write to disk */
WriteElement(index, element, m, updates, dup, entryPoint); WriteElement(index, procinfo, collation, element, m, neighbors, dup, entryPoint);
return true; return true;
} }

78
src/hnswutils.c
View File

@@ -372,8 +372,8 @@ LoadNeighborsFromPage(HnswElement element, Relation index, Page page)
/* /*
* Load neighbors * Load neighbors
*/ */
static void void
LoadNeighbors(HnswElement element, Relation index) HnswLoadNeighbors(HnswElement element, Relation index)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -571,7 +571,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
break; break;
if (c->element->neighbors == NULL) if (c->element->neighbors == NULL)
LoadNeighbors(c->element, index); HnswLoadNeighbors(c->element, index);
/* Get the neighborhood at layer lc */ /* Get the neighborhood at layer lc */
neighborhood = &c->element->neighbors[lc]; neighborhood = &c->element->neighbors[lc];
@@ -799,31 +799,12 @@ CompareCandidateDistances(const void *a, const void *b)
return 0; return 0;
} }
/*
* Create update
*/
static HnswUpdate *
CreateUpdate(HnswCandidate * hc, int level, int index)
{
HnswUpdate *update = palloc(sizeof(HnswUpdate));
update->hc = *hc;
update->level = level;
update->index = index;
return update;
}
/* /*
* Update connections * Update connections
*/ */
static void void
UpdateConnections(HnswElement element, List *neighbors, int m, int lc, List **updates, Relation index, FmgrInfo *procinfo, Oid collation) HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
{ {
ListCell *lc2;
foreach(lc2, neighbors)
{
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc]; HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc];
HnswCandidate hc2; HnswCandidate hc2;
@@ -835,9 +816,9 @@ UpdateConnections(HnswElement element, List *neighbors, int m, int lc, List **up
{ {
currentNeighbors->items[currentNeighbors->length++] = hc2; currentNeighbors->items[currentNeighbors->length++] = hc2;
/* Track updates */ /* Track update */
if (updates != NULL) if (updateIdx != NULL)
*updates = lappend(*updates, CreateUpdate(hc, lc, currentNeighbors->length - 1)); *updateIdx = currentNeighbors->length - 1;
} }
else else
{ {
@@ -880,7 +861,7 @@ UpdateConnections(HnswElement element, List *neighbors, int m, int lc, List **up
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
continue; return;
} }
/* Find and replace the pruned element */ /* Find and replace the pruned element */
@@ -890,28 +871,27 @@ UpdateConnections(HnswElement element, List *neighbors, int m, int lc, List **up
{ {
currentNeighbors->items[i] = hc2; currentNeighbors->items[i] = hc2;
/* Track updates */ /* Track update */
if (updates != NULL) if (updateIdx != NULL)
*updates = lappend(*updates, CreateUpdate(hc, lc, i)); *updateIdx = i;
break; break;
} }
} }
} }
} }
}
/* /*
* Algorithm 1 from paper * Algorithm 1 from paper
*/ */
HnswElement HnswElement
HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, List **updates, bool vacuuming) HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, List ***updateNeighbors, bool vacuuming)
{ {
List *ep = NIL; List *ep = NIL;
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
List **ws = palloc(sizeof(List *) * (level + 1)); List **neighbors = palloc(sizeof(List *) * (level + 1));
Datum q = PointerGetDatum(element->vec); Datum q = PointerGetDatum(element->vec);
HnswElement dup; HnswElement dup;
BlockNumber *skipPage = vacuuming ? &element->neighborPage : NULL; BlockNumber *skipPage = vacuuming ? &element->neighborPage : NULL;
@@ -940,20 +920,25 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
ep = w; ep = w;
} }
if (level > entryLevel) while (level > entryLevel)
level = entryLevel; {
neighbors[level] = NIL;
level--;
}
/* 2nd phase */ /* 2nd phase */
for (int lc = level; lc >= 0; lc--) for (int lc = level; lc >= 0; lc--)
{ {
int lm = HnswGetLayerM(m, lc);
w = HnswSearchLayer(q, ep, efConstruction, lc, index, procinfo, collation, true, skipPage, skipOffno); w = HnswSearchLayer(q, ep, efConstruction, lc, index, procinfo, collation, true, skipPage, skipOffno);
/* Remove entry point if it's being deleted */ /* Remove entry point if it's being deleted */
if (removeEntryPoint) if (removeEntryPoint)
w = list_delete_ptr(w, entryCandidate); w = list_delete_ptr(w, entryCandidate);
/* Save w for SelectNeighbors */ /* Always call on inserts since duplicate update can fail */
ws[lc] = w; neighbors[lc] = SelectNeighbors(w, lm, lc, procinfo, collation, NULL);
ep = w; ep = w;
} }
@@ -961,7 +946,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
/* Look for duplicate */ /* Look for duplicate */
if (level >= 0 && !vacuuming) if (level >= 0 && !vacuuming)
{ {
dup = HnswFindDuplicate(element, ws[0]); dup = HnswFindDuplicate(element, neighbors[0]);
if (dup != NULL) if (dup != NULL)
return dup; return dup;
} }
@@ -970,13 +955,20 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
for (int lc = level; lc >= 0; lc--) for (int lc = level; lc >= 0; lc--)
{ {
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
List *newNeighbors = SelectNeighbors(ws[lc], lm, lc, procinfo, collation, NULL);
AddConnections(element, newNeighbors, lm, lc); AddConnections(element, neighbors[lc], lm, lc);
if (!vacuuming) if (!vacuuming && updateNeighbors == NULL)
UpdateConnections(element, newNeighbors, lm, lc, updates, index, procinfo, collation); {
ListCell *lc2;
foreach(lc2, neighbors[lc])
HnswUpdateConnection(element, lfirst(lc2), lm, lc, NULL, index, procinfo, collation);
} }
}
if (updateNeighbors != NULL)
*updateNeighbors = neighbors;
return NULL; return NULL;
} }