Improved test for angular distance [skip ci]

CHANGELOG.md

@@ -1,13 +1,17 @@
+## 0.5.1 (unreleased)
+
+- Added `angular_distance` function
+
 ## 0.5.0 (2023-08-28)
 
 - Added HNSW index type
-- Added support for parallel index builds
+- Added support for parallel index builds for IVFFlat
 - Added `l1_distance` function
 - Added element-wise multiplication for vectors
 - Added `sum` aggregate
 - Improved performance of distance functions
 - Fixed out of range results for cosine distance
-- Fixed results for NULL and NaN distances
+- Fixed results for NULL and NaN distances for IVFFlat
 
 ## 0.4.4 (2023-06-12)
 

README.md

@@ -162,7 +162,7 @@ You can add an index to use approximate nearest neighbor search, which trades so
 Supported index types are:
 
 - [IVFFlat](#ivfflat)
-- [HNSW](#hnsw) - *added in 0.5.0*
+- [HNSW](#hnsw) - added in 0.5.0
 
 ## IVFFlat
 
@@ -282,8 +282,8 @@ SELECT phase, tuples_done, tuples_total FROM pg_stat_progress_create_index;
 The phases are:
 
 1. `initializing`
-2. `performing k-means` (IVFFlat only)
+2. `performing k-means` - IVFFlat only
-3. `assigning tuples` (IVFFlat only)
+3. `assigning tuples` - IVFFlat only
 4. `loading tuples`
 
 Note: `tuples_done` and `tuples_total` are only populated during the `loading tuples` phase

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/hnswscan.c

@@ -185,6 +185,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		scan->xs_ctup.t_self = *tid;
 #endif
 
+		/* Unpin buffer */
 		if (BufferIsValid(so->buf))
 			ReleaseBuffer(so->buf);
 

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/ivfscan.c

@@ -343,6 +343,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		scan->xs_ctup.t_self = *tid;
 #endif
 
+		/* Unpin buffer */
 		if (BufferIsValid(so->buf))
 			ReleaseBuffer(so->buf);
 

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/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();