Improved test for angular distance [skip ci]

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.5.1 (unreleased)
+
+- Added `angular_distance` function
+
 ## 0.5.0 (2023-08-28)
 
 - Added HNSW index type

sql/vector--0.5.0--0.5.1.sql

@@ -0,0 +1,5 @@
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "ALTER EXTENSION vector UPDATE TO '0.5.1'" to load this file. \quit
+
+CREATE FUNCTION angular_distance(vector, vector) RETURNS float8
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

sql/vector.sql

@@ -87,6 +87,9 @@ CREATE FUNCTION vector_l2_squared_distance(vector, vector) RETURNS float8
 CREATE FUNCTION vector_negative_inner_product(vector, vector) RETURNS float8
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
+CREATE FUNCTION angular_distance(vector, vector) RETURNS float8
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
 CREATE FUNCTION vector_spherical_distance(vector, vector) RETURNS float8
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 

src/hnsw.h

@@ -4,7 +4,6 @@
 #include "postgres.h"
 
 #include "access/generic_xlog.h"
-#include "access/parallel.h"
 #include "access/reloptions.h"
 #include "nodes/execnodes.h"
 #include "port.h"				/* for random() */
@@ -15,10 +14,6 @@
 #error "Requires PostgreSQL 11+"
 #endif
 
-#if PG_VERSION_NUM < 120000
-#include "access/relscan.h"
-#endif
-
 #define HNSW_MAX_DIM 2000
 
 /* Support functions */
@@ -137,49 +132,6 @@ typedef struct HnswOptions
 	int			efConstruction; /* size of dynamic candidate list */
 }			HnswOptions;
 
-typedef struct HnswSpool
-{
-	Relation	heap;
-	Relation	index;
-}			HnswSpool;
-
-typedef struct HnswShared
-{
-	/* Immutable state */
-	Oid			heaprelid;
-	Oid			indexrelid;
-	bool		isconcurrent;
-	int			scantuplesortstates;
-
-	/* Worker progress */
-	ConditionVariable workersdonecv;
-
-	/* Mutex for mutable state */
-	slock_t		mutex;
-
-	/* Mutable state */
-	int			nparticipantsdone;
-	double		reltuples;
-	double		indtuples;
-
-#if PG_VERSION_NUM < 120000
-	ParallelHeapScanDescData heapdesc;	/* must come last */
-#endif
-}			HnswShared;
-
-#if PG_VERSION_NUM >= 120000
-#define ParallelTableScanFromHnswShared(shared) \
-	(ParallelTableScanDesc) ((char *) (shared) + BUFFERALIGN(sizeof(HnswShared)))
-#endif
-
-typedef struct HnswLeader
-{
-	ParallelContext *pcxt;
-	int			nparticipanttuplesorts;
-	HnswShared *hnswshared;
-	Snapshot	snapshot;
-}			HnswLeader;
-
 typedef struct HnswBuildState
 {
 	/* Info */
@@ -213,9 +165,6 @@ typedef struct HnswBuildState
 
 	/* Memory */
 	MemoryContext tmpCtx;
-
-	/* Parallel builds */
-	HnswLeader *hnswleader;
 }			HnswBuildState;
 
 typedef struct HnswMetaPageData
@@ -336,7 +285,6 @@ void		HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation i
 void		HnswSetElementTuple(HnswElementTuple etup, HnswElement element);
 void		HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
 void		HnswLoadNeighbors(HnswElement element, Relation index);
-PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
 
 /* Index access methods */
 IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);

src/hnswbuild.c

@@ -2,15 +2,12 @@
 
 #include <math.h>
 
-#include "access/parallel.h"
-#include "access/xact.h"
 #include "catalog/index.h"
 #include "hnsw.h"
 #include "miscadmin.h"
 #include "lib/pairingheap.h"
 #include "nodes/pg_list.h"
 #include "storage/bufmgr.h"
-#include "tcop/tcopprot.h"
 #include "utils/memutils.h"
 
 #if PG_VERSION_NUM >= 140000
@@ -38,23 +35,6 @@
 #define UpdateProgress(index, val) ((void)val)
 #endif
 
