Merge branch 'master' into minibatch

.github/workflows/build.yml

@@ -33,6 +33,6 @@ jobs:
     - if: ${{ startsWith(matrix.os, 'macos') }}
       run: |
         brew install cpanm && cpanm IPC::Run
-        wget -q https://github.com/postgres/postgres/archive/refs/tags/REL_14_4.tar.gz
+        wget -q https://github.com/postgres/postgres/archive/refs/tags/REL_14_1.tar.gz
-        tar xf REL_14_4.tar.gz
+        tar xf REL_14_1.tar.gz
-        make prove_installcheck PROVE=prove PERL5LIB="postgres-REL_14_4/src/test/perl:/Users/runner/perl5/lib/perl5"
+        make prove_installcheck PROVE=prove PERL5LIB=postgres-REL_14_1/src/test/perl

CHANGELOG.md

@@ -1,9 +1,6 @@
-## 0.2.7 (2022-07-31)
+## 0.2.6 (unreleased)
-
-- Fixed `unexpected data beyond EOF` error
-
-## 0.2.6 (2022-05-22)
 
+- Switched to mini-batch k-means
 - Improved performance of index creation for Postgres < 12
 
 ## 0.2.5 (2022-02-11)

META.json

@@ -2,7 +2,7 @@
 	"name": "vector",
 	"abstract": "Open-source vector similarity search for Postgres",
 	"description": "Supports L2 distance, inner product, and cosine distance",
-	"version": "0.2.7",
+	"version": "0.2.5",
 	"maintainer": [
 		"Andrew Kane <andrew@ankane.org>"
 	],
@@ -20,7 +20,7 @@
 		"vector": {
 			"file": "sql/vector.sql",
 			"docfile": "README.md",
-			"version": "0.2.7",
+			"version": "0.2.5",
 			"abstract": "Open-source vector similarity search for Postgres"
 		}
 	},

Makefile

@@ -1,5 +1,5 @@
 EXTENSION = vector
-EXTVERSION = 0.2.7
+EXTVERSION = 0.2.5
 
 MODULE_big = vector
 DATA = $(wildcard sql/*--*.sql)

README.md

@@ -17,7 +17,7 @@ Supports L2 distance, inner product, and cosine distance
 Compile and install the extension (supports Postgres 9.6+)
 
 ```sh
-git clone --branch v0.2.7 https://github.com/pgvector/pgvector.git
+git clone --branch v0.2.5 https://github.com/pgvector/pgvector.git
 cd pgvector
 make
 make install # may need sudo
@@ -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
 
@@ -220,14 +219,14 @@ This adds pgvector to the [Postgres image](https://hub.docker.com/_/postgres).
 You can also build the image manually
 
 ```sh
-git clone --branch v0.2.7 https://github.com/pgvector/pgvector.git
+git clone --branch v0.2.5 https://github.com/pgvector/pgvector.git
 cd pgvector
 docker build -t pgvector .
 ```
 
 ### Homebrew
 
-With Homebrew Postgres, you can use:
+On Mac with Homebrew Postgres, you can use:
 
 ```sh
 brew install pgvector/brew/pgvector
@@ -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)
 

sql/vector--0.2.5--0.2.6.sql

@@ -1,2 +0,0 @@
--- complain if script is sourced in psql, rather than via CREATE EXTENSION
-\echo Use "ALTER EXTENSION vector UPDATE TO '0.2.6'" to load this file. \quit

sql/vector--0.2.6--0.2.7.sql

@@ -1,2 +0,0 @@
--- complain if script is sourced in psql, rather than via CREATE EXTENSION
-\echo Use "ALTER EXTENSION vector UPDATE TO '0.2.7'" to load this file. \quit

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	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))
 
@@ -200,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/ivfinsert.c

@@ -53,6 +53,18 @@ FindInsertPage(Relation rel, Datum *values, BlockNumber *insertPage, ListInfo *
 	}
 }
 
+/*
+ * Prepare to insert an index tuple
+ */
+static void
+LoadInsertPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state, BlockNumber insertPage)
+{
+	*buf = ReadBuffer(index, insertPage);
+	LockBuffer(*buf, BUFFER_LOCK_EXCLUSIVE);
+	*state = GenericXLogStart(index);
+	*page = GenericXLogRegisterBuffer(*state, *buf, 0);
+}
+
 /*
  * Insert a tuple into the index
  */
