Merge branch 'master' into minibatch

CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.2.6 (unreleased)
+
+- Switched to mini-batch k-means
+- Improved performance of index creation for Postgres < 12
+
 ## 0.2.5 (2022-02-11)
 
 - Reduced memory usage during index creation

README.md

@@ -119,10 +119,9 @@ SELECT phase, tuples_done, tuples_total FROM pg_stat_progress_create_index;
 The phases are:
 
 1. `initializing`
-2. `sampling table`
+2. `performing k-means`
-3. `performing k-means`
+3. `sorting tuples`
-4. `sorting tuples`
+4. `loading tuples`
-5. `loading tuples`
 
 Note: `tuples_done` and `tuples_total` are only populated during the `loading tuples` phase
 
@@ -264,7 +263,7 @@ Thanks to:
 
 - [PASE: PostgreSQL Ultra-High-Dimensional Approximate Nearest Neighbor Search Extension](https://dl.acm.org/doi/pdf/10.1145/3318464.3386131)
 - [Faiss: A Library for Efficient Similarity Search and Clustering of Dense Vectors](https://github.com/facebookresearch/faiss)
-- [Using the Triangle Inequality to Accelerate k-means](https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf)
+- [Web-Scale k-means Clustering](https://www.eecs.tufts.edu/~dsculley/papers/fastkmeans.pdf)
 - [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf)
 - [Concept Decompositions for Large Sparse Text Data using Clustering](https://www.cs.utexas.edu/users/inderjit/public_papers/concept_mlj.pdf)
 

src/ivfbuild.c

@@ -36,97 +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)
-#endif
-}
-
-/*
- * Callback for sampling
- */
-static void
-SampleCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
-			   bool *isnull, bool tupleIsAlive, void *state)
-{
-	IvfflatBuildState *buildstate = (IvfflatBuildState *) state;
-	VectorArray samples = buildstate->samples;
-	int			targsamples = samples->maxlen;
-	Datum		value = values[0];
-
-	/* Skip nulls */
-	if (isnull[0])
-		return;
-
-	/*
-	 * Normalize with KMEANS_NORM_PROC since spherical distance function
-	 * expects unit vectors
-	 */
-	if (buildstate->kmeansnormprocinfo != NULL)
-	{
-		if (!IvfflatNormValue(buildstate->kmeansnormprocinfo, buildstate->collation, &value, buildstate->normvec))
-			return;
-	}
-
-	if (samples->length < targsamples)
-	{
-		VectorArraySet(samples, samples->length, DatumGetVector(value));
-		samples->length++;
-	}
-	else
-	{
-		if (buildstate->rowstoskip < 0)
-			buildstate->rowstoskip = reservoir_get_next_S(&buildstate->rstate, samples->length, targsamples);
-
-		if (buildstate->rowstoskip <= 0)
-		{
-			int			k = (int) (targsamples * sampler_random_fract(buildstate->rstate.randstate));
-
-			Assert(k >= 0 && k < targsamples);
-			VectorArraySet(samples, k, DatumGetVector(value));
-		}
-
-		buildstate->rowstoskip -= 1;
-	}
-}
-
-/*
- * Sample rows with same logic as ANALYZE
- */
-static void
-SampleRows(IvfflatBuildState * buildstate)
-{
-	int			targsamples = buildstate->samples->maxlen;
-	BlockNumber totalblocks = RelationGetNumberOfBlocks(buildstate->heap);
-
-	UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_SAMPLE);
-
-	buildstate->rowstoskip = -1;
-
-	BlockSampler_Init(&buildstate->bs, totalblocks, targsamples, random());
-
-	reservoir_init_selection_state(&buildstate->rstate, targsamples);
-	while (BlockSampler_HasMore(&buildstate->bs))
-	{
-		BlockNumber targblock = BlockSampler_Next(&buildstate->bs);
-
-#if PG_VERSION_NUM >= 120000
-		table_index_build_range_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
-									 false, true, true, targblock, 1, SampleCallback, (void *) buildstate, NULL);
-#elif PG_VERSION_NUM >= 110000
-		IndexBuildHeapRangeScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
-								true, true, targblock, 1, SampleCallback, (void *) buildstate, NULL);
 #else
-		IndexBuildHeapRangeScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
+#define UpdateProgress(index, val) ((void)val)
-								true, true, targblock, 1, SampleCallback, (void *) buildstate);
 #endif
-	}
-}
 
 /*
  * Callback for table_index_build_scan
@@ -170,16 +84,20 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
 		}
 	}
 
+#ifdef IVFFLAT_KMEANS_DEBUG
+	buildstate->inertia += minDistance;
+	buildstate->listSums[closestCenter] += minDistance;
+	buildstate->listCounts[closestCenter]++;
+#endif
+
 	/* Create a virtual tuple */
 	ExecClearTuple(slot);
 	slot->tts_values[0] = Int32GetDatum(closestCenter);
 	slot->tts_isnull[0] = false;
-	slot->tts_values[1] = Int32GetDatum(ItemPointerGetBlockNumberNoCheck(tid));
+	slot->tts_values[1] = PointerGetDatum(tid);
 	slot->tts_isnull[1] = false;
-	slot->tts_values[2] = Int32GetDatum(ItemPointerGetOffsetNumberNoCheck(tid));
+	slot->tts_values[2] = value;
 	slot->tts_isnull[2] = false;
