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

Ensure centers and samples fit into maintenance_work_mem before allocating for IVFFlat index builds - closes #986

Browse Source
This commit is contained in:
Andrew Kane
2026-06-18 12:29:58 -07:00
parent 90cd2b4ee5
commit bdf19077db
4 changed files with 19 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/ivfbuild.c
View File

@@ -391,8 +391,8 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
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);
IvfflatCheckMemoryUsage(buildstate->memoryUsed);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->itemsize);
/* TODO Move allocation to page creation */
@@ -449,8 +449,8 @@ ComputeCenters(IvfflatBuildState * buildstate)
numSamples = 1;
/* Sample rows */
/* TODO Ensure within maintenance_work_mem */
buildstate->memoryUsed += VECTOR_ARRAY_SIZE(numSamples, buildstate->itemsize);
IvfflatCheckMemoryUsage(buildstate->memoryUsed);
buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions, buildstate->itemsize);
if (buildstate->heap != NULL)
{

1
src/ivfflat.h
View File

@@ -327,6 +327,7 @@ FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum);
Datum IvfflatNormValue(const IvfflatTypeInfo * typeInfo, Oid collation, Datum value);
bool IvfflatCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
void IvfflatNormVectors(const IvfflatTypeInfo * typeInfo, Oid collation, VectorArray arr, MemoryContext tmpCtx);
void IvfflatCheckMemoryUsage(Size totalSize);
int IvfflatGetLists(Relation index);
void IvfflatGetMetaPageInfo(Relation index, int *lists, int *dimensions);
void IvfflatUpdateList(Relation index, ListInfo listInfo, BlockNumber insertPage, BlockNumber originalInsertPage, BlockNumber startPage, ForkNumber forkNum);

7
src/ivfkmeans.c
View File

@@ -277,12 +277,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, const Ivff
Size totalSize = memoryUsed + newCentersSize + aggSize + centerCountsSize + closestCentersSize + lowerBoundSize + upperBoundSize + sSize + halfcdistSize + newcdistSize;
/* Check memory requirements */
/* Add one to error message to ceil */
if (totalSize > (Size) maintenance_work_mem * 1024L)
ereport(ERROR,
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("memory required is %zu MB, maintenance_work_mem is %d MB",
totalSize / (1024 * 1024) + 1, maintenance_work_mem / 1024)));
IvfflatCheckMemoryUsage(totalSize);
/* Ensure indexing does not overflow */
if (numCenters * numCenters > INT_MAX)

15
src/ivfutils.c
View File

@@ -6,6 +6,7 @@
#include "halfutils.h"
#include "halfvec.h"
#include "ivfflat.h"
#include "miscadmin.h"
#include "storage/bufmgr.h"
#include "utils/memutils.h"
#include "utils/relcache.h"
@@ -109,6 +110,20 @@ IvfflatNormVectors(const IvfflatTypeInfo * typeInfo, Oid collation, VectorArray
MemoryContextSwitchTo(oldCtx);
}
/*
* Check memory usage
*/
void
IvfflatCheckMemoryUsage(Size totalSize)
{
/* Add one to error message to ceil */
if (totalSize > (Size) maintenance_work_mem * 1024L)
ereport(ERROR,
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("memory required is %zu MB, maintenance_work_mem is %d MB",
totalSize / (1024 * 1024) + 1, maintenance_work_mem / 1024)));
}
/*
* New buffer
*/