Added subvector function

CHANGELOG.md

@@ -4,6 +4,7 @@
 - Added `hamming_distance` function
 - Added `jaccard_distance` function
 - Added `quantize_binary` function
+- Added `subvector` function
 - Updated comparison operators to support vectors with different dimensions
 
 ## 0.6.2 (2024-03-18)

README.md

@@ -738,6 +738,7 @@ inner_product(vector, vector) → double precision | inner product |
 l2_distance(vector, vector) → double precision | Euclidean distance |
 l1_distance(vector, vector) → double precision | taxicab distance | 0.5.0
 quantize_binary(vector) → bit | quantize | unreleased
+subvector(vector, integer, integer) → vector | subvector | unreleased
 vector_dims(vector) → integer | number of dimensions |
 vector_norm(vector) → double precision | Euclidean norm |
 

sql/vector--0.6.2--0.7.0.sql

@@ -4,6 +4,9 @@
 CREATE FUNCTION quantize_binary(vector) RETURNS bit
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
+CREATE FUNCTION subvector(vector, int, int) RETURNS vector
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
 CREATE FUNCTION hamming_distance(bit, bit) RETURNS float8
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 

sql/vector.sql

@@ -61,6 +61,9 @@ CREATE FUNCTION vector_mul(vector, vector) RETURNS vector
 CREATE FUNCTION quantize_binary(vector) RETURNS bit
 	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
 
+CREATE FUNCTION subvector(vector, int, int) RETURNS vector
+	AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
+
 -- vector private functions
 
 CREATE FUNCTION vector_lt(vector, vector) RETURNS bool

src/vector.c

@@ -877,6 +877,37 @@ quantize_binary(PG_FUNCTION_ARGS)
 	PG_RETURN_VARBIT_P(result);
 }
 
+/*
+ * Get a subvector
+ */
+PGDLLEXPORT PG_FUNCTION_INFO_V1(subvector);
+Datum
+subvector(PG_FUNCTION_ARGS)
+{
+	Vector	   *a = PG_GETARG_VECTOR_P(0);
+	int32		start = PG_GETARG_INT32(1);
+	int32		count = PG_GETARG_INT32(2);
+	int32		end = start + count;
+	float	   *ax = a->x;
+	Vector	   *result;
+	int			dim;
+
+	/* Indexing starts at 1, like substring */
+	if (start < 1)
+		start = 1;
+
+	if (end > a->dim)
+		end = a->dim + 1;
+
+	dim = end - start;
+	CheckDim(dim);
+	result = InitVector(dim);
+
+	for (int i = 0; i < dim; i++)
+		result->x[i] = ax[start - 1 + i];
+
+	PG_RETURN_POINTER(result);
+}
 
 /*
  * Internal helper to compare vectors

test/expected/vector_functions.out

@@ -284,6 +284,36 @@ SELECT quantize_binary('[0,0.1,-0.2,-0.3,0.4,0.5,0.6,-0.7,0.8,-0.9,1]');
  01001110101
 (1 row)
 
+SELECT subvector('[1,2,3,4,5]', 1, 3);
+ subvector 
+-----------
+ [1,2,3]
+(1 row)
+
+SELECT subvector('[1,2,3,4,5]', 3, 2);
+ subvector 
+-----------
+ [3,4]
+(1 row)
+
+SELECT subvector('[1,2,3,4,5]', -1, 3);
+ subvector 
+-----------
+ [1]
+(1 row)
+
+SELECT subvector('[1,2,3,4,5]', 3, 9);
+ subvector 
+-----------
+ [3,4,5]
+(1 row)
+
+SELECT subvector('[1,2,3,4,5]', 1, 0);
+ERROR:  vector must have at least 1 dimension
+SELECT subvector('[1,2,3,4,5]', 3, -1);
+ERROR:  vector must have at least 1 dimension
+SELECT subvector('[1,2,3,4,5]', -1, 2);
+ERROR:  vector must have at least 1 dimension
 SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
     avg    
 -----------

test/sql/vector_functions.sql

@@ -61,6 +61,14 @@ SELECT l1_distance('[3e38]'::vector, '[-3e38]');
 SELECT quantize_binary('[1,0,-1]');
 SELECT quantize_binary('[0,0.1,-0.2,-0.3,0.4,0.5,0.6,-0.7,0.8,-0.9,1]');
 
+SELECT subvector('[1,2,3,4,5]', 1, 3);
+SELECT subvector('[1,2,3,4,5]', 3, 2);
+SELECT subvector('[1,2,3,4,5]', -1, 3);
+SELECT subvector('[1,2,3,4,5]', 3, 9);
+SELECT subvector('[1,2,3,4,5]', 1, 0);
+SELECT subvector('[1,2,3,4,5]', 3, -1);
+SELECT subvector('[1,2,3,4,5]', -1, 2);
+
 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;