@@ -75,18 +87,11 @@ InsertTuple(Relation rel, IndexTuple itup, Relation heapRel, Datum *values)
 	itemsz = MAXALIGN(IndexTupleSize(itup));
 	Assert(itemsz <= BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(IvfflatPageOpaqueData)));
 
+	LoadInsertPage(rel, &buf, &page, &state, insertPage);
+
 	/* Find a page to insert the item */
-	for (;;)
+	while (PageGetFreeSpace(page) < itemsz)
 	{
-		buf = ReadBuffer(rel, insertPage);
-		LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
-
-		state = GenericXLogStart(rel);
-		page = GenericXLogRegisterBuffer(state, buf, 0);
-
-		if (PageGetFreeSpace(page) >= itemsz)
-			break;
-
 		insertPage = IvfflatPageGetOpaque(page)->nextblkno;
 
 		if (BlockNumberIsValid(insertPage))
@@ -94,31 +99,15 @@ InsertTuple(Relation rel, IndexTuple itup, Relation heapRel, Datum *values)
 			/* Move to next page */
 			GenericXLogAbort(state);
 			UnlockReleaseBuffer(buf);
+
+			LoadInsertPage(rel, &buf, &page, &state, insertPage);
 		}
 		else
 		{
 			/* Add a new page */
-			Buffer		newbuf = IvfflatNewBuffer(rel, MAIN_FORKNUM);
+			IvfflatAppendPage(rel, &buf, &page, &state, MAIN_FORKNUM);
-			Page		newpage = GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE);
 
-			insertPage = BufferGetBlockNumber(newbuf);
+			insertPage = BufferGetBlockNumber(buf);
-
-			/* Update previous buffer */
-			IvfflatPageGetOpaque(page)->nextblkno = insertPage;
-
-			/* Init page */
-			PageInit(newpage, BufferGetPageSize(newbuf), sizeof(IvfflatPageOpaqueData));
-			IvfflatPageGetOpaque(newpage)->nextblkno = InvalidBlockNumber;
-			IvfflatPageGetOpaque(newpage)->page_id = IVFFLAT_PAGE_ID;
-
-			/* Commit */
-			MarkBufferDirty(buf);
-			MarkBufferDirty(newbuf);
-			GenericXLogFinish(state);
-
-			/* Unlock */
-			UnlockReleaseBuffer(buf);
-			UnlockReleaseBuffer(newbuf);
 		}
 	}
 

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,106 +161,147 @@ 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);
+
+	for (int i = 0; i < centers->length; i++)
+		v[i] = 0;
+
+	for (i = 0; i < t; i++)
 	{
+		/* Can take a while, so ensure we can interrupt */
+		CHECK_FOR_INTERRUPTS();
+
+		/* Get b examples picked randomly from X */
+		SampleRows(buildstate);
+
+		/* Cache nearest center to x */
+		for (j = 0; j < m->length; j++)
+		{
+			/* compute closest */
 			minDistance = DBL_MAX;
 			closestCenter = -1;
 
-		vec = VectorArrayGet(samples, j);
+			x = VectorArrayGet(m, j);
 
 			/* Find closest center */
-		for (k = 0; k < numCenters; k++)
+			for (k = 0; k < centers->length; k++)
 			{
-			/* TODO Use Lemma 1 in k-means++ initialization */
+				distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(x), PointerGetDatum(VectorArrayGet(centers, k))));
-			distance = lowerBound[j * numCenters + k];
 
 				if (distance < minDistance)
 				{
@@ -264,191 +310,53 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
 				}
 			}
 
-		upperBound[j] = minDistance;
+			d[j] = closestCenter;
-		closestCenters[j] = closestCenter;
 		}
 
