Use normalize_l2 for ivfflat

CHANGELOG.md

@@ -3,6 +3,7 @@
 - Added HNSW index type
 - Added support for parallel index builds
 - Added `l1_distance` function
+- Added `normalize_l2` function
 - Added element-wise multiplication for vectors
 - Added `sum` aggregate
 - Improved performance of distance functions

README.md

@@ -392,6 +392,7 @@ cosine_distance(vector, vector) → double precision | cosine distance
 inner_product(vector, vector) → double precision | inner product
 l2_distance(vector, vector) → double precision | Euclidean distance
 l1_distance(vector, vector) → double precision | taxicab distance [unreleased]
+normalize_l2(vector) → vector | normalize with Euclidean norm [unreleased]
 vector_dims(vector) → integer | number of dimensions
 vector_norm(vector) → double precision | Euclidean norm
 

sql/vector--0.4.4--0.5.0.sql

@@ -4,6 +4,9 @@
 CREATE FUNCTION l1_distance(vector, vector) RETURNS float8
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
+CREATE FUNCTION normalize_l2(vector) RETURNS vector
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
 CREATE FUNCTION vector_mul(vector, vector) RETURNS vector
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 

sql/vector.sql

@@ -49,6 +49,9 @@ CREATE FUNCTION vector_dims(vector) RETURNS integer
 CREATE FUNCTION vector_norm(vector) RETURNS float8
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
+CREATE FUNCTION normalize_l2(vector) RETURNS vector
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
 CREATE FUNCTION vector_add(vector, vector) RETURNS vector
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
@@ -265,15 +268,15 @@ CREATE OPERATOR CLASS vector_ip_ops
 	OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops,
 	FUNCTION 1 vector_negative_inner_product(vector, vector),
 	FUNCTION 3 vector_spherical_distance(vector, vector),
-	FUNCTION 4 vector_norm(vector);
+	FUNCTION 6 normalize_l2(vector);
 
 CREATE OPERATOR CLASS vector_cosine_ops
 	FOR TYPE vector USING ivfflat AS
 	OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
 	FUNCTION 1 vector_negative_inner_product(vector, vector),
-	FUNCTION 2 vector_norm(vector),
 	FUNCTION 3 vector_spherical_distance(vector, vector),
-	FUNCTION 4 vector_norm(vector);
+	FUNCTION 5 normalize_l2(vector),
+	FUNCTION 6 normalize_l2(vector);
 
 CREATE OPERATOR CLASS vector_l2_ops
 	FOR TYPE vector USING hnsw AS
@@ -289,4 +292,4 @@ CREATE OPERATOR CLASS vector_cosine_ops
 	FOR TYPE vector USING hnsw AS
 	OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
 	FUNCTION 1 vector_negative_inner_product(vector, vector),
-	FUNCTION 2 vector_norm(vector);
+	FUNCTION 3 normalize_l2(vector);

src/hnsw.c

@@ -167,7 +167,7 @@ hnswhandler(PG_FUNCTION_ARGS)
 	IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
 
 	amroutine->amstrategies = 0;
-	amroutine->amsupport = 2;
+	amroutine->amsupport = 3;
 #if PG_VERSION_NUM >= 130000
 	amroutine->amoptsprocnum = 0;
 #endif

src/hnsw.h

@@ -19,6 +19,7 @@
 /* Support functions */
 #define HNSW_DISTANCE_PROC 1
 #define HNSW_NORM_PROC 2
+#define HNSW_NORMALIZE_PROC 3
 
 #define HNSW_VERSION	1
 #define HNSW_MAGIC_NUMBER 0xA953A953
@@ -147,6 +148,7 @@ typedef struct HnswBuildState
 	/* Support functions */
 	FmgrInfo   *procinfo;
 	FmgrInfo   *normprocinfo;
+	FmgrInfo   *normalizeprocinfo;
 	Oid			collation;
 
 	/* Variables */
@@ -220,6 +222,7 @@ typedef struct HnswScanOpaqueData
 	/* Support functions */
 	FmgrInfo   *procinfo;
 	FmgrInfo   *normprocinfo;
+	FmgrInfo   *normalizeprocinfo;
 	Oid			collation;
 }			HnswScanOpaqueData;
 
@@ -255,7 +258,7 @@ typedef struct HnswVacuumState
 int			HnswGetM(Relation index);
 int			HnswGetEfConstruction(Relation index);
 FmgrInfo   *HnswOptionalProcInfo(Relation rel, uint16 procnum);
-bool		HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
+void		HnswNormValue(FmgrInfo *procinfo, FmgrInfo *normalizeprocinfo, Oid collation, Datum *value, Vector * result);
 void		HnswCommitBuffer(Buffer buf, GenericXLogState *state);
 Buffer		HnswNewBuffer(Relation index, ForkNumber forkNum);
 void		HnswInitPage(Buffer buf, Page page);

