From bdf19077db1061b4139199c67ba21281f3ed3766 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Thu, 18 Jun 2026 12:29:58 -0700
Subject: [PATCH] Ensure centers and samples fit into maintenance_work_mem
 before allocating for IVFFlat index builds - closes #986

---
 src/ivfbuild.c  |  4 ++--
 src/ivfflat.h   |  1 +
 src/ivfkmeans.c |  7 +------
 src/ivfutils.c  | 15 +++++++++++++++
 4 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/ivfbuild.c b/src/ivfbuild.c
index d2427a4..e35c352 100644
--- a/src/ivfbuild.c
+++ b/src/ivfbuild.c
@@ -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)
 	{
diff --git a/src/ivfflat.h b/src/ivfflat.h
index ef39f3a..34bcb97 100644
--- a/src/ivfflat.h
+++ b/src/ivfflat.h
@@ -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);
diff --git a/src/ivfkmeans.c b/src/ivfkmeans.c
index ee647c1..83e534f 100644
--- a/src/ivfkmeans.c
+++ b/src/ivfkmeans.c
@@ -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)
diff --git a/src/ivfutils.c b/src/ivfutils.c
index b7ec566..610fa54 100644
--- a/src/ivfutils.c
+++ b/src/ivfutils.c
@@ -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
  */