Fixed alloc [skip ci]

CHANGELOG.md

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

README.md

@@ -119,9 +119,10 @@ SELECT phase, tuples_done, tuples_total FROM pg_stat_progress_create_index;
 The phases are:
 
 1. `initializing`
-2. `performing k-means`
-3. `sorting tuples`
-4. `loading tuples`
+2. `sampling table`
+3. `performing k-means`
+4. `sorting tuples`
+5. `loading tuples`
 
 Note: `tuples_done` and `tuples_total` are only populated during the `loading tuples` phase
 
@@ -263,7 +264,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)
-- [Web-Scale k-means Clustering](https://www.eecs.tufts.edu/~dsculley/papers/fastkmeans.pdf)
+- [Using the Triangle Inequality to Accelerate k-means](https://www.aaai.org/Papers/ICML/2003/ICML03-022.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,6 +42,87 @@
 #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
  */
@@ -287,6 +368,38 @@ 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
  */
@@ -460,9 +573,7 @@ BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo,
 {
 	InitBuildState(buildstate, heap, index, indexInfo);
 
-	/* Perform k-means clustering */
-	UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_KMEANS);
-	IvfflatBench("k-means", IvfflatKmeans(buildstate));
+	ComputeCenters(buildstate);
 
 	/* Create pages */
 	CreateMetaPage(index, buildstate->dimensions, buildstate->lists, forkNum);

src/ivfflat.c

@@ -45,6 +45,8 @@ 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,9 +37,10 @@
 
 /* Build phases */
 /* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
-#define PROGRESS_IVFFLAT_PHASE_KMEANS	2
-#define PROGRESS_IVFFLAT_PHASE_SORT		3
-#define PROGRESS_IVFFLAT_PHASE_LOAD		4
+#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 IVFFLAT_LIST_SIZE(_dim)	(offsetof(IvfflatListData, center) + VECTOR_SIZE(_dim))
 
@@ -166,12 +167,28 @@ 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;
-	bool		first;
+	int			stage;
 	Buffer		buf;
 
+	/* Items */
+	int			maxItems;
+	int			itemCount;
+	pairingheap *itemQueue;
+	IvfflatScanItem *items;
+	IvfflatScanItem **sortedItems;
+	bool		heapFull;
+
 	/* Sorting */
 	Tuplesortstate *sortstate;
 	TupleDesc	tupdesc;
@@ -185,6 +202,7 @@ typedef struct IvfflatScanOpaqueData
 
 	/* Lists */
 	pairingheap *listQueue;
+	IvfflatScanList **sortedLists;
 	IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER];	/* must come last */
 }			IvfflatScanOpaqueData;
 
@@ -199,7 +217,7 @@ typedef IvfflatScanOpaqueData * IvfflatScanOpaque;
 void		_PG_init(void);
 VectorArray VectorArrayInit(int maxlen, int dimensions);
 void		PrintVectorArray(char *msg, VectorArray arr);
-void		IvfflatKmeans(IvfflatBuildState * buildstate);
+void		IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);
 FmgrInfo   *IvfflatOptionalProcInfo(Relation rel, uint16 procnum);
 bool		IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
 int			IvfflatGetLists(Relation index);

src/ivfkmeans.c

@@ -2,20 +2,8 @@
 
 #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++
@@ -23,7 +11,7 @@
  * https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf
  */
 static void
-InitCenters(Relation index, VectorArray samples, VectorArray centers)
+InitCenters(Relation index, VectorArray samples, VectorArray centers, float *lowerBound)
 {
 	FmgrInfo   *procinfo;
 	Oid			collation;
@@ -47,7 +35,7 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers)
 	for (j = 0; j < numSamples; j++)
 		weight[j] = DBL_MAX;
 
-	for (i = 0; i < numCenters - 1; i++)
+	for (i = 0; i < numCenters; i++)
 	{
 		CHECK_FOR_INTERRUPTS();
 
@@ -61,6 +49,9 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers)
 			/* 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;
 
@@ -70,6 +61,10 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers)
 			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++)
@@ -161,202 +156,299 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
 }
 
 /*
- * 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
+ * Use Elkan for performance. This requires distance function to satisfy triangle inequality.
  *
  * We use L2 distance for L2 (not L2 squared like index scan)
  * and angular distance for inner product and cosine distance
  *
- * https://www.eecs.tufts.edu/~dsculley/papers/fastkmeans.pdf
+ * https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf
  */
 static void
-MiniBatchKmeans(IvfflatBuildState * buildstate)
+ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
 {
-	VectorArray centers = buildstate->centers;
-	VectorArray m = buildstate->samples;
-	int			b = m->maxlen;
-	int			t = 20;
-	double		distance;
-	double		minDistance;
-	int			closestCenter;
-	int			i;
+	FmgrInfo   *procinfo;
+	FmgrInfo   *normprocinfo;
+	Oid			collation;
+	Vector	   *vec;
+	Vector	   *newCenter;
+	int			iteration;
 	int			j;
 	int			k;
-	Vector	   *c;
-	Vector	   *x;
-	int		   *v;
-	int		   *d;
-	double		eta;
+	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;
+	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)));
 
 	/* Set support functions */
