Compare commits

..

30 Commits

Author SHA1 Message Date
Andrew Kane
2b125a1956 Moved bit functions to separate file 2024-03-27 13:42:24 -07:00
Andrew Kane
fcd655d2a3 Improved test [skip ci] 2024-03-26 19:26:52 -07:00
Andrew Kane
97fe28940d Updated readme [skip ci] 2024-03-25 23:33:59 -07:00
Andrew Kane
23c5bf6ef6 Fixed CI [skip ci] 2024-03-25 23:22:03 -07:00
Andrew Kane
9ed4303a5e Fixed test 2024-03-25 23:16:38 -07:00
Andrew Kane
acd066463a Added Jaccard distance to test 2024-03-25 23:15:14 -07:00
Andrew Kane
be936075eb Improved test [skip ci] 2024-03-25 23:13:30 -07:00
Andrew Kane
a9959fede2 Improved code [skip ci] 2024-03-25 23:03:12 -07:00
Andrew Kane
02c4f4884c Added support for indexing Jaccard distance 2024-03-25 22:44:51 -07:00
Andrew Kane
791fc2436f Added jaccard_distance function 2024-03-25 22:35:53 -07:00
Andrew Kane
e7a7936bb2 Improved code [skip ci] 2024-03-25 19:23:30 -07:00
Andrew Kane
ce2ba65906 Fixed CI 2024-03-25 19:07:09 -07:00
Andrew Kane
023633a274 Added test for build recall 2024-03-25 19:02:37 -07:00
Andrew Kane
9baa051b5b Fixed CI 2024-03-25 17:20:35 -07:00
Andrew Kane
ac94ac7cf1 Updated changelog [skip ci] 2024-03-25 17:18:09 -07:00
Andrew Kane
8b819dfdc2 Moved bit code to separate files 2024-03-25 17:10:01 -07:00
Andrew Kane
d9ca850faf Updated max dimensions [skip ci] 2024-03-25 16:56:32 -07:00
Andrew Kane
131782999b Updated changelog [skip ci] 2024-03-25 16:47:25 -07:00
Andrew Kane
d57ef873c2 Improved test [skip ci] 2024-03-25 16:43:15 -07:00
Andrew Kane
30c86fb05a Improved test [skip ci] 2024-03-25 16:41:40 -07:00
Andrew Kane
833f379ebe Added quantize_binary function 2024-03-25 16:39:55 -07:00
Andrew Kane
709fc75ce0 Updated changelog [skip ci] 2024-03-25 16:11:52 -07:00
Andrew Kane
2bc959b3eb Updated readme [skip ci] 2024-03-25 15:53:44 -07:00
Andrew Kane
ec9e13b5fb Improved test [skip ci] 2024-03-25 15:50:38 -07:00
Andrew Kane
95e476d570 Fixed CI 2024-03-25 15:49:18 -07:00
Andrew Kane
6bc0c47a0a Updated comments [skip ci] 2024-03-25 15:43:05 -07:00
Andrew Kane
58eeefeef4 Handle nulls 2024-03-25 15:40:15 -07:00
Andrew Kane
263e684824 Started indexing 2024-03-25 15:25:04 -07:00
Andrew Kane
21dfed5719 Fixed CI 2024-03-25 14:24:19 -07:00
Andrew Kane
f3aec9fd03 Added hamming_distance function 2024-03-25 14:22:23 -07:00
20 changed files with 510 additions and 62 deletions

View File

@@ -1,6 +1,9 @@
## 0.7.0 (unreleased)
- Added `subvector` function
- Added support for binary vectors to HNSW
- Added `hamming_distance` function
- Added `jaccard_distance` function
- Added `quantize_binary` function
## 0.6.2 (2024-03-18)

View File