-#if PG_VERSION_NUM >= 140000
-#include "utils/backend_status.h"
-#include "utils/wait_event.h"
-#endif
-
-#if PG_VERSION_NUM >= 120000
-#include "access/table.h"
-#include "optimizer/optimizer.h"
-#else
-#include "access/heapam.h"
-#include "optimizer/planner.h"
-#include "pgstat.h"
-#endif
-
-#define PARALLEL_KEY_HNSW_SHARED		UINT64CONST(0xA000000000000001)
-#define PARALLEL_KEY_QUERY_TEXT			UINT64CONST(0xA000000000000002)
-
 /*
  * Create the metapage
  */
@@ -371,7 +351,6 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
 
 		oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
 
-		/* TODO Fix progress for parallel builds */
 		if (HnswInsertTuple(buildstate->index, values, isnull, tid, buildstate->heap))
 			UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
 
@@ -469,8 +448,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
 	buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
 											   "Hnsw build temporary context",
 											   ALLOCSET_DEFAULT_SIZES);
-
-	buildstate->hnswleader = NULL;
 }
 
 /*
@@ -483,373 +460,21 @@ FreeBuildState(HnswBuildState * buildstate)
 	MemoryContextDelete(buildstate->tmpCtx);
 }
 
-/*
- * Within leader, wait for end of heap scan
- */
-static double
-ParallelHeapScan(HnswBuildState * buildstate)
-{
-	HnswShared *hnswshared = buildstate->hnswleader->hnswshared;
-	int			nparticipanttuplesorts;
-	double		reltuples;
-
-	nparticipanttuplesorts = buildstate->hnswleader->nparticipanttuplesorts;
-	for (;;)
-	{
-		SpinLockAcquire(&hnswshared->mutex);
-		if (hnswshared->nparticipantsdone == nparticipanttuplesorts)
-		{
-			buildstate->indtuples = hnswshared->indtuples;
-			reltuples = hnswshared->reltuples;
-			SpinLockRelease(&hnswshared->mutex);
-			break;
-		}
-		SpinLockRelease(&hnswshared->mutex);
-
-		ConditionVariableSleep(&hnswshared->workersdonecv,
-							   WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN);
-	}
-
-	ConditionVariableCancelSleep();
-
-	return reltuples;
-}
-
-/*
- * Perform a worker's portion of a parallel insert
- */
-static void
-HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool progress)
-{
-	HnswBuildState buildstate;
-#if PG_VERSION_NUM >= 120000
-	TableScanDesc scan;
-#else
-	HeapScanDesc scan;
-#endif
-	double		reltuples;
-	IndexInfo  *indexInfo;
-
-	/* Join parallel scan */
-	indexInfo = BuildIndexInfo(hnswspool->index);
-	indexInfo->ii_Concurrent = hnswshared->isconcurrent;
-	InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM);
-	/* TODO Support in-memory builds */
-	buildstate.maxInMemoryElements = 0;
-	buildstate.flushed = true;
-#if PG_VERSION_NUM >= 120000
-	scan = table_beginscan_parallel(hnswspool->heap,
-									ParallelTableScanFromHnswShared(hnswshared));
-	reltuples = table_index_build_scan(hnswspool->heap, hnswspool->index, indexInfo,
-									   true, progress, BuildCallback,
-									   (void *) &buildstate, scan);
-#else
-	scan = heap_beginscan_parallel(hnswspool->heap, &hnswshared->heapdesc);
-	reltuples = IndexBuildHeapScan(hnswspool->heap, hnswspool->index, indexInfo,
-								   true, BuildCallback,
-								   (void *) &buildstate, scan);
-#endif
-
-	/* Record statistics */
-	SpinLockAcquire(&hnswshared->mutex);
-	hnswshared->nparticipantsdone++;
-	hnswshared->reltuples += reltuples;
-	hnswshared->indtuples += buildstate.indtuples;
-	SpinLockRelease(&hnswshared->mutex);
-
-	/* Log statistics */
-	if (progress)
-		ereport(DEBUG1, (errmsg("leader processed " INT64_FORMAT " tuples", (int64) reltuples)));
-	else
-		ereport(DEBUG1, (errmsg("worker processed " INT64_FORMAT " tuples", (int64) reltuples)));
-
-	/* Notify leader */
-	ConditionVariableSignal(&hnswshared->workersdonecv);
-
-	FreeBuildState(&buildstate);
-}
-
-/*
- * Perform work within a launched parallel process
- */
-void
-HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
-{
-	char	   *sharedquery;
-	HnswSpool  *hnswspool;
-	HnswShared *hnswshared;
-	Relation	heapRel;
-	Relation	indexRel;
-	LOCKMODE	heapLockmode;
-	LOCKMODE	indexLockmode;
-
-	/* Set debug_query_string for individual workers first */
-	sharedquery = shm_toc_lookup(toc, PARALLEL_KEY_QUERY_TEXT, true);
-	debug_query_string = sharedquery;
-
-	/* Report the query string from leader */
-	pgstat_report_activity(STATE_RUNNING, debug_query_string);
-
-	/* Look up shared state */
-	hnswshared = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_SHARED, false);
-
-	/* Open relations using lock modes known to be obtained by index.c */
-	if (!hnswshared->isconcurrent)
-	{
-		heapLockmode = ShareLock;
-		indexLockmode = AccessExclusiveLock;
-	}
-	else
-	{
-		heapLockmode = ShareUpdateExclusiveLock;
-		indexLockmode = RowExclusiveLock;
-	}
-
-	/* Open relations within worker */
-#if PG_VERSION_NUM >= 120000
-	heapRel = table_open(hnswshared->heaprelid, heapLockmode);
-#else
-	heapRel = heap_open(hnswshared->heaprelid, heapLockmode);
-#endif
-	indexRel = index_open(hnswshared->indexrelid, indexLockmode);
-
-	/* Initialize worker's own spool */
-	hnswspool = (HnswSpool *) palloc0(sizeof(HnswSpool));
-	hnswspool->heap = heapRel;
-	hnswspool->index = indexRel;
-
-	/* Perform inserts */
-	HnswParallelScanAndInsert(hnswspool, hnswshared, false);
-
-	/* Close relations within worker */
-	index_close(indexRel, indexLockmode);
-#if PG_VERSION_NUM >= 120000
-	table_close(heapRel, heapLockmode);
-#else
-	heap_close(heapRel, heapLockmode);
-#endif
-}
-
-/*
- * End parallel build
- */
-static void
-HnswEndParallel(HnswLeader * hnswleader)
-{
-	/* Shutdown worker processes */
-	WaitForParallelWorkersToFinish(hnswleader->pcxt);
-
-	/* Free last reference to MVCC snapshot, if one was used */
-	if (IsMVCCSnapshot(hnswleader->snapshot))
-		UnregisterSnapshot(hnswleader->snapshot);
-	DestroyParallelContext(hnswleader->pcxt);
-	ExitParallelMode();
-}
-
-/*
- * Return size of shared memory required for parallel index build
- */
-static Size
-ParallelEstimateShared(Relation heap, Snapshot snapshot)
-{
-#if PG_VERSION_NUM >= 120000
-	return add_size(BUFFERALIGN(sizeof(HnswShared)), table_parallelscan_estimate(heap, snapshot));
-#else
-	if (!IsMVCCSnapshot(snapshot))
-	{
-		Assert(snapshot == SnapshotAny);
-		return sizeof(HnswShared);
-	}
-
-	return add_size(offsetof(HnswShared, heapdesc) +
-					offsetof(ParallelHeapScanDescData, phs_snapshot_data),
-					EstimateSnapshotSpace(snapshot));
-#endif
-}
-
-/*
- * Within leader, participate as a parallel worker
- */
-static void
-HnswLeaderParticipateAsWorker(HnswBuildState * buildstate)
-{
-	HnswLeader *hnswleader = buildstate->hnswleader;
-	HnswSpool  *leaderworker;
-
-	/* Allocate memory and initialize private spool */
-	leaderworker = (HnswSpool *) palloc0(sizeof(HnswSpool));
-	leaderworker->heap = buildstate->heap;
-	leaderworker->index = buildstate->index;
-
-	/* Perform work common to all participants */
-	HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, true);
-}
-
-/*
- * Begin parallel build
- */
-static void
-HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
-{
-	ParallelContext *pcxt;
-	int			scantuplesortstates;
-	Snapshot	snapshot;
-	Size		esthnswshared;
-	HnswShared *hnswshared;
-	HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader));
-	bool		leaderparticipates = true;
-	int			querylen;
-
-#ifdef DISABLE_LEADER_PARTICIPATION
-	leaderparticipates = false;
-#endif
-
-	/* Enter parallel mode and create context */
-	EnterParallelMode();
-	Assert(request > 0);
-#if PG_VERSION_NUM >= 120000
-	pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request);
-#else
-	pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request, true);
-#endif
-
-	scantuplesortstates = leaderparticipates ? request + 1 : request;
-
-	/* Get snapshot for table scan */
-	if (!isconcurrent)
-		snapshot = SnapshotAny;
-	else
-		snapshot = RegisterSnapshot(GetTransactionSnapshot());
-
-	/* Estimate size of workspaces */
-	esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot);
-	shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared);
-	shm_toc_estimate_keys(&pcxt->estimator, 1);
-
-	/* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */
-	if (debug_query_string)
-	{
-		querylen = strlen(debug_query_string);
-		shm_toc_estimate_chunk(&pcxt->estimator, querylen + 1);
-		shm_toc_estimate_keys(&pcxt->estimator, 1);
-	}
-	else
-		querylen = 0;			/* keep compiler quiet */
-
-	/* Everyone's had a chance to ask for space, so now create the DSM */
-	InitializeParallelDSM(pcxt);
-
-	/* If no DSM segment was available, back out (do serial build) */
-	if (pcxt->seg == NULL)
-	{
-		if (IsMVCCSnapshot(snapshot))
-			UnregisterSnapshot(snapshot);
-		DestroyParallelContext(pcxt);
-		ExitParallelMode();
-		return;
-	}
-
-	/* Store shared build state, for which we reserved space */
-	hnswshared = (HnswShared *) shm_toc_allocate(pcxt->toc, esthnswshared);
-	/* Initialize immutable state */
-	hnswshared->heaprelid = RelationGetRelid(buildstate->heap);
-	hnswshared->indexrelid = RelationGetRelid(buildstate->index);
-	hnswshared->isconcurrent = isconcurrent;
-	hnswshared->scantuplesortstates = scantuplesortstates;
-	ConditionVariableInit(&hnswshared->workersdonecv);
-	SpinLockInit(&hnswshared->mutex);
-	/* Initialize mutable state */
-	hnswshared->nparticipantsdone = 0;
-	hnswshared->reltuples = 0;
-	hnswshared->indtuples = 0;
-#if PG_VERSION_NUM >= 120000
-	table_parallelscan_initialize(buildstate->heap,
-								  ParallelTableScanFromHnswShared(hnswshared),
-								  snapshot);
-#else
-	heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot);
-#endif
-
-	shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_SHARED, hnswshared);
-
-	/* Store query string for workers */
-	if (debug_query_string)
-	{
-		char	   *sharedquery;
-
-		sharedquery = (char *) shm_toc_allocate(pcxt->toc, querylen + 1);
-		memcpy(sharedquery, debug_query_string, querylen + 1);
-		shm_toc_insert(pcxt->toc, PARALLEL_KEY_QUERY_TEXT, sharedquery);
-	}
-
-	/* Launch workers, saving status for leader/caller */
-	LaunchParallelWorkers(pcxt);
-	hnswleader->pcxt = pcxt;
-	hnswleader->nparticipanttuplesorts = pcxt->nworkers_launched;
-	if (leaderparticipates)
-		hnswleader->nparticipanttuplesorts++;
-	hnswleader->hnswshared = hnswshared;
-	hnswleader->snapshot = snapshot;
-
-	/* If no workers were successfully launched, back out (do serial build) */
-	if (pcxt->nworkers_launched == 0)
-	{
-		HnswEndParallel(hnswleader);
-		return;
-	}
-
-	/* Log participants */
-	ereport(DEBUG1, (errmsg("using %d parallel workers", pcxt->nworkers_launched)));
-
-	/* Save leader state now that it's clear build will be parallel */
-	buildstate->hnswleader = hnswleader;
-
-	/* Join heap scan ourselves */
-	if (leaderparticipates)
-		HnswLeaderParticipateAsWorker(buildstate);
-
-	/* Wait for all launched workers */
-	WaitForParallelWorkersToAttach(pcxt);
-}
-
 /*
  * Build graph
  */
 static void
 BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
 {
-	int			parallel_workers = 0;
-
 	UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD);
 