-	FmgrInfo   *procinfo = index_getprocinfo(buildstate->index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
-	FmgrInfo   *normprocinfo = buildstate->kmeansnormprocinfo;
-	Oid			collation = buildstate->index->rd_indcollation[0];
+	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;
+	}
 
 	/* Pick initial centers */
-	InitCenters(buildstate->index, buildstate->samples, buildstate->centers);
+	InitCenters(index, samples, centers, lowerBound);
 
-	v = palloc(sizeof(int) * centers->maxlen);
-	d = palloc(sizeof(int) * b);
+	/* 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;
 
-	for (int i = 0; i < centers->length; i++)
-		v[i] = 0;
+		vec = VectorArrayGet(samples, j);
 
-	for (i = 0; i < t; i++)
+		/* 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++)
 	{
 		/* Can take a while, so ensure we can interrupt */
 		CHECK_FOR_INTERRUPTS();
 
-		/* Get b examples picked randomly from X */
-		SampleRows(buildstate);
+		changes = 0;
 
-		/* Cache nearest center to x */
-		for (j = 0; j < m->length; j++)
+		/* 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++)
 		{
-			/* compute closest */
 			minDistance = DBL_MAX;
-			closestCenter = -1;
 
-			x = VectorArrayGet(m, j);
-
-			/* Find closest center */
-			for (k = 0; k < centers->length; k++)
+			for (k = 0; k < numCenters; k++)
 			{
-				distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(x), PointerGetDatum(VectorArrayGet(centers, k))));
+				if (j == k)
+					continue;
 
+				distance = halfcdist[j * numCenters + k];
 				if (distance < minDistance)
-				{
 					minDistance = distance;
-					closestCenter = k;
-				}
 			}
 
-			d[j] = closestCenter;
+			s[j] = minDistance;
 		}
 
-		for (j = 0; j < m->length; j++)
+		rjreset = iteration != 0;
+
+		for (j = 0; j < numSamples; j++)
 		{
-			x = VectorArrayGet(m, j);
+			/* Step 2: Identify all points x such that u(x) <= s(c(x)) */
+			if (upperBound[j] <= s[closestCenters[j]])
+				continue;
 
-			/* Get cached center for this x */
-			c = VectorArrayGet(centers, d[j]);
+			rj = rjreset;
 
-			/* 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++)
+			for (k = 0; k < numCenters; k++)
 			{
-				if (v[j] == 0)
-				{
-					c = VectorArrayGet(centers, 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)
+				{
+					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++;
+					}
 
-					/* TODO Handle empty centers properly */
-					for (k = 0; k < c->dim; k++)
-						c->x[k] = ((double) random()) / MAX_RANDOM_VALUE;
 				}
 			}
 		}
 
-		/* Normalize if needed */
-		if (normprocinfo != NULL)
+		/* Step 4: For each center c, let m(c) be mean of all points assigned */
+		for (j = 0; j < numCenters; j++)
 		{
-			for (j = 0; j < centers->length; j++)
-				ApplyNorm(normprocinfo, collation, VectorArrayGet(centers, j));
+			vec = VectorArrayGet(newCenters, j);
+			for (k = 0; k < dimensions; k++)
+				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(v);
-	pfree(d);
+	pfree(newCenters);
+	pfree(centerCounts);
+	pfree(closestCenters);
+	pfree(lowerBound);
+	pfree(upperBound);
+	pfree(s);
+	pfree(halfcdist);
+	pfree(newcdist);
 }
 
 /*
@@ -399,48 +491,16 @@ CheckCenters(Relation index, VectorArray centers)
 }
 
 /*
- * Perform k-means clustering
+ * Perform naive k-means centering
  * We use spherical k-means for inner product and cosine
  */
 void
-IvfflatKmeans(IvfflatBuildState * buildstate)
+IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers)
 {
-	int			numSamples;
-	Size		totalSize;
-
-	/* 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);
+	if (samples->length <= centers->maxlen)
+		QuickCenters(index, samples, centers);
 	else
-		MiniBatchKmeans(buildstate);
+		ElkanKmeans(index, samples, centers);
 
-	CheckCenters(buildstate->index, buildstate->centers);
-
-	/* Free samples before we allocate more memory */
-	pfree(buildstate->samples);
+	CheckCenters(index, centers);
 }

src/ivfscan.c

@@ -30,6 +30,21 @@ 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
  */
