Added bound option

src/ivfbuild.c

@@ -36,16 +36,11 @@
 #define CALLBACK_ITEM_POINTER HeapTuple hup
 #endif
 
-/*
- * Update build phase progress
- */
-static inline void
-UpdateProgress(int index, int64 val)
-{
 #if PG_VERSION_NUM >= 120000
-	pgstat_progress_update_param(index, val);
+#define UpdateProgress(index, val) pgstat_progress_update_param(index, val)
+#else
+#define UpdateProgress(index, val) ((void)val)
 #endif
-}
 
 /*
  * Callback for sampling
@@ -170,6 +165,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 +349,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 +505,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);
 }

src/ivfflat.c

@@ -13,6 +13,7 @@
 #endif
 
 int			ivfflat_probes;
+int			ivfflat_bound;
 static relopt_kind ivfflat_relopt_kind;
 
 /*
@@ -32,6 +33,10 @@ _PG_init(void)
 	DefineCustomIntVariable("ivfflat.probes", "Sets the number of probes",
 							"Valid range is 1..lists.", &ivfflat_probes,
 							1, 1, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
+
+	DefineCustomIntVariable("ivfflat.bound", "Sets the max results from index (experimental)",
+							NULL, &ivfflat_bound,
+							0, 0, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL);
 }
 
 /*

src/ivfflat.h

@@ -69,6 +69,7 @@
 
 /* Variables */
 extern int	ivfflat_probes;
+extern int	ivfflat_bound;
 
 typedef struct VectorArrayData
 {
@@ -120,6 +121,10 @@ typedef struct IvfflatBuildState
 	ListInfo   *listInfo;
 	Vector	   *normvec;
 
+#ifdef IVFFLAT_KMEANS_DEBUG
+	double		inertia;
+#endif
+
 	/* Sampling */
 	BlockSamplerData bs;
 	ReservoirStateData rstate;

src/ivfscan.c

@@ -111,6 +111,10 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	 */
 	BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
 
+	/* Set the max number of results */
+	if (ivfflat_bound > 0)
+		tuplesort_set_bound(so->sortstate, ivfflat_bound);
+
 	/* Search closest probes lists */
 	for (i = 0; i < so->probes; i++)
 	{

test/t/002_vacuum.pl

@@ -13,7 +13,7 @@ $node->start;
 $node->safe_psql("postgres", "CREATE EXTENSION vector;");
 $node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
 $node->safe_psql("postgres",
-  "INSERT INTO tst SELECT i%10, ARRAY[i%1000, i%333, i%55] FROM generate_series(1, 100000) i;"
+	"INSERT INTO tst SELECT i % 10, ARRAY[i % 1000, i % 333, i % 55] FROM generate_series(1, 100000) i;"
 );
 $node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
 
@@ -24,7 +24,7 @@ my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_id
 $node->safe_psql("postgres", "DELETE FROM tst;");
 $node->safe_psql("postgres", "VACUUM tst;");
 $node->safe_psql("postgres",
-  "INSERT INTO tst SELECT i % 10, ARRAY[i % 1000, i % 333, i % 55] FROM generate_series(1, 100000) i;"
+	"INSERT INTO tst SELECT i % 10, ARRAY[i % 1000, i % 333, i % 55] FROM generate_series(1, 100000) i;"
 );
 
 # Check size

test/t/004_centers.pl

@@ -0,0 +1,36 @@
+use strict;
+use warnings;
+use PostgresNode;
+use TestLib;
+use Test::More tests => 3;
+
+# Initialize node
+my $node = get_new_node('node');
+$node->init;
+$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",
+	"INSERT INTO tst SELECT i, '[1,2,3]' FROM generate_series(1, 10) i;"
+);
+
+sub test_centers
+{
+	my ($lists, $min) = @_;
+
+	my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING ivfflat (v) WITH (lists = $lists);");
+	is($ret, 0, $stderr);
+}
+
+# Test no error for duplicate centers
+test_centers(5);
+test_centers(10);
+
+$node->safe_psql("postgres",
+	"INSERT INTO tst SELECT i, '[4,5,6]' FROM generate_series(1, 10) i;"
+);
+
+# Test no error for duplicate centers
+test_centers(10);