From a508b120c167af5a2c55e8860ed751d411a89c2b Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Sun, 24 Dec 2023 09:27:09 -0500
Subject: [PATCH] Added IVFFLAT_MEMORY flag to show memory usage [skip ci]

---
 src/ivfkmeans.c | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/ivfkmeans.c b/src/ivfkmeans.c
index a87edcb..1c73225 100644
--- a/src/ivfkmeans.c
+++ b/src/ivfkmeans.c
@@ -6,6 +6,10 @@
 #include "ivfflat.h"
 #include "miscadmin.h"
 
+#ifdef IVFFLAT_MEMORY
+#include "utils/memutils.h"
+#endif
+
 /*
  * Initialize with kmeans++
  *
@@ -151,6 +155,23 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
 	}
 }
 
+#ifdef IVFFLAT_MEMORY
+/*
+ * Show memory usage
+ */
+static void
+ShowMemoryUsage(Size estimatedSize)
+{
+#if PG_VERSION_NUM >= 130000
+	elog(INFO, "total memory: %zu MB",
+		 MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
+#else
+	MemoryContextStats(CurrentMemoryContext);
+#endif
+	elog(INFO, "estimated memory: %zu MB", estimatedSize / (1024 * 1024));
+}
+#endif
+
 /*
  * Use Elkan for performance. This requires distance function to satisfy triangle inequality.
  *
@@ -231,6 +252,10 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
 		vec->dim = dimensions;
 	}
 
+#ifdef IVFFLAT_MEMORY
+	ShowMemoryUsage(totalSize);
+#endif
+
 	/* Pick initial centers */
 	InitCenters(index, samples, centers, lowerBound);