-	slot->tts_values[3] = value;
-	slot->tts_isnull[3] = false;
 	ExecStoreVirtualTuple(slot);
 
 	/*
@@ -201,8 +119,6 @@ GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot,
 {
 	Datum		value;
 	bool		isnull;
-	int			tupblk;
-	int			tupoff;
 
 #if PG_VERSION_NUM >= 100000
 	if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
@@ -211,13 +127,11 @@ GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot,
 #endif
 	{
 		*list = DatumGetInt32(slot_getattr(slot, 1, &isnull));
-		tupblk = DatumGetInt32(slot_getattr(slot, 2, &isnull));
+		value = slot_getattr(slot, 3, &isnull);
-		tupoff = DatumGetInt32(slot_getattr(slot, 3, &isnull));
-		value = slot_getattr(slot, 4, &isnull);
 
 		/* Form the index tuple */
 		*itup = index_form_tuple(tupdesc, &value, &isnull);
-		ItemPointerSet(&(*itup)->t_tid, tupblk, tupoff);
+		(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &isnull)));
 	}
 	else
 		*list = -1;
@@ -326,17 +240,16 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
 
 	/* Create tuple description for sorting */
 #if PG_VERSION_NUM >= 120000
-	buildstate->tupdesc = CreateTemplateTupleDesc(4);
+	buildstate->tupdesc = CreateTemplateTupleDesc(3);
 #else
-	buildstate->tupdesc = CreateTemplateTupleDesc(4, false);
+	buildstate->tupdesc = CreateTemplateTupleDesc(3, false);
 #endif
 	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "blkno", INT4OID, -1, 0);
+	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "offset", INT4OID, -1, 0);
 #if PG_VERSION_NUM >= 110000
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 4, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0);
+	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0);
 #else
-	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 4, "vector", RelationGetDescr(index)->attrs[0]->atttypid, -1, 0);
+	TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0]->atttypid, -1, 0);
 #endif
 
 #if PG_VERSION_NUM >= 120000
@@ -350,6 +263,12 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
 
 	/* Reuse for each tuple */
 	buildstate->normvec = InitVector(buildstate->dimensions);
+
+#ifdef IVFFLAT_KMEANS_DEBUG
+	buildstate->inertia = 0;
+	buildstate->listSums = palloc0(sizeof(double) * buildstate->lists);
+	buildstate->listCounts = palloc0(sizeof(int) * buildstate->lists);
+#endif
 }
 
 /*
@@ -361,38 +280,11 @@ FreeBuildState(IvfflatBuildState * buildstate)
 	pfree(buildstate->centers);
 	pfree(buildstate->listInfo);
 	pfree(buildstate->normvec);
-}
 
-/*
+#ifdef IVFFLAT_KMEANS_DEBUG
- * Compute centers
+	pfree(buildstate->listSums);
- */
+	pfree(buildstate->listCounts);
-static void
+#endif
-ComputeCenters(IvfflatBuildState * buildstate)
-{
-	int			numSamples;
-
-	/* Target 50 samples per list, with at least 10000 samples */
-	/* The number of samples has a large effect on index build time */
-	numSamples = buildstate->lists * 50;
-	if (numSamples < 10000)
-		numSamples = 10000;
-
-	/* Skip samples for unlogged table */
-	if (buildstate->heap == NULL)
-		numSamples = 1;
-
-	/* Sample rows */
-	/* TODO Ensure within maintenance_work_mem */
-	buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions);
-	if (buildstate->heap != NULL)
-		SampleRows(buildstate);
-
-	/* Calculate centers */
-	UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_KMEANS);
-	IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers));
-
-	/* Free samples before we allocate more memory */
-	pfree(buildstate->samples);
 }
 
 /*
@@ -468,6 +360,51 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
 	pfree(list);
 }
 
+/*
+ * Print k-means metrics
+ */
+#ifdef IVFFLAT_KMEANS_DEBUG
+static void
+PrintKmeansMetrics(IvfflatBuildState * buildstate)
+{
+	elog(INFO, "inertia: %.3e", buildstate->inertia);
+
+	/* Calculate Davies-Bouldin index */
+	if (buildstate->lists > 1)
+	{
+		double		db = 0.0;
+
+		/* Calculate average distance */
+		for (int i = 0; i < buildstate->lists; i++)
+		{
+			if (buildstate->listCounts[i] > 0)
+				buildstate->listSums[i] /= buildstate->listCounts[i];
+		}
+
+		for (int i = 0; i < buildstate->lists; i++)
+		{
+			double		max = 0.0;
+			double		distance;
+
+			for (int j = 0; j < buildstate->lists; j++)
+			{
+				if (j == i)
+					continue;
+
+				distance = DatumGetFloat8(FunctionCall2Coll(buildstate->procinfo, buildstate->collation, PointerGetDatum(VectorArrayGet(buildstate->centers, i)), PointerGetDatum(VectorArrayGet(buildstate->centers, j))));
+				distance = (buildstate->listSums[i] + buildstate->listSums[j]) / distance;
+
+				if (distance > max)
+					max = distance;
+			}
+			db += max;
+		}
+		db /= buildstate->lists;
+		elog(INFO, "davies-bouldin: %.3f", db);
+	}
+}
+#endif
+
 /*
  * Create entry pages
  */
@@ -502,8 +439,14 @@ CreateEntryPages(IvfflatBuildState * buildstate, ForkNumber forkNum)
 #endif
 	}
 
-	/* Sort and insert */
+	/* Sort */
 	tuplesort_performsort(buildstate->sortstate);
+
+#ifdef IVFFLAT_KMEANS_DEBUG
+	PrintKmeansMetrics(buildstate);
+#endif
+
+	/* Insert */
 	InsertTuples(buildstate->index, buildstate, forkNum);
 	tuplesort_end(buildstate->sortstate);
 }
@@ -517,7 +460,9 @@ BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo,
 {
 	InitBuildState(buildstate, heap, index, indexInfo);
 
-	ComputeCenters(buildstate);
+	/* Perform k-means clustering */
+	UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_KMEANS);
+	IvfflatBench("k-means", IvfflatKmeans(buildstate));
 
 	/* Create pages */
 	CreateMetaPage(index, buildstate->dimensions, buildstate->lists, forkNum);