src/hnswbuild.c

@@ -278,11 +278,7 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
 	Datum		value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
 
 	/* Normalize if needed */
-	if (buildstate->normprocinfo != NULL)
-	{
-		if (!HnswNormValue(buildstate->normprocinfo, collation, &value, buildstate->normvec))
-			return false;
-	}
+	HnswNormValue(buildstate->normprocinfo, buildstate->normalizeprocinfo, collation, &value, buildstate->normvec);
 
 	/* Copy value to element so accessible outside of memory context */
 	memcpy(element->vec, DatumGetVector(value), VECTOR_SIZE(buildstate->dimensions));
@@ -413,6 +409,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
 	/* Get support functions */
 	buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
 	buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
+	buildstate->normalizeprocinfo = HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC);
 	buildstate->collation = index->rd_indcollation[0];
 
 	buildstate->elements = NIL;

src/hnswinsert.c

@@ -417,6 +417,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
 {
 	Datum		value;
 	FmgrInfo   *normprocinfo;
+	FmgrInfo   *normalizeprocinfo;
 	HnswElement entryPoint;
 	HnswElement element;
 	int			m = HnswGetM(index);
@@ -432,11 +433,8 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
 
 	/* Normalize if needed */
 	normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
-	if (normprocinfo != NULL)
-	{
-		if (!HnswNormValue(normprocinfo, collation, &value, NULL))
-			return false;
-	}
+	normalizeprocinfo = HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC);
+	HnswNormValue(normprocinfo, normalizeprocinfo, collation, &value, NULL);
 
 	/* Create an element */
 	element = HnswInitElement(heap_tid, m, ml, HnswGetMaxLevel(m));

src/hnswscan.c

@@ -78,6 +78,7 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 	/* Set support functions */
 	so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
 	so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
+	so->normalizeprocinfo = HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC);
 	so->collation = index->rd_indcollation[0];
 
 	scan->opaque = so;
@@ -140,8 +141,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 			Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
 
 			/* Fine if normalization fails */
-			if (so->normprocinfo != NULL)
-				HnswNormValue(so->normprocinfo, so->collation, &value, NULL);
+			HnswNormValue(so->normprocinfo, so->normalizeprocinfo, so->collation, &value, NULL);
 		}
 
 		GetScanItems(scan, value);

src/hnswutils.c

@@ -47,17 +47,26 @@ HnswOptionalProcInfo(Relation rel, uint16 procnum)
 }
 
 /*
- * Divide by the norm
- *
- * Returns false if value should not be indexed
+ * Normalize a vector
  *
  * The caller needs to free the pointer stored in value
  * if it's different than the original value
  */
-bool
-HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
+void
+HnswNormValue(FmgrInfo *procinfo, FmgrInfo *normalizeprocinfo, Oid collation, Datum *value, Vector * result)
 {
-	double		norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
+	double		norm;
+
+	if (normalizeprocinfo != NULL)
+	{
+		*value = FunctionCall1Coll(normalizeprocinfo, collation, *value);
+		return;
+	}
+
+	if (procinfo == NULL)
+		return;
+
+	norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
 
 	if (norm > 0)
 	{
@@ -70,11 +79,7 @@ HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
 			result->x[i] = v->x[i] / norm;
 
 		*value = PointerGetDatum(result);
-
-		return true;
 	}
-
-	return false;
 }
 
 /*

src/ivfbuild.c

@@ -75,11 +75,7 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
 	 * 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;
-	}
+	IvfflatNormValue(buildstate->kmeansnormprocinfo, buildstate->kmeansnormalizeprocinfo, buildstate->collation, &value, buildstate->normvec);
 
 	if (samples->length < targsamples)
 	{
@@ -176,11 +172,7 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState
 	Datum		value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
 
 	/* Normalize if needed */
-	if (buildstate->normprocinfo != NULL)
-	{
-		if (!IvfflatNormValue(buildstate->normprocinfo, buildstate->collation, &value, buildstate->normvec))
-			return;
-	}
+	IvfflatNormValue(buildstate->normprocinfo, buildstate->normalizeprocinfo, buildstate->collation, &value, buildstate->normvec);
 
 	/* Find the list that minimizes the distance */
 	for (int i = 0; i < centers->length; i++)
@@ -368,6 +360,8 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
 	buildstate->procinfo = index_getprocinfo(index, 1, IVFFLAT_DISTANCE_PROC);
 	buildstate->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
 	buildstate->kmeansnormprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
+	buildstate->normalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORMALIZE_PROC);
+	buildstate->kmeansnormalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORMALIZE_PROC);
 	buildstate->collation = index->rd_indcollation[0];
 
 	/* Require more than one dimension for spherical k-means */

