Merge branch 'master' into minibatch

CHANGELOG.md

@@ -1,6 +1,6 @@
 ## 0.2.6 (unreleased)
 
-- Significantly improved index query performance
+- Switched to mini-batch k-means
 - Improved performance of index creation for Postgres < 12
 
 ## 0.2.5 (2022-02-11)

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`
-3. `performing k-means`
-4. `sorting tuples`
-5. `loading tuples`
+2. `performing k-means`
+3. `sorting tuples`
+4. `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

@@ -42,87 +42,6 @@
 #define UpdateProgress(index, val) ((void)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, 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
-	}
-}
-
 /*
  * Callback for table_index_build_scan
  */
@@ -368,38 +287,6 @@ FreeBuildState(IvfflatBuildState * buildstate)
 #endif
 }
 
-/*
- * Compute centers
- */
-static void
-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);
-}
-
 /*
  * Create the metapage
  */
@@ -573,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	3
-#define PROGRESS_IVFFLAT_PHASE_SORT		4
-#define PROGRESS_IVFFLAT_PHASE_LOAD		5
+#define PROGRESS_IVFFLAT_PHASE_KMEANS	2
+#define PROGRESS_IVFFLAT_PHASE_SORT		3
+#define PROGRESS_IVFFLAT_PHASE_LOAD		4
 
 #define IVFFLAT_LIST_SIZE(_dim)	(offsetof(IvfflatListData, center) + VECTOR_SIZE(_dim))
 
@@ -167,28 +166,12 @@ typedef struct IvfflatScanList
 	double		distance;
 }			IvfflatScanList;
 
-typedef struct IvfflatScanItem
-{
-	pairingheap_node ph_node;
-	BlockNumber searchPage;
-	double		distance;
-	ItemPointerData tid;
-}			IvfflatScanItem;
-
 typedef struct IvfflatScanOpaqueData
 {
 	int			probes;
-	int			stage;
+	bool		first;
 	Buffer		buf;
 
-	/* Items */
-	int			maxItems;
-	int			itemCount;
-	pairingheap *itemQueue;
-	IvfflatScanItem *items;
-	IvfflatScanItem **sortedItems;
-	bool		heapFull;
-
 	/* Sorting */
 	Tuplesortstate *sortstate;
 	TupleDesc	tupdesc;
@@ -202,7 +185,6 @@ typedef struct IvfflatScanOpaqueData
 
 	/* Lists */
 	pairingheap *listQueue;
-	IvfflatScanList **sortedLists;
 	IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER];	/* must come last */
 }			IvfflatScanOpaqueData;
 
@@ -217,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;
-	FmgrInfo   *normprocinfo;
-	Oid			collation;
-	Vector	   *vec;
-	Vector	   *newCenter;
-	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;
+	VectorArray centers = buildstate->centers;
+	VectorArray m = buildstate->samples;
+	int			b = m->maxlen;
+	int			t = 20;
+	double		distance;
 	double		minDistance;
 	int			closestCenter;
-	double		distance;
-	bool		rj;
-	bool		rjreset;
-	double		dxcx;
-	double		dxc;
-
-	/* Calculate allocation sizes */
-	Size		samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->dim);
-	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)));
+	int			i;
+	int			j;
+	int			k;
+	Vector	   *c;
+	Vector	   *x;
+	int		   *v;
+	int		   *d;
+	double		eta;
 
 	/* Set support functions */
-	procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
-	normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
-	collation = 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;
-	}
+	FmgrInfo   *procinfo = index_getprocinfo(buildstate->index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
+	FmgrInfo   *normprocinfo = buildstate->kmeansnormprocinfo;
+	Oid			collation = buildstate->index->rd_indcollation[0];
 
 	/* 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) */
-	for (j = 0; j < numSamples; j++)
-	{
-		minDistance = DBL_MAX;
-		closestCenter = -1;
+	v = palloc(sizeof(int) * centers->maxlen);
+	d = palloc(sizeof(int) * b);
 
-		vec = VectorArrayGet(samples, j);
+	for (int i = 0; i < centers->length; i++)
+		v[i] = 0;
 
-		/* Find closest center */
-		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++)
+	for (i = 0; i < t; i++)
 	{
 		/* 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 */
-		for (j = 0; j < numCenters; 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++)
+		/* Cache nearest center to x */
+		for (j = 0; j < m->length; 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)
-					continue;
+				distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(x), PointerGetDatum(VectorArrayGet(centers, k))));
 