src/ivfflat.c

@@ -45,8 +45,6 @@ ivfflatbuildphasename(int64 phasenum)
 	{
 		case PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE:
 			return "initializing";
-		case PROGRESS_IVFFLAT_PHASE_SAMPLE:
-			return "sampling table";
 		case PROGRESS_IVFFLAT_PHASE_KMEANS:
 			return "performing k-means";
 		case PROGRESS_IVFFLAT_PHASE_SORT:

src/ivfflat.h

@@ -37,10 +37,9 @@
 
 /* Build phases */
 /* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
-#define PROGRESS_IVFFLAT_PHASE_SAMPLE	2
+#define PROGRESS_IVFFLAT_PHASE_KMEANS	2
-#define PROGRESS_IVFFLAT_PHASE_KMEANS	3
+#define PROGRESS_IVFFLAT_PHASE_SORT		3
-#define PROGRESS_IVFFLAT_PHASE_SORT		4
+#define PROGRESS_IVFFLAT_PHASE_LOAD		4
-#define PROGRESS_IVFFLAT_PHASE_LOAD		5
 
 #define IVFFLAT_LIST_SIZE(_dim)	(offsetof(IvfflatListData, center) + VECTOR_SIZE(_dim))
 
@@ -62,11 +61,6 @@
 #define IvfflatBench(name, code) (code)
 #endif
 
-#if PG_VERSION_NUM < 100000
-#define ItemPointerGetBlockNumberNoCheck ItemPointerGetBlockNumber
-#define ItemPointerGetOffsetNumberNoCheck ItemPointerGetOffsetNumber
-#endif
-
 /* Variables */
 extern int	ivfflat_probes;
 
@@ -120,6 +114,12 @@ typedef struct IvfflatBuildState
 	ListInfo   *listInfo;
 	Vector	   *normvec;
 
+#ifdef IVFFLAT_KMEANS_DEBUG
+	double		inertia;
+	double	   *listSums;
+	int		   *listCounts;
+#endif
+
 	/* Sampling */
 	BlockSamplerData bs;
 	ReservoirStateData rstate;
@@ -161,6 +161,7 @@ typedef IvfflatListData * IvfflatList;
 
 typedef struct IvfflatScanList
 {
+	pairingheap_node ph_node;
 	BlockNumber startPage;
 	double		distance;
 }			IvfflatScanList;
@@ -182,6 +183,8 @@ typedef struct IvfflatScanOpaqueData
 	FmgrInfo   *normprocinfo;
 	Oid			collation;
 
+	/* Lists */
+	pairingheap *listQueue;
 	IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER];	/* must come last */
 }			IvfflatScanOpaqueData;
 
@@ -196,7 +199,7 @@ typedef IvfflatScanOpaqueData * IvfflatScanOpaque;
 void		_PG_init(void);
 VectorArray VectorArrayInit(int maxlen, int dimensions);
 void		PrintVectorArray(char *msg, VectorArray arr);
-void		IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);
+void		IvfflatKmeans(IvfflatBuildState * buildstate);
 FmgrInfo   *IvfflatOptionalProcInfo(Relation rel, uint16 procnum);
 bool		IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
 int			IvfflatGetLists(Relation index);

src/ivfkmeans.c

@@ -2,8 +2,20 @@
 
 #include <float.h>
 
+#include "catalog/index.h"
 #include "ivfflat.h"
 #include "miscadmin.h"
+#include "storage/bufmgr.h"
+
+#if PG_VERSION_NUM >= 120000
+#include "access/tableam.h"
+#endif
+
+#if PG_VERSION_NUM >= 130000
+#define CALLBACK_ITEM_POINTER ItemPointer tid
+#else
+#define CALLBACK_ITEM_POINTER HeapTuple hup
+#endif
 
 /*
  * Initialize with kmeans++
@@ -11,7 +23,7 @@
  * https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf
  */
 static void
-InitCenters(Relation index, VectorArray samples, VectorArray centers, float *lowerBound)
+InitCenters(Relation index, VectorArray samples, VectorArray centers)
 {
 	FmgrInfo   *procinfo;
 	Oid			collation;
@@ -35,7 +47,7 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
 	for (j = 0; j < numSamples; j++)
 		weight[j] = DBL_MAX;
 
-	for (i = 0; i < numCenters; i++)
+	for (i = 0; i < numCenters - 1; i++)
 	{
 		CHECK_FOR_INTERRUPTS();
 
@@ -49,9 +61,6 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
 			/* TODO Use triangle inequality to reduce distance calculations */
 			distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, i))));
 
-			/* Set lower bound */
-			lowerBound[j * numCenters + i] = distance;
-
 			/* Use distance squared for weighted probability distribution */
 			distance *= distance;
 
@@ -61,10 +70,6 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
 			sum += weight[j];
 		}
 
-		/* Only compute lower bound on last iteration */
-		if (i + 1 == numCenters)
-			break;
-
 		/* Choose new center using weighted probability distribution. */
 		choice = sum * (((double) random()) / MAX_RANDOM_VALUE);
 		for (j = 0; j < numSamples - 1; j++)
