mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-05 20:20:59 +08:00
Added support for bit to IVFFlat
This commit is contained in:
32
test/expected/ivfflat_bit_hamming.out
Normal file
32
test/expected/ivfflat_bit_hamming.out
Normal file
@@ -0,0 +1,32 @@
|
||||
SET enable_seqscan = off;
|
||||
CREATE TABLE t (val bit(3));
|
||||
INSERT INTO t (val) VALUES (B'000'), (B'100'), (B'111'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val bit_hamming_ops) WITH (lists = 1);
|
||||
INSERT INTO t (val) VALUES (B'110');
|
||||
SELECT * FROM t ORDER BY val <~> B'111';
|
||||
val
|
||||
-----
|
||||
111
|
||||
110
|
||||
100
|
||||
000
|
||||
(4 rows)
|
||||
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <~> (SELECT NULL::bit)) t2;
|
||||
count
|
||||
-------
|
||||
4
|
||||
(1 row)
|
||||
|
||||
DROP TABLE t;
|
||||
-- TODO move
|
||||
CREATE TABLE t (val varbit(3));
|
||||
CREATE INDEX ON t USING ivfflat (val bit_hamming_ops) WITH (lists = 1);
|
||||
ERROR: type not supported for ivfflat index
|
||||
CREATE INDEX ON t USING ivfflat ((val::bit(3)) bit_hamming_ops) WITH (lists = 1);
|
||||
NOTICE: ivfflat index created with little data
|
||||
DETAIL: This will cause low recall.
|
||||
HINT: Drop the index until the table has more data.
|
||||
CREATE INDEX ON t USING ivfflat ((val::bit(64001)) bit_hamming_ops) WITH (lists = 1);
|
||||
ERROR: column cannot have more than 64000 dimensions for ivfflat index
|
||||
DROP TABLE t;
|
||||
21
test/expected/ivfflat_bit_jaccard.out
Normal file
21
test/expected/ivfflat_bit_jaccard.out
Normal file
@@ -0,0 +1,21 @@
|
||||
SET enable_seqscan = off;
|
||||
CREATE TABLE t (val bit(4));
|
||||
INSERT INTO t (val) VALUES (B'0000'), (B'1100'), (B'1111'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val bit_jaccard_ops) WITH (lists = 1);
|
||||
INSERT INTO t (val) VALUES (B'1110');
|
||||
SELECT * FROM t ORDER BY val <%> B'1111';
|
||||
val
|
||||
------
|
||||
1111
|
||||
1110
|
||||
1100
|
||||
0000
|
||||
(4 rows)
|
||||
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <%> (SELECT NULL::bit)) t2;
|
||||
count
|
||||
-------
|
||||
4
|
||||
(1 row)
|
||||
|
||||
DROP TABLE t;
|
||||
19
test/sql/ivfflat_bit_hamming.sql
Normal file
19
test/sql/ivfflat_bit_hamming.sql
Normal file
@@ -0,0 +1,19 @@
|
||||
SET enable_seqscan = off;
|
||||
|
||||
CREATE TABLE t (val bit(3));
|
||||
INSERT INTO t (val) VALUES (B'000'), (B'100'), (B'111'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val bit_hamming_ops) WITH (lists = 1);
|
||||
|
||||
INSERT INTO t (val) VALUES (B'110');
|
||||
|
||||
SELECT * FROM t ORDER BY val <~> B'111';
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <~> (SELECT NULL::bit)) t2;
|
||||
|
||||
DROP TABLE t;
|
||||
|
||||
-- TODO move
|
||||
CREATE TABLE t (val varbit(3));
|
||||
CREATE INDEX ON t USING ivfflat (val bit_hamming_ops) WITH (lists = 1);
|
||||
CREATE INDEX ON t USING ivfflat ((val::bit(3)) bit_hamming_ops) WITH (lists = 1);
|
||||
CREATE INDEX ON t USING ivfflat ((val::bit(64001)) bit_hamming_ops) WITH (lists = 1);
|
||||
DROP TABLE t;
|
||||
12
test/sql/ivfflat_bit_jaccard.sql
Normal file
12
test/sql/ivfflat_bit_jaccard.sql
Normal file
@@ -0,0 +1,12 @@
|
||||
SET enable_seqscan = off;
|
||||
|
||||
CREATE TABLE t (val bit(4));
|
||||
INSERT INTO t (val) VALUES (B'0000'), (B'1100'), (B'1111'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val bit_jaccard_ops) WITH (lists = 1);
|
||||
|
||||
INSERT INTO t (val) VALUES (B'1110');
|
||||
|
||||
SELECT * FROM t ORDER BY val <%> B'1111';
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <%> (SELECT NULL::bit)) t2;
|
||||
|
||||
DROP TABLE t;
|
||||
128
test/t/035_ivfflat_bit_build_recall.pl
Normal file
128
test/t/035_ivfflat_bit_build_recall.pl
Normal file
@@ -0,0 +1,128 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
my $node;
|
||||
my @queries = ();
|
||||
my @expected;
|
||||
my $limit = 20;
|
||||
my $dim = 52;
|
||||
my $max = 2**$dim;
|
||||
|
||||
sub test_recall
|
||||
{
|
||||
my ($probes, $min, $operator) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET ivfflat.probes = $probes;
|
||||
EXPLAIN ANALYZE SELECT i FROM tst 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;
|
||||
SELECT i FROM tst 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);
|
||||
}
|
||||
|
||||
# Initialize node
|
||||
$node = get_new_node('node');
|
||||
$node->init;
|
||||
$node->start;
|
||||
|
||||
# Create table
|
||||
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
|
||||
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v bit($dim));");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, (random() * $max)::bigint::bit($dim) FROM generate_series(1, 100000) i;"
|
||||
);
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my $r = int(rand() * $max);
|
||||
push(@queries, "${r}::bigint::bit($dim)");
|
||||
}
|
||||
|
||||
# Check each index type
|
||||
my @operators = ("<~>", "<\%>");
|
||||
my @opclasses = ("bit_hamming_ops", "bit_jaccard_ops");
|
||||
|
||||
for my $i (0 .. $#operators)
|
||||
{
|
||||
my $operator = $operators[$i];
|
||||
my $opclass = $opclasses[$i];
|
||||
|
||||
# Get exact results
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
WITH top AS (
|
||||
SELECT v $operator $_ AS distance FROM tst ORDER BY distance LIMIT $limit
|
||||
)
|
||||
SELECT i FROM tst WHERE (v $operator $_) <= (SELECT MAX(distance) FROM top)
|
||||
));
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
# Build index serially
|
||||
$node->safe_psql("postgres", qq(
|
||||
SET max_parallel_maintenance_workers = 0;
|
||||
CREATE INDEX idx ON tst USING ivfflat (v $opclass);
|
||||
));
|
||||
|
||||
# Test approximate results
|
||||
test_recall(1, 0.10, $operator);
|
||||
test_recall(10, 0.55, $operator);
|
||||
|
||||
# Test probes equals lists
|
||||
test_recall(100, 1.00, $operator);
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
|
||||
# Build index in parallel
|
||||
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
|
||||
SET client_min_messages = DEBUG;
|
||||
SET min_parallel_table_scan_size = 1;
|
||||
CREATE INDEX idx ON tst USING ivfflat (v $opclass);
|
||||
));
|
||||
is($ret, 0, $stderr);
|
||||
like($stderr, qr/using \d+ parallel workers/);
|
||||
|
||||
# Test approximate results
|
||||
test_recall(1, 0.10, $operator);
|
||||
test_recall(10, 0.55, $operator);
|
||||
|
||||
# Test probes equals lists
|
||||
test_recall(100, 1.00, $operator);
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
}
|
||||
|
||||
done_testing();
|
||||
Reference in New Issue
Block a user