-	/* Give 500 iterations to converge */
+		for (j = 0; j < m->length; j++)
-	for (iteration = 0; iteration < 500; iteration++)
 		{
-		/* Can take a while, so ensure we can interrupt */
+			x = VectorArrayGet(m, j);
-		CHECK_FOR_INTERRUPTS();
 
-		changes = 0;
+			/* Get cached center for this x */
+			c = VectorArrayGet(centers, d[j]);
 
-		/* Step 1: For all centers, compute distance */
+			/* Update per-center counts */
-		for (j = 0; j < numCenters; j++)
+			v[d[j]]++;
-		{
-			vec = VectorArrayGet(centers, j);
 
-			for (k = j + 1; k < numCenters; k++)
+			/* Get per-center learning rate */
-			{
+			eta = 1.0 / v[d[j]];
-				distance = 0.5 * DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, k))));
+
-				halfcdist[j * numCenters + k] = distance;
+			/* Take gradient step */
-				halfcdist[k * numCenters + j] = distance;
+			for (k = 0; k < c->dim; k++)
-			}
+				c->x[k] = (1 - eta) * c->x[k] + eta * x->x[k];
 		}
 
-		/* For all centers c, compute s(c) */
+		/* Check for empty centers (likely duplicates) */
-		for (j = 0; j < numCenters; j++)
+		if (i == 0)
 		{
-			minDistance = DBL_MAX;
+			for (j = 0; j < centers->length; j++)
-
-			for (k = 0; k < numCenters; k++)
 			{
-				if (j == k)
+				if (v[j] == 0)
-					continue;
-
-				distance = halfcdist[j * numCenters + k];
-				if (distance < minDistance)
-					minDistance = distance;
-			}
-
-			s[j] = minDistance;
-		}
-
-		rjreset = iteration != 0;
-
-		for (j = 0; j < numSamples; j++)
 				{
-			/* Step 2: Identify all points x such that u(x) <= s(c(x)) */
+					c = VectorArrayGet(centers, j);
-			if (upperBound[j] <= s[closestCenters[j]])
-				continue;
 
-			rj = rjreset;
-
-			for (k = 0; k < numCenters; k++)
-			{
-				/* 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++;
-					}
-
-				}
-			}
-		}
-
-		/* Step 4: For each center c, let m(c) be mean of all points assigned */
-		for (j = 0; j < numCenters; 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++)
+					for (k = 0; k < c->dim; k++)
-					vec->x[k] = ((double) random()) / MAX_RANDOM_VALUE;
+						c->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++)
+			for (j = 0; j < centers->length; j++)
-			{
+				ApplyNorm(normprocinfo, collation, VectorArrayGet(centers, j));
-				distance = lowerBound[j * numCenters + k] - newcdist[k];
-
-				if (distance < 0)
-					distance = 0;
-
-				lowerBound[j * numCenters + k] = distance;
 		}
 	}
 
-		/* Step 6 */
+	pfree(v);
-		/* We reset r(x) before Step 3 in the next iteration */
+	pfree(d);
-		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);
 }
 
 /*
@@ -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/ivfutils.c

@@ -135,29 +135,17 @@ IvfflatCommitBuffer(Buffer buf, GenericXLogState *state)
 void
 IvfflatAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state, ForkNumber forkNum)
 {
-	/* Get new buffer */
+	Buffer		prevbuf = *buf;
-	Buffer		newbuf = IvfflatNewBuffer(index, forkNum);
-	Page		newpage = GenericXLogRegisterBuffer(*state, newbuf, GENERIC_XLOG_FULL_IMAGE);
 
-	/* Update the previous buffer */
+	/* Get new buffer */
-	IvfflatPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(newbuf);
+	*buf = IvfflatNewBuffer(index, forkNum);
+
+	/* Update and commit previous buffer */
+	IvfflatPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(*buf);
+	IvfflatCommitBuffer(prevbuf, *state);
 
 	/* Init new page */