@@ -156,299 +161,202 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
 }
 
 /*
- * Use Elkan for performance. This requires distance function to satisfy triangle inequality.
+ * Callback for sampling
+ */
+static void
+SampleCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
+			   bool *isnull, bool tupleIsAlive, void *state)
+{
+	IvfflatBuildState *buildstate = (IvfflatBuildState *) state;
+	VectorArray samples = buildstate->samples;
+	int			targsamples = samples->maxlen;
+	Datum		value = values[0];
+
+	/* Skip nulls */
+	if (isnull[0])
+		return;
+
+	/*
+	 * Normalize with KMEANS_NORM_PROC since spherical distance function
+	 * expects unit vectors
+	 */
+	if (buildstate->kmeansnormprocinfo != NULL)
+	{
+		if (!IvfflatNormValue(buildstate->kmeansnormprocinfo, buildstate->collation, &value, buildstate->normvec))
+			return;
+	}
+
+	if (samples->length < targsamples)
+	{
+		VectorArraySet(samples, samples->length, DatumGetVector(value));
+		samples->length++;
+	}
+	else
+	{
+		if (buildstate->rowstoskip < 0)
+			buildstate->rowstoskip = reservoir_get_next_S(&buildstate->rstate, samples->length, targsamples);
+
+		if (buildstate->rowstoskip <= 0)
+		{
+			int			k = (int) (targsamples * sampler_random_fract(buildstate->rstate.randstate));
+
+			Assert(k >= 0 && k < targsamples);
+			VectorArraySet(samples, k, DatumGetVector(value));
+		}
+
+		buildstate->rowstoskip -= 1;
+	}
+}
+
+/*
+ * Sample rows with same logic as ANALYZE
+ */
+static void
+SampleRows(IvfflatBuildState * buildstate)
+{
+	int			targsamples = buildstate->samples->maxlen;
+	BlockNumber totalblocks = RelationGetNumberOfBlocks(buildstate->heap);
+
+	buildstate->rowstoskip = -1;
+	buildstate->samples->length = 0;
+
+	BlockSampler_Init(&buildstate->bs, totalblocks, targsamples, random());
+
+	reservoir_init_selection_state(&buildstate->rstate, targsamples);
+	while (BlockSampler_HasMore(&buildstate->bs))
+	{
+		BlockNumber targblock = BlockSampler_Next(&buildstate->bs);
+
+#if PG_VERSION_NUM >= 120000
+		table_index_build_range_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
+									 false, true, false, targblock, 1, SampleCallback, (void *) buildstate, NULL);
+#elif PG_VERSION_NUM >= 110000
+		IndexBuildHeapRangeScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
+								false, true, targblock, 1, SampleCallback, (void *) buildstate, NULL);
+#else
+		IndexBuildHeapRangeScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
+								false, true, targblock, 1, SampleCallback, (void *) buildstate);
+#endif
+	}
+}
+
+/*
+ * Use mini-batch k-means
  *
  * We use L2 distance for L2 (not L2 squared like index scan)
  * and angular distance for inner product and cosine distance
  *
- * https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf
+ * https://www.eecs.tufts.edu/~dsculley/papers/fastkmeans.pdf
  */
 static void
-ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
+MiniBatchKmeans(IvfflatBuildState * buildstate)
 {
-	FmgrInfo   *procinfo;
+	VectorArray centers = buildstate->centers;
-	FmgrInfo   *normprocinfo;
+	VectorArray m = buildstate->samples;
-	Oid			collation;
+	int			b = m->maxlen;
-	Vector	   *vec;
+	int			t = 20;
-	Vector	   *newCenter;
+	double		distance;
-	int			iteration;
-	int			j;
-	int			k;
-	int			dimensions = centers->dim;
-	int			numCenters = centers->maxlen;
-	int			numSamples = samples->length;
-	VectorArray newCenters;
-	int		   *centerCounts;
-	int		   *closestCenters;
-	float	   *lowerBound;
-	float	   *upperBound;
-	float	   *s;
-	float	   *halfcdist;
-	float	   *newcdist;
-	int			changes;
 	double		minDistance;
 	int			closestCenter;
-	double		distance;
+	int			i;
-	bool		rj;
+	int			j;
-	bool		rjreset;
+	int			k;
-	double		dxcx;
+	Vector	   *c;
-	double		dxc;
+	Vector	   *x;
-
+	int		   *v;
-	/* Calculate allocation sizes */
+	int		   *d;
-	Size		samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->dim);
+	double		eta;
-	Size		centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->dim);
-	Size		newCentersSize = VECTOR_ARRAY_SIZE(numCenters, dimensions);
-	Size		centerCountsSize = sizeof(int) * numCenters;
-	Size		closestCentersSize = sizeof(int) * numSamples;
-	Size		lowerBoundSize = sizeof(float) * numSamples * numCenters;
-	Size		upperBoundSize = sizeof(float) * numSamples;
-	Size		sSize = sizeof(float) * numCenters;
-	Size		halfcdistSize = sizeof(float) * numCenters * numCenters;
-	Size		newcdistSize = sizeof(float) * numCenters;
-
-	/* Calculate total size */
-	Size		totalSize = samplesSize + centersSize + newCentersSize + centerCountsSize + closestCentersSize + lowerBoundSize + upperBoundSize + sSize + halfcdistSize + newcdistSize;
-
-	/* Check memory requirements */
-	/* Add one to error message to ceil */
-	if (totalSize / 1024 > maintenance_work_mem)
-		ereport(ERROR,
-				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
-				 errmsg("memory required is %zu MB, maintenance_work_mem is %d MB",
-						totalSize / (1024 * 1024) + 1, maintenance_work_mem / 1024)));
 
 	/* Set support functions */
