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

Improved memory tracking for IVFFlat index builds [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2026-06-18 11:57:03 -07:00
parent b44d1b4c5f
commit 90cd2b4ee5
3 changed files with 10 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/ivfbuild.c
View File

@@ -388,9 +388,11 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual);
buildstate->memoryUsed = 0;
buildstate->itemsize = buildstate->typeInfo->itemSize(buildstate->dimensions);
/* TODO Ensure within maintenance_work_mem */
buildstate->memoryUsed += VECTOR_ARRAY_SIZE(buildstate->lists, buildstate->itemsize);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->itemsize);
/* TODO Move allocation to page creation */
@@ -448,6 +450,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
/* Sample rows */
/* TODO Ensure within maintenance_work_mem */
buildstate->memoryUsed += VECTOR_ARRAY_SIZE(numSamples, buildstate->itemsize);
buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions, buildstate->itemsize);
if (buildstate->heap != NULL)
{
@@ -463,7 +466,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
}
/* Calculate centers */
IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers, buildstate->typeInfo));
IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers, buildstate->typeInfo, buildstate->memoryUsed));
/* Free samples before we allocate more memory */
VectorArrayFree(buildstate->samples);