Added support for uuid

sql/vector--0.6.1--0.7.0.sql

@@ -1,10 +1,34 @@
 -- complain if script is sourced in psql, rather than via CREATE EXTENSION
 \echo Use "ALTER EXTENSION vector UPDATE TO '0.7.0'" to load this file. \quit
 
-CREATE OPERATOR CLASS vector_integer_ops
-	DEFAULT FOR TYPE integer USING hnsw AS
-	OPERATOR 2 = (integer, integer);
-
 CREATE OPERATOR CLASS vector_bigint_ops
 	DEFAULT FOR TYPE bigint USING hnsw AS
-	OPERATOR 2 = (bigint, bigint);
+	OPERATOR 2 < ,
+	OPERATOR 3 <= ,
+	OPERATOR 4 = ,
+	OPERATOR 5 >= ,
+	OPERATOR 6 > ;
+
+CREATE OPERATOR CLASS vector_integer_ops
+	DEFAULT FOR TYPE integer USING hnsw AS
+	OPERATOR 2 < ,
+	OPERATOR 3 <= ,
+	OPERATOR 4 = ,
+	OPERATOR 5 >= ,
+	OPERATOR 6 > ;
+
+CREATE OPERATOR CLASS vector_text_ops
+	DEFAULT FOR TYPE text USING hnsw AS
+	OPERATOR 2 < ,
+	OPERATOR 3 <= ,
+	OPERATOR 4 = ,
+	OPERATOR 5 >= ,
+	OPERATOR 6 > ;
+
+CREATE OPERATOR CLASS vector_uuid_ops
+	DEFAULT FOR TYPE uuid USING hnsw AS
+	OPERATOR 2 < ,
+	OPERATOR 3 <= ,
+	OPERATOR 4 = ,
+	OPERATOR 5 >= ,
+	OPERATOR 6 > ;

sql/vector.sql

@@ -293,10 +293,34 @@ CREATE OPERATOR CLASS vector_cosine_ops
 
 -- hnsw attributes
 
-CREATE OPERATOR CLASS vector_integer_ops
-	DEFAULT FOR TYPE integer USING hnsw AS
-	OPERATOR 2 = (integer, integer);
-
 CREATE OPERATOR CLASS vector_bigint_ops
 	DEFAULT FOR TYPE bigint USING hnsw AS
-	OPERATOR 2 = (bigint, bigint);
+	OPERATOR 2 < ,
+	OPERATOR 3 <= ,
+	OPERATOR 4 = ,
+	OPERATOR 5 >= ,
+	OPERATOR 6 > ;
+
+CREATE OPERATOR CLASS vector_integer_ops
+	DEFAULT FOR TYPE integer USING hnsw AS
+	OPERATOR 2 < ,
+	OPERATOR 3 <= ,
+	OPERATOR 4 = ,
+	OPERATOR 5 >= ,
+	OPERATOR 6 > ;
+
+CREATE OPERATOR CLASS vector_text_ops
+	DEFAULT FOR TYPE text USING hnsw AS
+	OPERATOR 2 < ,
+	OPERATOR 3 <= ,
+	OPERATOR 4 = ,
+	OPERATOR 5 >= ,
+	OPERATOR 6 > ;
+
+CREATE OPERATOR CLASS vector_uuid_ops
+	DEFAULT FOR TYPE uuid USING hnsw AS
+	OPERATOR 2 < ,
+	OPERATOR 3 <= ,
+	OPERATOR 4 = ,
+	OPERATOR 5 >= ,
+	OPERATOR 6 > ;

src/hnswbuild.c

@@ -481,7 +481,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
 	bool		unused;
 
 	/* Detoast once for all calls */
-	Datum		value = PointerGetDatum(PG_DETOAST_DATUM(values[IndexRelationGetNumberOfKeyAttributes(index) - 1]));
+	Datum		value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
 
 	/* Normalize if needed */
 	if (buildstate->normprocinfo != NULL)
@@ -551,7 +551,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
 	/* Copy the index tuple */
 	memcpy(itupPtr, itup, itupSize);
 	HnswPtrStore(base, element->itup, itupPtr);
-	valuePtr = DatumGetPointer(index_getattr(itupPtr, IndexRelationGetNumberOfKeyAttributes(index), tupdesc, &unused));
+	valuePtr = DatumGetPointer(index_getattr(itupPtr, 1, tupdesc, &unused));
 	HnswPtrStore(base, element->value, valuePtr);
 
 	/* Create a lock for the element */