-	procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
+	FmgrInfo   *procinfo = index_getprocinfo(buildstate->index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
-	normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
+	FmgrInfo   *normprocinfo = buildstate->kmeansnormprocinfo;
-	collation = index->rd_indcollation[0];
+	Oid			collation = buildstate->index->rd_indcollation[0];
-
-	/* Allocate space */
-	/* Use float instead of double to save memory */
-	centerCounts = palloc(centerCountsSize);
-	closestCenters = palloc(closestCentersSize);
-	lowerBound = palloc_extended(lowerBoundSize, MCXT_ALLOC_HUGE);
-	upperBound = palloc(upperBoundSize);
-	s = palloc(sSize);
-	halfcdist = palloc(halfcdistSize);
-	newcdist = palloc(newcdistSize);
-
-	newCenters = VectorArrayInit(numCenters, dimensions);
-	for (j = 0; j < numCenters; j++)
-	{
-		vec = VectorArrayGet(newCenters, j);
-		SET_VARSIZE(vec, VECTOR_SIZE(dimensions));
-		vec->dim = dimensions;
-	}
 
 	/* Pick initial centers */
-	InitCenters(index, samples, centers, lowerBound);
+	InitCenters(buildstate->index, buildstate->samples, buildstate->centers);
 
-	/* Assign each x to its closest initial center c(x) = argmin d(x,c) */
+	v = palloc(sizeof(int) * centers->maxlen);
-	for (j = 0; j < numSamples; j++)
+	d = palloc(sizeof(int) * b);
-	{
-		minDistance = DBL_MAX;
-		closestCenter = -1;
 
-		vec = VectorArrayGet(samples, j);
+	for (int i = 0; i < centers->length; i++)
+		v[i] = 0;
 
-		/* Find closest center */
+	for (i = 0; i < t; i++)
-		for (k = 0; k < numCenters; k++)
-		{
-			/* TODO Use Lemma 1 in k-means++ initialization */
-			distance = lowerBound[j * numCenters + k];
-
-			if (distance < minDistance)
-			{
-				minDistance = distance;
-				closestCenter = k;
-			}
-		}
-
-		upperBound[j] = minDistance;
-		closestCenters[j] = closestCenter;
-	}
-
-	/* Give 500 iterations to converge */
-	for (iteration = 0; iteration < 500; iteration++)
 	{
 		/* Can take a while, so ensure we can interrupt */
 		CHECK_FOR_INTERRUPTS();
 
-		changes = 0;
+		/* Get b examples picked randomly from X */
+		SampleRows(buildstate);
 
-		/* Step 1: For all centers, compute distance */
+		/* Cache nearest center to x */
-		for (j = 0; j < numCenters; j++)
+		for (j = 0; j < m->length; j++)
-		{
-			vec = VectorArrayGet(centers, j);
-
-			for (k = j + 1; k < numCenters; k++)
-			{
-				distance = 0.5 * DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, k))));
-				halfcdist[j * numCenters + k] = distance;
-				halfcdist[k * numCenters + j] = distance;
-			}
-		}
-
-		/* For all centers c, compute s(c) */
-		for (j = 0; j < numCenters; j++)
 		{
+			/* compute closest */
 			minDistance = DBL_MAX;
+			closestCenter = -1;
 
-			for (k = 0; k < numCenters; k++)
+			x = VectorArrayGet(m, j);
+
+			/* Find closest center */
+			for (k = 0; k < centers->length; k++)
 			{
-				if (j == k)
+				distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(x), PointerGetDatum(VectorArrayGet(centers, k))));
-					continue;
 
-				distance = halfcdist[j * numCenters + k];
 				if (distance < minDistance)
+				{
 					minDistance = distance;
+					closestCenter = k;
+				}
 			}
 
-			s[j] = minDistance;
+			d[j] = closestCenter;
 		}
 
-		rjreset = iteration != 0;
+		for (j = 0; j < m->length; j++)
-
-		for (j = 0; j < numSamples; j++)
 		{
-			/* Step 2: Identify all points x such that u(x) <= s(c(x)) */
+			x = VectorArrayGet(m, j);
-			if (upperBound[j] <= s[closestCenters[j]])
-				continue;
 
-			rj = rjreset;
+			/* Get cached center for this x */
+			c = VectorArrayGet(centers, d[j]);
 
-			for (k = 0; k < numCenters; k++)
+			/* Update per-center counts */
+			v[d[j]]++;
+
+			/* Get per-center learning rate */
+			eta = 1.0 / v[d[j]];
+
+			/* Take gradient step */
+			for (k = 0; k < c->dim; k++)
+				c->x[k] = (1 - eta) * c->x[k] + eta * x->x[k];
+		}
+
+		/* Check for empty centers (likely duplicates) */
+		if (i == 0)
+		{
+			for (j = 0; j < centers->length; j++)
 			{
-				/* Step 3: For all remaining points x and centers c */
+				if (v[j] == 0)
-				if (k == closestCenters[j])
-					continue;
-
-				if (upperBound[j] <= lowerBound[j * numCenters + k])
-					continue;
-
-				if (upperBound[j] <= halfcdist[closestCenters[j] * numCenters + k])
-					continue;
-
-				vec = VectorArrayGet(samples, j);
-
-				/* Step 3a */
-				if (rj)
 				{
-					dxcx = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, closestCenters[j]))));
+					c = VectorArrayGet(centers, j);
-
-					/* d(x,c(x)) computed, which is a form of d(x,c) */
-					lowerBound[j * numCenters + closestCenters[j]] = dxcx;
-					upperBound[j] = dxcx;
-
-					rj = false;
-				}
-				else
-					dxcx = upperBound[j];
-
-				/* Step 3b */
-				if (dxcx > lowerBound[j * numCenters + k] || dxcx > halfcdist[closestCenters[j] * numCenters + k])
-				{
-					dxc = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, k))));
-
-					/* d(x,c) calculated */
-					lowerBound[j * numCenters + k] = dxc;
-
-					if (dxc < dxcx)
-					{
-						closestCenters[j] = k;
-
-						/* c(x) changed */
-						upperBound[j] = dxc;
-
-						changes++;
-					}
 