@@ -3,7 +3,7 @@ EXTVERSION = 0.6.2
MODULE_big = vector
DATA = $(wildcard sql/*--*.sql)
OBJS = src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/vector.o
OBJS = src/bitvector.o src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/vector.o
HEADERS = src/vector.h
TESTS = $(wildcard test/sql/*.sql)

View File

@@ -1,7 +1,7 @@
EXTENSION = vector
EXTVERSION = 0.6.2
OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj
OBJS = src\bitvector.obj src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj
HEADERS = src\vector.h
REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged

View File

@@ -5,7 +5,7 @@ Open-source vector similarity search for Postgres
Store your vectors with the rest of your data. Supports:
- exact and approximate nearest neighbor search
- L2 distance, inner product, and cosine distance
- L2 distance, inner product, cosine distance, and more
- any [language](#languages) with a Postgres client
Plus [ACID](https://en.wikipedia.org/wiki/ACID) compliance, point-in-time recovery, JOINs, and all of the other [great features](https://www.postgresql.org/about/) of Postgres
@@ -221,7 +221,19 @@ Cosine distance
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);
```
Vectors with up to 2,000 dimensions can be indexed.
Hamming distance - added in 0.7.0
```sql
CREATE INDEX ON items USING hnsw (embedding bit_hamming_ops);
```
Jaccard distance - added in 0.7.0
```sql
CREATE INDEX ON items USING hnsw (embedding bit_jaccard_ops);
```
Vectors with up to 2,000 dimensions can be indexed, or bit vectors with up to 64,000 dimensions.
### Index Options
@@ -706,7 +718,7 @@ cosine_distance(vector, vector) → double precision | cosine distance |
inner_product(vector, vector) → double precision | inner product |
l2_distance(vector, vector) → double precision | Euclidean distance |
l1_distance(vector, vector) → double precision | taxicab distance | 0.5.0
subvector(vector, integer, integer) → vector | subvector | 0.7.0 [unreleased]
quantize_binary(vector) → bit | quantize | 0.7.0
vector_dims(vector) → integer | number of dimensions |
vector_norm(vector) → double precision | Euclidean norm |
@@ -717,6 +729,20 @@ Function | Description | Added
avg(vector) → vector | average |
sum(vector) → vector | sum | 0.5.0
### Bit Operators
Operator | Description | Added
--- | --- | ---
<~> | Hamming distance | 0.7.0
<%> | Jaccard distance | 0.7.0
### Bit Functions
Function | Description | Added
--- | --- | ---
hamming_distance(bit, bit) → double precision | Hamming distance | 0.7.0
jaccard_distance(bit, bit) → double precision | Jaccard distance | 0.7.0
## Installation Notes - Linux and Mac
### Postgres Location

View File

@@ -1,5 +1,31 @@
-- 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 FUNCTION subvector(vector, int, int) RETURNS vector
CREATE FUNCTION quantize_binary(vector) RETURNS bit
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION hamming_distance(bit, bit) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION jaccard_distance(bit, bit) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR <~> (
LEFTARG = bit, RIGHTARG = bit, PROCEDURE = hamming_distance,
COMMUTATOR = '<~>'
);
CREATE OPERATOR <%> (
LEFTARG = bit, RIGHTARG = bit, PROCEDURE = jaccard_distance,
COMMUTATOR = '<%>'
);
CREATE OPERATOR CLASS bit_hamming_ops
FOR TYPE bit USING hnsw AS
OPERATOR 1 <~> (bit, bit) FOR ORDER BY float_ops,
FUNCTION 1 hamming_distance(bit, bit);
CREATE OPERATOR CLASS bit_jaccard_ops
FOR TYPE bit USING hnsw AS
OPERATOR 1 <%> (bit, bit) FOR ORDER BY float_ops,
FUNCTION 1 jaccard_distance(bit, bit);

View File

@@ -58,7 +58,7 @@ CREATE FUNCTION vector_sub(vector, vector) RETURNS vector
CREATE FUNCTION vector_mul(vector, vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION subvector(vector, int, int) RETURNS vector
CREATE FUNCTION quantize_binary(vector) RETURNS bit
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- private functions
@@ -290,3 +290,31 @@ CREATE OPERATOR CLASS vector_cosine_ops
OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector),
FUNCTION 2 vector_norm(vector);
-- bit functions
CREATE FUNCTION hamming_distance(bit, bit) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION jaccard_distance(bit, bit) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR <~> (
LEFTARG = bit, RIGHTARG = bit, PROCEDURE = hamming_distance,
COMMUTATOR = '<~>'
);
CREATE OPERATOR <%> (
LEFTARG = bit, RIGHTARG = bit, PROCEDURE = jaccard_distance,
COMMUTATOR = '<%>'
);
CREATE OPERATOR CLASS bit_hamming_ops
FOR TYPE bit USING hnsw AS
OPERATOR 1 <~> (bit, bit) FOR ORDER BY float_ops,
FUNCTION 1 hamming_distance(bit, bit);
CREATE OPERATOR CLASS bit_jaccard_ops
FOR TYPE bit USING hnsw AS
OPERATOR 1 <%> (bit, bit) FOR ORDER BY float_ops,
FUNCTION 1 jaccard_distance(bit, bit);

90
src/bitvector.c Normal file
View File

@@ -0,0 +1,90 @@
#include "postgres.h"
#include "bitvector.h"
#include "port/pg_bitutils.h"
#include "utils/varbit.h"
#if PG_VERSION_NUM >= 160000
#include "varatt.h"
#endif
/*
* Allocate and initialize a new bit vector
*/
VarBit *
InitBitVector(int dim)
{
VarBit *result;
int size;
size = VARBITTOTALLEN(dim);
result = (VarBit *) palloc0(size);
SET_VARSIZE(result, size);
VARBITLEN(result) = dim;
return result;
}
/*
* Ensure same number of bits
*/
static inline void
CheckBitLengths(uint32 aLen, uint32 bLen)
{
if (aLen != bLen)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("different bit lengths %u and %u", aLen, bLen)));
}
/*
* Get the Hamming distance between two bit strings
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(hamming_distance);
Datum
hamming_distance(PG_FUNCTION_ARGS)
{
VarBit *a = PG_GETARG_VARBIT_P(0);
VarBit *b = PG_GETARG_VARBIT_P(1);
unsigned char *ax = VARBITS(a);
unsigned char *bx = VARBITS(b);
uint64 distance = 0;
CheckBitLengths(VARBITLEN(a), VARBITLEN(b));
/* TODO Improve performance */
for (uint32 i = 0; i < VARBITBYTES(a); i++)
distance += pg_number_of_ones[ax[i] ^ bx[i]];
PG_RETURN_FLOAT8((double) distance);
}
/*
* Get the Jaccard distance between two bit strings
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(jaccard_distance);
Datum
jaccard_distance(PG_FUNCTION_ARGS)
{
VarBit *a = PG_GETARG_VARBIT_P(0);
VarBit *b = PG_GETARG_VARBIT_P(1);
unsigned char *ax = VARBITS(a);
unsigned char *bx = VARBITS(b);
uint64 ab = 0;
uint64 aa;
uint64 bb;
CheckBitLengths(VARBITLEN(a), VARBITLEN(b));
/* TODO Improve performance */
for (uint32 i = 0; i < VARBITBYTES(a); i++)
ab += pg_number_of_ones[ax[i] & bx[i]];
if (ab == 0)
PG_RETURN_FLOAT8(1);
aa = pg_popcount((char *) ax, VARBITBYTES(a));
bb = pg_popcount((char *) bx, VARBITBYTES(b));
PG_RETURN_FLOAT8(1 - (ab / ((double) (aa + bb - ab))));
}