-				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 < numSamples; j++)
+		for (j = 0; j < m->length; j++)
 		{
-			/* Step 2: Identify all points x such that u(x) <= s(c(x)) */
-			if (upperBound[j] <= s[closestCenters[j]])
-				continue;
+			x = VectorArrayGet(m, j);
 
-			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 (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)
+				if (v[j] == 0)
 				{
-					dxcx = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, closestCenters[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++;
-					}
+					c = VectorArrayGet(centers, j);
 
+					/* 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 */
-		for (j = 0; j < numCenters; j++)
+		/* Normalize if needed */
+		if (normprocinfo != NULL)
 		{
-			vec = VectorArrayGet(newCenters, j);
-			for (k = 0; k < dimensions; k++)
-				vec->x[k] = 0.0;
-
-			centerCounts[j] = 0;
+			for (j = 0; j < centers->length; j++)
+				ApplyNorm(normprocinfo, collation, VectorArrayGet(centers, j));
 		}
-
-		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(centerCounts);
-	pfree(closestCenters);
-	pfree(lowerBound);
-	pfree(upperBound);
-	pfree(s);
-	pfree(halfcdist);
-	pfree(newcdist);
+	pfree(v);
+	pfree(d);
 }
 
 /*
@@ -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)
-		QuickCenters(index, samples, centers);
-	else
-		ElkanKmeans(index, samples, centers);
+	int			numSamples;
+	Size		totalSize;
 
-	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

@@ -30,21 +30,6 @@ CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
 	return 0;
 }
 
-/*
- * Compare item distances
- */
-static int
-CompareItems(const pairingheap_node *a, const pairingheap_node *b, void *arg)
-{
-	if (((const IvfflatScanItem *) a)->distance > ((const IvfflatScanItem *) b)->distance)
-		return 1;
-
-	if (((const IvfflatScanItem *) a)->distance < ((const IvfflatScanItem *) b)->distance)
-		return -1;
-
-	return ItemPointerCompare(&((IvfflatScanItem *) a)->tid, &((IvfflatScanItem *) b)->tid);
-}
-
 /*
  * Get lists and sort by distance
  */
@@ -59,7 +44,6 @@ GetScanLists(IndexScanDesc scan, Datum value)
 	BlockNumber nextblkno = IVFFLAT_HEAD_BLKNO;
 	int			listCount = 0;
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
-	int			i;
 	double		distance;
 	IvfflatScanList *scanlist;
 	double		maxDistance = DBL_MAX;
@@ -113,140 +97,6 @@ GetScanLists(IndexScanDesc scan, Datum value)
 
 		UnlockReleaseBuffer(cbuf);
 	}
