liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-06-06 05:51:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into intvec

Browse Source
This commit is contained in:
Andrew Kane
2024-10-13 19:21:59 -07:00
parent 67de46487e 049972a4a3
commit a344ecb60d
33 changed files with 1609 additions and 572 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
test/expected/hnsw_vector.out
View File

@@ -99,6 +99,32 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <+> (SELECT NULL::vector)) t2
4
(1 row)
DROP TABLE t;
-- iterative
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops);
SET hnsw.iterative_search = strict_order;
SET hnsw.ef_search = 1;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,1,1]
[0,0,0]
(3 rows)
SET hnsw.iterative_search = relaxed_order;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,1,1]
[0,0,0]
(3 rows)
RESET hnsw.iterative_search;
RESET hnsw.ef_search;
DROP TABLE t;
-- unlogged
CREATE UNLOGGED TABLE t (val vector(3));
@@ -139,4 +165,21 @@ SET hnsw.ef_search = 0;
ERROR: 0 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000)
SET hnsw.ef_search = 1001;
ERROR: 1001 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000)
SHOW hnsw.iterative_search;
hnsw.iterative_search
-----------------------
off
(1 row)
SET hnsw.iterative_search = on;
ERROR: invalid value for parameter "hnsw.iterative_search": "on"
HINT: Available values: off, relaxed_order, strict_order.
SHOW hnsw.max_search_tuples;
hnsw.max_search_tuples
------------------------
-1
(1 row)
SET hnsw.max_search_tuples = -2;
ERROR: -2 is outside the valid range for parameter "hnsw.max_search_tuples" (-1 .. 2147483647)
DROP TABLE t;

61
test/expected/ivfflat_vector.out
View File

@@ -81,6 +81,44 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2
3
(1 row)
DROP TABLE t;
-- iterative
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 3);
SET ivfflat.iterative_search = relaxed_order;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,1,1]
[0,0,0]
(3 rows)
SET ivfflat.max_probes = 0;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
(1 row)
SET ivfflat.max_probes = 1;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
(1 row)
SET ivfflat.max_probes = 2;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,1,1]
(2 rows)
RESET ivfflat.iterative_search;
RESET ivfflat.max_probes;
DROP TABLE t;
-- unlogged
CREATE UNLOGGED TABLE t (val vector(3));
@@ -109,4 +147,27 @@ SHOW ivfflat.probes;
1
(1 row)
SET ivfflat.probes = 0;
ERROR: 0 is outside the valid range for parameter "ivfflat.probes" (1 .. 32768)
SET ivfflat.probes = 32769;
ERROR: 32769 is outside the valid range for parameter "ivfflat.probes" (1 .. 32768)
SHOW ivfflat.iterative_search;
ivfflat.iterative_search
--------------------------
off
(1 row)
SET ivfflat.iterative_search = on;
ERROR: invalid value for parameter "ivfflat.iterative_search": "on"
HINT: Available values: off, relaxed_order.
SHOW ivfflat.max_probes;
ivfflat.max_probes
--------------------
-1
(1 row)
SET ivfflat.max_probes = -2;
ERROR: -2 is outside the valid range for parameter "ivfflat.max_probes" (-1 .. 32768)
SET ivfflat.max_probes = 32769;
ERROR: 32769 is outside the valid range for parameter "ivfflat.max_probes" (-1 .. 32768)
DROP TABLE t;

25
test/sql/hnsw_vector.sql
View File

@@ -57,6 +57,23 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <+> (SELECT NULL::vector)) t2
DROP TABLE t;
-- iterative
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops);
SET hnsw.iterative_search = strict_order;
SET hnsw.ef_search = 1;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SET hnsw.iterative_search = relaxed_order;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
RESET hnsw.iterative_search;
RESET hnsw.ef_search;
DROP TABLE t;
-- unlogged
CREATE UNLOGGED TABLE t (val vector(3));
@@ -81,4 +98,12 @@ SHOW hnsw.ef_search;
SET hnsw.ef_search = 0;
SET hnsw.ef_search = 1001;
SHOW hnsw.iterative_search;
SET hnsw.iterative_search = on;
SHOW hnsw.max_search_tuples;
SET hnsw.max_search_tuples = -2;
DROP TABLE t;

34
test/sql/ivfflat_vector.sql
View File

@@ -44,6 +44,28 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2
DROP TABLE t;
-- iterative
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 3);
SET ivfflat.iterative_search = relaxed_order;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SET ivfflat.max_probes = 0;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SET ivfflat.max_probes = 1;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SET ivfflat.max_probes = 2;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
RESET ivfflat.iterative_search;
RESET ivfflat.max_probes;
DROP TABLE t;
-- unlogged
CREATE UNLOGGED TABLE t (val vector(3));
@@ -62,4 +84,16 @@ CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 32769);
SHOW ivfflat.probes;
SET ivfflat.probes = 0;
SET ivfflat.probes = 32769;
SHOW ivfflat.iterative_search;
SET ivfflat.iterative_search = on;
SHOW ivfflat.max_probes;
SET ivfflat.max_probes = -2;
SET ivfflat.max_probes = 32769;
DROP TABLE t;