+					/* TODO Handle empty centers properly */
+					for (k = 0; k < c->dim; k++)
+						c->x[k] = ((double) random()) / MAX_RANDOM_VALUE;
 				}
 			}
 		}
 
-		/* Step 4: For each center c, let m(c) be mean of all points assigned */
+		/* Normalize if needed */
-		for (j = 0; j < numCenters; j++)
+		if (normprocinfo != NULL)
 		{
-			vec = VectorArrayGet(newCenters, j);
+			for (j = 0; j < centers->length; j++)
-			for (k = 0; k < dimensions; k++)
+				ApplyNorm(normprocinfo, collation, VectorArrayGet(centers, j));
-				vec->x[k] = 0.0;
-
-			centerCounts[j] = 0;
 		}
-
-		for (j = 0; j < numSamples; j++)
-		{
-			vec = VectorArrayGet(samples, j);
-			closestCenter = closestCenters[j];
-
-			/* Increment sum and count of closest center */
-			newCenter = VectorArrayGet(newCenters, closestCenter);
-			for (k = 0; k < dimensions; k++)
-				newCenter->x[k] += vec->x[k];
-
-			centerCounts[closestCenter] += 1;
-		}
-
-		for (j = 0; j < numCenters; j++)
-		{
-			vec = VectorArrayGet(newCenters, j);
-
-			if (centerCounts[j] > 0)
-			{
-				for (k = 0; k < dimensions; k++)
-					vec->x[k] /= centerCounts[j];
-			}
-			else
-			{
-				/* TODO Handle empty centers properly */
-				for (k = 0; k < dimensions; k++)
-					vec->x[k] = ((double) random()) / MAX_RANDOM_VALUE;
-			}
-
-			/* Normalize if needed */
-			if (normprocinfo != NULL)
-				ApplyNorm(normprocinfo, collation, vec);
-		}
-
-		/* Step 5 */
-		for (j = 0; j < numCenters; j++)
-			newcdist[j] = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(VectorArrayGet(centers, j)), PointerGetDatum(VectorArrayGet(newCenters, j))));
-
-		for (j = 0; j < numSamples; j++)
-		{
-			for (k = 0; k < numCenters; k++)
-			{
-				distance = lowerBound[j * numCenters + k] - newcdist[k];
-
-				if (distance < 0)
-					distance = 0;
-
-				lowerBound[j * numCenters + k] = distance;
-			}
-		}
-
-		/* Step 6 */
-		/* We reset r(x) before Step 3 in the next iteration */
-		for (j = 0; j < numSamples; j++)
-			upperBound[j] += newcdist[closestCenters[j]];
-
-		/* Step 7 */
-		for (j = 0; j < numCenters; j++)
-			memcpy(VectorArrayGet(centers, j), VectorArrayGet(newCenters, j), VECTOR_SIZE(dimensions));
-
-		if (changes == 0 && iteration != 0)
-			break;
 	}
 
-	pfree(newCenters);
+	pfree(v);
-	pfree(centerCounts);
+	pfree(d);
-	pfree(closestCenters);
-	pfree(lowerBound);
-	pfree(upperBound);
-	pfree(s);
-	pfree(halfcdist);
-	pfree(newcdist);
 }
 
 /*
@@ -491,16 +399,48 @@ CheckCenters(Relation index, VectorArray centers)
 }
 
 /*
- * Perform naive k-means centering
+ * Perform k-means clustering
  * We use spherical k-means for inner product and cosine
  */
 void
-IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers)
+IvfflatKmeans(IvfflatBuildState * buildstate)
 {
-	if (samples->length <= centers->maxlen)
+	int			numSamples;
-		QuickCenters(index, samples, centers);
+	Size		totalSize;
-	else
-		ElkanKmeans(index, samples, centers);
 
-	CheckCenters(index, centers);
+	/* Target 10 samples per list, with at least 10000 samples */
+	/* The number of samples has a large effect on index build time */
+	numSamples = buildstate->lists * 10;
+	if (numSamples < 10000)
+		numSamples = 10000;
+
+	/* Skip samples for unlogged table */
+	if (buildstate->heap == NULL)
+		numSamples = 1;
+
+	/* Calculate total size */
+	totalSize = VECTOR_ARRAY_SIZE(numSamples, buildstate->dimensions);
+
+	/* Check memory requirements */
+	/* Add one to error message to ceil */
+	if (totalSize / 1024 > maintenance_work_mem)
+		ereport(ERROR,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				 errmsg("memory required is %zu MB, maintenance_work_mem is %d MB",
+						totalSize / (1024 * 1024) + 1, maintenance_work_mem / 1024)));
+
+	/* Sample rows */
+	buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions);
+	if (buildstate->heap != NULL)
+		SampleRows(buildstate);
+
+	if (buildstate->samples->length <= buildstate->centers->maxlen)
+		QuickCenters(buildstate->index, buildstate->samples, buildstate->centers);
+	else
+		MiniBatchKmeans(buildstate);
+
+	CheckCenters(buildstate->index, buildstate->centers);
+
+	/* Free samples before we allocate more memory */
+	pfree(buildstate->samples);
 }

src/ivfscan.c

@@ -1,5 +1,7 @@
 #include "postgres.h"
 
+#include <float.h>
+
 #include "access/relscan.h"
 #include "ivfflat.h"
 #include "miscadmin.h"
@@ -17,14 +19,12 @@
  * Compare list distances
  */
 static int