8
src/bitvector.h Normal file
View File

@@ -0,0 +1,8 @@
#ifndef BITVECTOR_H
#define BITVECTOR_H
#include "utils/varbit.h"
VarBit *InitBitVector(int dim);
#endif

View File

@@ -44,6 +44,7 @@
#include "access/xact.h"
#include "access/xloginsert.h"
#include "catalog/index.h"
#include "catalog/pg_type_d.h"
#include "commands/progress.h"
#include "hnsw.h"
#include "miscadmin.h"
@@ -671,6 +672,12 @@ HnswSharedMemoryAlloc(Size size, void *state)
static void
InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, IndexInfo *indexInfo, ForkNumber forkNum)
{
int maxDimensions = HNSW_MAX_DIM;
Oid typid = TupleDescAttr(index->rd_att, 0)->atttypid;
if (typid == BITOID || typid == VARBITOID)
maxDimensions *= 32;
buildstate->heap = heap;
buildstate->index = index;
buildstate->indexInfo = indexInfo;
@@ -684,8 +691,8 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
if (buildstate->dimensions < 0)
elog(ERROR, "column does not have dimensions");
if (buildstate->dimensions > HNSW_MAX_DIM)
elog(ERROR, "column cannot have more than %d dimensions for hnsw index", HNSW_MAX_DIM);
if (buildstate->dimensions > maxDimensions)
elog(ERROR, "column cannot have more than %d dimensions for hnsw index", maxDimensions);
if (buildstate->efConstruction < 2 * buildstate->m)
elog(ERROR, "ef_construction must be greater than or equal to 2 * m");

View File

@@ -1,6 +1,8 @@
#include "postgres.h"
#include "access/relscan.h"
#include "bitvector.h"
#include "catalog/pg_type_d.h"
#include "hnsw.h"
#include "pgstat.h"
#include "storage/bufmgr.h"
@@ -73,7 +75,15 @@ GetScanValue(IndexScanDesc scan)
Datum value;
if (scan->orderByData->sk_flags & SK_ISNULL)
value = PointerGetDatum(InitVector(GetDimensions(scan->indexRelation)));
{
Oid typid = TupleDescAttr(scan->indexRelation->rd_att, 0)->atttypid;
int dimensions = GetDimensions(scan->indexRelation);
if (typid == BITOID || typid == VARBITOID)
value = PointerGetDatum(InitBitVector(dimensions));
else
value = PointerGetDatum(InitVector(dimensions));
}
else
{
value = scan->orderByData->sk_argument;

View File

@@ -2,6 +2,7 @@
#include <math.h>
#include "bitvector.h"
#include "catalog/pg_type.h"
#include "common/shortest_dec.h"
#include "fmgr.h"
@@ -861,36 +862,25 @@ vector_mul(PG_FUNCTION_ARGS)
}
/*
* Get a subset of a vector
* Quantize a vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(subvector);
PGDLLEXPORT PG_FUNCTION_INFO_V1(quantize_binary);
Datum
subvector(PG_FUNCTION_ARGS)
quantize_binary(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
int32 start = PG_GETARG_INT32(1);
int32 count = PG_GETARG_INT32(2);
int32 end = start + count;
float *ax = a->x;
Vector *result;
int dim;
VarBit *result = InitBitVector(a->dim);
unsigned char *rx = VARBITS(result);
if (start < 1)
start = 1;
/* TODO Improve */
for (int i = 0; i < a->dim; i++)
rx[i / 8] |= (ax[i] > 0) << (7 - (i % 8));
if (end > a->dim)
end = a->dim + 1;
dim = end - start;
CheckDim(dim);
result = InitVector(dim);
for (int i = 0; i < dim; i++)
result->x[i] = ax[start - 1 + i];
PG_RETURN_POINTER(result);
PG_RETURN_VARBIT_P(result);
}
/*
* Internal helper to compare vectors
*/

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

