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

@@ -155,15 +155,6 @@ hnswvalidate(Oid opclassoid)
 	return true;
 }
 
-/*
- * Checks if index-only scan is supported
- */
-static bool
-hnswcanreturn(Relation indexRelation, int attno)
-{
-	return attno == 1 && !OidIsValid(index_getprocid(indexRelation, 1, HNSW_NORM_PROC));
-}
-
 /*
  * Define index handler
  *
@@ -205,7 +196,7 @@ hnswhandler(PG_FUNCTION_ARGS)
 	amroutine->aminsert = hnswinsert;
 	amroutine->ambulkdelete = hnswbulkdelete;
 	amroutine->amvacuumcleanup = hnswvacuumcleanup;
-	amroutine->amcanreturn = hnswcanreturn;
+	amroutine->amcanreturn = NULL;	/* tuple not included in heapsort */
 	amroutine->amcostestimate = hnswcostestimate;
 	amroutine->amoptions = hnswoptions;
 	amroutine->amproperty = NULL;	/* TODO AMPROP_DISTANCE_ORDERABLE */

src/hnsw.h

@@ -266,7 +266,7 @@ Buffer		HnswNewBuffer(Relation index, ForkNumber forkNum);
 void		HnswInitPage(Buffer buf, Page page);
 void		HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
 void		HnswInit(void);
-List	   *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, HnswElement skipElement);
+List	   *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, HnswElement skipElement);
 HnswElement HnswGetEntryPoint(Relation index);
 HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
 void		HnswFreeElement(HnswElement element);

src/hnswscan.c

@@ -17,7 +17,6 @@ GetScanItems(IndexScanDesc scan, Datum q)
 	Relation	index = scan->indexRelation;
 	FmgrInfo   *procinfo = so->procinfo;
 	Oid			collation = so->collation;
-	bool		loadVec = scan->xs_want_itup;
 	List	   *ep;
 	List	   *w;
 	HnswElement entryPoint = HnswGetEntryPoint(index);
@@ -25,15 +24,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
 	if (entryPoint == NULL)
 		return NIL;
 
-	ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, loadVec));
+	ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, false));
 
 	for (int lc = entryPoint->level; lc >= 1; lc--)
 	{
-		w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, loadVec, NULL);
+		w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, false, NULL);
 		ep = w;
 	}
 
-	return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, loadVec, NULL);
+	return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, false, NULL);
 }
 
 /*
@@ -84,9 +83,6 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 
 	scan->opaque = so;
 
-	/* OK to always set since cheap */
-	scan->xs_itupdesc = RelationGetDescr(index);
-
 	return scan;
 }
 
@@ -169,7 +165,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		ItemPointer tid;
 		BlockNumber indexblkno;
 
-		/* Move to next element if no valid heap TIDs */
+		/* Move to next element if no valid heap tids */
 		if (list_length(hc->element->heaptids) == 0)
 		{
 			so->w = list_delete_last(so->w);
@@ -181,15 +177,6 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 
 		hc->element->heaptids = list_delete_last(hc->element->heaptids);
 
-		if (scan->xs_want_itup)
-		{
-			Datum		value = PointerGetDatum(hc->element->vec);
-			bool		isnull = false;
-
-			scan->xs_itup = index_form_tuple(scan->xs_itupdesc, &value, &isnull);
-			scan->xs_itup->t_tid = *tid;
-		}
-
 		MemoryContextSwitchTo(oldCtx);
 
 #if PG_VERSION_NUM >= 120000

src/hnswutils.c

@@ -543,7 +543,7 @@ AddToVisited(HTAB *v, HnswCandidate * hc, Relation index, bool *found)
  * Algorithm 2 from paper
  */
 List *
-HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, HnswElement skipElement)
+HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, HnswElement skipElement)
 {
 	ListCell   *lc2;
 
@@ -619,7 +619,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
 				if (index == NULL)
 					eDistance = GetCandidateDistance(e, q, procinfo, collation);
 				else
-					HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, loadVec);
+					HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, inserting);
 
 				Assert(!e->element->deleted);
 

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

test/expected/functions.out

@@ -106,12 +106,6 @@ SELECT cosine_distance('[1,1]', '[1,1]');
                0
 (1 row)
 
-SELECT cosine_distance('[1,0]', '[0,2]');
- cosine_distance 
------------------
-               1
-(1 row)
-
 SELECT cosine_distance('[1,1]', '[-1,-1]');
  cosine_distance 
 -----------------
@@ -158,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

@@ -25,7 +25,6 @@ SELECT inner_product('[3e38]', '[3e38]');
 SELECT cosine_distance('[1,2]', '[2,4]');
 SELECT cosine_distance('[1,2]', '[0,0]');
 SELECT cosine_distance('[1,1]', '[1,1]');
-SELECT cosine_distance('[1,0]', '[0,2]');
 SELECT cosine_distance('[1,1]', '[-1,-1]');
 SELECT cosine_distance('[1,2]', '[3]');
 SELECT cosine_distance('[1,1]', '[1.1,1.1]');
@@ -37,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;