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
\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
DEFAULT FOR TYPE integer USING hnsw AS
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
CREATE OPERATOR CLASS vector_bigint_ops
DEFAULT FOR TYPE bigint USING hnsw AS
OPERATOR 2 = (bigint, bigint);
CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING hnsw AS
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 $array_sql = join(",", ('random()') x $dim);
my $nc = 50;
my $type = "int4"; # int4, int8, text, varchar
sub test_recall
{
my ($min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $cast = $type eq "int8" ? "::int8" : "";
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 = '$cs[0]'$cast ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond/);
@@ -29,7 +31,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 = '$cs[$i]'$cast ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
@@ -58,7 +60,7 @@ $node->start;
# Create table
$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",
"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(
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);
}
@@ -94,7 +96,7 @@ for my $i (0 .. $#queries)
test_recall(0.99, '<->');
# 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;");
# Test columns