-	/* Calculate parallel workers */
-	parallel_workers = plan_create_index_workers(RelationGetRelid(buildstate->heap), RelationGetRelid(buildstate->index));
-
-	/* Attempt to launch parallel worker scan when required */
-	if (parallel_workers > 0)
-	{
-		/* TODO Support in-memory builds */
-		FlushPages(buildstate);
-		HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
-	}
-
-	/* Add tuples to sort */
-	if (buildstate->hnswleader)
-		buildstate->reltuples = ParallelHeapScan(buildstate);
-	else
-	{
 #if PG_VERSION_NUM >= 120000
-		buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
-													   true, true, BuildCallback, (void *) buildstate, NULL);
+	buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
+												   true, true, BuildCallback, (void *) buildstate, NULL);
 #else
-		buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
-												   true, BuildCallback, (void *) buildstate, NULL);
+	buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
+											   true, BuildCallback, (void *) buildstate, NULL);
 #endif
-	}
-
-	/* End parallel build */
-	if (buildstate->hnswleader)
-		HnswEndParallel(buildstate->hnswleader);
 }
 
 /*

src/ivfbuild.c

@@ -10,8 +10,8 @@
 #include "ivfflat.h"
 #include "miscadmin.h"
 #include "storage/bufmgr.h"
-#include "utils/memutils.h"
 #include "tcop/tcopprot.h"
+#include "utils/memutils.h"
 
 #if PG_VERSION_NUM >= 140000
 #include "utils/backend_progress.h"

src/vector.c

@@ -684,6 +684,49 @@ cosine_distance(PG_FUNCTION_ARGS)
 	PG_RETURN_FLOAT8(1.0 - similarity);
 }
 
+/*
+ * Get the angular distance between two vectors
+ */
+PGDLLEXPORT PG_FUNCTION_INFO_V1(angular_distance);
+Datum
+angular_distance(PG_FUNCTION_ARGS)
+{
+	Vector	   *a = PG_GETARG_VECTOR_P(0);
+	Vector	   *b = PG_GETARG_VECTOR_P(1);
+	float	   *ax = a->x;
+	float	   *bx = b->x;
+	float		distance = 0.0;
+	float		norma = 0.0;
+	float		normb = 0.0;
+	double		similarity;
+
+	CheckDims(a, b);
+
+	/* Auto-vectorized */
+	for (int i = 0; i < a->dim; i++)
+	{
+		distance += ax[i] * bx[i];
+		norma += ax[i] * ax[i];
+		normb += bx[i] * bx[i];
+	}
+
+	similarity = (double) distance / sqrt((double) norma * (double) normb);
+
+#ifdef _MSC_VER
+	/* /fp:fast may not propagate NaN */
+	if (isnan(similarity))
+		PG_RETURN_FLOAT8(NAN);
+#endif
+
+	/* Prevent NaN with acos with loss of precision */
+	if (similarity > 1)
+		similarity = 1;
+	else if (similarity < -1)
+		similarity = -1;
+
+	PG_RETURN_FLOAT8(acos(similarity) / M_PI);
+}
+
 /*
  * Get the distance for spherical k-means
  * Currently uses angular distance since needs to satisfy triangle inequality
@@ -695,6 +738,8 @@ vector_spherical_distance(PG_FUNCTION_ARGS)
 {
 	Vector	   *a = PG_GETARG_VECTOR_P(0);
 	Vector	   *b = PG_GETARG_VECTOR_P(1);
+	float	   *ax = a->x;
+	float	   *bx = b->x;
 	float		dp = 0.0;
 	double		distance;
 
@@ -702,7 +747,7 @@ vector_spherical_distance(PG_FUNCTION_ARGS)
 
 	/* Auto-vectorized */
 	for (int i = 0; i < a->dim; i++)
