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

@@ -219,9 +219,6 @@ typedef struct HnswScanOpaqueData
 {
 	bool		first;
 	Buffer		buf;
-	ItemPointerData heaptid;
-	OffsetNumber offno;
-	int			removedCount;
 	List	   *w;
 	MemoryContext tmpCtx;
 

src/hnswscan.c

@@ -58,75 +58,6 @@ GetDimensions(Relation index)
 	return dimensions;
 }
 
-/*
- * Remove deleted heap TID
- */
-static void
-RemoveHeapTid(IndexScanDesc scan)
-{
-	HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
-	Relation	index = scan->indexRelation;
-	Buffer		buf = so->buf;
-	Page		page;
-	GenericXLogState *state;
-	ItemId		itemid;
-	HnswElementTuple etup;
-	Size		etupSize;
-	int			idx = -1;
-
-	/* Safety check */
-	if (!BufferIsValid(buf) || !OffsetNumberIsValid(so->offno) || !ItemPointerIsValid(&so->heaptid))
-		return;
-
-	/* Use WAL rather than hint */
-	LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
-	state = GenericXLogStart(index);
-	page = GenericXLogRegisterBuffer(state, buf, 0);
-	itemid = PageGetItemId(page, so->offno);
-	etup = (HnswElementTuple) PageGetItem(page, itemid);
-	etupSize = ItemIdGetLength(itemid);
-
-	Assert(HnswIsElementTuple(etup));
-
-	/* Find index */
-	for (int i = 0; i < HNSW_HEAPTIDS; i++)
-	{
-		if (!ItemPointerIsValid(&etup->heaptids[i]))
-			break;
-
-		if (ItemPointerEquals(&etup->heaptids[i], &so->heaptid))
-		{
-			idx = i;
-			break;
-		}
-	}
-
-	if (idx == -1)
-		GenericXLogAbort(state);
-	else
-	{
-		/* Move pointers forward */
-		for (int i = idx; i < HNSW_HEAPTIDS; i++)
-		{
-			if (i + 1 == HNSW_HEAPTIDS || !ItemPointerIsValid(&etup->heaptids[i + 1]))
-				ItemPointerSetInvalid(&etup->heaptids[i]);
-			else
-				ItemPointerCopy(&etup->heaptids[i + 1], &etup->heaptids[i]);
-		}
-
-		/* Overwrite tuple */
-		if (!PageIndexTupleOverwrite(page, so->offno, (Item) etup, etupSize))
-			elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
-
-		/* Commit */
-		MarkBufferDirty(buf);
-		GenericXLogFinish(state);
-	}
-
-	/* Unlock buffer */
-	LockBuffer(buf, BUFFER_LOCK_UNLOCK);
-}
-
 /*
  * Prepare for an index scan
  */
@@ -140,9 +71,6 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 
 	so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
 	so->buf = InvalidBuffer;
-	ItemPointerSetInvalid(&so->heaptid);
-	so->offno = InvalidOffsetNumber;
-	so->removedCount = 0;
 	so->first = true;
 	so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
 									   "Hnsw scan temporary context",
@@ -167,7 +95,6 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
 	HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
 
 	so->first = true;
-	ItemPointerSetInvalid(&so->heaptid);
 	MemoryContextReset(so->tmpCtx);
 
 	if (keys && scan->numberOfKeys > 0)
@@ -231,25 +158,12 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 
 		so->first = false;
 	}
-	else
-	{
-		/*
-		 * Remove dead tuples. kill_prior_tuple will only be true if not in
-		 * recovery. Limit the number removed per scan for performance.
-		 */
-		if (scan->kill_prior_tuple && so->removedCount < 3)
-		{
-			RemoveHeapTid(scan);
-			so->removedCount++;
-		}
-	}
 
 	while (list_length(so->w) > 0)
 	{
 		HnswCandidate *hc = llast(so->w);
 		ItemPointer tid;
 		BlockNumber indexblkno;
-		OffsetNumber indexoffno;
 
 		/* Move to next element if no valid heap tids */
 		if (list_length(hc->element->heaptids) == 0)
@@ -260,7 +174,6 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 
 		tid = llast(hc->element->heaptids);
 		indexblkno = hc->element->blkno;
-		indexoffno = hc->element->offno;
 
 		hc->element->heaptids = list_delete_last(hc->element->heaptids);
 
@@ -272,10 +185,6 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		scan->xs_ctup.t_self = *tid;
 #endif
 
-		/* Keep track of info needed to remove dead tuples */
-		so->heaptid = *tid;
-		so->offno = indexoffno;
-
 		/* Unpin buffer */
 		if (BufferIsValid(so->buf))
 			ReleaseBuffer(so->buf);

src/hnswutils.c

@@ -551,8 +551,6 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
 	pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
 	pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
 	int			wlen = 0;
-	uint64		dead = 0;
-	uint64		maxAdditional = skipElement == NULL ? ef : PG_UINT64_MAX;
 	HASHCTL		hash_ctl;
 	HTAB	   *v;
 
@@ -581,14 +579,13 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
 		pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node));
 		pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node));
 
