From fa5e90585d96f50f06ffbf03db8535f55eb73732 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Sat, 12 Feb 2022 21:23:53 -0800
Subject: [PATCH] Added inertia

---
 src/ivfbuild.c | 16 +++++++++++++++-
 src/ivfflat.h  |  4 ++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/ivfbuild.c b/src/ivfbuild.c
index 51c6a55..65e4f38 100644
--- a/src/ivfbuild.c
+++ b/src/ivfbuild.c
@@ -170,6 +170,10 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
 		}
 	}
 
+#ifdef IVFFLAT_KMEANS_DEBUG
+	buildstate->inertia += minDistance;
+#endif
+
 	/* Create a virtual tuple */
 	ExecClearTuple(slot);
 	slot->tts_values[0] = Int32GetDatum(closestCenter);
@@ -350,6 +354,10 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
 
 	/* Reuse for each tuple */
 	buildstate->normvec = InitVector(buildstate->dimensions);
+
+#ifdef IVFFLAT_KMEANS_DEBUG
+	buildstate->inertia = 0;
+#endif
 }
 
 /*
@@ -502,8 +510,14 @@ CreateEntryPages(IvfflatBuildState * buildstate, ForkNumber forkNum)
 #endif
 	}
 
-	/* Sort and insert */
+	/* Sort */
 	tuplesort_performsort(buildstate->sortstate);
+
+#ifdef IVFFLAT_KMEANS_DEBUG
+	elog(INFO, "inertia: %.3e", buildstate->inertia);
+#endif
+
+	/* Insert */
 	InsertTuples(buildstate->index, buildstate, forkNum);
 	tuplesort_end(buildstate->sortstate);
 }
diff --git a/src/ivfflat.h b/src/ivfflat.h
index 0533a83..9ec14f1 100644
--- a/src/ivfflat.h
+++ b/src/ivfflat.h
@@ -120,6 +120,10 @@ typedef struct IvfflatBuildState
 	ListInfo   *listInfo;
 	Vector	   *normvec;
 
+#ifdef IVFFLAT_KMEANS_DEBUG
+	double inertia;
+#endif
+
 	/* Sampling */
 	BlockSamplerData bs;
 	ReservoirStateData rstate;