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_bigint_ops
+	DEFAULT FOR TYPE bigint USING hnsw AS
+	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 = (integer, integer);
+	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 = (text, text);
+	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_bigint_ops
+	DEFAULT FOR TYPE bigint USING hnsw AS
+	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 = (integer, integer);
+	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 = (text, text);
+	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

@@ -677,8 +677,8 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
 	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, 1, HNSW_DISTANCE_PROC)))
 		elog(ERROR, "first column must be a vector");

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 text);");
-$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 (v vector_l2_ops, c);");
-$node->safe_psql("postgres",
-	"INSERT INTO tst SELECT i, ARRAY[$array_sql], (i % $nc)::text 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,8 +110,32 @@ 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
@@ -104,8 +145,8 @@ like($stderr, qr/first 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 (v vector_l2_ops, c, c);");
-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 2 cannot be a vector/);