-		/* Do not count certain number of dead elements towards ef */
-		if (list_length(hc->element->heaptids) == 0)
-		{
-			if ((++dead) <= maxAdditional)
-				continue;
-		}
-
-		wlen++;
+		/*
+		 * Do not count elements being deleted towards ef when vacuuming. It
+		 * would be ideal to do this for inserts as well, but this could
+		 * affect insert performance.
+		 */
+		if (skipElement == NULL || list_length(hc->element->heaptids) != 0)
+			wlen++;
 	}
 
 	while (!pairingheap_is_empty(C))
@@ -641,18 +638,19 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
 					pairingheap_add(C, &(CreatePairingHeapNode(ec)->ph_node));
 					pairingheap_add(W, &(CreatePairingHeapNode(ec)->ph_node));
 
-					/* Do not count certain number of dead elements towards ef */
-					if (list_length(e->element->heaptids) == 0)
+					/*
+					 * Do not count elements being deleted towards ef when
+					 * vacuuming. It would be ideal to do this for inserts as
+					 * well, but this could affect insert performance.
+					 */
+					if (skipElement == NULL || list_length(e->element->heaptids) != 0)
 					{
-						if ((++dead) <= maxAdditional)
-							continue;
+						wlen++;
+
+						/* No need to decrement wlen */
+						if (wlen > ef)
+							pairingheap_remove_first(W);
 					}
-
-					wlen++;
-
-					/* No need to decrement wlen */
-					if (wlen > ef)
-						pairingheap_remove_first(W);
 				}
 			}
 		}

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/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";
@@ -38,9 +36,8 @@ sub test_index_replay
 	);
 
 	# Run test queries and compare their result
-	# Query replica first since index scan on primary can generate WAL removing tuples
-	my $replica_result = $node_replica->safe_psql("postgres", $queries);
 	my $primary_result = $node_primary->safe_psql("postgres", $queries);
+	my $replica_result = $node_replica->safe_psql("postgres", $queries);
 
 	is($primary_result, $replica_result, "$test_name: query result matches");
 	return;

test/t/015_hnsw_duplicates.pl

@@ -23,27 +23,25 @@ sub insert_vectors
 
 sub test_duplicates
 {
-	my ($exp) = @_;
-
 	my $res = $node->safe_psql("postgres", qq(
 		SET enable_seqscan = off;
 		SET hnsw.ef_search = 1;
 		SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '[1,1,1]') t;
 	));
-	is($res, $exp);
+	is($res, 10);
 }
 
 # Test duplicates with build
 insert_vectors();
 $node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
-test_duplicates(10);
+test_duplicates();
 
 # Reset
 $node->safe_psql("postgres", "TRUNCATE tst;");
 
 # Test duplicates with inserts
 insert_vectors();
-test_duplicates(10);
+test_duplicates();
 
 # Test fallback path for inserts
 $node->pgbench(
@@ -57,15 +55,4 @@ $node->pgbench(
 	}
 );
 
-# Reset
-$node->safe_psql("postgres", "TRUNCATE tst;");
-
-# Test deletes with index scan
-$node->safe_psql("postgres", "INSERT INTO tst SELECT '[1,1,1]' FROM generate_series(1, 10) i;");
-$node->safe_psql("postgres", "DELETE FROM tst WHERE ctid IN (SELECT ctid FROM tst ORDER BY random() LIMIT 5);");
-for (1 .. 3)
-{
-	test_duplicates(5);
-}
-
 done_testing();