From 63d5d121a3cb0141fc5f73e3cbd73f3562adb478 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Fri, 20 Sep 2024 16:20:27 -0700
Subject: [PATCH] Added HNSW_BENCH define [skip ci]

---
 src/hnsw.h                          | 19 +++++++++++++++++++
 src/hnswscan.c                      |  4 ++--
 test/t/040_hnsw_streaming_recall.pl | 16 ++++++++++------
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/src/hnsw.h b/src/hnsw.h
index 2ecd369..e6a5be5 100644
--- a/src/hnsw.h
+++ b/src/hnsw.h
@@ -12,6 +12,10 @@
 #include "utils/sampling.h"
 #include "vector.h"
 
+#ifdef HNSW_BENCH
+#include "portability/instr_time.h"
+#endif
+
 #define HNSW_MAX_DIM 2000
 #define HNSW_MAX_NNZ 1000
 
@@ -69,6 +73,21 @@
 #define HnswPageGetOpaque(page)	((HnswPageOpaque) PageGetSpecialPointer(page))
 #define HnswPageGetMeta(page)	((HnswMetaPageData *) PageGetContents(page))
 
+#ifdef HNSW_BENCH
+#define HnswBench(name, code) \
+	do { \
+		instr_time	start; \
+		instr_time	duration; \
+		INSTR_TIME_SET_CURRENT(start); \
+		(code); \
+		INSTR_TIME_SET_CURRENT(duration); \
+		INSTR_TIME_SUBTRACT(duration, start); \
+		elog(INFO, "%s: %.3f ms", name, INSTR_TIME_GET_MILLISEC(duration)); \
+	} while (0)
+#else
+#define HnswBench(name, code) (code)
+#endif
+
 #if PG_VERSION_NUM >= 150000
 #define RandomDouble() pg_prng_double(&pg_global_prng_state)
 #define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
diff --git a/src/hnswscan.c b/src/hnswscan.c
index b269f06..87337f3 100644
--- a/src/hnswscan.c
+++ b/src/hnswscan.c
@@ -180,7 +180,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		 */
 		LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
 
-		so->w = GetScanItems(scan, value);
+		HnswBench("scan iteration", so->w = GetScanItems(scan, value));
 
 		/* Release shared lock */
 		UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
@@ -206,7 +206,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 
 			LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
 
-			so->w = ResumeScanItems(scan);
+			HnswBench("scan iteration", so->w = ResumeScanItems(scan));
 
 			UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
 
diff --git a/test/t/040_hnsw_streaming_recall.pl b/test/t/040_hnsw_streaming_recall.pl
index dfef942..2c2eca4 100644
--- a/test/t/040_hnsw_streaming_recall.pl
+++ b/test/t/040_hnsw_streaming_recall.pl
@@ -8,6 +8,8 @@ my $node;
 my @queries = ();
 my @expected;
 my $limit = 20;
+my $dim = 3;
+my $array_sql = join(",", ('random()') x $dim);
 my @cs = (100, 1000);
 
 sub test_recall
@@ -58,18 +60,20 @@ $node->start;
 
 # Create table
 $node->safe_psql("postgres", "CREATE EXTENSION vector;");
-$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
+$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
 $node->safe_psql("postgres",
-	"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
+	"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
 );
 
 # Generate queries
 for (1 .. 20)
 {
-	my $r1 = rand();
-	my $r2 = rand();
-	my $r3 = rand();
-	push(@queries, "[$r1,$r2,$r3]");
+	my @r = ();
+	for (1 .. $dim)
+	{
+		push(@r, rand());
+	}
+	push(@queries, "[" . join(",", @r) . "]");
 }
 
 # Check each index type