-		dp += a->x[i] * b->x[i];
+		dp += ax[i] * bx[i];
 
 	distance = (double) dp;
 

test/expected/functions.out

@@ -152,6 +152,56 @@ SELECT l1_distance('[3e38]', '[-3e38]');
     Infinity
 (1 row)
 
+SELECT angular_distance('[1,2]', '[2,4]');
+ angular_distance 
+------------------
+                0
+(1 row)
+
+SELECT angular_distance('[1,2]', '[0,0]');
+ angular_distance 
+------------------
+              NaN
+(1 row)
+
+SELECT angular_distance('[1,1]', '[1,1]');
+ angular_distance 
+------------------
+                0
+(1 row)
+
+SELECT angular_distance('[1,0]', '[0,2]');
+ angular_distance 
+------------------
+              0.5
+(1 row)
+
+SELECT angular_distance('[1,1]', '[-1,-1]');
+ angular_distance 
+------------------
+                1
+(1 row)
+
+SELECT angular_distance('[1,2]', '[3]');
+ERROR:  different vector dimensions 2 and 1
+SELECT angular_distance('[1,1]', '[1.1,1.1]');
+ angular_distance 
+------------------
+                0
+(1 row)
+
+SELECT angular_distance('[1,1]', '[-1.1,-1.1]');
+ angular_distance 
+------------------
+                1
+(1 row)
+
+SELECT angular_distance('[3e38]', '[3e38]');
+ angular_distance 
+------------------
+              NaN
+(1 row)
+
 SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
     avg    
 -----------