-
-	for (i = 0; i < so->probes; i++)
-		so->sortedLists[i] = (IvfflatScanList *) pairingheap_remove_first(so->listQueue);
-
-	Assert(pairingheap_is_empty(so->listQueue));
-}
-
-/*
- * Get items
- */
-static void
-GetScanItemsQuick(IndexScanDesc scan, Datum value)
-{
-	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
-	Buffer		buf;
-	Page		page;
-	IndexTuple	itup;
-	BlockNumber searchPage;
-	OffsetNumber offno;
-	OffsetNumber maxoffno;
-	Datum		datum;
-	bool		isnull;
-	TupleDesc	tupdesc = RelationGetDescr(scan->indexRelation);
-	int			i;
-	double		distance;
-	IvfflatScanItem *scanitem;
-	double		maxDistance = DBL_MAX;
-
-	/*
-	 * Reuse same set of shared buffers for scan
-	 *
-	 * See postgres/src/backend/storage/buffer/README for description
-	 */
-	BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
-
-	/* Search closest probes lists */
-	for (i = 0; i < so->probes; i++)
-	{
-		/* Read closest lists first for performance */
-		searchPage = so->sortedLists[i]->startPage;
-
-		/* Search all entry pages for list */
-		while (BlockNumberIsValid(searchPage))
-		{
-			buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas);
-			LockBuffer(buf, BUFFER_LOCK_SHARE);
-			page = BufferGetPage(buf);
-			maxoffno = PageGetMaxOffsetNumber(page);
-
-			for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
-			{
-				itup = (IndexTuple) PageGetItem(page, PageGetItemId(page, offno));
-				datum = index_getattr(itup, 1, tupdesc, &isnull);
-				distance = DatumGetFloat8(FunctionCall2Coll(so->procinfo, so->collation, datum, value));
-
-				if (so->itemCount < so->maxItems)
-				{
-					scanitem = &so->items[so->itemCount];
-					scanitem->searchPage = searchPage;
-					scanitem->tid = itup->t_tid;
-					scanitem->distance = distance;
-					so->itemCount++;
-
-					/* Add to heap */
-					pairingheap_add(so->itemQueue, &scanitem->ph_node);
-
-					/* Calculate max distance */
-					if (so->itemCount == so->maxItems)
-					{
-						maxDistance = ((IvfflatScanItem *) pairingheap_first(so->itemQueue))->distance;
-						scanitem = &so->items[so->itemCount];
-					}
-				}
-				else if (distance <= maxDistance)
-				{
-					/* Reuse */
-					scanitem->searchPage = searchPage;
-					scanitem->tid = itup->t_tid;
-					scanitem->distance = distance;
-					pairingheap_add(so->itemQueue, &scanitem->ph_node);
-
-					/* Remove */
-					scanitem = (IvfflatScanItem *) pairingheap_remove_first(so->itemQueue);
-
-					/* Update max distance */
-					maxDistance = ((IvfflatScanItem *) pairingheap_first(so->itemQueue))->distance;
-				}
-			}
-
-			searchPage = IvfflatPageGetOpaque(page)->nextblkno;
-
-			UnlockReleaseBuffer(buf);
-		}
-	}
-
-	for (i = 0; i < so->itemCount; i++)
-		so->sortedItems[i] = (IvfflatScanItem *) pairingheap_remove_first(so->itemQueue);
-
-	Assert(pairingheap_is_empty(so->itemQueue));
-}
-
-/*
- * Initialize sort
- */
-static void
-InitSort(IvfflatScanOpaque so)
-{
-	AttrNumber	attNums[] = {1, 2};
-	Oid			sortOperators[] = {Float8LessOperator, TIDLessOperator};
-	Oid			sortCollations[] = {InvalidOid, InvalidOid};
-	bool		nullsFirstFlags[] = {false, false};
-
-	/* Create tuple description for sorting */
-#if PG_VERSION_NUM >= 120000
-	so->tupdesc = CreateTemplateTupleDesc(3);
-#else
-	so->tupdesc = CreateTemplateTupleDesc(3, false);
-#endif
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0);
-
-	/* Prep sort */
-#if PG_VERSION_NUM >= 110000
-	so->sortstate = tuplesort_begin_heap(so->tupdesc, sizeof(attNums) / sizeof(attNums[0]), attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
-#else
-	so->sortstate = tuplesort_begin_heap(so->tupdesc, sizeof(attNums) / sizeof(attNums[0]), attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, false);
-#endif
-
-#if PG_VERSION_NUM >= 120000
-	so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
-#else
-	so->slot = MakeSingleTupleTableSlot(so->tupdesc);
-#endif
 }
 
 /*
@@ -265,7 +115,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	Datum		datum;
 	bool		isnull;
 	TupleDesc	tupdesc = RelationGetDescr(scan->indexRelation);
-	int			i;
 
 #if PG_VERSION_NUM >= 120000
 	TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
@@ -281,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->sortedLists[i]->startPage;
+		searchPage = ((IvfflatScanList *) pairingheap_remove_first(so->listQueue))->startPage;
 
 		/* Search all entry pages for list */
 		while (BlockNumberIsValid(searchPage))
@@ -323,7 +172,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	}
 
 	tuplesort_performsort(so->sortstate);
