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

Added support for bit vectors to HNSW

Browse Source
This commit is contained in:
Andrew Kane
2024-04-01 20:30:55 -07:00
parent 7ee9074a9c
commit 94a444f029
21 changed files with 541 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
test/expected/bit_functions.out Normal file
View File

@@ -0,0 +1,64 @@
SELECT hamming_distance(B'111', B'111');
hamming_distance
------------------
0
(1 row)
SELECT hamming_distance(B'111', B'110');
hamming_distance
------------------
1
(1 row)
SELECT hamming_distance(B'111', B'100');
hamming_distance
------------------
2
(1 row)
SELECT hamming_distance(B'111', B'000');
hamming_distance
------------------
3
(1 row)
SELECT hamming_distance(B'111', B'00');
ERROR: different bit lengths 3 and 2
SELECT jaccard_distance(B'1111', B'1111');
jaccard_distance
------------------
0
(1 row)
SELECT jaccard_distance(B'1111', B'1110');
jaccard_distance
------------------
0.25
(1 row)
SELECT jaccard_distance(B'1111', B'1100');
jaccard_distance
------------------
0.5
(1 row)
SELECT jaccard_distance(B'1111', B'1000');
jaccard_distance
------------------
0.75
(1 row)
SELECT jaccard_distance(B'1111', B'0000');
jaccard_distance
------------------
1
(1 row)
SELECT jaccard_distance(B'1100', B'1000');
jaccard_distance
------------------
0.5
(1 row)
SELECT jaccard_distance(B'1111', B'000');
ERROR: different bit lengths 4 and 3

12
test/expected/functions.out
View File

@@ -208,6 +208,18 @@ SELECT l1_distance('[3e38]'::vector, '[-3e38]');
Infinity
(1 row)
SELECT quantize_binary('[1,0,-1]');
quantize_binary
-----------------
100
(1 row)
SELECT quantize_binary('[0,0.1,-0.2,-0.3,0.4,0.5,0.6,-0.7,0.8,-0.9,1]');
quantize_binary
-----------------
01001110101
(1 row)
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
avg
-----------

21
test/expected/hnsw_hamming.out Normal file
View File

@@ -0,0 +1,21 @@
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 hnsw (val bit_hamming_ops);
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;

21
test/expected/hnsw_jaccard.out Normal file
View 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 hnsw (val bit_jaccard_ops);
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;

13
test/sql/bit_functions.sql Normal file
View File

@@ -0,0 +1,13 @@
SELECT hamming_distance(B'111', B'111');
SELECT hamming_distance(B'111', B'110');
SELECT hamming_distance(B'111', B'100');
SELECT hamming_distance(B'111', B'000');
SELECT hamming_distance(B'111', B'00');
SELECT jaccard_distance(B'1111', B'1111');
SELECT jaccard_distance(B'1111', B'1110');
SELECT jaccard_distance(B'1111', B'1100');
SELECT jaccard_distance(B'1111', B'1000');
SELECT jaccard_distance(B'1111', B'0000');
SELECT jaccard_distance(B'1100', B'1000');
SELECT jaccard_distance(B'1111', B'000');

3
test/sql/functions.sql
View File

@@ -48,6 +48,9 @@ SELECT l1_distance('[0,0]'::vector, '[0,1]');
SELECT l1_distance('[1,2]'::vector, '[3]');
SELECT l1_distance('[3e38]'::vector, '[-3e38]');
SELECT quantize_binary('[1,0,-1]');
SELECT quantize_binary('[0,0.1,-0.2,-0.3,0.4,0.5,0.6,-0.7,0.8,-0.9,1]');
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v;
SELECT avg(v) FROM unnest(ARRAY[]::vector[]) v;

12
test/sql/hnsw_hamming.sql Normal file
View File

@@ -0,0 +1,12 @@
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 hnsw (val bit_hamming_ops);
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;

12
test/sql/hnsw_jaccard.sql Normal file
View 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 hnsw (val bit_jaccard_ops);
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;

137
test/t/020_hnsw_bit_build_recall.pl Normal file
View File

@@ -0,0 +1,137 @@
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 ($min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = 100;
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v $operator $queries[0] LIMIT $limit;
));
like($explain, qr/Index Scan/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = 100;
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, 10000) 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)
{
# Handle ties
my $res = $node->safe_psql("postgres", qq(
WITH top AS (
SELECT v $operator $_ AS distance FROM tst ORDER BY v $operator $_ 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 hnsw (v $opclass);
));
# Test approximate results
my $min = $operator eq "<\%>" ? 0.96 : 0.99;
test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
# Build index in parallel in memory
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 hnsw (v $opclass);
));
is($ret, 0, $stderr);
like($stderr, qr/using \d+ parallel workers/);
# Test approximate results
test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
# Build index in parallel on disk
# Set parallel_workers on table to use workers with low maintenance_work_mem
($ret, $stdout, $stderr) = $node->psql("postgres", qq(
ALTER TABLE tst SET (parallel_workers = 2);
SET client_min_messages = DEBUG;
SET maintenance_work_mem = '4MB';
CREATE INDEX idx ON tst USING hnsw (v $opclass);
ALTER TABLE tst RESET (parallel_workers);
));
is($ret, 0, $stderr);
like($stderr, qr/using \d+ parallel workers/);
like($stderr, qr/hnsw graph no longer fits into maintenance_work_mem/);
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();