src/ivfflat.c

@@ -194,7 +194,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
 	IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
 
 	amroutine->amstrategies = 0;
-	amroutine->amsupport = 4;
+	amroutine->amsupport = 6;
 #if PG_VERSION_NUM >= 130000
 	amroutine->amoptsprocnum = 0;
 #endif

src/ivfflat.h

@@ -31,6 +31,8 @@
 #define IVFFLAT_NORM_PROC 2
 #define IVFFLAT_KMEANS_DISTANCE_PROC 3
 #define IVFFLAT_KMEANS_NORM_PROC 4
+#define IVFFLAT_NORMALIZE_PROC 5
+#define IVFFLAT_KMEANS_NORMALIZE_PROC 6
 
 #define IVFFLAT_VERSION	1
 #define IVFFLAT_MAGIC_NUMBER 0x14FF1A7
@@ -172,6 +174,8 @@ typedef struct IvfflatBuildState
 	FmgrInfo   *procinfo;
 	FmgrInfo   *normprocinfo;
 	FmgrInfo   *kmeansnormprocinfo;
+	FmgrInfo   *normalizeprocinfo;
+	FmgrInfo   *kmeansnormalizeprocinfo;
 	Oid			collation;
 
 	/* Variables */
@@ -253,6 +257,7 @@ typedef struct IvfflatScanOpaqueData
 	/* Support functions */
 	FmgrInfo   *procinfo;
 	FmgrInfo   *normprocinfo;
+	FmgrInfo   *normalizeprocinfo;
 	Oid			collation;
 
 	/* Lists */
@@ -273,7 +278,7 @@ void		VectorArrayFree(VectorArray arr);
 void		PrintVectorArray(char *msg, VectorArray arr);
 void		IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);
 FmgrInfo   *IvfflatOptionalProcInfo(Relation rel, uint16 procnum);
-bool		IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
+void		IvfflatNormValue(FmgrInfo *procinfo, FmgrInfo *normalizeprocinfo, Oid collation, Datum *value, Vector * result);
 int			IvfflatGetLists(Relation index);
 void		IvfflatUpdateList(Relation index, ListInfo listInfo, BlockNumber insertPage, BlockNumber originalInsertPage, BlockNumber startPage, ForkNumber forkNum);
 void		IvfflatCommitBuffer(Buffer buf, GenericXLogState *state);

src/ivfinsert.c

@@ -68,6 +68,7 @@ InsertTuple(Relation rel, Datum *values, bool *isnull, ItemPointer heap_tid, Rel
 	IndexTuple	itup;
 	Datum		value;
 	FmgrInfo   *normprocinfo;
+	FmgrInfo   *normalizeprocinfo;
 	Buffer		buf;
 	Page		page;
 	GenericXLogState *state;
@@ -81,11 +82,8 @@ InsertTuple(Relation rel, Datum *values, bool *isnull, ItemPointer heap_tid, Rel
 
 	/* Normalize if needed */
 	normprocinfo = IvfflatOptionalProcInfo(rel, IVFFLAT_NORM_PROC);
-	if (normprocinfo != NULL)
-	{
-		if (!IvfflatNormValue(normprocinfo, rel->rd_indcollation[0], &value, NULL))
-			return;
-	}
+	normalizeprocinfo = IvfflatOptionalProcInfo(rel, IVFFLAT_NORMALIZE_PROC);
+	IvfflatNormValue(normprocinfo, normalizeprocinfo, rel->rd_indcollation[0], &value, NULL);
 
 	/* Find the insert page - sets the page and list info */
 	FindInsertPage(rel, values, &insertPage, &listInfo);

src/ivfscan.c

@@ -232,6 +232,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	/* Set support functions */
 	so->procinfo = index_getprocinfo(index, 1, IVFFLAT_DISTANCE_PROC);
 	so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
+	so->normalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORMALIZE_PROC);
 	so->collation = index->rd_indcollation[0];
 
 	/* Create tuple description for sorting */
@@ -319,8 +320,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 			Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
 
 			/* Fine if normalization fails */
-			if (so->normprocinfo != NULL)
-				IvfflatNormValue(so->normprocinfo, so->collation, &value, NULL);
+			IvfflatNormValue(so->normprocinfo, so->normalizeprocinfo, so->collation, &value, NULL);
 		}
 
 		IvfflatBench("GetScanLists", GetScanLists(scan, value));

src/ivfutils.c

@@ -66,17 +66,26 @@ IvfflatOptionalProcInfo(Relation rel, uint16 procnum)
 }
 
 /*
- * Divide by the norm
- *
- * Returns false if value should not be indexed
+ * Normalize a vector
  *
  * The caller needs to free the pointer stored in value
  * if it's different than the original value
  */