-	tuplesort_skiptuples(so->sortstate, so->maxItems, true);
 }
 
 /*
@@ -332,17 +180,24 @@ GetScanItems(IndexScanDesc scan, Datum value)
 IndexScanDesc
 ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 {
-	IndexScanDesc scan = RelationGetIndexScan(index, nkeys, norderbys);
-	int			lists = IvfflatGetLists(scan->indexRelation);
-	int			probes = ivfflat_probes;
+	IndexScanDesc scan;
 	IvfflatScanOpaque so;
+	int			lists;
+	AttrNumber	attNums[] = {1};
+	Oid			sortOperators[] = {Float8LessOperator};
+	Oid			sortCollations[] = {InvalidOid};
+	bool		nullsFirstFlags[] = {false};
+	int			probes = ivfflat_probes;
+
+	scan = RelationGetIndexScan(index, nkeys, norderbys);
+	lists = IvfflatGetLists(scan->indexRelation);
 
 	if (probes > lists)
 		probes = lists;
 
 	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
 	so->buf = InvalidBuffer;
-	so->stage = 0;
+	so->first = true;
 	so->probes = probes;
 
 	/* Set support functions */
@@ -350,16 +205,30 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
 	so->collation = index->rd_indcollation[0];
 
+	/* Create tuple description for sorting */
+#if PG_VERSION_NUM >= 120000
+	so->tupdesc = CreateTemplateTupleDesc(3);
+#else
+	so->tupdesc = CreateTemplateTupleDesc(3, false);
+#endif
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0);
+
+	/* Prep sort */
+#if PG_VERSION_NUM >= 110000
+	so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
+#else
+	so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, false);
+#endif
+
+#if PG_VERSION_NUM >= 120000
+	so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
+#else
+	so->slot = MakeSingleTupleTableSlot(so->tupdesc);
+#endif
+
 	so->listQueue = pairingheap_allocate(CompareLists, scan);
-	so->sortedLists = palloc(sizeof(IvfflatScanItem *) * probes);
-
-	so->maxItems = 1024;
-	so->itemCount = 0;
-	so->itemQueue = pairingheap_allocate(CompareItems, scan);
-	so->items = palloc(sizeof(IvfflatScanItem) * (so->maxItems + 1));
-	so->sortedItems = palloc(sizeof(IvfflatScanItem *) * so->maxItems);
-
-	so->sortstate = NULL;
 
 	scan->opaque = so;
 
@@ -375,14 +244,12 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
 
 #if PG_VERSION_NUM >= 130000
-	if (so->sortstate != NULL)
+	if (!so->first)
 		tuplesort_reset(so->sortstate);
 #endif
 
-	so->stage = 0;
+	so->first = true;
 	pairingheap_reset(so->listQueue);
-	pairingheap_reset(so->itemQueue);
-	so->itemCount = 0;
 
 	if (keys && scan->numberOfKeys > 0)
 		memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -405,7 +272,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 	 */
 	Assert(ScanDirectionIsForward(dir));
 
-	if (so->stage == 0)
+	if (so->first)
 	{
 		Datum		value;
 
@@ -427,101 +294,42 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		}
 
 		IvfflatBench("GetScanLists", GetScanLists(scan, value));
-		IvfflatBench("GetScanItemsQuick", GetScanItemsQuick(scan, value));
-		so->heapFull = so->itemCount == so->maxItems;
-		so->stage++;
+		IvfflatBench("GetScanItems", GetScanItems(scan, value));
+		so->first = false;
 
 		/* Clean up if we allocated a new value */
 		if (value != scan->orderByData->sk_argument)
 			pfree(DatumGetPointer(value));
 	}
 
