From 1bc29ac697de5bd3c5eedb49ae42c59ac2fae3a6 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Sat, 13 Jul 2024 08:47:53 -0700
Subject: [PATCH] Added tuple stats for HNSW [skip ci]

---
 src/hnsw.c      |  8 ++++++++
 src/hnsw.h      |  4 ++++
 src/hnswscan.c  | 12 +++++++++++-
 src/hnswutils.c |  5 +++++
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/hnsw.c b/src/hnsw.c
index 0850a69..6d9cd7c 100644
--- a/src/hnsw.c
+++ b/src/hnsw.c
@@ -20,6 +20,10 @@ int			hnsw_ef_search;
 int			hnsw_lock_tranche_id;
 static relopt_kind hnsw_relopt_kind;
 
+#ifdef HNSW_STATS
+int			hnsw_tuples;
+#endif
+
 /*
  * Assign a tranche ID for our LWLocks. This only needs to be done by one
  * backend, as the tranche ID is remembered in shared memory.
@@ -135,6 +139,10 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
 
 	genericcostestimate(root, path, loop_count, &costs);
 
+#ifdef HNSW_STATS
+	elog(INFO, "estimated tuples = %0.f", costs.numIndexTuples);
+#endif
+
 	/* Use total cost since most work happens before first tuple is returned */
 	*indexStartupCost = costs.indexTotalCost;
 	*indexTotalCost = costs.indexTotalCost;
diff --git a/src/hnsw.h b/src/hnsw.h
index 480ad9f..ef514a1 100644
--- a/src/hnsw.h
+++ b/src/hnsw.h
@@ -113,6 +113,10 @@
 extern int	hnsw_ef_search;
 extern int	hnsw_lock_tranche_id;
 
+#ifdef HNSW_STATS
+extern int	hnsw_tuples;
+#endif
+
 typedef struct HnswElementData HnswElementData;
 typedef struct HnswNeighborArray HnswNeighborArray;
 
diff --git a/src/hnswscan.c b/src/hnswscan.c
index 0efbaa1..7b31f6f 100644
--- a/src/hnswscan.c
+++ b/src/hnswscan.c
@@ -31,13 +31,23 @@ GetScanItems(IndexScanDesc scan, Datum q)
 
 	ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, false));
 
+#ifdef HNSW_STATS
+	hnsw_tuples = 1;
+#endif
+
 	for (int lc = entryPoint->level; lc >= 1; lc--)
 	{
 		w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL);
 		ep = w;
 	}
 
-	return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
+	w = HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
+
+#ifdef HNSW_STATS
+	elog(INFO, "total tuples = %d", hnsw_tuples);
+#endif
+
+	return w;
 }
 
 /*
diff --git a/src/hnswutils.c b/src/hnswutils.c
index f7ce99a..e41585d 100644
--- a/src/hnswutils.c
+++ b/src/hnswutils.c
@@ -801,6 +801,11 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
 				HnswElement eElement = HnswPtrAccess(base, e->element);
 				bool		alwaysAdd = wlen < ef;
 
+#ifdef HNSW_STATS
+				if (!inserting)
+					hnsw_tuples++;
+#endif
+
 				f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
 
 				if (index == NULL)