@@ -44,6 +59,7 @@ 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;
@@ -97,6 +113,140 @@ 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
 }
 
 /*
@@ -115,6 +265,7 @@ 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);
@@ -130,9 +281,9 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
 
 	/* Search closest probes lists */
-	while (!pairingheap_is_empty(so->listQueue))
+	for (i = 0; i < so->probes; i++)
 	{
-		searchPage = ((IvfflatScanList *) pairingheap_remove_first(so->listQueue))->startPage;
+		searchPage = so->sortedLists[i]->startPage;
 
 		/* Search all entry pages for list */
 		while (BlockNumberIsValid(searchPage))
@@ -172,6 +323,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
 	}
 
 	tuplesort_performsort(so->sortstate);
+	tuplesort_skiptuples(so->sortstate, so->maxItems, true);
 }
 
 /*
@@ -180,24 +332,17 @@ GetScanItems(IndexScanDesc scan, Datum value)
 IndexScanDesc
 ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 {
-	IndexScanDesc scan;
-	IvfflatScanOpaque so;
-	int			lists;
-	AttrNumber	attNums[] = {1};
-	Oid			sortOperators[] = {Float8LessOperator};
-	Oid			sortCollations[] = {InvalidOid};
-	bool		nullsFirstFlags[] = {false};
+	IndexScanDesc scan = RelationGetIndexScan(index, nkeys, norderbys);
+	int			lists = IvfflatGetLists(scan->indexRelation);
 	int			probes = ivfflat_probes;
-
-	scan = RelationGetIndexScan(index, nkeys, norderbys);
-	lists = IvfflatGetLists(scan->indexRelation);
+	IvfflatScanOpaque so;
 
 	if (probes > lists)
 		probes = lists;
 
 	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
 	so->buf = InvalidBuffer;
-	so->first = true;
+	so->stage = 0;
 	so->probes = probes;
 
 	/* Set support functions */
@@ -205,30 +350,16 @@ 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;
 
@@ -244,12 +375,14 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
 	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
 
 #if PG_VERSION_NUM >= 130000
-	if (!so->first)
+	if (so->sortstate != NULL)
 		tuplesort_reset(so->sortstate);
 #endif
 
-	so->first = true;
+	so->stage = 0;
 	pairingheap_reset(so->listQueue);
+	pairingheap_reset(so->itemQueue);
+	so->itemCount = 0;
 
 	if (keys && scan->numberOfKeys > 0)
 		memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -272,7 +405,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 	 */
 	Assert(ScanDirectionIsForward(dir));
 
-	if (so->first)
+	if (so->stage == 0)
 	{
 		Datum		value;
 
@@ -294,42 +427,101 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		}
 
 		IvfflatBench("GetScanLists", GetScanLists(scan, value));
-		IvfflatBench("GetScanItems", GetScanItems(scan, value));
-		so->first = false;
+		IvfflatBench("GetScanItemsQuick", GetScanItemsQuick(scan, value));
+		so->heapFull = so->itemCount == so->maxItems;
+		so->stage++;
 
 		/* Clean up if we allocated a new value */
 		if (value != scan->orderByData->sk_argument)
 			pfree(DatumGetPointer(value));
 	}
 
-#if PG_VERSION_NUM >= 100000
-	if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
-#else
-	if (tuplesort_gettupleslot(so->sortstate, true, so->slot, NULL))
-#endif
+	if (so->stage == 1)
 	{
-		ItemPointer tid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
-		BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
+		if (so->itemCount > 0)
+		{
+			IvfflatScanItem *scanitem;
+
+			so->itemCount--;
+
+			scanitem = so->sortedItems[so->itemCount];
 
 #if PG_VERSION_NUM >= 120000
-		scan->xs_heaptid = *tid;
+			scan->xs_heaptid = scanitem->tid;
 #else
-		scan->xs_ctup.t_self = *tid;
+			scan->xs_ctup.t_sef = scanitem->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, scanitem->searchPage);
 
-		scan->xs_recheckorderby = false;
-		return true;
+			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))
+#else
+		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));
+
+#if PG_VERSION_NUM >= 120000
+			scan->xs_heaptid = *tid;
+#else
+			scan->xs_ctup.t_self = *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, indexblkno);
+
+			scan->xs_recheckorderby = false;
+			return true;
+		}
 	}
 
 	return false;
@@ -348,7 +540,14 @@ ivfflatendscan(IndexScanDesc scan)
 		ReleaseBuffer(so->buf);
 
 	pairingheap_free(so->listQueue);
-	tuplesort_end(so->sortstate);
+	pfree(so->sortedLists);
+
+	if (so->sortstate != NULL)
+		tuplesort_end(so->sortstate);
+
+	pairingheap_free(so->itemQueue);
+	pfree(so->items);
+	pfree(so->sortedItems);
 
 	pfree(so);
 	scan->opaque = NULL;

test/t/007_stages.pl

@@ -0,0 +1,38 @@
+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]);
+}