View File

@@ -208,34 +208,18 @@ SELECT l1_distance('[3e38]', '[-3e38]');
Infinity
(1 row)
SELECT subvector('[1,2,3,4,5]', 1, 3);
subvector
-----------
[1,2,3]
SELECT quantize_binary('[1,0,-1]');
quantize_binary
-----------------
100
(1 row)
SELECT subvector('[1,2,3,4,5]', 3, 2);
subvector
-----------
[3,4]
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 subvector('[1,2,3,4,5]', -1, 3);
subvector
-----------
[1]
(1 row)
SELECT subvector('[1,2,3,4,5]', 3, 9);
subvector
-----------
[3,4,5]
(1 row)
SELECT subvector('[1,2,3,4,5]', 1, 0);
ERROR: vector must have at least 1 dimension
SELECT subvector('[1,2,3,4,5]', -1, 2);
ERROR: vector must have at least 1 dimension
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
avg
-----------

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;

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;

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');

View File

@@ -48,12 +48,8 @@ SELECT l1_distance('[0,0]', '[0,1]');
SELECT l1_distance('[1,2]', '[3]');
SELECT l1_distance('[3e38]', '[-3e38]');
SELECT subvector('[1,2,3,4,5]', 1, 3);
SELECT subvector('[1,2,3,4,5]', 3, 2);
SELECT subvector('[1,2,3,4,5]', -1, 3);
SELECT subvector('[1,2,3,4,5]', 3, 9);
SELECT subvector('[1,2,3,4,5]', 1, 0);
SELECT subvector('[1,2,3,4,5]', -1, 2);
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;

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;

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();