test/sql/functions.sql

@@ -36,6 +36,16 @@ SELECT l1_distance('[0,0]', '[0,1]');
 SELECT l1_distance('[1,2]', '[3]');
 SELECT l1_distance('[3e38]', '[-3e38]');
 
+SELECT angular_distance('[1,2]', '[2,4]');
+SELECT angular_distance('[1,2]', '[0,0]');
+SELECT angular_distance('[1,1]', '[1,1]');
+SELECT angular_distance('[1,0]', '[0,2]');
+SELECT angular_distance('[1,1]', '[-1,-1]');
+SELECT angular_distance('[1,2]', '[3]');
+SELECT angular_distance('[1,1]', '[1.1,1.1]');
+SELECT angular_distance('[1,1]', '[-1.1,-1.1]');
+SELECT angular_distance('[3e38]', '[3e38]');
+
 SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
 SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v;
 SELECT avg(v) FROM unnest(ARRAY[]::vector[]) v;

test/t/001_ivfflat_wal.pl

@@ -19,8 +19,6 @@ sub test_index_replay
 
 	# Wait for replica to catch up
 	my $applname = $node_replica->name;
-
-	my $server_version_num = $node_primary->safe_psql("postgres", "SHOW server_version_num");
 	my $caughtup_query = "SELECT pg_current_wal_lsn() <= replay_lsn FROM pg_stat_replication WHERE application_name = '$applname';";
 	$node_primary->poll_query_until('postgres', $caughtup_query)
 	  or die "Timed out while waiting for replica 1 to catch up";