-	PageInit(newpage, BufferGetPageSize(newbuf), sizeof(IvfflatPageOpaqueData));
+	IvfflatInitPage(index, buf, page, state);
-	IvfflatPageGetOpaque(newpage)->nextblkno = InvalidBlockNumber;
-	IvfflatPageGetOpaque(newpage)->page_id = IVFFLAT_PAGE_ID;
-
-	/* Commit */
-	MarkBufferDirty(*buf);
-	MarkBufferDirty(newbuf);
-	GenericXLogFinish(*state);
-
-	/* Unlock */
-	UnlockReleaseBuffer(*buf);
-
-	*state = GenericXLogStart(index);
-	*page = GenericXLogRegisterBuffer(*state, newbuf, GENERIC_XLOG_FULL_IMAGE);
-	*buf = newbuf;
 }
 
 /*

test/t/001_wal.pl

@@ -7,8 +7,6 @@ use PostgresNode;
 use TestLib;
 use Test::More tests => 31;
 
-my $dim = 32;
-
 my $node_primary;
 my $node_replica;
 
@@ -32,15 +30,13 @@ sub test_index_replay
 	$node_primary->poll_query_until('postgres', $caughtup_query)
 	  or die "Timed out while waiting for replica 1 to catch up";
 
-	my @r = ();
+	my $r1 = rand();
-	for (1 .. $dim) {
+	my $r2 = rand();
-		push(@r, rand());
+	my $r3 = rand();
-	}
-	my $sql = join(",", @r);
 
 	my $queries = qq(
 		SET enable_seqscan = off;
-		SELECT * FROM tst ORDER BY v <-> '[$sql]' LIMIT 10;
+		SELECT * FROM tst ORDER BY v <-> '[$r1,$r2,$r3]' LIMIT 10;
 	);
 
 	# Run test queries and compare their result
@@ -54,10 +50,6 @@ sub test_index_replay
 # Initialize primary node
 $node_primary = get_new_node('primary');
 $node_primary->init(allows_streaming => 1);
-if ($dim > 32) {
-	# TODO use wal_keep_segments for Postgres < 13
-	$node_primary->append_conf('postgresql.conf', qq(wal_keep_size = 1GB));
-}
 $node_primary->start;
 my $backup_name = 'my_backup';
 
@@ -72,9 +64,9 @@ $node_replica->start;
 
 # Create ivfflat index on primary
 $node_primary->safe_psql("postgres", "CREATE EXTENSION vector;");
-$node_primary->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
+$node_primary->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
 $node_primary->safe_psql("postgres",
-	"INSERT INTO tst SELECT i % 10, (SELECT array_agg(random()) FROM generate_series(1, $dim)) FROM generate_series(1, 100000) i;"
+	"INSERT INTO tst SELECT i % 10, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
 );
 $node_primary->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
 
@@ -90,7 +82,7 @@ for my $i (1 .. 10)
 	test_index_replay("vacuum $i");
 	my ($start, $end) = (100001 + ($i - 1) * 10000, 100000 + $i * 10000);
 	$node_primary->safe_psql("postgres",
-		"INSERT INTO tst SELECT i % 10, (SELECT array_agg(random()) FROM generate_series(1, $dim)) FROM generate_series($start, $end) i;"
+		"INSERT INTO tst SELECT i % 10, ARRAY[random(), random(), random()] FROM generate_series($start, $end) i;"
 	);
 	test_index_replay("insert $i");
 }

test/t/007_inserts.pl

@@ -1,33 +0,0 @@
-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 and index
-$node->safe_psql("postgres", "CREATE EXTENSION vector;");
-$node->safe_psql("postgres", "CREATE TABLE tst (v vector(768));");
-$node->safe_psql("postgres",
-	"INSERT INTO tst SELECT (SELECT array_agg(random()) FROM generate_series(1, 768)) FROM generate_series(1, 10000) i;"
-);
-$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
-
-$node->pgbench(
-	"--no-vacuum --client=5 --transactions=100",
-	0,
-	[qr{actually processed}],
-	[qr{^$}],
-	"concurrent INSERTs",
-	{
-		"007_concurrent" => q(
-			BEGIN;
-			INSERT INTO tst SELECT (SELECT array_agg(random()) FROM generate_series(1, 768)) FROM generate_series(1, 10) i;
-			COMMIT;
-		),
-	}
-);

vector.control

@@ -1,4 +1,4 @@
 comment = 'vector data type and ivfflat access method'
-default_version = '0.2.7'
+default_version = '0.2.5'
 module_pathname = '$libdir/vector'
 relocatable = true