-	if (so->stage == 1)
-	{
-		if (so->itemCount > 0)
-		{
-			IvfflatScanItem *scanitem;
-
-			so->itemCount--;
-
-			scanitem = so->sortedItems[so->itemCount];
-
-#if PG_VERSION_NUM >= 120000
-			scan->xs_heaptid = scanitem->tid;
-#else
-			scan->xs_ctup.t_sef = scanitem->tid;
-#endif
-
-			if (BufferIsValid(so->buf))
-				ReleaseBuffer(so->buf);
-
-			/*
-			 * An index scan must maintain a pin on the index page holding the
-			 * item last returned by amgettuple
-			 *
-			 * https://www.postgresql.org/docs/current/index-locking.html
-			 */
-			so->buf = ReadBuffer(scan->indexRelation, scanitem->searchPage);
-
-			scan->xs_recheckorderby = false;
-			return true;
-		}
-		else if (so->heapFull)
-		{
-			Datum value = scan->orderByData->sk_argument;
-
-			if (so->normprocinfo != NULL)
-			{
-				/* No items will match if normalization fails */
-				if (!IvfflatNormValue(so->normprocinfo, so->collation, &value, NULL))
-					return false;
-			}
-
-			if (so->sortstate == NULL)
-				InitSort(so);
-
-			IvfflatBench("GetScanItems", GetScanItems(scan, value));
-			so->stage++;
-
-			/* Clean up if we allocated a new value */
-			if (value != scan->orderByData->sk_argument)
-				pfree(DatumGetPointer(value));
-		}
-		else
-			so->stage = 3;
-	}
-
-	if (so->stage == 2)
-	{
 #if PG_VERSION_NUM >= 100000
-		if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
+	if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
 #else
-		if (tuplesort_gettupleslot(so->sortstate, true, so->slot, NULL))
+	if (tuplesort_gettupleslot(so->sortstate, true, so->slot, NULL))
 #endif
-		{
-			ItemPointer tid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
-			BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
+	{
+		ItemPointer tid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
+		BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
 
 #if PG_VERSION_NUM >= 120000
-			scan->xs_heaptid = *tid;
+		scan->xs_heaptid = *tid;
 #else
-			scan->xs_ctup.t_self = *tid;
+		scan->xs_ctup.t_self = *tid;
 #endif
 
-			if (BufferIsValid(so->buf))
-				ReleaseBuffer(so->buf);
+		if (BufferIsValid(so->buf))
+			ReleaseBuffer(so->buf);
 
-			/*
-			 * An index scan must maintain a pin on the index page holding the
-			 * item last returned by amgettuple
-			 *
-			 * https://www.postgresql.org/docs/current/index-locking.html
-			 */
-			so->buf = ReadBuffer(scan->indexRelation, indexblkno);
+		/*
+		 * An index scan must maintain a pin on the index page holding the
+		 * item last returned by amgettuple
+		 *
+		 * https://www.postgresql.org/docs/current/index-locking.html
+		 */
+		so->buf = ReadBuffer(scan->indexRelation, indexblkno);
 
-			scan->xs_recheckorderby = false;
-			return true;
-		}
+		scan->xs_recheckorderby = false;
+		return true;
 	}
 
 	return false;
@@ -540,14 +348,7 @@ ivfflatendscan(IndexScanDesc scan)
 		ReleaseBuffer(so->buf);
 
 	pairingheap_free(so->listQueue);
-	pfree(so->sortedLists);
-
-	if (so->sortstate != NULL)
-		tuplesort_end(so->sortstate);
-
-	pairingheap_free(so->itemQueue);
-	pfree(so->items);
-	pfree(so->sortedItems);
+	tuplesort_end(so->sortstate);
 
 	pfree(so);
 	scan->opaque = NULL;

test/t/007_stages.pl

@@ -1,38 +0,0 @@
-use strict;
-use warnings;
-use PostgresNode;
-use TestLib;
-use Test::More tests => 2;
-
-# 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, ARRAY[i % 1000, i % 1000, i % 1000] FROM generate_series(1, 10000) i;"
-);
-
-my @limits = (128, 2048);
-my @expected = ();
-
-foreach (@limits) {
-	my $res = $node->safe_psql("postgres", "SELECT i, v FROM tst ORDER BY v <-> '[0,0,0]', i LIMIT $_;");
-	push(@expected, $res);
-}
-
-$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v) WITH (lists = 5);");
-
-for my $i (0 .. $#limits) {
-	my $res = $node->safe_psql("postgres", qq(
-		SET enable_seqscan = off;
-		SET ivfflat.probes = 5;
-		WITH tmp AS (
-			SELECT *, v <-> '[0,0,0]' AS d FROM tst ORDER BY v <-> '[0,0,0]' LIMIT $limits[$i]
-		) SELECT i, v FROM tmp ORDER BY d, i;
-	));
-	is($res, $expected[$i]);
-}