-CompareLists(const void *a, const void *b)
+CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
 {
-	double		diff = (((IvfflatScanList *) a)->distance - ((IvfflatScanList *) b)->distance);
+	if (((const IvfflatScanList *) a)->distance > ((const IvfflatScanList *) b)->distance)
-
-	if (diff > 0)
 		return 1;
 
-	if (diff < 0)
+	if (((const IvfflatScanList *) a)->distance < ((const IvfflatScanList *) b)->distance)
 		return -1;
 
 	return 0;
@@ -45,6 +45,8 @@ GetScanLists(IndexScanDesc scan, Datum value)
 	int			listCount = 0;
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
 	double		distance;
+	IvfflatScanList *scanlist;
+	double		maxDistance = DBL_MAX;
 
 	/* Search all list pages */
 	while (BlockNumberIsValid(nextblkno))
@@ -62,22 +64,39 @@ GetScanLists(IndexScanDesc scan, Datum value)
 			/* Use procinfo from the index instead of scan key for performance */
 			distance = DatumGetFloat8(FunctionCall2Coll(so->procinfo, so->collation, PointerGetDatum(&list->center), value));
 
-			so->lists[listCount].startPage = list->startPage;
+			if (listCount < so->probes)
-			so->lists[listCount].distance = distance;
+			{
-			listCount++;
+				scanlist = &so->lists[listCount];
+				scanlist->startPage = list->startPage;
+				scanlist->distance = distance;
+				listCount++;
+
+				/* Add to heap */
+				pairingheap_add(so->listQueue, &scanlist->ph_node);
+
+				/* Calculate max distance */
+				if (listCount == so->probes)
+					maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
+			}
+			else if (distance < maxDistance)
+			{
+				/* Remove */
+				scanlist = (IvfflatScanList *) pairingheap_remove_first(so->listQueue);
+
+				/* Reuse */
+				scanlist->startPage = list->startPage;
+				scanlist->distance = distance;
+				pairingheap_add(so->listQueue, &scanlist->ph_node);
+
+				/* Update max distance */
+				maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
+			}
 		}
 
 		nextblkno = IvfflatPageGetOpaque(cpage)->nextblkno;
 
 		UnlockReleaseBuffer(cbuf);
 	}
-
-	/* Sort by distance */
-	/* TODO Use heap for performance */
-	qsort(so->lists, listCount, sizeof(IvfflatScanList), CompareLists);
-
-	if (so->probes > listCount)
-		so->probes = listCount;
 }
 
 /*
@@ -95,7 +114,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	OffsetNumber maxoffno;
 	Datum		datum;
 	bool		isnull;
-	int			i;
 	TupleDesc	tupdesc = RelationGetDescr(scan->indexRelation);
 
 #if PG_VERSION_NUM >= 120000
@@ -112,9 +130,9 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
 
 	/* Search closest probes lists */
-	for (i = 0; i < so->probes; i++)
+	while (!pairingheap_is_empty(so->listQueue))
 	{
-		searchPage = so->lists[i].startPage;
+		searchPage = ((IvfflatScanList *) pairingheap_remove_first(so->listQueue))->startPage;
 
 		/* Search all entry pages for list */
 		while (BlockNumberIsValid(searchPage))
@@ -138,12 +156,10 @@ GetScanItems(IndexScanDesc scan, Datum value)
 				ExecClearTuple(slot);
 				slot->tts_values[0] = FunctionCall2Coll(so->procinfo, so->collation, datum, value);
 				slot->tts_isnull[0] = false;
-				slot->tts_values[1] = Int32GetDatum((int) ItemPointerGetBlockNumberNoCheck(&itup->t_tid));
+				slot->tts_values[1] = PointerGetDatum(&itup->t_tid);
 				slot->tts_isnull[1] = false;
-				slot->tts_values[2] = Int32GetDatum((int) ItemPointerGetOffsetNumberNoCheck(&itup->t_tid));
+				slot->tts_values[2] = Int32GetDatum((int) searchPage);
 				slot->tts_isnull[2] = false;
-				slot->tts_values[3] = Int32GetDatum((int) searchPage);
-				slot->tts_isnull[3] = false;
 				ExecStoreVirtualTuple(slot);
 
 				tuplesort_puttupleslot(so->sortstate, slot);
@@ -171,13 +187,18 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	Oid			sortOperators[] = {Float8LessOperator};
 	Oid			sortCollations[] = {InvalidOid};
 	bool		nullsFirstFlags[] = {false};
+	int			probes = ivfflat_probes;
 
 	scan = RelationGetIndexScan(index, nkeys, norderbys);
 	lists = IvfflatGetLists(scan->indexRelation);
 
-	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + lists * sizeof(IvfflatScanList));
+	if (probes > lists)
+		probes = lists;
+
+	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
 	so->buf = InvalidBuffer;
 	so->first = true;
+	so->probes = probes;
 
 	/* Set support functions */
 	so->procinfo = index_getprocinfo(index, 1, IVFFLAT_DISTANCE_PROC);
@@ -186,14 +207,13 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 
 	/* Create tuple description for sorting */
 #if PG_VERSION_NUM >= 120000
-	so->tupdesc = CreateTemplateTupleDesc(4);
+	so->tupdesc = CreateTemplateTupleDesc(3);
 #else
-	so->tupdesc = CreateTemplateTupleDesc(4, false);
+	so->tupdesc = CreateTemplateTupleDesc(3, false);
 #endif
 	TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "blkno", INT4OID, -1, 0);
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "offset", INT4OID, -1, 0);
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0);
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 4, "indexblkno", INT4OID, -1, 0);
 
 	/* Prep sort */
 #if PG_VERSION_NUM >= 110000
@@ -208,6 +228,8 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	so->slot = MakeSingleTupleTableSlot(so->tupdesc);
 #endif
 
+	so->listQueue = pairingheap_allocate(CompareLists, scan);
+
 	scan->opaque = so;
 
 	return scan;
@@ -227,7 +249,7 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
 #endif
 
 	so->first = true;