test/t/003_ivfflat_build_recall.pl

@@ -94,7 +94,7 @@ for my $i (0 .. $#operators)
 	# Test approximate results
 	if ($operator ne "<#>")
 	{
-		# TODO fix test
+		# TODO Fix test (uniform random vectors all have similar inner product)
 		test_recall(1, 0.71, $operator);
 		test_recall(10, 0.95, $operator);
 	}
@@ -115,7 +115,7 @@ for my $i (0 .. $#operators)
 	# Test approximate results
 	if ($operator ne "<#>")
 	{
-		# TODO fix test
+		# TODO Fix test (uniform random vectors all have similar inner product)
 		test_recall(1, 0.71, $operator);
 		test_recall(10, 0.95, $operator);
 	}

test/t/010_hnsw_wal.pl

@@ -19,8 +19,6 @@ sub test_index_replay
 
 	# Wait for replica to catch up
 	my $applname = $node_replica->name;
-
-	my $server_version_num = $node_primary->safe_psql("postgres", "SHOW server_version_num");
 	my $caughtup_query = "SELECT pg_current_wal_lsn() <= replay_lsn FROM pg_stat_replication WHERE application_name = '$applname';";
 	$node_primary->poll_query_until('postgres', $caughtup_query)
 	  or die "Timed out while waiting for replica 1 to catch up";

test/t/012_hnsw_build_recall.pl

@@ -83,31 +83,11 @@ for my $i (0 .. $#operators)
 		push(@expected, $res);
 	}
 
-	# Build index serially
-	$node->safe_psql("postgres", qq(
-		SET max_parallel_maintenance_workers = 0;
-		CREATE INDEX idx ON tst USING hnsw (v $opclass);
-	));
+	# Add index
+	$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v $opclass);");
 
-	# Test approximate results
 	my $min = $operator eq "<#>" ? 0.80 : 0.99;
 	test_recall($min, $operator);
-
-	$node->safe_psql("postgres", "DROP INDEX idx;");
-
-	# Build index in parallel
-	my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
-		SET client_min_messages = DEBUG;
-		SET min_parallel_table_scan_size = 1;
-		CREATE INDEX idx ON tst USING hnsw (v $opclass);
-	));
-	is($ret, 0, $stderr);
-	like($stderr, qr/using \d+ parallel workers/);
-
-	# Test approximate results
-	test_recall($min, $operator);
-
-	$node->safe_psql("postgres", "DROP INDEX idx;");
 }
 
 done_testing();