@@ -582,7 +582,7 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
 #endif
 
 	/* Skip nulls */
-	if (isnull[IndexRelationGetNumberOfKeyAttributes(index) - 1])
+	if (isnull[0])
 		return;
 
 	/* Use memory context */
@@ -674,16 +674,16 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
 
 	buildstate->m = HnswGetM(index);
 	buildstate->efConstruction = HnswGetEfConstruction(index);
-	buildstate->dimensions = TupleDescAttr(index->rd_att, IndexRelationGetNumberOfKeyAttributes(index) - 1)->atttypmod;
+	buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
 
 	/* For now */
-	if (IndexRelationGetNumberOfKeyAttributes(index) > 2)
-		elog(ERROR, "index cannot have more than two columns");
+	if (IndexRelationGetNumberOfKeyAttributes(index) > 3)
+		elog(ERROR, "index cannot have more than three columns");
 
-	if (!OidIsValid(index_getprocid(index, IndexRelationGetNumberOfKeyAttributes(index), HNSW_DISTANCE_PROC)))
-		elog(ERROR, "last column must be a vector");
+	if (!OidIsValid(index_getprocid(index, 1, HNSW_DISTANCE_PROC)))
+		elog(ERROR, "first column must be a vector");
 
-	for (int i = 0; i < IndexRelationGetNumberOfKeyAttributes(index) - 1; i++)
+	for (int i = 1; i < IndexRelationGetNumberOfKeyAttributes(index); i++)
 	{
 		if (OidIsValid(index_getprocid(index, i + 1, HNSW_DISTANCE_PROC)))
 			elog(ERROR, "column %d cannot be a vector", i + 1);
@@ -703,9 +703,9 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
 	buildstate->indtuples = 0;
 
 	/* Get support functions */
-	buildstate->procinfo = index_getprocinfo(index, IndexRelationGetNumberOfKeyAttributes(index), HNSW_DISTANCE_PROC);
+	buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
 	buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
-	buildstate->collation = index->rd_indcollation[IndexRelationGetNumberOfKeyAttributes(index) - 1];
+	buildstate->collation = index->rd_indcollation[0];
 
 	InitGraph(&buildstate->graphData, NULL, maintenance_work_mem * 1024L);
 	buildstate->graph = &buildstate->graphData;

src/hnswinsert.c

@@ -561,8 +561,8 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
 	HnswElement element;
 	int			m;
 	int			efConstruction = HnswGetEfConstruction(index);
-	FmgrInfo   *procinfo = index_getprocinfo(index, IndexRelationGetNumberOfKeyAttributes(index), HNSW_DISTANCE_PROC);
-	Oid			collation = index->rd_indcollation[IndexRelationGetNumberOfKeyAttributes(index) - 1];
+	FmgrInfo   *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
+	Oid			collation = index->rd_indcollation[0];
 	LOCKMODE	lockmode = ShareLock;
 	char	   *base = NULL;
 	TupleDesc	tupdesc = HnswTupleDesc(index);
@@ -584,7 +584,7 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
 	element = HnswInitElement(base, heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
 	itup = HnswFormIndexTuple(index, tupdesc, value, values, isnull);
 	HnswPtrStore(base, element->itup, itup);
-	valuePtr = DatumGetPointer(index_getattr(itup, IndexRelationGetNumberOfKeyAttributes(index), tupdesc, &unused));
+	valuePtr = DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused));
 	HnswPtrStore(base, element->value, valuePtr);
 
 	/* Prevent concurrent inserts when likely updating entry point */
