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

Moved graph to separate struct

Browse Source
This commit is contained in:
Andrew Kane
2024-01-07 20:15:30 -08:00
parent c7fe1571ee
commit 19a0e1b341
2 changed files with 60 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/hnsw.h
View File

@@ -148,6 +148,15 @@ typedef struct HnswOptions
int efConstruction; /* size of dynamic candidate list */ int efConstruction; /* size of dynamic candidate list */
} HnswOptions; } HnswOptions;
typedef struct HnswGraph
{
slist_head elements;
HnswElement entryPoint;
long memoryUsed;
long memoryTotal;
bool flushed;
} HnswGraph;
typedef struct HnswSpool typedef struct HnswSpool
{ {
Relation heap; Relation heap;
@@ -172,6 +181,7 @@ typedef struct HnswShared
int nparticipantsdone; int nparticipantsdone;
double reltuples; double reltuples;
double indtuples; double indtuples;
HnswGraph graphData;
#if PG_VERSION_NUM < 120000 #if PG_VERSION_NUM < 120000
ParallelHeapScanDescData heapdesc; /* must come last */ ParallelHeapScanDescData heapdesc; /* must come last */
@@ -214,13 +224,10 @@ typedef struct HnswBuildState
Oid collation; Oid collation;
/* Variables */ /* Variables */
slist_head elements; HnswGraph graphData;
HnswElement entryPoint; HnswGraph *graph;
double ml; double ml;
int maxLevel; int maxLevel;
long memoryUsed;
long memoryTotal;
bool flushed;
Vector *normvec; Vector *normvec;
/* Memory */ /* Memory */

79
src/hnswbuild.c
View File

@@ -152,7 +152,7 @@ CreateElementPages(HnswBuildState * buildstate)
page = BufferGetPage(buf); page = BufferGetPage(buf);
HnswInitPage(buf, page); HnswInitPage(buf, page);
slist_foreach(iter, &buildstate->elements) slist_foreach(iter, &buildstate->graph->elements)
{ {
HnswElement element = slist_container(HnswElementData, next, iter.cur); HnswElement element = slist_container(HnswElementData, next, iter.cur);
Size etupSize; Size etupSize;
@@ -212,7 +212,7 @@ CreateElementPages(HnswBuildState * buildstate)
MarkBufferDirty(buf); MarkBufferDirty(buf);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->entryPoint, insertPage, forkNum, true); HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->graph->entryPoint, insertPage, forkNum, true);
pfree(etup); pfree(etup);
pfree(ntup); pfree(ntup);
@@ -233,7 +233,7 @@ CreateNeighborPages(HnswBuildState * buildstate)
/* Allocate once */ /* Allocate once */
ntup = palloc0(BLCKSZ); ntup = palloc0(BLCKSZ);
slist_foreach(iter, &buildstate->elements) slist_foreach(iter, &buildstate->graph->elements)
{ {
HnswElement e = slist_container(HnswElementData, next, iter.cur); HnswElement e = slist_container(HnswElementData, next, iter.cur);
Buffer buf; Buffer buf;
@@ -293,7 +293,7 @@ FlushPages(HnswBuildState * buildstate)
CreateElementPages(buildstate); CreateElementPages(buildstate);
CreateNeighborPages(buildstate); CreateNeighborPages(buildstate);
buildstate->flushed = true; buildstate->graph->flushed = true;
MemoryContextReset(buildstate->graphCtx); MemoryContextReset(buildstate->graphCtx);
} }
@@ -324,7 +324,8 @@ InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuil
{ {
FmgrInfo *procinfo = buildstate->procinfo; FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation; Oid collation = buildstate->collation;
HnswElement entryPoint = buildstate->entryPoint; HnswGraph *graph = buildstate->graph;
HnswElement entryPoint = graph->entryPoint;
int efConstruction = buildstate->efConstruction; int efConstruction = buildstate->efConstruction;
int m = buildstate->m; int m = buildstate->m;
MemoryContext oldCtx; MemoryContext oldCtx;
@@ -356,7 +357,7 @@ InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuil
if (dup == NULL) if (dup == NULL)
{ {
/* Add element */ /* Add element */
slist_push_head(&buildstate->elements, &element->next); slist_push_head(&graph->elements, &element->next);
/* Update neighbors */ /* Update neighbors */
for (int lc = element->level; lc >= 0; lc--) for (int lc = element->level; lc >= 0; lc--)
@@ -370,7 +371,7 @@ InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuil
/* Update entry point if needed */ /* Update entry point if needed */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
buildstate->entryPoint = element; graph->entryPoint = element;
} }
else else
{ {
@@ -381,9 +382,9 @@ InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuil
/* Update memory usage */ /* Update memory usage */
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
buildstate->memoryUsed = MemoryContextMemAllocated(buildstate->graphCtx, false); graph->memoryUsed = MemoryContextMemAllocated(buildstate->graphCtx, false);
#else #else
buildstate->memoryUsed += HnswElementMemory(element, buildstate->m); graph->memoryUsed += HnswElementMemory(element, buildstate->m);
#endif #endif
return true; return true;
@@ -397,6 +398,8 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
bool *isnull, bool tupleIsAlive, void *state) bool *isnull, bool tupleIsAlive, void *state)
{ {
HnswBuildState *buildstate = (HnswBuildState *) state; HnswBuildState *buildstate = (HnswBuildState *) state;
HnswGraph *graph = buildstate->graph;
HnswShared *hnswshared = buildstate->hnswshared;
MemoryContext oldCtx; MemoryContext oldCtx;
bool inserted; bool inserted;
@@ -409,20 +412,28 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
return; return;
/* Flush pages if needed */ /* Flush pages if needed */
if (!buildstate->flushed && buildstate->memoryUsed >= buildstate->memoryTotal) if (!graph->flushed && graph->memoryUsed >= graph->memoryTotal)
{ {
ereport(NOTICE, if (hnswshared)
(errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) buildstate->indtuples), SpinLockAcquire(&hnswshared->mutex);
errdetail("Building will take significantly more time."),
errhint("Increase maintenance_work_mem to speed up builds.")));
FlushPages(buildstate); if (!hnswshared)
ereport(NOTICE,
(errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) buildstate->indtuples),
errdetail("Building will take significantly more time."),
errhint("Increase maintenance_work_mem to speed up builds.")));
if (!graph->flushed)
FlushPages(buildstate);
if (hnswshared)
SpinLockRelease(&hnswshared->mutex);
} }
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Insert tuple */ /* Insert tuple */
if (buildstate->flushed) if (graph->flushed)
inserted = HnswInsertTuple(index, values, isnull, tid, buildstate->heap, true); inserted = HnswInsertTuple(index, values, isnull, tid, buildstate->heap, true);
else else
inserted = InsertTupleInMemory(index, values, tid, buildstate); inserted = InsertTupleInMemory(index, values, tid, buildstate);
@@ -430,8 +441,6 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
/* Update progress */ /* Update progress */
if (inserted) if (inserted)
{ {
HnswShared *hnswshared = buildstate->hnswshared;
if (hnswshared) if (hnswshared)
{ {
SpinLockAcquire(&hnswshared->mutex); SpinLockAcquire(&hnswshared->mutex);
@@ -447,6 +456,19 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
MemoryContextReset(buildstate->tmpCtx); MemoryContextReset(buildstate->tmpCtx);
} }
/*
* Initialize the graph
*/
static void
InitGraph(HnswGraph * graph)
{
slist_init(&graph->elements);
graph->entryPoint = NULL;
graph->memoryUsed = 0;
graph->memoryTotal = maintenance_work_mem * 1024L;
graph->flushed = false;
}
/* /*
* Initialize the build state * Initialize the build state
*/ */
@@ -480,13 +502,10 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->collation = index->rd_indcollation[0]; buildstate->collation = index->rd_indcollation[0];
slist_init(&buildstate->elements); InitGraph(&buildstate->graphData);
buildstate->entryPoint = NULL; buildstate->graph = &buildstate->graphData;
buildstate->ml = HnswGetMl(buildstate->m); buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m); buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->memoryUsed = 0;
buildstate->memoryTotal = maintenance_work_mem * 1024L;
buildstate->flushed = false;
/* Reuse for each tuple */ /* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions); buildstate->normvec = InitVector(buildstate->dimensions);
@@ -532,6 +551,7 @@ ParallelHeapScan(HnswBuildState * buildstate)
SpinLockAcquire(&hnswshared->mutex); SpinLockAcquire(&hnswshared->mutex);
if (hnswshared->nparticipantsdone == nparticipanttuplesorts) if (hnswshared->nparticipantsdone == nparticipanttuplesorts)
{ {
buildstate->graph = &hnswshared->graphData;
buildstate->indtuples = hnswshared->indtuples; buildstate->indtuples = hnswshared->indtuples;
reltuples = hnswshared->reltuples; reltuples = hnswshared->reltuples;
SpinLockRelease(&hnswshared->mutex); SpinLockRelease(&hnswshared->mutex);
@@ -567,9 +587,7 @@ HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool p
indexInfo = BuildIndexInfo(hnswspool->index); indexInfo = BuildIndexInfo(hnswspool->index);
indexInfo->ii_Concurrent = hnswshared->isconcurrent; indexInfo->ii_Concurrent = hnswshared->isconcurrent;
InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM); InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM);
/* TODO Support in-memory builds */ buildstate.graph = &hnswshared->graphData;
buildstate.memoryTotal = 0;
buildstate.flushed = true;
buildstate.hnswshared = hnswshared; buildstate.hnswshared = hnswshared;
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
scan = table_beginscan_parallel(hnswspool->heap, scan = table_beginscan_parallel(hnswspool->heap,
@@ -795,6 +813,9 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
hnswshared->nparticipantsdone = 0; hnswshared->nparticipantsdone = 0;
hnswshared->reltuples = 0; hnswshared->reltuples = 0;
hnswshared->indtuples = 0; hnswshared->indtuples = 0;
InitGraph(&hnswshared->graphData);
/* TODO Support in-memory builds */
hnswshared->graphData.memoryTotal = 0;
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
table_parallelscan_initialize(buildstate->heap, table_parallelscan_initialize(buildstate->heap,
ParallelTableScanFromHnswShared(hnswshared), ParallelTableScanFromHnswShared(hnswshared),
@@ -882,11 +903,7 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
/* Attempt to launch parallel worker scan when required */ /* Attempt to launch parallel worker scan when required */
if (parallel_workers > 0) if (parallel_workers > 0)
{
/* TODO Support in-memory builds */
FlushPages(buildstate);
HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers); HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
}
/* Add tuples to graph */ /* Add tuples to graph */
if (buildstate->heap != NULL) if (buildstate->heap != NULL)
@@ -906,7 +923,7 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
} }
/* Flush pages */ /* Flush pages */
if (!buildstate->flushed) if (!buildstate->graph->flushed)
FlushPages(buildstate); FlushPages(buildstate);
/* End parallel build */ /* End parallel build */