test/t/016_hnsw_vacuum_recall.pl

@@ -89,7 +89,7 @@ foreach (@queries)
 test_recall(0.20, $limit, "before vacuum");
 test_recall(0.95, 100, "before vacuum");
 
-# TODO test concurrent inserts with vacuum
+# TODO Test concurrent inserts with vacuum
 $node->safe_psql("postgres", "VACUUM tst;");
 
 test_recall(0.95, $limit, "after vacuum");

test/t/017_ivfflat_insert_recall.pl

@@ -0,0 +1,117 @@
+use strict;
+use warnings;
+use PostgresNode;
+use TestLib;
+use Test::More;
+
+my $node;
+my @queries = ();
+my @expected;
+my $limit = 20;
+
+sub test_recall
+{
+	my ($probes, $min, $operator) = @_;
+	my $correct = 0;
+	my $total = 0;
+
+	my $explain = $node->safe_psql("postgres", qq(
+		SET enable_seqscan = off;
+		SET ivfflat.probes = $probes;
+		EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v $operator '$queries[0]' LIMIT $limit;
+	));
+	like($explain, qr/Index Scan using idx on tst/);
+
+	for my $i (0 .. $#queries)
+	{
+		my $actual = $node->safe_psql("postgres", qq(
+			SET enable_seqscan = off;
+			SET ivfflat.probes = $probes;
+			SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
+		));
+		my @actual_ids = split("\n", $actual);
+		my %actual_set = map { $_ => 1 } @actual_ids;
+
+		my @expected_ids = split("\n", $expected[$i]);
+
+		foreach (@expected_ids)
+		{
+			if (exists($actual_set{$_}))
+			{
+				$correct++;
+			}
+			$total++;
+		}
+	}
+
+	cmp_ok($correct / $total, ">=", $min, $operator);
+}
+
+# Initialize node
+$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 serial, v vector(3));");
+
+# Generate queries
+for (1 .. 20)
+{
+	my $r1 = rand();
+	my $r2 = rand();
+	my $r3 = rand();
+	push(@queries, "[$r1,$r2,$r3]");
+}
+
+# Check each index type
+my @operators = ("<->", "<#>", "<=>");
+my @opclasses = ("vector_l2_ops", "vector_ip_ops", "vector_cosine_ops");
+
+for my $i (0 .. $#operators)
+{
+	my $operator = $operators[$i];
+	my $opclass = $opclasses[$i];
+
+	# Add index
+	$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v $opclass);");
+
+	# Use concurrent inserts
+	$node->pgbench(
+		"--no-vacuum --client=10 --transactions=1000",
+		0,
+		[qr{actually processed}],
+		[qr{^$}],
+		"concurrent INSERTs",
+		{
+			"017_ivfflat_insert_recall_$opclass" => "INSERT INTO tst (v) SELECT ARRAY[random(), random(), random()] FROM generate_series(1, 10) i;"
+		}
+	);
+
+	# Get exact results
+	@expected = ();
+	foreach (@queries)
+	{
+		my $res = $node->safe_psql("postgres", qq(
+			SET enable_indexscan = off;
+			SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;
+		));
+		push(@expected, $res);
+	}
+
+	# Test approximate results
+	if ($operator ne "<#>")
+	{
+		# TODO Fix test (uniform random vectors all have similar inner product)
+		test_recall(1, 0.71, $operator);
+		test_recall(10, 0.95, $operator);
+	}
+	# Account for equal distances
+	test_recall(100, 0.9925, $operator);
+
+	$node->safe_psql("postgres", "DROP INDEX idx;");
+	$node->safe_psql("postgres", "TRUNCATE tst;");
+}
+
+done_testing();