pgvector.git / summary / log / commit / refs

commit    f59b1184fcd1fbc261058124eaeacbb24496f35d
Author:   Andrew Kane <andrew@ankane.org>
Date:     Fri Feb 11 11:33:43 2022 +0000

    Calculate all allocation sizes


 src/ivfkmeans.c | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/ivfkmeans.c b/src/ivfkmeans.c
index 66f08b59..763dc76d 100644
--- a/src/ivfkmeans.c
+++ b/src/ivfkmeans.c
@@ -190,10 +190,18 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
 	bool		rjreset;
 	double		dxcx;
 	double		dxc;
+
+	/* Calculate allocation sizes */
+	Size		centerCountsSize = sizeof(int) * numCenters;
+	Size		closestCentersSize = sizeof(int) * numSamples;
 	Size		lowerBoundSize = sizeof(float) * numSamples * numCenters;
+	Size		upperBoundSize = sizeof(float) * numSamples;
+	Size		sSize = sizeof(float) * numCenters;
+	Size		halfcdistSize = sizeof(float) * numCenters * numCenters;
+	Size		newcdistSize = sizeof(float) * numCenters;
 
-	/* TODO Add other allocations */
-	Size		totalSize = lowerBoundSize;
+	/* Calculate total size */
+	Size		totalSize = centerCountsSize + closestCentersSize + lowerBoundSize + upperBoundSize + sSize + halfcdistSize + newcdistSize;
 
 	/* Check memory requirements */
 	if (totalSize / 1024 > maintenance_work_mem)
@@ -209,13 +217,13 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
 
 	/* Allocate space */
 	/* Use float instead of double to save memory */
-	centerCounts = palloc(sizeof(int) * numCenters);
-	closestCenters = palloc(sizeof(int) * numSamples);
+	centerCounts = palloc(centerCountsSize);
+	closestCenters = palloc(closestCentersSize);
 	lowerBound = palloc_extended(lowerBoundSize, MCXT_ALLOC_HUGE);
-	upperBound = palloc(sizeof(float) * numSamples);
-	s = palloc(sizeof(float) * numCenters);
-	halfcdist = palloc(sizeof(float) * numCenters * numCenters);
-	newcdist = palloc(sizeof(float) * numCenters);
+	upperBound = palloc(upperBoundSize);
+	s = palloc(sSize);
+	halfcdist = palloc(halfcdistSize);
+	newcdist = palloc(newcdistSize);
 
 	newCenters = VectorArrayInit(numCenters, dimensions);
 	for (j = 0; j < numCenters; j++)


[parent: cab103a2e74e]