@@ -621,10 +621,10 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
 {
 	Datum		value;
 	FmgrInfo   *normprocinfo;
-	Oid			collation = index->rd_indcollation[IndexRelationGetNumberOfKeyAttributes(index) - 1];
+	Oid			collation = index->rd_indcollation[0];
 
 	/* Detoast once for all calls */
-	value = PointerGetDatum(PG_DETOAST_DATUM(values[IndexRelationGetNumberOfKeyAttributes(index) - 1]));
+	value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
 
 	/* Normalize if needed */
 	normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
@@ -653,7 +653,7 @@ hnswinsert(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid,
 	MemoryContext insertCtx;
 
 	/* Skip nulls */
-	if (isnull[IndexRelationGetNumberOfKeyAttributes(index) - 1])
+	if (isnull[0])
 		return false;
 
 	/* Create memory context */

src/hnswscan.c

@@ -109,9 +109,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 									   ALLOCSET_DEFAULT_SIZES);
 
 	/* Set support functions */
-	so->procinfo = index_getprocinfo(index, IndexRelationGetNumberOfKeyAttributes(index), HNSW_DISTANCE_PROC);
+	so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
 	so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
-	so->collation = index->rd_indcollation[IndexRelationGetNumberOfKeyAttributes(index) - 1];
+	so->collation = index->rd_indcollation[0];
 
 	scan->opaque = so;
 

src/hnswutils.c

@@ -148,10 +148,10 @@ HnswGetEfConstruction(Relation index)
 FmgrInfo *
 HnswOptionalProcInfo(Relation index, uint16 procnum)
 {
-	if (!OidIsValid(index_getprocid(index, IndexRelationGetNumberOfKeyAttributes(index), procnum)))
+	if (!OidIsValid(index_getprocid(index, 1, procnum)))
 		return NULL;
 
-	return index_getprocinfo(index, IndexRelationGetNumberOfKeyAttributes(index), procnum);
+	return index_getprocinfo(index, 1, procnum);
 }
 
 /*
@@ -324,7 +324,7 @@ HnswFormIndexTuple(Relation index, TupleDesc tupdesc, Datum value, Datum *values
 	Datum	   *newValues = palloc(size);
 
 	memcpy(newValues, values, size);
-	newValues[IndexRelationGetNumberOfKeyAttributes(index) - 1] = value;
+	newValues[0] = value;
 
 	return index_form_tuple(tupdesc, newValues, isnull);
 }
@@ -642,7 +642,7 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
 			TupleDesc	tupdesc = RelationGetDescr(index);
 			bool		unused;
 			IndexTuple	itup = CopyIndexTuple((IndexTuple) &etup->data);
-			Datum		value = index_getattr(itup, IndexRelationGetNumberOfKeyAttributes(index), tupdesc, &unused);
+			Datum		value = index_getattr(itup, 1, tupdesc, &unused);
 
 			HnswPtrStore(base, element->itup, itup);
 			HnswPtrStore(base, element->value, DatumGetPointer(value));
@@ -722,7 +722,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
 			TupleDesc	tupdesc = RelationGetDescr(index);
 			bool		unused;
 
-			value = index_getattr(itup, IndexRelationGetNumberOfKeyAttributes(index), tupdesc, &unused);
+			value = index_getattr(itup, 1, tupdesc, &unused);
 		}
 		else
 			value = PointerGetDatum(&etup->data);

src/hnswvacuum.c

@@ -585,8 +585,8 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
 	vacuumstate->callback_state = callback_state;
 	vacuumstate->efConstruction = HnswGetEfConstruction(index);
 	vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD);
-	vacuumstate->procinfo = index_getprocinfo(index, IndexRelationGetNumberOfKeyAttributes(index), HNSW_DISTANCE_PROC);
-	vacuumstate->collation = index->rd_indcollation[IndexRelationGetNumberOfKeyAttributes(index) - 1];
+	vacuumstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
+	vacuumstate->collation = index->rd_indcollation[0];
 	vacuumstate->ntup = palloc0(HNSW_TUPLE_ALLOC_SIZE);
 	vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
 												"Hnsw vacuum temporary context",

test/t/020_hnsw_filtering.pl

@@ -12,6 +12,22 @@ my $limit = 20;
 my $dim = 3;
 my $array_sql = join(",", ('random()') x $dim);
 my $nc = 50;
+my @types = ("int4", "int8", "text", "uuid", "varchar");
+my $type = $types[rand(@types)];
+
+sub cast_c
+{
+	my ($c) = @_;
+
+	if ($type eq "uuid")
+	{
+		return "('00000000-0000-0000-0000-0000000000' || LPAD(($c)::text, 2, '0'))::uuid";
+	}
+	else
+	{
+		return "($c)::$type";
+	}
+}
 
 sub test_recall
 {
@@ -21,7 +37,7 @@ sub test_recall
 
 	my $explain = $node->safe_psql("postgres", qq(
 		SET enable_seqscan = off;
-		EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $cs[0] ORDER BY v $operator '$queries[0]' LIMIT $limit;
+		EXPLAIN ANALYZE SELECT i FROM tst WHERE c = @{[cast_c($cs[0])]} ORDER BY v $operator '$queries[0]' LIMIT $limit;
 	));
 	like($explain, qr/Index Cond/);
 
@@ -29,7 +45,7 @@ sub test_recall
 	{
 		my $actual = $node->safe_psql("postgres", qq(
 			SET enable_seqscan = off;
-			SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v $operator '$queries[$i]' LIMIT $limit;
+			SELECT i FROM tst WHERE c = @{[cast_c($cs[$i])]} ORDER BY v $operator '$queries[$i]' LIMIT $limit;
 		));
 		my @actual_ids = split("\n", $actual);
 		my %actual_set = map { $_ => 1 } @actual_ids;
@@ -58,14 +74,15 @@ $node->start;
 
 # Create table
 $node->safe_psql("postgres", "CREATE EXTENSION vector;");
-$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int8);");
-$node->safe_psql("postgres",
-	"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc FROM generate_series(1, 10000) i;"
-);
-$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (c, v vector_l2_ops);");
-$node->safe_psql("postgres",
-	"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc FROM generate_series(1, 10000) i;"
-);
+$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c $type, c2 $type);");
+$node->safe_psql("postgres", qq(
+	INSERT INTO tst SELECT i, ARRAY[$array_sql], @{[cast_c("i % $nc")]}, @{[cast_c("i % $nc")]} FROM generate_series(1, 10000) i;
+));
+$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, c, c2);");
+$node->safe_psql("postgres", qq(
+	INSERT INTO tst SELECT i, ARRAY[$array_sql], @{[cast_c("i % $nc")]}, @{[cast_c("i % $nc")]} FROM generate_series(1, 10000) i;
+));
+$node->safe_psql("postgres", "ANALYZE tst;");
 
 # Generate queries
 for (1 .. 20)
@@ -85,7 +102,7 @@ for my $i (0 .. $#queries)
 {
 	my $res = $node->safe_psql("postgres", qq(
 		SET enable_indexscan = off;
-		SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v <-> '$queries[$i]' LIMIT $limit;
+		SELECT i FROM tst WHERE c = @{[cast_c($cs[$i])]} ORDER BY v <-> '$queries[$i]' LIMIT $limit;
 	));
 	push(@expected, $res);
 }
@@ -93,21 +110,45 @@ for my $i (0 .. $#queries)
 # Test recall
 test_recall(0.99, '<->');
 
+# Test no conditions
+my $explain = $node->safe_psql("postgres", qq(
+	EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$queries[0]' LIMIT $limit;
+));
+like($explain, qr/Index Scan/);
+
+# Test range
+$explain = $node->safe_psql("postgres", qq(
+	EXPLAIN ANALYZE SELECT i FROM tst WHERE c >= @{[cast_c(1)]} AND c <= @{[cast_c(3)]} ORDER BY v <-> '$queries[0]' LIMIT $limit;
+));
+like($explain, qr/Index Cond: \(\(\S+ >= \S+\) AND \(\S+ <= \S+\)\)/);
+
+# Test multiple conditions
+$explain = $node->safe_psql("postgres", qq(
+	EXPLAIN ANALYZE SELECT i FROM tst WHERE c = @{[cast_c($cs[0])]} AND c2 = @{[cast_c($cs[0])]} ORDER BY v <-> '$queries[0]' LIMIT $limit;
+));
+like($explain, qr/Index Cond: \(\(\S+ = \S+\) AND \(\S+ = \S+\)\)/);
+
+# Test no order
+$explain = $node->safe_psql("postgres", qq(
+	EXPLAIN ANALYZE SELECT i FROM tst WHERE c = @{[cast_c($cs[0])]} LIMIT $limit;
+));
+like($explain, qr/Seq Scan/);
+
 # Test vacuum
-$node->safe_psql("postgres", "DELETE FROM tst WHERE c > 5;");
+$node->safe_psql("postgres", "DELETE FROM tst WHERE c > @{[cast_c(5)]};");
 $node->safe_psql("postgres", "VACUUM tst;");
 
 # Test columns
 my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c);");
-like($stderr, qr/last column must be a vector/);
+like($stderr, qr/first column must be a vector/);
 
-($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c, c);");
-like($stderr, qr/last column must be a vector/);
+($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c, v vector_l2_ops);");
+like($stderr, qr/first column must be a vector/);
 
-($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c, c, v vector_l2_ops);");
-like($stderr, qr/index cannot have more than two columns/);
+($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, c, c, c);");
+like($stderr, qr/index cannot have more than three columns/);
 
 ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, v vector_l2_ops);");
-like($stderr, qr/column 1 cannot be a vector/);
+like($stderr, qr/column 2 cannot be a vector/);
 
 done_testing();