17
test/t/002_ivfflat_vacuum.pl
View File

@@ -6,13 +6,7 @@ use Test::More;
my $dim = 3;
my @r = ();
for (1 .. $dim)
{
my $v = int(rand(1000)) + 1;
push(@r, "i % $v");
}
my $array_sql = join(", ", @r);
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
@@ -23,19 +17,20 @@ $node->start;
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
# Get size
my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
# Store values
$node->safe_psql("postgres", "CREATE TABLE tmp AS SELECT * FROM tst;");
# Delete all, vacuum, and insert same data
$node->safe_psql("postgres", "DELETE FROM tst;");
$node->safe_psql("postgres", "VACUUM tst;");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", "INSERT INTO tst SELECT * FROM tmp;");
# Check size
my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");

6
test/t/009_ivfflat_filtering.pl
View File

@@ -94,8 +94,7 @@ like($explain, qr/Seq Scan/);
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query';
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
like($explain, qr/Seq Scan/);
# Test attribute index
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
@@ -110,7 +109,6 @@ $node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING ivfflat (v v
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Use partial index
like($explain, qr/Index Scan using idx/);
like($explain, qr/Index Scan using partial_idx/);
done_testing();

26
test/t/017_hnsw_filtering.pl
View File

@@ -18,9 +18,13 @@ $node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int4, t text);");
$node->safe_psql("postgres", "CREATE TABLE cat (i int4 PRIMARY KEY, t text, b boolean);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc, 'test ' || i FROM generate_series(1, 10000) i;"
);
$node->safe_psql("postgres",
"INSERT INTO cat SELECT i, 'cat ' || i, i % 5 = 0 FROM generate_series(1, $nc) i;"
);
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
$node->safe_psql("postgres", "ANALYZE tst;");
@@ -37,8 +41,7 @@ my $c = int(rand() * $nc);
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
like($explain, qr/Seq Scan/);
# Test attribute filtering with few rows removed
$explain = $node->safe_psql("postgres", qq(
@@ -56,8 +59,7 @@ like($explain, qr/Index Scan using idx/);
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
like($explain, qr/Seq Scan/);
# Test attribute filtering with few rows removed like
$explain = $node->safe_psql("postgres", qq(
@@ -96,13 +98,25 @@ $explain = $node->safe_psql("postgres", qq(
));
like($explain, qr/Seq Scan/);
# Test join
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT cat.t FROM cat INNER JOIN tst ON cat.i = tst.c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test join with attribute filtering
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT cat.t FROM cat INNER JOIN tst ON cat.i = tst.c WHERE cat.b = 't' ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute index
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Use attribute index
like($explain, qr/Index Scan using idx/);
# Use attribute index
like($explain, qr/Bitmap Index Scan on attribute_idx/);
# Test partial index
$node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING hnsw (v vector_l2_ops) WHERE (c = $c);");

60
test/t/039_hnsw_cost.pl Normal file
View File

@@ -0,0 +1,60 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my @dims = (384, 1536);
my $limit = 10;
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
for my $dim (@dims)
{
my $array_sql = join(",", ('random()') x $dim);
# Create table and index
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 2000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
$node->safe_psql("postgres", "ANALYZE tst;");
# Generate query
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
my $query = "[" . join(",", @r) . "]";
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# 3x the rows are needed for distance filters
# since the planner uses DEFAULT_INEQ_SEL for the selectivity (should be 1)
# Recreate index for performance
$node->safe_psql("postgres", "DROP INDEX idx;");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(2001, 6000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
$node->safe_psql("postgres", "ANALYZE tst;");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
$node->safe_psql("postgres", "DROP TABLE tst;");
}
done_testing();

50
test/t/040_ivfflat_cost.pl Normal file
View File

@@ -0,0 +1,50 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my @dims = (384, 1536);
my $limit = 10;
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
for my $dim (@dims)
{
my $array_sql = join(",", ('random()') x $dim);
# Create table and index
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 5000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 5);");
$node->safe_psql("postgres", "ANALYZE tst;");
# Generate query
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
my $query = "[" . join(",", @r) . "]";
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
$node->safe_psql("postgres", "DROP TABLE tst;");
}
done_testing();

54
test/t/041_ivfflat_iterative_search.pl Normal file
View File

@@ -0,0 +1,54 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4 PRIMARY KEY, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = 10;
SET ivfflat.iterative_search = relaxed_order;
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
is($count, 10);
foreach ((30, 50, 70))
{
my $max_probes = $_;
my $expected = $max_probes / 10;
my $sum = 0;
for my $i (1 .. 20)
{
$count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = 10;
SET ivfflat.iterative_search = relaxed_order;
SET ivfflat.max_probes = $max_probes;
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
));
$sum += $count;
}
my $avg = $sum / 20;
cmp_ok($avg, '>', $expected - 2);
cmp_ok($avg, '<', $expected + 2);
}
done_testing();

125
test/t/042_ivfflat_iterative_search_recall.pl Normal file
View File

@@ -0,0 +1,125 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
my @cs = (100, 1000);
sub test_recall
{
my ($c, $probes, $min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
SET ivfflat.iterative_search = relaxed_order;
EXPLAIN ANALYZE SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan using idx on tst/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
SET ivfflat.iterative_search = relaxed_order;
SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my @expected_ids = split("\n", $expected[$i]);
my %expected_set = map { $_ => 1 } @expected_ids;
foreach (@actual_ids)
{
if (exists($expected_set{$_}))
{
$correct++;
}
}
$total += $limit;
}
cmp_ok($correct / $total, ">=", $min, "$operator $c");
}
# Initialize node
$node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
);
# Generate queries
for (1 .. 20)
{
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
push(@queries, "[$r1,$r2,$r3]");
}
# Check each index type
my @operators = ("<->", "<=>");
my @opclasses = ("vector_l2_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v $opclass);");
foreach (@cs)
{
my $c = $_;
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
WITH top AS (
SELECT v $operator '$_' AS distance FROM tst WHERE i % $c = 0 ORDER BY distance LIMIT $limit
)
SELECT i FROM tst WHERE (v $operator '$_') <= (SELECT MAX(distance) FROM top)
));
push(@expected, $res);
}
if ($c == 100)
{
test_recall($c, 1, 0.57, $operator);
test_recall($c, 10, 0.98, $operator);
}
else
{
if ($operator eq "<->")
{
test_recall($c, 1, 0.80, $operator);
}
else
{
test_recall($c, 1, 0.88, $operator);
}
}
}
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();

