From d1ebb8db73c8d33f18cd77c11871fd0c7b6874b2 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Fri, 11 Oct 2024 11:43:32 -0700
Subject: [PATCH] Use -1 for no limit for ivfflat.max_probes [skip ci]

---
 src/hnsw.c                       |  1 -
 src/ivfflat.c                    |  4 ++--
 src/ivfscan.c                    | 20 ++++++++++++++------
 test/expected/ivfflat_vector.out | 31 +++++++++++++++++++++++++++----
 test/sql/ivfflat_vector.sql      | 12 +++++++++++-
 5 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/src/hnsw.c b/src/hnsw.c
index fdd35b3..d1bc233 100644
--- a/src/hnsw.c
+++ b/src/hnsw.c
@@ -82,7 +82,6 @@ HnswInit(void)
 							 NULL, &hnsw_iterative_search,
 							 HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
 
-	/* TODO Ensure ivfflat.max_probes uses same value for no limit */
 	DefineCustomIntVariable("hnsw.max_search_tuples", "Sets the max number of candidates to visit for iterative search",
 							"-1 means no limit", &hnsw_max_search_tuples,
 							-1, -1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL);
diff --git a/src/ivfflat.c b/src/ivfflat.c
index 4c02a5a..531bc3a 100644
--- a/src/ivfflat.c
+++ b/src/ivfflat.c
@@ -46,8 +46,8 @@ IvfflatInit(void)
 							 IVFFLAT_ITERATIVE_SEARCH_OFF, ivfflat_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
 
 	DefineCustomIntVariable("ivfflat.max_probes", "Sets the max number of probes for iterative search",
-							"Zero sets to the number of lists", &ivfflat_max_probes,
-							0, 0, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
+							"-1 means no limit", &ivfflat_max_probes,
+							-1, -1, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
 
 	MarkGUCPrefixReserved("ivfflat");
 }
diff --git a/src/ivfscan.c b/src/ivfscan.c
index 548af81..bc795e3 100644
--- a/src/ivfscan.c
+++ b/src/ivfscan.c
@@ -259,19 +259,27 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	/* Get lists and dimensions from metapage */
 	IvfflatGetMetaPageInfo(index, &lists, &dimensions);
 
-	if (probes > lists)
-		probes = lists;
-
 	if (ivfflat_iterative_search != IVFFLAT_ITERATIVE_SEARCH_OFF)
 	{
-		if (ivfflat_max_probes == 0)
+		maxProbes = ivfflat_max_probes;
+
+		if (maxProbes < 0)
 			maxProbes = lists;
-		else
-			maxProbes = Min(ivfflat_max_probes, lists);
+		else if (maxProbes < probes)
+		{
+			/* TODO Show notice */
+			maxProbes = probes;
+		}
 	}
 	else
 		maxProbes = probes;
 
+	if (probes > lists)
+		probes = lists;
+
+	if (maxProbes > lists)
+		maxProbes = lists;
+
 	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + maxProbes * sizeof(IvfflatScanList));
 	so->typeInfo = IvfflatGetTypeInfo(index);
 	so->first = true;
diff --git a/test/expected/ivfflat_vector.out b/test/expected/ivfflat_vector.out
index 9bb7b13..8a80ea3 100644
--- a/test/expected/ivfflat_vector.out
+++ b/test/expected/ivfflat_vector.out
@@ -95,7 +95,30 @@ SELECT * FROM t ORDER BY val <-> '[3,3,3]';
  [0,0,0]
 (3 rows)
 
+SET ivfflat.max_probes = 0;
+SELECT * FROM t ORDER BY val <-> '[3,3,3]';
+   val   
+---------
+ [1,2,3]
+(1 row)
+
+SET ivfflat.max_probes = 1;
+SELECT * FROM t ORDER BY val <-> '[3,3,3]';
+   val   
+---------
+ [1,2,3]
+(1 row)
+
+SET ivfflat.max_probes = 2;
+SELECT * FROM t ORDER BY val <-> '[3,3,3]';
+   val   
+---------
+ [1,2,3]
+ [1,1,1]
+(2 rows)
+
 RESET ivfflat.iterative_search;
+RESET ivfflat.max_probes;
 DROP TABLE t;
 -- unlogged
 CREATE UNLOGGED TABLE t (val vector(3));
@@ -140,11 +163,11 @@ HINT:  Available values: off, relaxed_order.
 SHOW ivfflat.max_probes;
  ivfflat.max_probes 
 --------------------
- 0
+ -1
 (1 row)
 
-SET ivfflat.max_probes = -1;
-ERROR:  -1 is outside the valid range for parameter "ivfflat.max_probes" (0 .. 32768)
+SET ivfflat.max_probes = -2;
+ERROR:  -2 is outside the valid range for parameter "ivfflat.max_probes" (-1 .. 32768)
 SET ivfflat.max_probes = 32769;
-ERROR:  32769 is outside the valid range for parameter "ivfflat.max_probes" (0 .. 32768)
+ERROR:  32769 is outside the valid range for parameter "ivfflat.max_probes" (-1 .. 32768)
 DROP TABLE t;
diff --git a/test/sql/ivfflat_vector.sql b/test/sql/ivfflat_vector.sql
index 7e83915..9e20060 100644
--- a/test/sql/ivfflat_vector.sql
+++ b/test/sql/ivfflat_vector.sql
@@ -53,7 +53,17 @@ CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 3);
 SET ivfflat.iterative_search = relaxed_order;
 SELECT * FROM t ORDER BY val <-> '[3,3,3]';
 
+SET ivfflat.max_probes = 0;
+SELECT * FROM t ORDER BY val <-> '[3,3,3]';
+
+SET ivfflat.max_probes = 1;
+SELECT * FROM t ORDER BY val <-> '[3,3,3]';
+
+SET ivfflat.max_probes = 2;
+SELECT * FROM t ORDER BY val <-> '[3,3,3]';
+
 RESET ivfflat.iterative_search;
+RESET ivfflat.max_probes;
 DROP TABLE t;
 
 -- unlogged
@@ -83,7 +93,7 @@ SET ivfflat.iterative_search = on;
 
 SHOW ivfflat.max_probes;
 
-SET ivfflat.max_probes = -1;
+SET ivfflat.max_probes = -2;
 SET ivfflat.max_probes = 32769;
 
 DROP TABLE t;