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

Reduced memory usage for HNSW index builds - #384

Browse Source 
Co-authored-by: Heikki Linnakangas <heikki.linnakangas@iki.fi>
This commit is contained in:
Andrew Kane
2023-12-22 12:41:47 -05:00
parent f61d4087b5
commit dcbe0b6f0d
3 changed files with 18 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGELOG.md
View File

@@ -2,6 +2,7 @@
- Improved performance of HNSW - Improved performance of HNSW
- Added support for on-disk parallel index builds for HNSW - Added support for on-disk parallel index builds for HNSW
- Reduced memory usage for HNSW index builds
- Reduced WAL generation for HNSW index builds - Reduced WAL generation for HNSW index builds
- Fixed `invalid memory alloc request size` error with HNSW index build - Fixed `invalid memory alloc request size` error with HNSW index build

1
src/hnsw.h
View File

@@ -220,6 +220,7 @@ typedef struct HnswBuildState
Vector *normvec; Vector *normvec;
/* Memory */ /* Memory */
MemoryContext graphCtx;
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Parallel builds */ /* Parallel builds */

38
src/hnswbuild.c
View File

@@ -259,20 +259,6 @@ CreateNeighborPages(HnswBuildState * buildstate)
pfree(ntup); pfree(ntup);
} }
/*
* Free elements
*/
static void
FreeElements(HnswBuildState * buildstate)
{
ListCell *lc;
foreach(lc, buildstate->elements)
HnswFreeElement(lfirst(lc));
list_free(buildstate->elements);
}
/* /*
* Flush pages * Flush pages
*/ */
@@ -281,9 +267,9 @@ FlushPages(HnswBuildState * buildstate)
{ {
#ifdef HNSW_MEMORY #ifdef HNSW_MEMORY
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, false) / (1024 * 1024)); elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(buildstate->graphCtx, false) / (1024 * 1024));
#else #else
MemoryContextStats(CurrentMemoryContext); MemoryContextStats(buildstate->graphCtx);
#endif #endif
#endif #endif
@@ -292,14 +278,14 @@ FlushPages(HnswBuildState * buildstate)
CreateNeighborPages(buildstate); CreateNeighborPages(buildstate);
buildstate->flushed = true; buildstate->flushed = true;
FreeElements(buildstate); MemoryContextReset(buildstate->graphCtx);
} }
/* /*
* Insert tuple * Insert tuple
*/ */
static bool static bool
InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup, MemoryContext outerCtx) InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup)
{ {
FmgrInfo *procinfo = buildstate->procinfo; FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation; Oid collation = buildstate->collation;
@@ -319,7 +305,7 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
} }
/* Copy value to element so accessible outside of memory context */ /* Copy value to element so accessible outside of memory context */
oldCtx = MemoryContextSwitchTo(outerCtx); oldCtx = MemoryContextSwitchTo(buildstate->graphCtx);
element->value = datumCopy(value, false, -1); element->value = datumCopy(value, false, -1);
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -428,18 +414,22 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
} }
/* Allocate necessary memory outside of memory context */ /* Allocate necessary memory outside of memory context */
oldCtx = MemoryContextSwitchTo(buildstate->graphCtx);
element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel); element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel);
MemoryContextSwitchTo(oldCtx);
/* Use memory context since detoast can allocate */ /* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Insert tuple */ /* Insert tuple */
inserted = InsertTuple(index, values, element, buildstate, &dup, oldCtx); inserted = InsertTuple(index, values, element, buildstate, &dup);
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx); MemoryContextReset(buildstate->tmpCtx);
oldCtx = MemoryContextSwitchTo(buildstate->graphCtx);
/* Add outside memory context */ /* Add outside memory context */
if (dup != NULL) if (dup != NULL)
{ {
@@ -453,8 +443,8 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
buildstate->elements = lappend(buildstate->elements, element); buildstate->elements = lappend(buildstate->elements, element);
buildstate->memoryLeft -= HnswElementMemory(element, buildstate->m); buildstate->memoryLeft -= HnswElementMemory(element, buildstate->m);
} }
else
HnswFreeElement(element); MemoryContextSwitchTo(oldCtx);
} }
/* /*
@@ -500,6 +490,9 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
/* Reuse for each tuple */ /* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions); buildstate->normvec = InitVector(buildstate->dimensions);
buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw build graph context",
1024 * 1024, 1024 * 1024, 1024 * 1024);
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw build temporary context", "Hnsw build temporary context",
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
@@ -515,6 +508,7 @@ static void
FreeBuildState(HnswBuildState * buildstate) FreeBuildState(HnswBuildState * buildstate)
{ {
pfree(buildstate->normvec); pfree(buildstate->normvec);
MemoryContextDelete(buildstate->graphCtx);
MemoryContextDelete(buildstate->tmpCtx); MemoryContextDelete(buildstate->tmpCtx);
} }