67
test/t/043_hnsw_iterative_search.pl Normal file
View File

@@ -0,0 +1,67 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4 PRIMARY KEY, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
SET max_parallel_maintenance_workers = 2;
CREATE INDEX ON tst USING hnsw (v vector_l2_ops)
));
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.iterative_search = relaxed_order;
SET work_mem = '8MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
is($count, 10);
foreach ((30000, 50000, 70000))
{
my $max_tuples = $_;
my $expected = $max_tuples / 10000;
my $sum = 0;
for my $i (1 .. 20)
{
$count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.iterative_search = relaxed_order;
SET hnsw.max_search_tuples = $max_tuples;
SET work_mem = '8MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
));
$sum += $count;
}
my $avg = $sum / 20;
cmp_ok($avg, '>', $expected - 2);
cmp_ok($avg, '<', $expected + 2);
}
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.iterative_search = relaxed_order;
SET client_min_messages = debug1;
SET work_mem = '2MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
like($stderr, qr/hnsw index scan exceeded work_mem after \d+ tuples/);
done_testing();

118
test/t/044_hnsw_iterative_search_recall.pl Normal file
View File

@@ -0,0 +1,118 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
my @cs = (50, 500);
sub test_recall
{
my ($c, $ef_search, $min, $operator, $mode) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SET hnsw.iterative_search = $mode;
EXPLAIN ANALYZE SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan using idx on tst/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SET hnsw.iterative_search = $mode;
SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my @expected_ids = split("\n", $expected[$i]);
my %expected_set = map { $_ => 1 } @expected_ids;
foreach (@actual_ids)
{
if (exists($expected_set{$_}))
{
$correct++;
}
}
$total += $limit;
}
cmp_ok($correct / $total, ">=", $min, "$operator $mode $c");
}
# Initialize node
$node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 50000) i;"
);
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
push(@queries, "[" . join(",", @r) . "]");
}
# Check each index type
my @operators = ("<->", "<=>");
my @opclasses = ("vector_l2_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
CREATE INDEX idx ON tst USING hnsw (v $opclass);
));
foreach (@cs)
{
my $c = $_;
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
WITH top AS (
SELECT v $operator '$_' AS distance FROM tst WHERE i % $c = 0 ORDER BY distance LIMIT $limit
)
SELECT i FROM tst WHERE (v $operator '$_') <= (SELECT MAX(distance) FROM top)
));
push(@expected, $res);
}
test_recall($c, 40, 0.99, $operator, "strict_order");
test_recall($c, 40, 0.99, $operator, "relaxed_order");
}
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();