Added support for bigint, text, and varchar [skip ci]

This commit is contained in:
Andrew Kane
2024-02-13 21:07:39 -08:00
parent f712847781
commit 07f2cbdee4
3 changed files with 23 additions and 5 deletions

View File

@@ -1,6 +1,14 @@
-- complain if script is sourced in psql, rather than via CREATE EXTENSION -- 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 \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 = (bigint, bigint);
CREATE OPERATOR CLASS vector_integer_ops CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING hnsw AS DEFAULT FOR TYPE integer USING hnsw AS
OPERATOR 2 = (integer, integer); OPERATOR 2 = (integer, integer);
CREATE OPERATOR CLASS vector_text_ops
DEFAULT FOR TYPE text USING hnsw AS
OPERATOR 2 = (text, text);

View File

@@ -293,6 +293,14 @@ CREATE OPERATOR CLASS vector_cosine_ops
-- hnsw attributes -- hnsw attributes
CREATE OPERATOR CLASS vector_bigint_ops
DEFAULT FOR TYPE bigint USING hnsw AS
OPERATOR 2 = (bigint, bigint);
CREATE OPERATOR CLASS vector_integer_ops CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING hnsw AS DEFAULT FOR TYPE integer USING hnsw AS
OPERATOR 2 = (integer, integer); OPERATOR 2 = (integer, integer);
CREATE OPERATOR CLASS vector_text_ops
DEFAULT FOR TYPE text USING hnsw AS
OPERATOR 2 = (text, text);

View File

@@ -12,16 +12,18 @@ my $limit = 20;
my $dim = 3; my $dim = 3;
my $array_sql = join(",", ('random()') x $dim); my $array_sql = join(",", ('random()') x $dim);
my $nc = 50; my $nc = 50;
my $type = "int4"; # int4, int8, text, varchar
sub test_recall sub test_recall
{ {
my ($min, $operator) = @_; my ($min, $operator) = @_;
my $correct = 0; my $correct = 0;
my $total = 0; my $total = 0;
my $cast = $type eq "int8" ? "::int8" : "";
my $explain = $node->safe_psql("postgres", qq( my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; 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 = '$cs[0]'$cast ORDER BY v $operator '$queries[0]' LIMIT $limit;
)); ));
like($explain, qr/Index Cond/); like($explain, qr/Index Cond/);
@@ -29,7 +31,7 @@ sub test_recall
{ {
my $actual = $node->safe_psql("postgres", qq( my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; 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 = '$cs[$i]'$cast ORDER BY v $operator '$queries[$i]' LIMIT $limit;
)); ));
my @actual_ids = split("\n", $actual); my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids; my %actual_set = map { $_ => 1 } @actual_ids;
@@ -58,7 +60,7 @@ $node->start;
# Create table # Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;"); $node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int4);"); $node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c $type);");
$node->safe_psql("postgres", $node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc FROM generate_series(1, 10000) i;" "INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc FROM generate_series(1, 10000) i;"
); );
@@ -85,7 +87,7 @@ for my $i (0 .. $#queries)
{ {
my $res = $node->safe_psql("postgres", qq( my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off; 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 = '$cs[$i]' ORDER BY v <-> '$queries[$i]' LIMIT $limit;
)); ));
push(@expected, $res); push(@expected, $res);
} }
@@ -94,7 +96,7 @@ for my $i (0 .. $#queries)
test_recall(0.99, '<->'); test_recall(0.99, '<->');
# Test vacuum # Test vacuum
$node->safe_psql("postgres", "DELETE FROM tst WHERE c > 5;"); $node->safe_psql("postgres", "DELETE FROM tst WHERE c > '5';");
$node->safe_psql("postgres", "VACUUM tst;"); $node->safe_psql("postgres", "VACUUM tst;");
# Test columns # Test columns