-bool
-IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
+void
+IvfflatNormValue(FmgrInfo *procinfo, FmgrInfo *normalizeprocinfo, Oid collation, Datum *value, Vector * result)
 {
-	double		norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
+	double		norm;
+
+	if (normalizeprocinfo != NULL)
+	{
+		*value = FunctionCall1Coll(normalizeprocinfo, collation, *value);
+		return;
+	}
+
+	if (procinfo == NULL)
+		return;
+
+	norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
 
 	if (norm > 0)
 	{
@@ -89,11 +98,7 @@ IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * resul
 			result->x[i] = v->x[i] / norm;
 
 		*value = PointerGetDatum(result);
-
-		return true;
 	}
-
-	return false;
 }
 
 /*

src/vector.c

@@ -745,6 +745,52 @@ vector_norm(PG_FUNCTION_ARGS)
 	PG_RETURN_FLOAT8(sqrt((double) norm));
 }
 
+/*
+ * Normalize a vector with the L2 norm
+ */
+PGDLLEXPORT PG_FUNCTION_INFO_V1(normalize_l2);
+Datum
+normalize_l2(PG_FUNCTION_ARGS)
+{
+	Vector	   *a = PG_GETARG_VECTOR_P(0);
+	float	   *ax = a->x;
+	double		norm = 0.0;
+#ifdef _MSC_VER
+	/* Fix precision on Windows */
+	double		normf;
+#else
+	float		normf;
+#endif
+	Vector	   *result;
+	float	   *rx;
+
+	result = InitVector(a->dim);
+	rx = result->x;
+
+	/* Auto-vectorized */
+	for (int i = 0; i < a->dim; i++)
+		norm += (double) ax[i] * (double) ax[i];
+
+	norm = sqrt(norm);
+	normf = norm;
+
+	if (normf > 0)
+	{
+		/* Auto-vectorized */
+		for (int i = 0, imax = a->dim; i < imax; i++)
+			rx[i] = ax[i] / normf;
+
+		/* Check for overflow */
+		for (int i = 0, imax = a->dim; i < imax; i++)
+		{
+			if (isinf(rx[i]))
+				float_overflow_error();
+		}
+	}
+
+	PG_RETURN_POINTER(result);
+}
+
 /*
  * Add vectors
  */

test/expected/functions.out

@@ -48,6 +48,36 @@ SELECT vector_norm('[0,1]');
            1
 (1 row)
 
+SELECT normalize_l2('[3,4]');
+ normalize_l2 
+--------------
+ [0.6,0.8]
+(1 row)
+
+SELECT normalize_l2('[3,0]');
+ normalize_l2 
+--------------
+ [1,0]
+(1 row)
+
+SELECT normalize_l2('[0,0.1]');
+ normalize_l2 
+--------------
+ [0,1]
+(1 row)
+
+SELECT normalize_l2('[0,0]');
+ normalize_l2 
+--------------
+ [0,0]
+(1 row)
+
+SELECT normalize_l2('[3e38]');
+ normalize_l2 
+--------------
+ [1]
+(1 row)
+
 SELECT l2_distance('[0,0]', '[3,4]');
  l2_distance 
 -------------

test/expected/hnsw_cosine.out

@@ -9,18 +9,19 @@ SELECT * FROM t ORDER BY val <=> '[3,3,3]';
  [1,1,1]
  [1,2,3]
  [1,2,4]
-(3 rows)
+ [0,0,0]
+(4 rows)
 
 SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
  count 
 -------
-     3
+     4
 (1 row)
 
 SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2;
  count 
 -------
-     3
+     4
 (1 row)
 
 DROP TABLE t;

test/expected/ivfflat_cosine.out

@@ -9,18 +9,19 @@ SELECT * FROM t ORDER BY val <=> '[3,3,3]';
  [1,1,1]
  [1,2,3]
  [1,2,4]
-(3 rows)
+ [0,0,0]
+(4 rows)
 
 SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
  count 
 -------
-     3
+     4
 (1 row)
 
 SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2;
  count 
 -------
-     3
+     4
 (1 row)
 
 DROP TABLE t;

test/sql/functions.sql

@@ -12,6 +12,12 @@ SELECT round(vector_norm('[1,1]')::numeric, 5);
 SELECT vector_norm('[3,4]');
 SELECT vector_norm('[0,1]');
 
+SELECT normalize_l2('[3,4]');
+SELECT normalize_l2('[3,0]');
+SELECT normalize_l2('[0,0.1]');
+SELECT normalize_l2('[0,0]');
+SELECT normalize_l2('[3e38]');
+
 SELECT l2_distance('[0,0]', '[3,4]');
 SELECT l2_distance('[0,0]', '[0,1]');
 SELECT l2_distance('[1,2]', '[3]');