From fd65bcfb10bc9d4a74a08cc8f353bdf49a182118 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Mon, 23 Sep 2024 17:01:32 -0700
Subject: [PATCH] Added indexing [skip ci]

---
 sql/vector--0.7.4--0.8.0.sql | 12 ++++++++++++
 sql/vector.sql               | 16 ++++++++++++++++
 src/hnswutils.c              | 14 ++++++++++++++
 src/minivec.c                | 15 +++++++++++++++
 4 files changed, 57 insertions(+)

diff --git a/sql/vector--0.7.4--0.8.0.sql b/sql/vector--0.7.4--0.8.0.sql
index 38058c8..0c67ec8 100644
--- a/sql/vector--0.7.4--0.8.0.sql
+++ b/sql/vector--0.7.4--0.8.0.sql
@@ -1,6 +1,9 @@
 -- complain if script is sourced in psql, rather than via CREATE EXTENSION
 \echo Use "ALTER EXTENSION vector UPDATE TO '0.8.0'" to load this file. \quit
 
+CREATE FUNCTION hnsw_minivec_support(internal) RETURNS internal
+	AS 'MODULE_PATHNAME' LANGUAGE C;
+
 CREATE TYPE minivec;
 
 CREATE FUNCTION minivec_in(cstring, oid, integer) RETURNS minivec
@@ -30,11 +33,20 @@ CREATE TYPE minivec (
 CREATE FUNCTION l2_distance(minivec, minivec) RETURNS float8
 	AS 'MODULE_PATHNAME', 'minivec_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
+CREATE FUNCTION minivec_l2_squared_distance(minivec, minivec) RETURNS float8
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
 CREATE OPERATOR <-> (
 	LEFTARG = minivec, RIGHTARG = minivec, PROCEDURE = l2_distance,
 	COMMUTATOR = '<->'
 );
 
+CREATE OPERATOR CLASS minivec_l2_ops
+	FOR TYPE minivec USING hnsw AS
+	OPERATOR 1 <-> (minivec, minivec) FOR ORDER BY float_ops,
+	FUNCTION 1 minivec_l2_squared_distance(minivec, minivec),
+	FUNCTION 3 hnsw_minivec_support(internal);
+
 CREATE FUNCTION array_to_sparsevec(integer[], integer, boolean) RETURNS sparsevec
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
diff --git a/sql/vector.sql b/sql/vector.sql
index a3c54f1..1502293 100644
--- a/sql/vector.sql
+++ b/sql/vector.sql
@@ -272,6 +272,9 @@ CREATE FUNCTION ivfflat_bit_support(internal) RETURNS internal
 CREATE FUNCTION hnsw_halfvec_support(internal) RETURNS internal
 	AS 'MODULE_PATHNAME' LANGUAGE C;
 
+CREATE FUNCTION hnsw_minivec_support(internal) RETURNS internal
+	AS 'MODULE_PATHNAME' LANGUAGE C;
+
 CREATE FUNCTION hnsw_bit_support(internal) RETURNS internal
 	AS 'MODULE_PATHNAME' LANGUAGE C;
 
@@ -680,6 +683,11 @@ CREATE TYPE minivec (
 CREATE FUNCTION l2_distance(minivec, minivec) RETURNS float8
 	AS 'MODULE_PATHNAME', 'minivec_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
+-- minivec private functions
+
+CREATE FUNCTION minivec_l2_squared_distance(minivec, minivec) RETURNS float8
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
 -- minivec operators
 
 CREATE OPERATOR <-> (
@@ -687,6 +695,14 @@ CREATE OPERATOR <-> (
 	COMMUTATOR = '<->'
 );
 
+-- minivec op classes
+
+CREATE OPERATOR CLASS minivec_l2_ops
+	FOR TYPE minivec USING hnsw AS
+	OPERATOR 1 <-> (minivec, minivec) FOR ORDER BY float_ops,
+	FUNCTION 1 minivec_l2_squared_distance(minivec, minivec),
+	FUNCTION 3 hnsw_minivec_support(internal);
+
 -- bit functions
 
 CREATE FUNCTION hamming_distance(bit, bit) RETURNS float8
diff --git a/src/hnswutils.c b/src/hnswutils.c
index ac1e7de..7311f17 100644
--- a/src/hnswutils.c
+++ b/src/hnswutils.c
@@ -1375,6 +1375,20 @@ hnsw_halfvec_support(PG_FUNCTION_ARGS)
 	PG_RETURN_POINTER(&typeInfo);
 };
 
+FUNCTION_PREFIX PG_FUNCTION_INFO_V1(hnsw_minivec_support);
+Datum
+hnsw_minivec_support(PG_FUNCTION_ARGS)
+{
+	static const HnswTypeInfo typeInfo = {
+		.maxDimensions = HNSW_MAX_DIM * 4,
+		/* TODO */
+		.normalize = NULL,
+		.checkValue = NULL
+	};
+
+	PG_RETURN_POINTER(&typeInfo);
+};
+
 FUNCTION_PREFIX PG_FUNCTION_INFO_V1(hnsw_bit_support);
 Datum
 hnsw_bit_support(PG_FUNCTION_ARGS)
diff --git a/src/minivec.c b/src/minivec.c
index cdab86d..d7ab626 100644
--- a/src/minivec.c
+++ b/src/minivec.c
@@ -377,3 +377,18 @@ minivec_l2_distance(PG_FUNCTION_ARGS)
 
 	PG_RETURN_FLOAT8(sqrt((double) MinivecL2SquaredDistance(a->dim, a->x, b->x)));
 }
+
+/*
+ * Get the L2 squared distance between fp8 vectors
+ */
+FUNCTION_PREFIX PG_FUNCTION_INFO_V1(minivec_l2_squared_distance);
+Datum
+minivec_l2_squared_distance(PG_FUNCTION_ARGS)
+{
+	MiniVector *a = PG_GETARG_MINIVEC_P(0);
+	MiniVector *b = PG_GETARG_MINIVEC_P(1);
+
+	CheckDims(a, b);
+
+	PG_RETURN_FLOAT8((double) MinivecL2SquaredDistance(a->dim, a->x, b->x));
+}