-	so->probes = ivfflat_probes;
+	pairingheap_reset(so->listQueue);
 
 	if (keys && scan->numberOfKeys > 0)
 		memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -286,14 +308,13 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 	if (tuplesort_gettupleslot(so->sortstate, true, so->slot, NULL))
 #endif
 	{
-		BlockNumber blkno = DatumGetInt32(slot_getattr(so->slot, 2, &so->isnull));
+		ItemPointer tid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
-		OffsetNumber offset = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
+		BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
-		BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 4, &so->isnull));
 
 #if PG_VERSION_NUM >= 120000
-		ItemPointerSet(&scan->xs_heaptid, blkno, offset);
+		scan->xs_heaptid = *tid;
 #else
-		ItemPointerSet(&scan->xs_ctup.t_self, blkno, offset);
+		scan->xs_ctup.t_self = *tid;
 #endif
 
 		if (BufferIsValid(so->buf))
@@ -326,6 +347,7 @@ ivfflatendscan(IndexScanDesc scan)
 	if (BufferIsValid(so->buf))
 		ReleaseBuffer(so->buf);
 
+	pairingheap_free(so->listQueue);
 	tuplesort_end(so->sortstate);
 
 	pfree(so);

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/003_recall.pl

@@ -2,15 +2,16 @@ use strict;
 use warnings;
 use PostgresNode;
 use TestLib;
-use Test::More tests => 2;
+use Test::More tests => 9;
 
 my $node;
 my @queries = ();
-my @expected = ();
+my @expected;
+my $limit = 20;
 
 sub test_recall
 {
-	my ($probes, $min) = @_;
+	my ($probes, $min, $operator) = @_;
 	my $correct = 0;
 	my $total = 0;
 
@@ -18,7 +19,7 @@ sub test_recall
 		my $actual = $node->safe_psql("postgres", qq(
 			SET enable_seqscan = off;
 			SET ivfflat.probes = $probes;
-			SELECT i FROM tst ORDER BY v <-> '$queries[$i]' LIMIT 10;
+			SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
 		));
 		my @actual_ids = split("\n", $actual);
 		my %actual_set = map { $_ => 1 } @actual_ids;
@@ -33,7 +34,7 @@ sub test_recall
 		}
 	}
 
-	cmp_ok($correct / $total, ">=", $min);
+	cmp_ok($correct / $total, ">=", $min, $operator);
 }
 
 # Initialize node
@@ -56,17 +57,32 @@ for (1..20) {
 	push(@queries, "[$r1,$r2,$r3]");
 }
 
-# Get exact results
+# Check each index type
-foreach (@queries) {
+my @operators = ("<->", "<#>", "<=>");
-	my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v <-> '$_' LIMIT 10;");
+
-	push(@expected, $res);
+foreach (@operators) {
+	my $operator = $_;
+
+	# Get exact results
+	@expected = ();
+	foreach (@queries) {
+		my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
+		push(@expected, $res);
+	}
+
+	# Add index
+	my $opclass;
+	if ($operator == "<->") {
+		$opclass = "vector_l2_ops";
+	} elsif ($operator == "<#>") {
+		$opclass = "vector_ip_ops";
+	} else {
+		$opclass = "vector_cosine_ops";
+	}
+	$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v $opclass);");
+
+	# Test approximate results
+	test_recall(1, 0.75, $operator);
+	test_recall(10, 0.95, $operator);
+	test_recall(100, 1.0, $operator);
 }
-
-# Add index
-$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
-
-# Test approximate results
-test_recall(1, 0.8);
-
-# Test probes
-test_recall(100, 1.0);

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);

test/t/005_query_recall.pl

@@ -0,0 +1,45 @@
+use strict;
+use warnings;
+use PostgresNode;
+use TestLib;
+use Test::More tests => 60;
+
+# 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 primary key, v vector(3));");
+$node->safe_psql("postgres",
+	"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
+);
+
+# Check each index type
+my @operators = ("<->", "<#>", "<=>");
+foreach (@operators) {
+	my $operator = $_;
+
+	# Add index
+	my $opclass;
+	if ($operator == "<->") {
+		$opclass = "vector_l2_ops";
+	} elsif ($operator == "<#>") {
+		$opclass = "vector_ip_ops";
+	} else {
+		$opclass = "vector_cosine_ops";
+	}
+	$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v $opclass);");
+
+	# Test 100% recall
+	for (1..20) {
+		my $i = int(rand() * 100000);
+		my $query = $node->safe_psql("postgres", "SELECT v FROM tst WHERE i = $i;");
+		my $res = $node->safe_psql("postgres", qq(
+			SET enable_seqscan = off;
+			SELECT v FROM tst ORDER BY v <-> '$query' LIMIT 1;
+		));
+		is($res, $query);
+	}
+}

test/t/006_lists.pl

@@ -0,0 +1,31 @@
+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 (v vector(3));");
+$node->safe_psql("postgres",
+	"INSERT INTO tst SELECT ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
+);
+
+$node->safe_psql("postgres", "CREATE INDEX lists50 ON tst USING ivfflat (v) WITH (lists = 50);");
+$node->safe_psql("postgres", "CREATE INDEX lists100 ON tst USING ivfflat (v) WITH (lists = 100);");
+
+# Test prefers more lists
+my $res = $node->safe_psql("postgres", "EXPLAIN SELECT v FROM tst ORDER BY v <-> '[0.5,0.5,0.5]' LIMIT 10;");
+like($res, qr/lists100/);
+unlike($res, qr/lists50/);
+
+# Test errors with too much memory
+my ($ret, $stdout, $stderr) = $node->psql("postgres",
+	"CREATE INDEX lists10000 ON tst USING ivfflat (v) WITH (lists = 10000);"
+);
+like($stderr, qr/memory required is/);