Compare commits

..

56 Commits

Author SHA1 Message Date
Andrew Kane
8d5631da72 Use integer[] for consistency [skip ci] 2024-10-13 21:15:46 -07:00
Andrew Kane
54d91e1116 Fixed CI 2024-10-13 19:52:17 -07:00
Andrew Kane
f298a72ada Added casting [skip ci] 2024-10-13 19:46:16 -07:00
Andrew Kane
f5b225530f Removed unneeded cast [skip ci] 2024-10-13 19:45:04 -07:00
Andrew Kane
05944f5a6f Improved distance functions [skip ci] 2024-10-13 19:33:43 -07:00
Andrew Kane
a344ecb60d Merge branch 'master' into intvec 2024-10-13 19:21:59 -07:00
Andrew Kane
67de46487e Fixed flaky test [skip ci] 2024-10-13 19:20:02 -07:00
Andrew Kane
f52abfc26c Added more TAP tests [skip ci] 2024-10-13 19:14:09 -07:00
Andrew Kane
1066b4be60 Added intvec to distance functions test [skip ci] 2024-10-13 19:01:04 -07:00
Andrew Kane
6221779f74 Added casting to int[] [skip ci] 2024-10-13 18:55:12 -07:00
Andrew Kane
4765df1c50 Updated upgrade script [skip ci] 2024-10-13 18:46:45 -07:00
Andrew Kane
9514a7ffcd Added vector_dims [skip ci] 2024-10-13 18:46:09 -07:00
Andrew Kane
98deaa7407 Added intvec to inputs test [skip ci] 2024-10-13 18:41:57 -07:00
Andrew Kane
faa25d31f7 Added cast tests [skip ci] 2024-10-13 18:39:46 -07:00
Andrew Kane
39af767617 Fixed casting [skip ci] 2024-10-13 18:36:25 -07:00
Andrew Kane
3424f46298 Updated upgrade script [skip ci] 2024-10-13 18:27:27 -07:00
Andrew Kane
ab12f1622d Added more intvec functions [skip ci] 2024-10-13 18:14:31 -07:00
Andrew Kane
2fb71a6183 Updated readme [skip ci] 2024-10-13 18:02:43 -07:00
Andrew Kane
049972a4a3 Improved test output [skip ci] 2024-10-13 17:22:49 -07:00
Andrew Kane
61027645e9 Improved test output [skip ci] 2024-10-13 17:21:38 -07:00
Andrew Kane
a41b327b33 Speed up test [skip ci] 2024-10-13 17:12:12 -07:00
Andrew Kane
7f735ebd9b Added test for strict order [skip ci] 2024-10-13 17:04:03 -07:00
Andrew Kane
02b01e1ca9 Show tuples with memory usage [skip ci] 2024-10-12 22:05:13 -07:00
Andrew Kane
388e42f6e6 Fixed flaky test [skip ci] 2024-10-11 15:48:19 -07:00
Andrew Kane
bf379eed86 Use a memory context for IVFFlat index scans 2024-10-11 15:46:38 -07:00
Andrew Kane
e1bc929429 Simplify lists for IvfflatScanOpaque [skip ci] 2024-10-11 15:29:23 -07:00
Andrew Kane
38285aacc7 Revert "Updated IVFFlat to support multiple attributes (not enabled yet)"
This reverts commit 772ab69de6.
2024-10-11 15:01:54 -07:00
Andrew Kane
a2408e60fa Revert "Added IndexTuple to HNSW elements (first step to support multiple attributes)"
This reverts commit 53a8734bac.
2024-10-11 14:57:57 -07:00
Andrew Kane
53a8734bac Added IndexTuple to HNSW elements (first step to support multiple attributes) 2024-10-11 14:12:01 -07:00
Andrew Kane
7484625227 Added comments [skip ci] 2024-10-11 11:59:36 -07:00
Andrew Kane
d1ebb8db73 Use -1 for no limit for ivfflat.max_probes [skip ci] 2024-10-11 11:43:32 -07:00
Andrew Kane
42af8aa1d1 Updated GUC descriptions [skip ci] 2024-10-11 11:26:27 -07:00
Andrew Kane
9d15a76b60 Improved enum naming [skip ci] 2024-10-11 11:20:36 -07:00
Andrew Kane
a3a20f9816 Simplified GUC names [skip ci] 2024-10-11 11:18:01 -07:00
Andrew Kane
b26a21b848 Added regression tests for iterative search [skip ci] 2024-10-11 11:07:11 -07:00
Andrew Kane
2dc392ed6c Updated GUC names [skip ci] 2024-10-10 23:50:11 -07:00
Andrew Kane
960d2848cb Updated comment [skip ci] 2024-10-10 21:02:33 -07:00
Andrew Kane
8e88b481a6 Use a lower max allocation size than default to allow scanning more tuples for iterative search before exceeding work_mem [skip ci] 2024-10-10 20:57:57 -07:00
Andrew Kane
124018b8dd Added HnswInitSearchCandidate function 2024-10-10 19:30:47 -07:00
Andrew Kane
35b252a3e3 Switched to macos-13 on CI [skip ci] 2024-10-10 18:33:35 -07:00
Andrew Kane
2832e746f0 Use get_float8_infinity instead of INFINITY
Co-authored-by: "Jonathan S. Katz" <jkatz@amazon.com>
2024-10-10 18:16:39 -07:00
Andrew Kane
961cb17d80 Added iterative search for HNSW [skip ci] 2024-10-10 18:14:39 -07:00
Andrew Kane
c91ed7b2c3 Added iterative search for IVFFlat [skip ci] 2024-10-10 18:12:27 -07:00
Andrew Kane
48fe70c219 Fixed sortdesc for multiple attributes [skip ci] 2024-10-10 15:22:23 -07:00
Andrew Kane
29908405ab List assignment does not change with scaling [skip ci] 2024-10-10 15:16:08 -07:00
Andrew Kane
08d0340655 Improved IVFFlat vacuum test [skip ci] 2024-10-10 14:24:26 -07:00
Andrew Kane
7d2eb49c2a Fixed list assignment for cosine distance and non-normalized vectors for IVFFlat 2024-10-10 12:54:13 -07:00
Andrew Kane
772ab69de6 Updated IVFFlat to support multiple attributes (not enabled yet) 2024-10-10 12:34:03 -07:00
Andrew Kane
e13e9a9614 Improved scoping [skip ci] 2024-10-10 12:19:08 -07:00
Andrew Kane
e2fab306ac Added tupdesc to buildstate [skip ci] 2024-10-10 12:18:41 -07:00
Andrew Kane
edc2126a4a Renamed tupdesc to sortdesc [skip ci] 2024-10-10 12:16:57 -07:00
Andrew Kane
67b4ba5a14 Updated readme [skip ci] 2024-09-27 21:30:53 -07:00
Andrew Kane
8597eb79aa Restored isnan check [skip ci] 2024-09-23 13:44:29 -07:00
Andrew Kane
ffca8e720c Exclude zero vectors for cosine distance to be consistent with other types [skip ci] 2024-09-23 13:35:36 -07:00
Andrew Kane
0296a08ccf Added L1 distance to test [skip ci] 2024-09-23 13:25:22 -07:00
Andrew Kane
b067b5f073 Added intvec type 2024-09-23 13:20:49 -07:00
43 changed files with 2802 additions and 203 deletions

View File

@@ -49,7 +49,7 @@ jobs:
- postgres: 16 - postgres: 16
os: macos-14 os: macos-14
- postgres: 14 - postgres: 14
os: macos-12 os: macos-13
steps: steps:
- uses: actions/checkout@v4 - uses: actions/checkout@v4
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1

View File

@@ -1,6 +1,7 @@
## 0.8.0 (unreleased) ## 0.8.0 (unreleased)
- Added support for iterative index scans - Added support for iterative index scans
- Added `intvec` type
- Added casts for arrays to `sparsevec` - Added casts for arrays to `sparsevec`
- Improved cost estimation - Improved cost estimation
- Improved performance of HNSW inserts and on-disk index builds - Improved performance of HNSW inserts and on-disk index builds

View File

@@ -4,8 +4,8 @@ EXTVERSION = 0.7.4
MODULE_big = vector MODULE_big = vector
DATA = $(wildcard sql/*--*--*.sql) DATA = $(wildcard sql/*--*--*.sql)
DATA_built = sql/$(EXTENSION)--$(EXTVERSION).sql DATA_built = sql/$(EXTENSION)--$(EXTVERSION).sql
OBJS = src/bitutils.o src/bitvec.o src/halfutils.o src/halfvec.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/sparsevec.o src/vector.o OBJS = src/bitutils.o src/bitvec.o src/halfutils.o src/halfvec.o src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/intvec.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/sparsevec.o src/vector.o
HEADERS = src/halfvec.h src/sparsevec.h src/vector.h HEADERS = src/halfvec.h src/intvec.h src/sparsevec.h src/vector.h
TESTS = $(wildcard test/sql/*.sql) TESTS = $(wildcard test/sql/*.sql)
REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS)) REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS))

View File

@@ -2,8 +2,8 @@ EXTENSION = vector
EXTVERSION = 0.7.4 EXTVERSION = 0.7.4
DATA_built = sql\$(EXTENSION)--$(EXTVERSION).sql DATA_built = sql\$(EXTENSION)--$(EXTVERSION).sql
OBJS = src\bitutils.obj src\bitvec.obj src\halfutils.obj src\halfvec.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\sparsevec.obj src\vector.obj OBJS = src\bitutils.obj src\bitvec.obj src\halfutils.obj src\halfvec.obj src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\intvec.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\sparsevec.obj src\vector.obj
HEADERS = src\halfvec.h src\sparsevec.h src\vector.h HEADERS = src\halfvec.h src\intvec.h src\sparsevec.h src\vector.h
REGRESS = bit btree cast copy halfvec hnsw_bit hnsw_halfvec hnsw_sparsevec hnsw_vector ivfflat_bit ivfflat_halfvec ivfflat_vector sparsevec vector_type REGRESS = bit btree cast copy halfvec hnsw_bit hnsw_halfvec hnsw_sparsevec hnsw_vector ivfflat_bit ivfflat_halfvec ivfflat_vector sparsevec vector_type
REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION) REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION)

108
README.md
View File

@@ -223,7 +223,7 @@ L2 distance
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops); CREATE INDEX ON items USING hnsw (embedding vector_l2_ops);
``` ```
Note: Use `halfvec_l2_ops` for `halfvec` and `sparsevec_l2_ops` for `sparsevec` (and similar with the other distance functions) Note: Use `halfvec_l2_ops` for `halfvec`, `intvec_l2_ops` for `intvec`, and `sparsevec_l2_ops` for `sparsevec` (and similar with the other distance functions)
Inner product Inner product
@@ -259,6 +259,7 @@ Supported types are:
- `vector` - up to 2,000 dimensions - `vector` - up to 2,000 dimensions
- `halfvec` - up to 4,000 dimensions (added in 0.7.0) - `halfvec` - up to 4,000 dimensions (added in 0.7.0)
- `intvec` - up to 8,000 dimensions (added in 0.8.0)
- `bit` - up to 64,000 dimensions (added in 0.7.0) - `bit` - up to 64,000 dimensions (added in 0.7.0)
- `sparsevec` - up to 1,000 non-zero elements (added in 0.7.0) - `sparsevec` - up to 1,000 non-zero elements (added in 0.7.0)
@@ -451,77 +452,6 @@ Use [partitioning](https://www.postgresql.org/docs/current/ddl-partitioning.html
CREATE TABLE items (embedding vector(3), category_id int) PARTITION BY LIST(category_id); CREATE TABLE items (embedding vector(3), category_id int) PARTITION BY LIST(category_id);
``` ```
## Iterative Search [unreleased]
*Added in 0.8.0*
With approximate indexes, you can end up with less results than expected due to filtering conditions in the query.
Starting with 0.8.0, you can enable iterative search. If too few results from the initial index scan match the query filters, it will resume scanning until enough results are found. This can significantly improve recall (at the cost of speed).
```tsql
SET hnsw.streaming = on;
-- or
SET ivfflat.streaming = on;
```
However, there are some important caveats.
### Iterative Caveats
With iterative search, its possible for rows to be slightly out of order by distance. For strict ordering, use:
```sql
WITH approx_order AS MATERIALIZED (
SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5
) SELECT * FROM approx_order ORDER BY distance;
```
For distance filters, use a CTE and place the filter outside it.
```sql
WITH approx_order AS MATERIALIZED (
SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5
) SELECT * FROM approx_order WHERE distance < 0.1 ORDER BY distance;
```
### Iterative Options
Since scanning a large portion of the index is expensive, there are options to control when the scan ends.
#### HNSW
Specify the max number of additional tuples visited
```sql
SET hnsw.ef_stream = 10000;
```
The scan will also end if reaches `work_mem`. You can see when this happens by enabling debug messages.
```sql
SET client_min_messages = debug1;
```
```text
DEBUG: hnsw index scan exceeded work_mem after 10000 tuples
HINT: Increase work_mem to scan more tuples.
```
If the server has enough memory, you can adjust this with:
```sql
SET work_mem = '8MB';
```
#### IVFFlat
Specify the max number of probes
```sql
SET ivfflat.max_probes = 100;
```
## Half-Precision Vectors ## Half-Precision Vectors
*Added in 0.7.0* *Added in 0.7.0*
@@ -548,6 +478,16 @@ Get the nearest neighbors
SELECT * FROM items ORDER BY embedding::halfvec(3) <-> '[1,2,3]' LIMIT 5; SELECT * FROM items ORDER BY embedding::halfvec(3) <-> '[1,2,3]' LIMIT 5;
``` ```
## Integer Vectors
*Added in 0.8.0*
Use the `intvec` type to store 8-bit integer vectors
```sql
CREATE TABLE items (id bigserial PRIMARY KEY, embedding intvec(3));
```
## Binary Vectors ## Binary Vectors
Use the `bit` type to store binary vectors ([example](https://github.com/pgvector/pgvector-python/blob/master/examples/imagehash/example.py)) Use the `bit` type to store binary vectors ([example](https://github.com/pgvector/pgvector-python/blob/master/examples/imagehash/example.py))
@@ -1011,6 +951,30 @@ Function | Description | Added
avg(halfvec) → halfvec | average | 0.7.0 avg(halfvec) → halfvec | average | 0.7.0
sum(halfvec) → halfvec | sum | 0.7.0 sum(halfvec) → halfvec | sum | 0.7.0
### Intvec Type
Each int vector takes `dimensions + 8` bytes of storage. Each element is a single byte signed integer. Int vectors can have up to 16,000 dimensions.
### Intvec Operators
Operator | Description | Added
--- | --- | ---
<-> | Euclidean distance | 0.8.0
<#> | negative inner product | 0.8.0
<=> | cosine distance | 0.8.0
<+> | taxicab distance | 0.8.0
### Intvec Functions
Function | Description | Added
--- | --- | ---
cosine_distance(intvec, intvec) → double precision | cosine distance | 0.8.0
inner_product(intvec, intvec) → double precision | inner product | 0.8.0
l1_distance(intvec, intvec) → double precision | taxicab distance | 0.8.0
l2_distance(intvec, intvec) → double precision | Euclidean distance | 0.8.0
l2_norm(intvec) → double precision | Euclidean norm | 0.8.0
vector_dims(intvec) → integer | number of dimensions | 0.8.0
### Bit Type ### Bit Type
Each bit vector takes `dimensions / 8 + 8` bytes of storage. See the [Postgres docs](https://www.postgresql.org/docs/current/datatype-bit.html) for more info. Each bit vector takes `dimensions / 8 + 8` bytes of storage. See the [Postgres docs](https://www.postgresql.org/docs/current/datatype-bit.html) for more info.

View File

@@ -1,6 +1,188 @@
-- complain if script is sourced in psql, rather than via CREATE EXTENSION -- complain if script is sourced in psql, rather than via CREATE EXTENSION
\echo Use "ALTER EXTENSION vector UPDATE TO '0.8.0'" to load this file. \quit \echo Use "ALTER EXTENSION vector UPDATE TO '0.8.0'" to load this file. \quit
CREATE FUNCTION hnsw_intvec_support(internal) RETURNS internal
AS 'MODULE_PATHNAME' LANGUAGE C;
CREATE TYPE intvec;
CREATE FUNCTION intvec_in(cstring, oid, integer) RETURNS intvec
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_out(intvec) RETURNS cstring
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_typmod_in(cstring[]) RETURNS integer
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_recv(internal, oid, integer) RETURNS intvec
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_send(intvec) RETURNS bytea
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE TYPE intvec (
INPUT = intvec_in,
OUTPUT = intvec_out,
TYPMOD_IN = intvec_typmod_in,
RECEIVE = intvec_recv,
SEND = intvec_send,
STORAGE = external
);
CREATE FUNCTION l2_distance(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION inner_product(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_inner_product' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION cosine_distance(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_cosine_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION l1_distance(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_l1_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_dims(intvec) RETURNS integer
AS 'MODULE_PATHNAME', 'intvec_vector_dims' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION l2_norm(intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_l2_norm' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_lt(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_le(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_eq(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_ne(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_ge(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_gt(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_cmp(intvec, intvec) RETURNS int4
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_l2_squared_distance(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_negative_inner_product(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec(intvec, integer, boolean) RETURNS intvec
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION array_to_intvec(integer[], integer, boolean) RETURNS intvec
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_to_integer(intvec, integer, boolean) RETURNS integer[]
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE CAST (intvec AS intvec)
WITH FUNCTION intvec(intvec, integer, boolean) AS IMPLICIT;
CREATE CAST (intvec AS integer[])
WITH FUNCTION intvec_to_integer(intvec, integer, boolean) AS ASSIGNMENT;
CREATE CAST (integer[] AS intvec)
WITH FUNCTION array_to_intvec(integer[], integer, boolean) AS ASSIGNMENT;
CREATE OPERATOR <-> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = l2_distance,
COMMUTATOR = '<->'
);
CREATE OPERATOR <#> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_negative_inner_product,
COMMUTATOR = '<#>'
);
CREATE OPERATOR <=> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = cosine_distance,
COMMUTATOR = '<=>'
);
CREATE OPERATOR <+> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = l1_distance,
COMMUTATOR = '<+>'
);
CREATE OPERATOR < (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_lt,
COMMUTATOR = > , NEGATOR = >= ,
RESTRICT = scalarltsel, JOIN = scalarltjoinsel
);
CREATE OPERATOR <= (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_le,
COMMUTATOR = >= , NEGATOR = > ,
RESTRICT = scalarlesel, JOIN = scalarlejoinsel
);
CREATE OPERATOR = (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_eq,
COMMUTATOR = = , NEGATOR = <> ,
RESTRICT = eqsel, JOIN = eqjoinsel
);
CREATE OPERATOR <> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_ne,
COMMUTATOR = <> , NEGATOR = = ,
RESTRICT = eqsel, JOIN = eqjoinsel
);
CREATE OPERATOR >= (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_ge,
COMMUTATOR = <= , NEGATOR = < ,
RESTRICT = scalargesel, JOIN = scalargejoinsel
);
CREATE OPERATOR > (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_gt,
COMMUTATOR = < , NEGATOR = <= ,
RESTRICT = scalargtsel, JOIN = scalargtjoinsel
);
CREATE OPERATOR CLASS intvec_ops
DEFAULT FOR TYPE intvec USING btree AS
OPERATOR 1 < ,
OPERATOR 2 <= ,
OPERATOR 3 = ,
OPERATOR 4 >= ,
OPERATOR 5 > ,
FUNCTION 1 intvec_cmp(intvec, intvec);
CREATE OPERATOR CLASS intvec_l2_ops
FOR TYPE intvec USING hnsw AS
OPERATOR 1 <-> (intvec, intvec) FOR ORDER BY float_ops,
FUNCTION 1 intvec_l2_squared_distance(intvec, intvec),
FUNCTION 3 hnsw_intvec_support(internal);
CREATE OPERATOR CLASS intvec_ip_ops
FOR TYPE intvec USING hnsw AS
OPERATOR 1 <#> (intvec, intvec) FOR ORDER BY float_ops,
FUNCTION 1 intvec_negative_inner_product(intvec, intvec),
FUNCTION 3 hnsw_intvec_support(internal);
CREATE OPERATOR CLASS intvec_cosine_ops
FOR TYPE intvec USING hnsw AS
OPERATOR 1 <=> (intvec, intvec) FOR ORDER BY float_ops,
FUNCTION 1 cosine_distance(intvec, intvec),
FUNCTION 2 l2_norm(intvec),
FUNCTION 3 hnsw_intvec_support(internal);
CREATE OPERATOR CLASS intvec_l1_ops
FOR TYPE intvec USING hnsw AS
OPERATOR 1 <+> (intvec, intvec) FOR ORDER BY float_ops,
FUNCTION 1 l1_distance(intvec, intvec),
FUNCTION 3 hnsw_intvec_support(internal);
CREATE FUNCTION array_to_sparsevec(integer[], integer, boolean) RETURNS sparsevec CREATE FUNCTION array_to_sparsevec(integer[], integer, boolean) RETURNS sparsevec
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

View File

@@ -272,6 +272,9 @@ CREATE FUNCTION ivfflat_bit_support(internal) RETURNS internal
CREATE FUNCTION hnsw_halfvec_support(internal) RETURNS internal CREATE FUNCTION hnsw_halfvec_support(internal) RETURNS internal
AS 'MODULE_PATHNAME' LANGUAGE C; AS 'MODULE_PATHNAME' LANGUAGE C;
CREATE FUNCTION hnsw_intvec_support(internal) RETURNS internal
AS 'MODULE_PATHNAME' LANGUAGE C;
CREATE FUNCTION hnsw_bit_support(internal) RETURNS internal CREATE FUNCTION hnsw_bit_support(internal) RETURNS internal
AS 'MODULE_PATHNAME' LANGUAGE C; AS 'MODULE_PATHNAME' LANGUAGE C;
@@ -647,6 +650,199 @@ CREATE OPERATOR CLASS halfvec_l1_ops
FUNCTION 1 l1_distance(halfvec, halfvec), FUNCTION 1 l1_distance(halfvec, halfvec),
FUNCTION 3 hnsw_halfvec_support(internal); FUNCTION 3 hnsw_halfvec_support(internal);
-- intvec type
CREATE TYPE intvec;
CREATE FUNCTION intvec_in(cstring, oid, integer) RETURNS intvec
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_out(intvec) RETURNS cstring
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_typmod_in(cstring[]) RETURNS integer
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_recv(internal, oid, integer) RETURNS intvec
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_send(intvec) RETURNS bytea
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE TYPE intvec (
INPUT = intvec_in,
OUTPUT = intvec_out,
TYPMOD_IN = intvec_typmod_in,
RECEIVE = intvec_recv,
SEND = intvec_send,
STORAGE = external
);
-- intvec functions
CREATE FUNCTION l2_distance(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION inner_product(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_inner_product' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION cosine_distance(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_cosine_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION l1_distance(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_l1_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_dims(intvec) RETURNS integer
AS 'MODULE_PATHNAME', 'intvec_vector_dims' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION l2_norm(intvec) RETURNS float8
AS 'MODULE_PATHNAME', 'intvec_l2_norm' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- intvec private functions
CREATE FUNCTION intvec_lt(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_le(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_eq(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_ne(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_ge(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_gt(intvec, intvec) RETURNS bool
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_cmp(intvec, intvec) RETURNS int4
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_l2_squared_distance(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_negative_inner_product(intvec, intvec) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- intvec cast functions
CREATE FUNCTION intvec(intvec, integer, boolean) RETURNS intvec
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION array_to_intvec(integer[], integer, boolean) RETURNS intvec
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION intvec_to_integer(intvec, integer, boolean) RETURNS integer[]
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- intvec casts
CREATE CAST (intvec AS intvec)
WITH FUNCTION intvec(intvec, integer, boolean) AS IMPLICIT;
CREATE CAST (intvec AS integer[])
WITH FUNCTION intvec_to_integer(intvec, integer, boolean) AS ASSIGNMENT;
CREATE CAST (integer[] AS intvec)
WITH FUNCTION array_to_intvec(integer[], integer, boolean) AS ASSIGNMENT;
-- intvec operators
CREATE OPERATOR <-> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = l2_distance,
COMMUTATOR = '<->'
);
CREATE OPERATOR <#> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_negative_inner_product,
COMMUTATOR = '<#>'
);
CREATE OPERATOR <=> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = cosine_distance,
COMMUTATOR = '<=>'
);
CREATE OPERATOR <+> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = l1_distance,
COMMUTATOR = '<+>'
);
CREATE OPERATOR < (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_lt,
COMMUTATOR = > , NEGATOR = >= ,
RESTRICT = scalarltsel, JOIN = scalarltjoinsel
);
CREATE OPERATOR <= (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_le,
COMMUTATOR = >= , NEGATOR = > ,
RESTRICT = scalarlesel, JOIN = scalarlejoinsel
);
CREATE OPERATOR = (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_eq,
COMMUTATOR = = , NEGATOR = <> ,
RESTRICT = eqsel, JOIN = eqjoinsel
);
CREATE OPERATOR <> (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_ne,
COMMUTATOR = <> , NEGATOR = = ,
RESTRICT = eqsel, JOIN = eqjoinsel
);
CREATE OPERATOR >= (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_ge,
COMMUTATOR = <= , NEGATOR = < ,
RESTRICT = scalargesel, JOIN = scalargejoinsel
);
CREATE OPERATOR > (
LEFTARG = intvec, RIGHTARG = intvec, PROCEDURE = intvec_gt,
COMMUTATOR = < , NEGATOR = <= ,
RESTRICT = scalargtsel, JOIN = scalargtjoinsel
);
-- intvec opclasses
CREATE OPERATOR CLASS intvec_ops
DEFAULT FOR TYPE intvec USING btree AS
OPERATOR 1 < ,
OPERATOR 2 <= ,
OPERATOR 3 = ,
OPERATOR 4 >= ,
OPERATOR 5 > ,
FUNCTION 1 intvec_cmp(intvec, intvec);
CREATE OPERATOR CLASS intvec_l2_ops
FOR TYPE intvec USING hnsw AS
OPERATOR 1 <-> (intvec, intvec) FOR ORDER BY float_ops,
FUNCTION 1 intvec_l2_squared_distance(intvec, intvec),
FUNCTION 3 hnsw_intvec_support(internal);
CREATE OPERATOR CLASS intvec_ip_ops
FOR TYPE intvec USING hnsw AS
OPERATOR 1 <#> (intvec, intvec) FOR ORDER BY float_ops,
FUNCTION 1 intvec_negative_inner_product(intvec, intvec),
FUNCTION 3 hnsw_intvec_support(internal);
CREATE OPERATOR CLASS intvec_cosine_ops
FOR TYPE intvec USING hnsw AS
OPERATOR 1 <=> (intvec, intvec) FOR ORDER BY float_ops,
FUNCTION 1 cosine_distance(intvec, intvec),
FUNCTION 2 l2_norm(intvec),
FUNCTION 3 hnsw_intvec_support(internal);
CREATE OPERATOR CLASS intvec_l1_ops
FOR TYPE intvec USING hnsw AS
OPERATOR 1 <+> (intvec, intvec) FOR ORDER BY float_ops,
FUNCTION 1 l1_distance(intvec, intvec),
FUNCTION 3 hnsw_intvec_support(internal);
-- bit functions -- bit functions
CREATE FUNCTION hamming_distance(bit, bit) RETURNS float8 CREATE FUNCTION hamming_distance(bit, bit) RETURNS float8

View File

@@ -20,15 +20,13 @@
static const struct config_enum_entry hnsw_iterative_search_options[] = { static const struct config_enum_entry hnsw_iterative_search_options[] = {
{"off", HNSW_ITERATIVE_SEARCH_OFF, false}, {"off", HNSW_ITERATIVE_SEARCH_OFF, false},
{"strict", HNSW_ITERATIVE_SEARCH_STRICT, false}, {"relaxed_order", HNSW_ITERATIVE_SEARCH_RELAXED, false},
{"relaxed", HNSW_ITERATIVE_SEARCH_RELAXED, false}, {"strict_order", HNSW_ITERATIVE_SEARCH_STRICT, false},
/* TODO Change to strict before merging */
{"on", HNSW_ITERATIVE_SEARCH_RELAXED, false},
{NULL, 0, false} {NULL, 0, false}
}; };
int hnsw_ef_search; int hnsw_ef_search;
int hnsw_max_iterative_tuples; int hnsw_max_search_tuples;
int hnsw_iterative_search; int hnsw_iterative_search;
int hnsw_lock_tranche_id; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
@@ -80,16 +78,14 @@ HnswInit(void)
"Valid range is 1..1000.", &hnsw_ef_search, "Valid range is 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Change name */ DefineCustomEnumVariable("hnsw.iterative_search", "Sets the iterative search mode",
DefineCustomEnumVariable("hnsw.streaming", "Iterative search mode",
NULL, &hnsw_iterative_search, NULL, &hnsw_iterative_search,
HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Change name */ /* This is approximate and does not apply to the initial scan */
/* TODO Ensure ivfflat.max_probes uses same value for "all" */ DefineCustomIntVariable("hnsw.max_search_tuples", "Sets the max number of candidates to visit for iterative search",
DefineCustomIntVariable("hnsw.ef_stream", "Sets the max number of additional candidates to visit for streaming search", "-1 means no limit", &hnsw_max_search_tuples,
"-1 means all", &hnsw_max_iterative_tuples, -1, -1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL);
HNSW_DEFAULT_EF_STREAM, HNSW_MIN_EF_STREAM, HNSW_MAX_EF_STREAM, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("hnsw"); MarkGUCPrefixReserved("hnsw");
} }

View File

@@ -42,9 +42,6 @@
#define HNSW_DEFAULT_EF_SEARCH 40 #define HNSW_DEFAULT_EF_SEARCH 40
#define HNSW_MIN_EF_SEARCH 1 #define HNSW_MIN_EF_SEARCH 1
#define HNSW_MAX_EF_SEARCH 1000 #define HNSW_MAX_EF_SEARCH 1000
#define HNSW_DEFAULT_EF_STREAM -1
#define HNSW_MIN_EF_STREAM -1
#define HNSW_MAX_EF_STREAM INT_MAX
/* Tuple types */ /* Tuple types */
#define HNSW_ELEMENT_TUPLE_TYPE 1 #define HNSW_ELEMENT_TUPLE_TYPE 1
@@ -112,16 +109,16 @@
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_max_iterative_tuples;
extern int hnsw_iterative_search; extern int hnsw_iterative_search;
extern int hnsw_max_search_tuples;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef enum HnswIterativeSearchType typedef enum HnswIterativeSearchMode
{ {
HNSW_ITERATIVE_SEARCH_OFF, HNSW_ITERATIVE_SEARCH_OFF,
HNSW_ITERATIVE_SEARCH_STRICT, HNSW_ITERATIVE_SEARCH_RELAXED,
HNSW_ITERATIVE_SEARCH_RELAXED HNSW_ITERATIVE_SEARCH_STRICT
} HnswIterativeSearchType; } HnswIterativeSearchMode;
typedef struct HnswElementData HnswElementData; typedef struct HnswElementData HnswElementData;
typedef struct HnswNeighborArray HnswNeighborArray; typedef struct HnswNeighborArray HnswNeighborArray;

View File

@@ -102,6 +102,17 @@ GetScanValue(IndexScanDesc scan)
return value; return value;
} }
#if defined(HNSW_MEMORY)
/*
* Show memory usage
*/
static void
ShowMemoryUsage(HnswScanOpaque so)
{
elog(INFO, "memory: %zu KB, tuples: " INT64_FORMAT, MemoryContextMemAllocated(so->tmpCtx, false) / 1024, so->tuples);
}
#endif
/* /*
* Prepare for an index scan * Prepare for an index scan
*/ */
@@ -118,9 +129,14 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
so->first = true; so->first = true;
so->v.tids = NULL; so->v.tids = NULL;
so->discarded = NULL; so->discarded = NULL;
/*
* Use a lower max allocation size than default to allow scanning more
* tuples for iterative search before exceeding work_mem
*/
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context", "Hnsw scan temporary context",
ALLOCSET_DEFAULT_SIZES); 0, 8 * 1024, 512 * 1024);
/* Set support functions */ /* Set support functions */
HnswInitSupport(&so->support, index); HnswInitSupport(&so->support, index);
@@ -204,7 +220,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->first = false; so->first = false;
#if defined(HNSW_MEMORY) #if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024); ShowMemoryUsage(so);
#endif #endif
} }
@@ -224,8 +240,8 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
if (so->discarded == NULL) if (so->discarded == NULL)
break; break;
/* Reached max number of additional tuples */ /* Reached max number of tuples */
if (hnsw_max_iterative_tuples != -1 && so->tuples >= hnsw_ef_search + hnsw_max_iterative_tuples) if (hnsw_max_search_tuples != -1 && so->tuples >= hnsw_max_search_tuples)
{ {
if (pairingheap_is_empty(so->discarded)) if (pairingheap_is_empty(so->discarded))
break; break;
@@ -266,7 +282,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock); UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
#if defined(HNSW_MEMORY) #if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024); ShowMemoryUsage(so);
#endif #endif
} }

View File

@@ -1,5 +1,6 @@
#include "postgres.h" #include "postgres.h"
#include <float.h>
#include <math.h> #include <math.h>
#include "access/generic_xlog.h" #include "access/generic_xlog.h"
@@ -157,6 +158,9 @@ HnswInitSupport(HnswSupport * support, Relation index)
Datum Datum
HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value) HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value)
{ {
if (!typeInfo->normalize)
return value;
return DirectFunctionCall1Coll(typeInfo->normalize, collation, value); return DirectFunctionCall1Coll(typeInfo->normalize, collation, value);
} }
@@ -546,7 +550,14 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Hns
if (DatumGetPointer(q->value) == NULL) if (DatumGetPointer(q->value) == NULL)
*distance = 0; *distance = 0;
else else
{
*distance = HnswGetDistance(q->value, PointerGetDatum(&etup->data), support); *distance = HnswGetDistance(q->value, PointerGetDatum(&etup->data), support);
/* Needed for intvec cosine distance */
/* TODO Improve */
if (isnan(*distance))
*distance = DBL_MAX;
}
} }
/* Load element */ /* Load element */
@@ -581,21 +592,34 @@ GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport *
return HnswGetDistance(q->value, value, support); return HnswGetDistance(q->value, value, support);
} }
/*
* Allocate a search candidate
*/
static HnswSearchCandidate *
HnswInitSearchCandidate(char *base, HnswElement element, double distance)
{
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, sc->element, element);
sc->distance = distance;
return sc;
}
/* /*
* Create a candidate for the entry point * Create a candidate for the entry point
*/ */
HnswSearchCandidate * HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec) HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec)
{ {
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
bool inMemory = index == NULL; bool inMemory = index == NULL;
double distance;
HnswPtrStore(base, sc->element, entryPoint);
if (inMemory) if (inMemory)
sc->distance = GetElementDistance(base, entryPoint, q, support); distance = GetElementDistance(base, entryPoint, q, support);
else else
HnswLoadElement(entryPoint, &sc->distance, q, index, support, loadVec, NULL); HnswLoadElement(entryPoint, &distance, q, index, support, loadVec, NULL);
return sc;
return HnswInitSearchCandidate(base, entryPoint, distance);
} }
/* /*
@@ -912,9 +936,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
if (discarded != NULL) if (discarded != NULL)
{ {
/* Create a new candidate */ /* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate)); e = HnswInitSearchCandidate(base, eElement, eDistance);
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
pairingheap_add(*discarded, &e->w_node); pairingheap_add(*discarded, &e->w_node);
} }
@@ -926,9 +948,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
continue; continue;
/* Create a new candidate */ /* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate)); e = HnswInitSearchCandidate(base, eElement, eDistance);
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
pairingheap_add(C, &e->c_node); pairingheap_add(C, &e->c_node);
pairingheap_add(W, &e->w_node); pairingheap_add(W, &e->w_node);
@@ -1384,6 +1404,19 @@ hnsw_halfvec_support(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(&typeInfo); PG_RETURN_POINTER(&typeInfo);
}; };
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(hnsw_intvec_support);
Datum
hnsw_intvec_support(PG_FUNCTION_ARGS)
{
static const HnswTypeInfo typeInfo = {
.maxDimensions = HNSW_MAX_DIM * 4,
.normalize = NULL,
.checkValue = NULL
};
PG_RETURN_POINTER(&typeInfo);
};
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(hnsw_bit_support); FUNCTION_PREFIX PG_FUNCTION_INFO_V1(hnsw_bit_support);
Datum Datum
hnsw_bit_support(PG_FUNCTION_ARGS) hnsw_bit_support(PG_FUNCTION_ARGS)

754
src/intvec.c Normal file
View File

@@ -0,0 +1,754 @@
#include "postgres.h"
#include <limits.h>
#include <math.h>
#include "catalog/pg_type.h"
#include "fmgr.h"
#include "intvec.h"
#include "lib/stringinfo.h"
#include "libpq/pqformat.h"
#include "utils/array.h"
#include "utils/builtins.h"
#include "utils/lsyscache.h"
/*
* Ensure same dimensions
*/
static inline void
CheckDims(IntVector * a, IntVector * b)
{
if (a->dim != b->dim)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("different intvec dimensions %d and %d", a->dim, b->dim)));
}
/*
* Ensure expected dimensions
*/
static inline void
CheckExpectedDim(int32 typmod, int dim)
{
if (typmod != -1 && typmod != dim)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("expected %d dimensions, not %d", typmod, dim)));
}
/*
* Ensure valid dimensions
*/
static inline void
CheckDim(int dim)
{
if (dim < 1)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("intvec must have at least 1 dimension")));
if (dim > INTVEC_MAX_DIM)
ereport(ERROR,
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("intvec cannot have more than %d dimensions", INTVEC_MAX_DIM)));
}
/*
* Ensure element in range
*/
static inline void
CheckElement(long value)
{
if (value < SCHAR_MIN || value > SCHAR_MAX)
ereport(ERROR,
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("value \"%ld\" is out of range for type intvec", value)));
}
/*
* Allocate and initialize a new int vector
*/
IntVector *
InitIntVector(int dim)
{
IntVector *result;
int size;
size = INTVEC_SIZE(dim);
result = (IntVector *) palloc0(size);
SET_VARSIZE(result, size);
result->dim = dim;
return result;
}
/*
* Check for whitespace, since array_isspace() is static
*/
static inline bool
intvec_isspace(char ch)
{
if (ch == ' ' ||
ch == '\t' ||
ch == '\n' ||
ch == '\r' ||
ch == '\v' ||
ch == '\f')
return true;
return false;
}
/*
* Convert textual representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_in);
Datum
intvec_in(PG_FUNCTION_ARGS)
{
char *lit = PG_GETARG_CSTRING(0);
int32 typmod = PG_GETARG_INT32(2);
int8 x[INTVEC_MAX_DIM];
int dim = 0;
char *pt = lit;
IntVector *result;
while (intvec_isspace(*pt))
pt++;
if (*pt != '[')
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type intvec: \"%s\"", lit),
errdetail("Vector contents must start with \"[\".")));
pt++;
while (intvec_isspace(*pt))
pt++;
if (*pt == ']')
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("intvec must have at least 1 dimension")));
for (;;)
{
long val;
char *stringEnd;
if (dim == INTVEC_MAX_DIM)
ereport(ERROR,
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("intvec cannot have more than %d dimensions", VECTOR_MAX_DIM)));
while (intvec_isspace(*pt))
pt++;
/* Check for empty string like float4in */
if (*pt == '\0')
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type intvec: \"%s\"", lit)));
errno = 0;
/* Use similar logic as int2vectorin */
val = strtol(pt, &stringEnd, 10);
if (stringEnd == pt)
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type intvec: \"%s\"", lit)));
/* Check for range error like float4in */
if (errno == ERANGE || val < SCHAR_MIN || val > SCHAR_MAX)
ereport(ERROR,
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("\"%s\" is out of range for type intvec", pnstrdup(pt, stringEnd - pt))));
CheckElement(val);
x[dim++] = val;
pt = stringEnd;
while (intvec_isspace(*pt))
pt++;
if (*pt == ',')
pt++;
else if (*pt == ']')
{
pt++;
break;
}
else
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type intvec: \"%s\"", lit)));
}
/* Only whitespace is allowed after the closing brace */
while (intvec_isspace(*pt))
pt++;
if (*pt != '\0')
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type intvec: \"%s\"", lit),
errdetail("Junk after closing right brace.")));
CheckDim(dim);
CheckExpectedDim(typmod, dim);
result = InitIntVector(dim);
for (int i = 0; i < dim; i++)
result->x[i] = x[i];
PG_RETURN_POINTER(result);
}
/*
* Convert internal representation to textual representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_out);
Datum
intvec_out(PG_FUNCTION_ARGS)
{
IntVector *vector = PG_GETARG_INTVEC_P(0);
int dim = vector->dim;
char *buf;
char *ptr;
/*
* Need:
*
* dim * 4 bytes for elements (-128 to 127)
*
* dim - 1 bytes for separator
*
* 3 bytes for [, ], and \0
*/
buf = (char *) palloc(5 * dim + 2);
ptr = buf;
*ptr = '[';
ptr++;
for (int i = 0; i < dim; i++)
{
if (i > 0)
{
*ptr = ',';
ptr++;
}
#if PG_VERSION_NUM >= 140000
ptr += pg_ltoa(vector->x[i], ptr);
#else
pg_ltoa(vector->x[i], ptr);
while (*ptr != '\0')
ptr++;
#endif
}
*ptr = ']';
ptr++;
*ptr = '\0';
PG_FREE_IF_COPY(vector, 0);
PG_RETURN_CSTRING(buf);
}
/*
* Convert type modifier
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_typmod_in);
Datum
intvec_typmod_in(PG_FUNCTION_ARGS)
{
ArrayType *ta = PG_GETARG_ARRAYTYPE_P(0);
int32 *tl;
int n;
tl = ArrayGetIntegerTypmods(ta, &n);
if (n != 1)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("invalid type modifier")));
if (*tl < 1)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("dimensions for type intvec must be at least 1")));
if (*tl > INTVEC_MAX_DIM)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("dimensions for type intvec cannot exceed %d", INTVEC_MAX_DIM)));
PG_RETURN_INT32(*tl);
}
/*
* Convert external binary representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_recv);
Datum
intvec_recv(PG_FUNCTION_ARGS)
{
StringInfo buf = (StringInfo) PG_GETARG_POINTER(0);
int32 typmod = PG_GETARG_INT32(2);
IntVector *result;
int16 dim;
int16 unused;
dim = pq_getmsgint(buf, sizeof(int16));
unused = pq_getmsgint(buf, sizeof(int16));
CheckDim(dim);
CheckExpectedDim(typmod, dim);
if (unused != 0)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("expected unused to be 0, not %d", unused)));
result = InitIntVector(dim);
for (int i = 0; i < dim; i++)
result->x[i] = pq_getmsgint(buf, sizeof(int8));
PG_RETURN_POINTER(result);
}
/*
* Convert internal representation to the external binary representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_send);
Datum
intvec_send(PG_FUNCTION_ARGS)
{
IntVector *vec = PG_GETARG_INTVEC_P(0);
StringInfoData buf;
pq_begintypsend(&buf);
pq_sendint(&buf, vec->dim, sizeof(int16));
pq_sendint(&buf, vec->unused, sizeof(int16));
for (int i = 0; i < vec->dim; i++)
pq_sendint8(&buf, vec->x[i]);
PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
}
/*
* Convert int vector to int vector
* This is needed to check the type modifier
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec);
Datum
intvec(PG_FUNCTION_ARGS)
{
IntVector *vec = PG_GETARG_INTVEC_P(0);
int32 typmod = PG_GETARG_INT32(1);
CheckExpectedDim(typmod, vec->dim);
PG_RETURN_POINTER(vec);
}
/*
* Convert array to intvec vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(array_to_intvec);
Datum
array_to_intvec(PG_FUNCTION_ARGS)
{
ArrayType *array = PG_GETARG_ARRAYTYPE_P(0);
int32 typmod = PG_GETARG_INT32(1);
IntVector *result;
int16 typlen;
bool typbyval;
char typalign;
Datum *elemsp;
int nelemsp;
if (ARR_NDIM(array) > 1)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("array must be 1-D")));
if (ARR_HASNULL(array) && array_contains_nulls(array))
ereport(ERROR,
(errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
errmsg("array must not contain nulls")));
get_typlenbyvalalign(ARR_ELEMTYPE(array), &typlen, &typbyval, &typalign);
deconstruct_array(array, ARR_ELEMTYPE(array), typlen, typbyval, typalign, &elemsp, NULL, &nelemsp);
CheckDim(nelemsp);
CheckExpectedDim(typmod, nelemsp);
result = InitIntVector(nelemsp);
if (ARR_ELEMTYPE(array) == INT4OID)
{
for (int i = 0; i < nelemsp; i++)
{
long l = DatumGetInt32(elemsp[i]);
CheckElement(l);
result->x[i] = l;
}
}
else
{
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("unsupported array type")));
}
/*
* Free allocation from deconstruct_array. Do not free individual elements
* when pass-by-reference since they point to original array.
*/
pfree(elemsp);
PG_RETURN_POINTER(result);
}
/*
* Convert int vector to integer[]
*/
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(intvec_to_integer);
Datum
intvec_to_integer(PG_FUNCTION_ARGS)
{
IntVector *vec = PG_GETARG_INTVEC_P(0);
Datum *datums;
ArrayType *result;
datums = (Datum *) palloc(sizeof(Datum) * vec->dim);
for (int i = 0; i < vec->dim; i++)
datums[i] = Int32GetDatum((int) vec->x[i]);
result = construct_array(datums, vec->dim, INT4OID, sizeof(int32), true, TYPALIGN_INT);
pfree(datums);
PG_RETURN_POINTER(result);
}
static int
IntvecL2SquaredDistance(int dim, int8 *ax, int8 *bx)
{
int distance = 0;
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
{
int diff = (int) ax[i] - (int) bx[i];
distance += diff * diff;
}
return distance;
}
/*
* Get the L2 distance between int vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_l2_distance);
Datum
intvec_l2_distance(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8(sqrt((double) IntvecL2SquaredDistance(a->dim, a->x, b->x)));
}
/*
* Get the L2 squared distance between int vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_l2_squared_distance);
Datum
intvec_l2_squared_distance(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8((double) IntvecL2SquaredDistance(a->dim, a->x, b->x));
}
static int
IntvecInnerProduct(int dim, int8 *ax, int8 *bx)
{
int distance = 0;
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
distance += (int) ax[i] * (int) bx[i];
return distance;
}
/*
* Get the inner product of two int vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_inner_product);
Datum
intvec_inner_product(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8((double) IntvecInnerProduct(a->dim, a->x, b->x));
}
/*
* Get the negative inner product of two int vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_negative_inner_product);
Datum
intvec_negative_inner_product(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8((double) -IntvecInnerProduct(a->dim, a->x, b->x));
}
static double
IntvecCosineSimilarity(int dim, int8 *ax, int8 *bx)
{
int similarity = 0;
int norma = 0;
int normb = 0;
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
{
int axi = ax[i];
int bxi = bx[i];
similarity += axi * bxi;
norma += axi * axi;
normb += bxi * bxi;
}
/* Use sqrt(a * b) over sqrt(a) * sqrt(b) */
return (double) similarity / sqrt((double) norma * (double) normb);
}
/*
* Get the cosine distance between two int vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_cosine_distance);
Datum
intvec_cosine_distance(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
double similarity;
CheckDims(a, b);
similarity = IntvecCosineSimilarity(a->dim, a->x, b->x);
#ifdef _MSC_VER
/* /fp:fast may not propagate NaN */
if (isnan(similarity))
PG_RETURN_FLOAT8(NAN);
#endif
/* Keep in range */
if (similarity > 1)
similarity = 1;
else if (similarity < -1)
similarity = -1;
PG_RETURN_FLOAT8(1 - similarity);
}
static int
IntvecL1Distance(int dim, int8 *ax, int8 *bx)
{
int distance = 0;
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
distance += abs((int) ax[i] - (int) bx[i]);
return distance;
}
/*
* Get the L1 distance between two int vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_l1_distance);
Datum
intvec_l1_distance(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8((double) IntvecL1Distance(a->dim, a->x, b->x));
}
/*
* Get the dimensions of an int vector
*/
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(intvec_vector_dims);
Datum
intvec_vector_dims(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
PG_RETURN_INT32(a->dim);
}
/*
* Get the L2 norm of an int vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(intvec_l2_norm);
Datum
intvec_l2_norm(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
int8 *ax = a->x;
int norm = 0;
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++)
norm += (int) ax[i] * (int) ax[i];
PG_RETURN_FLOAT8(sqrt((double) norm));
}
/*
* Internal helper to compare int vectors
*/
static int
intvec_cmp_internal(IntVector * a, IntVector * b)
{
int dim = Min(a->dim, b->dim);
/* Check values before dimensions to be consistent with Postgres arrays */
for (int i = 0; i < dim; i++)
{
if (a->x[i] < b->x[i])
return -1;
if (a->x[i] > b->x[i])
return 1;
}
if (a->dim < b->dim)
return -1;
if (a->dim > b->dim)
return 1;
return 0;
}
/*
* Less than
*/
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(intvec_lt);
Datum
intvec_lt(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
PG_RETURN_BOOL(intvec_cmp_internal(a, b) < 0);
}
/*
* Less than or equal
*/
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(intvec_le);
Datum
intvec_le(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
PG_RETURN_BOOL(intvec_cmp_internal(a, b) <= 0);
}
/*
* Equal
*/
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(intvec_eq);
Datum
intvec_eq(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
PG_RETURN_BOOL(intvec_cmp_internal(a, b) == 0);
}
/*
* Not equal
*/
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(intvec_ne);
Datum
intvec_ne(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
PG_RETURN_BOOL(intvec_cmp_internal(a, b) != 0);
}
/*
* Greater than or equal
*/
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(intvec_ge);
Datum
intvec_ge(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
PG_RETURN_BOOL(intvec_cmp_internal(a, b) >= 0);
}
/*
* Greater than
*/
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(intvec_gt);
Datum
intvec_gt(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
PG_RETURN_BOOL(intvec_cmp_internal(a, b) > 0);
}
/*
* Compare int vectors
*/
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(intvec_cmp);
Datum
intvec_cmp(PG_FUNCTION_ARGS)
{
IntVector *a = PG_GETARG_INTVEC_P(0);
IntVector *b = PG_GETARG_INTVEC_P(1);
PG_RETURN_INT32(intvec_cmp_internal(a, b));
}

23
src/intvec.h Normal file
View File

@@ -0,0 +1,23 @@
#ifndef INTVEC_H
#define INTVEC_H
#include "vector.h"
#define INTVEC_MAX_DIM VECTOR_MAX_DIM
#define INTVEC_SIZE(_dim) (offsetof(IntVector, x) + sizeof(int8)*(_dim))
#define DatumGetIntVector(x) ((IntVector *) PG_DETOAST_DATUM(x))
#define PG_GETARG_INTVEC_P(x) DatumGetIntVector(PG_GETARG_DATUM(x))
#define PG_RETURN_INTVEC_P(x) PG_RETURN_POINTER(x)
typedef struct IntVector
{
int32 vl_len_; /* varlena header (do not touch directly!) */
int16 dim; /* number of dimensions */
int16 unused;
int8 x[FLEXIBLE_ARRAY_MEMBER];
} IntVector;
IntVector *InitIntVector(int dim);
#endif

View File

@@ -228,11 +228,11 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
static inline void static inline void
GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, IndexTuple *itup, int *list) GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, IndexTuple *itup, int *list)
{ {
if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
{
Datum value; Datum value;
bool isnull; bool isnull;
if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
{
*list = DatumGetInt32(slot_getattr(slot, 1, &isnull)); *list = DatumGetInt32(slot_getattr(slot, 1, &isnull));
value = slot_getattr(slot, 3, &isnull); value = slot_getattr(slot, 3, &isnull);
@@ -254,8 +254,8 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
IndexTuple itup = NULL; /* silence compiler warning */ IndexTuple itup = NULL; /* silence compiler warning */
int64 inserted = 0; int64 inserted = 0;
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsMinimalTuple); TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple);
TupleDesc tupdesc = RelationGetDescr(index); TupleDesc tupdesc = buildstate->tupdesc;
pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD); pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD);
@@ -319,6 +319,7 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->index = index; buildstate->index = index;
buildstate->indexInfo = indexInfo; buildstate->indexInfo = indexInfo;
buildstate->typeInfo = IvfflatGetTypeInfo(index); buildstate->typeInfo = IvfflatGetTypeInfo(index);
buildstate->tupdesc = RelationGetDescr(index);
buildstate->lists = IvfflatGetLists(index); buildstate->lists = IvfflatGetLists(index);
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod; buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
@@ -356,12 +357,12 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
errmsg("dimensions must be greater than one for this opclass"))); errmsg("dimensions must be greater than one for this opclass")));
/* Create tuple description for sorting */ /* Create tuple description for sorting */
buildstate->tupdesc = CreateTemplateTupleDesc(3); buildstate->sortdesc = CreateTemplateTupleDesc(3);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 3, "vector", buildstate->tupdesc->attrs[0].atttypid, -1, 0);
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual); buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions)); buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions));
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists); buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
@@ -633,7 +634,7 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo); InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo);
memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen); memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen);
buildstate.centers->length = buildstate.centers->maxlen; buildstate.centers->length = buildstate.centers->maxlen;
ivfspool->sortstate = InitBuildSortState(buildstate.tupdesc, sortmem, coordinate); ivfspool->sortstate = InitBuildSortState(buildstate.sortdesc, sortmem, coordinate);
buildstate.sortstate = ivfspool->sortstate; buildstate.sortstate = ivfspool->sortstate;
scan = table_beginscan_parallel(ivfspool->heap, scan = table_beginscan_parallel(ivfspool->heap,
ParallelTableScanFromIvfflatShared(ivfshared)); ParallelTableScanFromIvfflatShared(ivfshared));
@@ -950,7 +951,7 @@ AssignTuples(IvfflatBuildState * buildstate)
} }
/* Begin serial/leader tuplesort */ /* Begin serial/leader tuplesort */
buildstate->sortstate = InitBuildSortState(buildstate->tupdesc, maintenance_work_mem, coordinate); buildstate->sortstate = InitBuildSortState(buildstate->sortdesc, maintenance_work_mem, coordinate);
/* Add tuples to sort */ /* Add tuples to sort */
if (buildstate->heap != NULL) if (buildstate->heap != NULL)

View File

@@ -17,8 +17,16 @@
#endif #endif
int ivfflat_probes; int ivfflat_probes;
int ivfflat_iterative_search;
int ivfflat_max_probes;
static relopt_kind ivfflat_relopt_kind; static relopt_kind ivfflat_relopt_kind;
static const struct config_enum_entry ivfflat_iterative_search_options[] = {
{"off", IVFFLAT_ITERATIVE_SEARCH_OFF, false},
{"relaxed_order", IVFFLAT_ITERATIVE_SEARCH_RELAXED, false},
{NULL, 0, false}
};
/* /*
* Initialize index options and variables * Initialize index options and variables
*/ */
@@ -33,6 +41,15 @@ IvfflatInit(void)
"Valid range is 1..lists.", &ivfflat_probes, "Valid range is 1..lists.", &ivfflat_probes,
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL); IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
DefineCustomEnumVariable("ivfflat.iterative_search", "Sets the iterative search mode",
NULL, &ivfflat_iterative_search,
IVFFLAT_ITERATIVE_SEARCH_OFF, ivfflat_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
/* If this is less than probes, probes is used */
DefineCustomIntVariable("ivfflat.max_probes", "Sets the max number of probes for iterative search",
"-1 means no limit", &ivfflat_max_probes,
-1, -1, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("ivfflat"); MarkGUCPrefixReserved("ivfflat");
} }

View File

@@ -80,6 +80,14 @@
/* Variables */ /* Variables */
extern int ivfflat_probes; extern int ivfflat_probes;
extern int ivfflat_iterative_search;
extern int ivfflat_max_probes;
typedef enum IvfflatIterativeSearchMode
{
IVFFLAT_ITERATIVE_SEARCH_OFF,
IVFFLAT_ITERATIVE_SEARCH_RELAXED
} IvfflatIterativeSearchMode;
typedef struct VectorArrayData typedef struct VectorArrayData
{ {
@@ -165,6 +173,7 @@ typedef struct IvfflatBuildState
Relation index; Relation index;
IndexInfo *indexInfo; IndexInfo *indexInfo;
const IvfflatTypeInfo *typeInfo; const IvfflatTypeInfo *typeInfo;
TupleDesc tupdesc;
/* Settings */ /* Settings */
int dimensions; int dimensions;
@@ -198,7 +207,7 @@ typedef struct IvfflatBuildState
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
TupleDesc tupdesc; TupleDesc sortdesc;
TupleTableSlot *slot; TupleTableSlot *slot;
/* Memory */ /* Memory */
@@ -247,8 +256,11 @@ typedef struct IvfflatScanOpaqueData
{ {
const IvfflatTypeInfo *typeInfo; const IvfflatTypeInfo *typeInfo;
int probes; int probes;
int maxProbes;
int dimensions; int dimensions;
bool first; bool first;
Datum value;
MemoryContext tmpCtx;
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
@@ -265,7 +277,9 @@ typedef struct IvfflatScanOpaqueData
/* Lists */ /* Lists */
pairingheap *listQueue; pairingheap *listQueue;
IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */ BlockNumber *listPages;
int listIndex;
IvfflatScanList *lists;
} IvfflatScanOpaqueData; } IvfflatScanOpaqueData;
typedef IvfflatScanOpaqueData * IvfflatScanOpaque; typedef IvfflatScanOpaqueData * IvfflatScanOpaque;

View File

@@ -98,7 +98,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
IvfflatGetMetaPageInfo(index, NULL, NULL); IvfflatGetMetaPageInfo(index, NULL, NULL);
/* Find the insert page - sets the page and list info */ /* Find the insert page - sets the page and list info */
FindInsertPage(index, values, &insertPage, &listInfo); FindInsertPage(index, &value, &insertPage, &listInfo);
Assert(BlockNumberIsValid(insertPage)); Assert(BlockNumberIsValid(insertPage));
originalInsertPage = insertPage; originalInsertPage = insertPage;

View File

@@ -10,10 +10,7 @@
#include "miscadmin.h" #include "miscadmin.h"
#include "pgstat.h" #include "pgstat.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#ifdef IVFFLAT_MEMORY
#include "utils/memutils.h" #include "utils/memutils.h"
#endif
#define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr) #define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr)
#define GetScanListConst(ptr) pairingheap_const_container(IvfflatScanList, ph_node, ptr) #define GetScanListConst(ptr) pairingheap_const_container(IvfflatScanList, ph_node, ptr)
@@ -65,7 +62,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
/* Use procinfo from the index instead of scan key for performance */ /* Use procinfo from the index instead of scan key for performance */
distance = DatumGetFloat8(so->distfunc(so->procinfo, so->collation, PointerGetDatum(&list->center), value)); distance = DatumGetFloat8(so->distfunc(so->procinfo, so->collation, PointerGetDatum(&list->center), value));
if (listCount < so->probes) if (listCount < so->maxProbes)
{ {
IvfflatScanList *scanlist; IvfflatScanList *scanlist;
@@ -78,7 +75,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
pairingheap_add(so->listQueue, &scanlist->ph_node); pairingheap_add(so->listQueue, &scanlist->ph_node);
/* Calculate max distance */ /* Calculate max distance */
if (listCount == so->probes) if (listCount == so->maxProbes)
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance; maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
} }
else if (distance < maxDistance) else if (distance < maxDistance)
@@ -102,6 +99,11 @@ GetScanLists(IndexScanDesc scan, Datum value)
UnlockReleaseBuffer(cbuf); UnlockReleaseBuffer(cbuf);
} }
for (int i = listCount - 1; i >= 0; i--)
so->listPages[i] = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
Assert(pairingheap_is_empty(so->listQueue));
} }
/* /*
@@ -114,11 +116,14 @@ GetScanItems(IndexScanDesc scan, Datum value)
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation); TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
double tuples = 0; double tuples = 0;
TupleTableSlot *slot = so->vslot; TupleTableSlot *slot = so->vslot;
int batchProbes = 0;
tuplesort_reset(so->sortstate);
/* Search closest probes lists */ /* Search closest probes lists */
while (!pairingheap_is_empty(so->listQueue)) while (so->listIndex < so->maxProbes && (++batchProbes) <= so->probes)
{ {
BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage; BlockNumber searchPage = so->listPages[so->listIndex++];
/* Search all entry pages for list */ /* Search all entry pages for list */
while (BlockNumberIsValid(searchPage)) while (BlockNumberIsValid(searchPage))
@@ -166,13 +171,17 @@ GetScanItems(IndexScanDesc scan, Datum value)
} }
} }
if (tuples < 100) if (tuples < 100 && ivfflat_iterative_search == IVFFLAT_ITERATIVE_SEARCH_OFF)
ereport(DEBUG1, ereport(DEBUG1,
(errmsg("index scan found few tuples"), (errmsg("index scan found few tuples"),
errdetail("Index may have been created with little data."), errdetail("Index may have been created with little data."),
errhint("Recreate the index and possibly decrease lists."))); errhint("Recreate the index and possibly decrease lists.")));
tuplesort_performsort(so->sortstate); tuplesort_performsort(so->sortstate);
#if defined(IVFFLAT_MEMORY)
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#endif
} }
/* /*
@@ -209,7 +218,13 @@ GetScanValue(IndexScanDesc scan)
/* Normalize if needed */ /* Normalize if needed */
if (so->normprocinfo != NULL) if (so->normprocinfo != NULL)
{
MemoryContext oldCtx = MemoryContextSwitchTo(so->tmpCtx);
value = IvfflatNormValue(so->typeInfo, so->collation, value); value = IvfflatNormValue(so->typeInfo, so->collation, value);
MemoryContextSwitchTo(oldCtx);
}
} }
return value; return value;
@@ -240,19 +255,40 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
int lists; int lists;
int dimensions; int dimensions;
int probes = ivfflat_probes; int probes = ivfflat_probes;
int maxProbes;
MemoryContext oldCtx;
scan = RelationGetIndexScan(index, nkeys, norderbys); scan = RelationGetIndexScan(index, nkeys, norderbys);
/* Get lists and dimensions from metapage */ /* Get lists and dimensions from metapage */
IvfflatGetMetaPageInfo(index, &lists, &dimensions); IvfflatGetMetaPageInfo(index, &lists, &dimensions);
if (ivfflat_iterative_search != IVFFLAT_ITERATIVE_SEARCH_OFF)
{
maxProbes = ivfflat_max_probes;
if (maxProbes < 0)
maxProbes = lists;
else if (maxProbes < probes)
{
/* TODO Show notice */
maxProbes = probes;
}
}
else
maxProbes = probes;
if (probes > lists) if (probes > lists)
probes = lists; probes = lists;
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList)); if (maxProbes > lists)
maxProbes = lists;
so = (IvfflatScanOpaque) palloc(sizeof(IvfflatScanOpaqueData));
so->typeInfo = IvfflatGetTypeInfo(index); so->typeInfo = IvfflatGetTypeInfo(index);
so->first = true; so->first = true;
so->probes = probes; so->probes = probes;
so->maxProbes = maxProbes;
so->dimensions = dimensions; so->dimensions = dimensions;
/* Set support functions */ /* Set support functions */
@@ -260,6 +296,12 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC); so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
so->collation = index->rd_indcollation[0]; so->collation = index->rd_indcollation[0];
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Ivfflat scan temporary context",
ALLOCSET_DEFAULT_SIZES);
oldCtx = MemoryContextSwitchTo(so->tmpCtx);
/* Create tuple description for sorting */ /* Create tuple description for sorting */
so->tupdesc = CreateTemplateTupleDesc(2); so->tupdesc = CreateTemplateTupleDesc(2);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
@@ -280,6 +322,11 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->bas = GetAccessStrategy(BAS_BULKREAD); so->bas = GetAccessStrategy(BAS_BULKREAD);
so->listQueue = pairingheap_allocate(CompareLists, scan); so->listQueue = pairingheap_allocate(CompareLists, scan);
so->listPages = palloc(maxProbes * sizeof(BlockNumber));
so->listIndex = 0;
so->lists = palloc(maxProbes * sizeof(IvfflatScanList));
MemoryContextSwitchTo(oldCtx);
scan->opaque = so; scan->opaque = so;
@@ -294,11 +341,9 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
if (!so->first)
tuplesort_reset(so->sortstate);
so->first = true; so->first = true;
pairingheap_reset(so->listQueue); pairingheap_reset(so->listQueue);
so->listIndex = 0;
if (keys && scan->numberOfKeys > 0) if (keys && scan->numberOfKeys > 0)
memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData)); memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -314,6 +359,8 @@ bool
ivfflatgettuple(IndexScanDesc scan, ScanDirection dir) ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
ItemPointer heaptid;
bool isnull;
/* /*
* Index can be used to scan backward, but Postgres doesn't support * Index can be used to scan backward, but Postgres doesn't support
@@ -341,28 +388,23 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
IvfflatBench("GetScanLists", GetScanLists(scan, value)); IvfflatBench("GetScanLists", GetScanLists(scan, value));
IvfflatBench("GetScanItems", GetScanItems(scan, value)); IvfflatBench("GetScanItems", GetScanItems(scan, value));
so->first = false; so->first = false;
so->value = value;
#if defined(IVFFLAT_MEMORY)
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#endif
/* Clean up if we allocated a new value */
if (value != scan->orderByData->sk_argument)
pfree(DatumGetPointer(value));
} }
if (tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL)) while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
{ {
bool isnull; if (so->listIndex == so->maxProbes)
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull)); return false;
IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
}
heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;
scan->xs_recheck = false; scan->xs_recheck = false;
scan->xs_recheckorderby = false; scan->xs_recheckorderby = false;
return true; return true;
}
return false;
} }
/* /*
@@ -373,12 +415,10 @@ ivfflatendscan(IndexScanDesc scan)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
pairingheap_free(so->listQueue); /* Free any temporary files */
tuplesort_end(so->sortstate); tuplesort_end(so->sortstate);
FreeAccessStrategy(so->bas);
FreeTupleDesc(so->tupdesc);
/* TODO Free vslot and mslot without freeing TupleDesc */ MemoryContextDelete(so->tmpCtx);
pfree(so); pfree(so);
scan->opaque = NULL; scan->opaque = NULL;

View File

@@ -26,7 +26,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
Page cpage; Page cpage;
OffsetNumber coffno; OffsetNumber coffno;
OffsetNumber cmaxoffno; OffsetNumber cmaxoffno;
BlockNumber startPages[MaxOffsetNumber]; BlockNumber listPages[MaxOffsetNumber];
ListInfo listInfo; ListInfo listInfo;
cbuf = ReadBuffer(index, blkno); cbuf = ReadBuffer(index, blkno);
@@ -40,7 +40,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
{ {
IvfflatList list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, coffno)); IvfflatList list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, coffno));
startPages[coffno - FirstOffsetNumber] = list->startPage; listPages[coffno - FirstOffsetNumber] = list->startPage;
} }
listInfo.blkno = blkno; listInfo.blkno = blkno;
@@ -50,7 +50,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
for (coffno = FirstOffsetNumber; coffno <= cmaxoffno; coffno = OffsetNumberNext(coffno)) for (coffno = FirstOffsetNumber; coffno <= cmaxoffno; coffno = OffsetNumberNext(coffno))
{ {
BlockNumber searchPage = startPages[coffno - FirstOffsetNumber]; BlockNumber searchPage = listPages[coffno - FirstOffsetNumber];
BlockNumber insertPage = InvalidBlockNumber; BlockNumber insertPage = InvalidBlockNumber;
/* Iterate over entry pages */ /* Iterate over entry pages */

View File

@@ -140,6 +140,34 @@ SELECT '{1e-8,-1e-8}'::real[]::halfvec;
[0,-0] [0,-0]
(1 row) (1 row)
SELECT '[1,2,3]'::intvec::integer[];
int4
---------
{1,2,3}
(1 row)
SELECT '{1,2,3}'::integer[]::intvec;
intvec
---------
[1,2,3]
(1 row)
SELECT '{1,2,3}'::integer[]::intvec(3);
intvec
---------
[1,2,3]
(1 row)
SELECT '{1,2,3}'::integer[]::intvec(2);
ERROR: expected 2 dimensions, not 3
SELECT '{127,-128}'::integer[]::intvec;
intvec
------------
[127,-128]
(1 row)
SELECT '{128,-129}'::integer[]::intvec;
ERROR: value "128" is out of range for type intvec
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec; SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec;
sparsevec sparsevec
----------------- -----------------

View File

@@ -30,6 +30,23 @@ SELECT * FROM t2 ORDER BY val;
(4 rows) (4 rows)
DROP TABLE t;
DROP TABLE t2;
-- intvec
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE TABLE t2 (val intvec(3));
\copy t TO 'results/intvec.bin' WITH (FORMAT binary)
\copy t2 FROM 'results/intvec.bin' WITH (FORMAT binary)
SELECT * FROM t2 ORDER BY val;
val
---------
[0,0,0]
[1,1,1]
[1,2,3]
(4 rows)
DROP TABLE t; DROP TABLE t;
DROP TABLE t2; DROP TABLE t2;
-- sparsevec -- sparsevec

View File

@@ -0,0 +1,102 @@
SET enable_seqscan = off;
-- L2
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val intvec_l2_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,2,4]
[1,1,1]
[0,0,0]
(4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <-> (SELECT NULL::intvec)) t2;
count
-------
4
(1 row)
SELECT COUNT(*) FROM t;
count
-------
5
(1 row)
TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
-----
(0 rows)
DROP TABLE t;
-- inner product
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val intvec_ip_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]';
val
---------
[1,2,4]
[1,2,3]
[1,1,1]
[0,0,0]
(4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::intvec)) t2;
count
-------
4
(1 row)
DROP TABLE t;
-- cosine
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val intvec_cosine_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]';
val
---------
[1,1,1]
[1,2,3]
[1,2,4]
(3 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
count
-------
3
(1 row)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::intvec)) t2;
count
-------
3
(1 row)
DROP TABLE t;
-- L1
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val intvec_l1_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <+> '[3,3,3]';
val
---------
[1,2,3]
[1,2,4]
[1,1,1]
[0,0,0]
(4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <+> (SELECT NULL::intvec)) t2;
count
-------
4
(1 row)
DROP TABLE t;

View File

@@ -99,6 +99,32 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <+> (SELECT NULL::vector)) t2
4 4
(1 row) (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; DROP TABLE t;
-- unlogged -- unlogged
CREATE UNLOGGED TABLE t (val vector(3)); 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) ERROR: 0 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000)
SET hnsw.ef_search = 1001; SET hnsw.ef_search = 1001;
ERROR: 1001 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000) 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; DROP TABLE t;

328
test/expected/intvec.out Normal file
View File

@@ -0,0 +1,328 @@
SELECT '[1,2,3]'::intvec;
intvec
---------
[1,2,3]
(1 row)
SELECT '[-1,-2,-3]'::intvec;
intvec
------------
[-1,-2,-3]
(1 row)
SELECT ' [ 1, 2 , 3 ] '::intvec;
intvec
---------
[1,2,3]
(1 row)
SELECT '[1.23456]'::intvec;
ERROR: invalid input syntax for type intvec: "[1.23456]"
LINE 1: SELECT '[1.23456]'::intvec;
^
SELECT '[hello,1]'::intvec;
ERROR: invalid input syntax for type intvec: "[hello,1]"
LINE 1: SELECT '[hello,1]'::intvec;
^
SELECT '[127,-128]'::intvec;
intvec
------------
[127,-128]
(1 row)
SELECT '[128,-129]'::intvec;
ERROR: "128" is out of range for type intvec
LINE 1: SELECT '[128,-129]'::intvec;
^
SELECT '[1,2,3'::intvec;
ERROR: invalid input syntax for type intvec: "[1,2,3"
LINE 1: SELECT '[1,2,3'::intvec;
^
SELECT '[1,2,3]9'::intvec;
ERROR: invalid input syntax for type intvec: "[1,2,3]9"
LINE 1: SELECT '[1,2,3]9'::intvec;
^
DETAIL: Junk after closing right brace.
SELECT '1,2,3'::intvec;
ERROR: invalid input syntax for type intvec: "1,2,3"
LINE 1: SELECT '1,2,3'::intvec;
^
DETAIL: Vector contents must start with "[".
SELECT ''::intvec;
ERROR: invalid input syntax for type intvec: ""
LINE 1: SELECT ''::intvec;
^
DETAIL: Vector contents must start with "[".
SELECT '['::intvec;
ERROR: invalid input syntax for type intvec: "["
LINE 1: SELECT '['::intvec;
^
SELECT '[,'::intvec;
ERROR: invalid input syntax for type intvec: "[,"
LINE 1: SELECT '[,'::intvec;
^
SELECT '[]'::intvec;
ERROR: intvec must have at least 1 dimension
LINE 1: SELECT '[]'::intvec;
^
SELECT '[1,]'::intvec;
ERROR: invalid input syntax for type intvec: "[1,]"
LINE 1: SELECT '[1,]'::intvec;
^
SELECT '[1a]'::intvec;
ERROR: invalid input syntax for type intvec: "[1a]"
LINE 1: SELECT '[1a]'::intvec;
^
SELECT '[1,,3]'::intvec;
ERROR: invalid input syntax for type intvec: "[1,,3]"
LINE 1: SELECT '[1,,3]'::intvec;
^
SELECT '[1, ,3]'::intvec;
ERROR: invalid input syntax for type intvec: "[1, ,3]"
LINE 1: SELECT '[1, ,3]'::intvec;
^
SELECT '[1,2,3]'::intvec(3);
intvec
---------
[1,2,3]
(1 row)
SELECT '[1,2,3]'::intvec(2);
ERROR: expected 2 dimensions, not 3
SELECT '[1,2,3]'::intvec(3, 2);
ERROR: invalid type modifier
LINE 1: SELECT '[1,2,3]'::intvec(3, 2);
^
SELECT '[1,2,3]'::intvec('a');
ERROR: invalid input syntax for type integer: "a"
LINE 1: SELECT '[1,2,3]'::intvec('a');
^
SELECT '[1,2,3]'::intvec(0);
ERROR: dimensions for type intvec must be at least 1
LINE 1: SELECT '[1,2,3]'::intvec(0);
^
SELECT '[1,2,3]'::intvec(16001);
ERROR: dimensions for type intvec cannot exceed 16000
LINE 1: SELECT '[1,2,3]'::intvec(16001);
^
SELECT unnest('{"[1,2,3]", "[4,5,6]"}'::intvec[]);
unnest
---------
[1,2,3]
[4,5,6]
(2 rows)
SELECT '{"[1,2,3]"}'::intvec(2)[];
ERROR: expected 2 dimensions, not 3
SELECT '[1,2,3]'::intvec < '[1,2,3]';
?column?
----------
f
(1 row)
SELECT '[1,2,3]'::intvec < '[1,2]';
?column?
----------
f
(1 row)
SELECT '[1,2,3]'::intvec <= '[1,2,3]';
?column?
----------
t
(1 row)
SELECT '[1,2,3]'::intvec <= '[1,2]';
?column?
----------
f
(1 row)
SELECT '[1,2,3]'::intvec = '[1,2,3]';
?column?
----------
t
(1 row)
SELECT '[1,2,3]'::intvec = '[1,2]';
?column?
----------
f
(1 row)
SELECT '[1,2,3]'::intvec != '[1,2,3]';
?column?
----------
f
(1 row)
SELECT '[1,2,3]'::intvec != '[1,2]';
?column?
----------
t
(1 row)
SELECT '[1,2,3]'::intvec >= '[1,2,3]';
?column?
----------
t
(1 row)
SELECT '[1,2,3]'::intvec >= '[1,2]';
?column?
----------
t
(1 row)
SELECT '[1,2,3]'::intvec > '[1,2,3]';
?column?
----------
f
(1 row)
SELECT '[1,2,3]'::intvec > '[1,2]';
?column?
----------
t
(1 row)
SELECT intvec_cmp('[1,2,3]', '[1,2,3]');
intvec_cmp
------------
0
(1 row)
SELECT intvec_cmp('[1,2,3]', '[0,0,0]');
intvec_cmp
------------
1
(1 row)
SELECT intvec_cmp('[0,0,0]', '[1,2,3]');
intvec_cmp
------------
-1
(1 row)
SELECT intvec_cmp('[1,2]', '[1,2,3]');
intvec_cmp
------------
-1
(1 row)
SELECT intvec_cmp('[1,2,3]', '[1,2]');
intvec_cmp
------------
1
(1 row)
SELECT intvec_cmp('[1,2]', '[2,3,4]');
intvec_cmp
------------
-1
(1 row)
SELECT intvec_cmp('[2,3]', '[1,2,3]');
intvec_cmp
------------
1
(1 row)
SELECT vector_dims('[1,2,3]'::intvec);
vector_dims
-------------
3
(1 row)
SELECT l2_distance('[0,0]'::intvec, '[3,4]');
l2_distance
-------------
5
(1 row)
SELECT l2_distance('[0,0]'::intvec, '[0,1]');
l2_distance
-------------
1
(1 row)
SELECT l2_distance('[1,2]'::intvec, '[3]');
ERROR: different intvec dimensions 2 and 1
SELECT '[0,0]'::intvec <-> '[3,4]';
?column?
----------
5
(1 row)
SELECT inner_product('[1,2]'::intvec, '[3,4]');
inner_product
---------------
11
(1 row)
SELECT inner_product('[1,2]'::intvec, '[3]');
ERROR: different intvec dimensions 2 and 1
SELECT inner_product('[127]'::intvec, '[127]');
inner_product
---------------
16129
(1 row)
SELECT '[1,2]'::intvec <#> '[3,4]';
?column?
----------
-11
(1 row)
SELECT cosine_distance('[1,2]'::intvec, '[2,4]');
cosine_distance
-----------------
0
(1 row)
SELECT cosine_distance('[1,2]'::intvec, '[0,0]');
cosine_distance
-----------------
NaN
(1 row)
SELECT cosine_distance('[1,1]'::intvec, '[1,1]');
cosine_distance
-----------------
0
(1 row)
SELECT cosine_distance('[1,0]'::intvec, '[0,2]');
cosine_distance
-----------------
1
(1 row)
SELECT cosine_distance('[1,1]'::intvec, '[-1,-1]');
cosine_distance
-----------------
2
(1 row)
SELECT cosine_distance('[1,2]'::intvec, '[3]');
ERROR: different intvec dimensions 2 and 1
SELECT '[1,2]'::intvec <=> '[2,4]';
?column?
----------
0
(1 row)
SELECT l1_distance('[0,0]'::intvec, '[3,4]');
l1_distance
-------------
7
(1 row)
SELECT l1_distance('[0,0]'::intvec, '[0,1]');
l1_distance
-------------
1
(1 row)
SELECT l1_distance('[1,2]'::intvec, '[3]');
ERROR: different intvec dimensions 2 and 1

View File

@@ -81,6 +81,44 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2
3 3
(1 row) (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; DROP TABLE t;
-- unlogged -- unlogged
CREATE UNLOGGED TABLE t (val vector(3)); CREATE UNLOGGED TABLE t (val vector(3));
@@ -109,4 +147,27 @@ SHOW ivfflat.probes;
1 1
(1 row) (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; DROP TABLE t;

View File

@@ -38,6 +38,14 @@ SELECT '{1,2,3}'::real[]::halfvec(2);
SELECT '{65520,-65520}'::real[]::halfvec; SELECT '{65520,-65520}'::real[]::halfvec;
SELECT '{1e-8,-1e-8}'::real[]::halfvec; SELECT '{1e-8,-1e-8}'::real[]::halfvec;
SELECT '[1,2,3]'::intvec::integer[];
SELECT '{1,2,3}'::integer[]::intvec;
SELECT '{1,2,3}'::integer[]::intvec(3);
SELECT '{1,2,3}'::integer[]::intvec(2);
SELECT '{127,-128}'::integer[]::intvec;
SELECT '{128,-129}'::integer[]::intvec;
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec; SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec;
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec(5); SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec(5);
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec(4); SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec(4);

View File

@@ -28,6 +28,21 @@ SELECT * FROM t2 ORDER BY val;
DROP TABLE t; DROP TABLE t;
DROP TABLE t2; DROP TABLE t2;
-- intvec
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE TABLE t2 (val intvec(3));
\copy t TO 'results/intvec.bin' WITH (FORMAT binary)
\copy t2 FROM 'results/intvec.bin' WITH (FORMAT binary)
SELECT * FROM t2 ORDER BY val;
DROP TABLE t;
DROP TABLE t2;
-- sparsevec -- sparsevec
CREATE TABLE t (val sparsevec(3)); CREATE TABLE t (val sparsevec(3));

58
test/sql/hnsw_intvec.sql Normal file
View File

@@ -0,0 +1,58 @@
SET enable_seqscan = off;
-- L2
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val intvec_l2_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <-> (SELECT NULL::intvec)) t2;
SELECT COUNT(*) FROM t;
TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
DROP TABLE t;
-- inner product
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val intvec_ip_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]';
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::intvec)) t2;
DROP TABLE t;
-- cosine
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val intvec_cosine_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]';
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::intvec)) t2;
DROP TABLE t;
-- L1
CREATE TABLE t (val intvec(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val intvec_l1_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <+> '[3,3,3]';
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <+> (SELECT NULL::intvec)) t2;
DROP TABLE t;

View File

@@ -57,6 +57,23 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <+> (SELECT NULL::vector)) t2
DROP TABLE t; 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 -- unlogged
CREATE UNLOGGED TABLE t (val vector(3)); CREATE UNLOGGED TABLE t (val vector(3));
@@ -81,4 +98,12 @@ SHOW hnsw.ef_search;
SET hnsw.ef_search = 0; SET hnsw.ef_search = 0;
SET hnsw.ef_search = 1001; 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; DROP TABLE t;

73
test/sql/intvec.sql Normal file
View File

@@ -0,0 +1,73 @@
SELECT '[1,2,3]'::intvec;
SELECT '[-1,-2,-3]'::intvec;
SELECT ' [ 1, 2 , 3 ] '::intvec;
SELECT '[1.23456]'::intvec;
SELECT '[hello,1]'::intvec;
SELECT '[127,-128]'::intvec;
SELECT '[128,-129]'::intvec;
SELECT '[1,2,3'::intvec;
SELECT '[1,2,3]9'::intvec;
SELECT '1,2,3'::intvec;
SELECT ''::intvec;
SELECT '['::intvec;
SELECT '[,'::intvec;
SELECT '[]'::intvec;
SELECT '[1,]'::intvec;
SELECT '[1a]'::intvec;
SELECT '[1,,3]'::intvec;
SELECT '[1, ,3]'::intvec;
SELECT '[1,2,3]'::intvec(3);
SELECT '[1,2,3]'::intvec(2);
SELECT '[1,2,3]'::intvec(3, 2);
SELECT '[1,2,3]'::intvec('a');
SELECT '[1,2,3]'::intvec(0);
SELECT '[1,2,3]'::intvec(16001);
SELECT unnest('{"[1,2,3]", "[4,5,6]"}'::intvec[]);
SELECT '{"[1,2,3]"}'::intvec(2)[];
SELECT '[1,2,3]'::intvec < '[1,2,3]';
SELECT '[1,2,3]'::intvec < '[1,2]';
SELECT '[1,2,3]'::intvec <= '[1,2,3]';
SELECT '[1,2,3]'::intvec <= '[1,2]';
SELECT '[1,2,3]'::intvec = '[1,2,3]';
SELECT '[1,2,3]'::intvec = '[1,2]';
SELECT '[1,2,3]'::intvec != '[1,2,3]';
SELECT '[1,2,3]'::intvec != '[1,2]';
SELECT '[1,2,3]'::intvec >= '[1,2,3]';
SELECT '[1,2,3]'::intvec >= '[1,2]';
SELECT '[1,2,3]'::intvec > '[1,2,3]';
SELECT '[1,2,3]'::intvec > '[1,2]';
SELECT intvec_cmp('[1,2,3]', '[1,2,3]');
SELECT intvec_cmp('[1,2,3]', '[0,0,0]');
SELECT intvec_cmp('[0,0,0]', '[1,2,3]');
SELECT intvec_cmp('[1,2]', '[1,2,3]');
SELECT intvec_cmp('[1,2,3]', '[1,2]');
SELECT intvec_cmp('[1,2]', '[2,3,4]');
SELECT intvec_cmp('[2,3]', '[1,2,3]');
SELECT vector_dims('[1,2,3]'::intvec);
SELECT l2_distance('[0,0]'::intvec, '[3,4]');
SELECT l2_distance('[0,0]'::intvec, '[0,1]');
SELECT l2_distance('[1,2]'::intvec, '[3]');
SELECT '[0,0]'::intvec <-> '[3,4]';
SELECT inner_product('[1,2]'::intvec, '[3,4]');
SELECT inner_product('[1,2]'::intvec, '[3]');
SELECT inner_product('[127]'::intvec, '[127]');
SELECT '[1,2]'::intvec <#> '[3,4]';
SELECT cosine_distance('[1,2]'::intvec, '[2,4]');
SELECT cosine_distance('[1,2]'::intvec, '[0,0]');
SELECT cosine_distance('[1,1]'::intvec, '[1,1]');
SELECT cosine_distance('[1,0]'::intvec, '[0,2]');
SELECT cosine_distance('[1,1]'::intvec, '[-1,-1]');
SELECT cosine_distance('[1,2]'::intvec, '[3]');
SELECT '[1,2]'::intvec <=> '[2,4]';
SELECT l1_distance('[0,0]'::intvec, '[3,4]');
SELECT l1_distance('[0,0]'::intvec, '[0,1]');
SELECT l1_distance('[1,2]'::intvec, '[3]');

View File

@@ -44,6 +44,28 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2
DROP TABLE t; 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 -- unlogged
CREATE UNLOGGED TABLE t (val vector(3)); 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; 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; DROP TABLE t;

View File

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

View File

@@ -40,6 +40,10 @@ for (1 .. 50)
$actual = $node->safe_psql("postgres", "SELECT halfvec_cmp(v::halfvec, '$query'::real[]::halfvec) FROM tst"); $actual = $node->safe_psql("postgres", "SELECT halfvec_cmp(v::halfvec, '$query'::real[]::halfvec) FROM tst");
is($expected, $actual); is($expected, $actual);
# Test intvec
$actual = $node->safe_psql("postgres", "SELECT intvec_cmp(v::integer[]::intvec, '$query'::integer[]::intvec) FROM tst");
is($expected, $actual);
# Test sparsevec # Test sparsevec
$actual = $node->safe_psql("postgres", "SELECT sparsevec_cmp(v::vector::sparsevec, '$query'::real[]::vector::sparsevec) FROM tst"); $actual = $node->safe_psql("postgres", "SELECT sparsevec_cmp(v::vector::sparsevec, '$query'::real[]::vector::sparsevec) FROM tst");
is($expected, $actual); is($expected, $actual);

View File

@@ -45,6 +45,10 @@ for my $function (@functions)
my $actual = $node->safe_psql("postgres", "SELECT $function(v::halfvec, '$query'::vector::halfvec) FROM tst"); my $actual = $node->safe_psql("postgres", "SELECT $function(v::halfvec, '$query'::vector::halfvec) FROM tst");
is($expected, $actual, "halfvec $function"); is($expected, $actual, "halfvec $function");
# Test intvec
$actual = $node->safe_psql("postgres", "SELECT $function(v::real[]::integer[]::intvec, '$query'::vector::real[]::integer[]::intvec) FROM tst");
is($expected, $actual, "intvec $function");
# Test sparsevec # Test sparsevec
$actual = $node->safe_psql("postgres", "SELECT $function(v::sparsevec, '$query'::vector::sparsevec) FROM tst"); $actual = $node->safe_psql("postgres", "SELECT $function(v::sparsevec, '$query'::vector::sparsevec) FROM tst");
is($expected, $actual, "sparsevec $function"); is($expected, $actual, "sparsevec $function");

View File

@@ -12,8 +12,8 @@ $node->start;
# Create extension # Create extension
$node->safe_psql("postgres", "CREATE EXTENSION vector;"); $node->safe_psql("postgres", "CREATE EXTENSION vector;");
my @types = ("vector", "halfvec", "sparsevec"); my @types = ("vector", "halfvec", "intvec", "sparsevec");
my @inputs = ("[1.23,4.56,7.89]", "[1.23,4.56,7.89]", "{1:1.23,2:4.56,3:7.89}/3"); my @inputs = ("[1.23,4.56,7.89]", "[1.23,4.56,7.89]", "[1,2,3]", "{1:1.23,2:4.56,3:7.89}/3");
my @subs = (" ", " ", ",", ":", "-", "1", "9", "\0", "2147483648", "-2147483649"); my @subs = (" ", " ", ",", ":", "-", "1", "9", "\0", "2147483648", "-2147483649");
for my $i (0 .. $#types) for my $i (0 .. $#types)

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

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

View File

@@ -26,7 +26,7 @@ $node->safe_psql("postgres", qq(
my $count = $node->safe_psql("postgres", qq( my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.streaming = on; SET hnsw.iterative_search = relaxed_order;
SET work_mem = '8MB'; 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; SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
)); ));
@@ -34,16 +34,16 @@ is($count, 10);
foreach ((30000, 50000, 70000)) foreach ((30000, 50000, 70000))
{ {
my $ef_stream = $_; my $max_tuples = $_;
my $expected = $ef_stream / 10000; my $expected = $max_tuples / 10000;
my $sum = 0; my $sum = 0;
for my $i (1 .. 20) for my $i (1 .. 20)
{ {
$count = $node->safe_psql("postgres", qq( $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.streaming = on; SET hnsw.iterative_search = relaxed_order;
SET hnsw.ef_stream = $ef_stream; SET hnsw.max_search_tuples = $max_tuples;
SET work_mem = '8MB'; 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; SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
)); ));
@@ -57,7 +57,7 @@ foreach ((30000, 50000, 70000))
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq( my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.streaming = on; SET hnsw.iterative_search = relaxed_order;
SET client_min_messages = debug1; SET client_min_messages = debug1;
SET work_mem = '2MB'; 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; SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;

View File

@@ -10,18 +10,18 @@ my @expected;
my $limit = 20; my $limit = 20;
my $dim = 3; my $dim = 3;
my $array_sql = join(",", ('random()') x $dim); my $array_sql = join(",", ('random()') x $dim);
my @cs = (100, 1000); my @cs = (50, 500);
sub test_recall sub test_recall
{ {
my ($c, $ef_search, $min, $operator) = @_; my ($c, $ef_search, $min, $operator, $mode) = @_;
my $correct = 0; my $correct = 0;
my $total = 0; my $total = 0;
my $explain = $node->safe_psql("postgres", qq( my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search; SET hnsw.ef_search = $ef_search;
SET hnsw.streaming = on; SET hnsw.iterative_search = $mode;
EXPLAIN ANALYZE SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[0]' LIMIT $limit; 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/); like($explain, qr/Index Scan using idx on tst/);
@@ -31,7 +31,7 @@ sub test_recall
my $actual = $node->safe_psql("postgres", qq( my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search; SET hnsw.ef_search = $ef_search;
SET hnsw.streaming = on; SET hnsw.iterative_search = $mode;
SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit; SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
)); ));
my @actual_ids = split("\n", $actual); my @actual_ids = split("\n", $actual);
@@ -50,7 +50,7 @@ sub test_recall
$total += $limit; $total += $limit;
} }
cmp_ok($correct / $total, ">=", $min, $operator); cmp_ok($correct / $total, ">=", $min, "$operator $mode $c");
} }
# Initialize node # Initialize node
@@ -62,7 +62,7 @@ $node->start;
$node->safe_psql("postgres", "CREATE EXTENSION vector;"); $node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));"); $node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres", $node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;" "INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 50000) i;"
); );
# Generate queries # Generate queries
@@ -108,21 +108,8 @@ for my $i (0 .. $#operators)
push(@expected, $res); push(@expected, $res);
} }
if ($c == 100) test_recall($c, 40, 0.99, $operator, "strict_order");
{ test_recall($c, 40, 0.99, $operator, "relaxed_order");
test_recall($c, 40, 0.99, $operator);
}
else
{
if ($operator eq "<->")
{
test_recall($c, 40, 0.99, $operator);
}
else
{
test_recall($c, 40, 0.99, $operator);
}
}
} }
$node->safe_psql("postgres", "DROP INDEX idx;"); $node->safe_psql("postgres", "DROP INDEX idx;");

View File

@@ -0,0 +1,132 @@
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 = 10;
my $array_sql = join(",", ('(random() * 255)::int - 128') x $dim);
sub test_recall
{
my ($min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
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;
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
$correct++;
}
$total++;
}
}
cmp_ok($correct / $total, ">=", $min, $operator);
}
# 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 intvec($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
);
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, int(rand(256)) - 128);
}
push(@queries, "[" . join(",", @r) . "]");
}
# Check each index type
my @operators = ("<->", "<#>", "<=>", "<+>");
my @opclasses = ("intvec_l2_ops", "intvec_ip_ops", "intvec_cosine_ops", "intvec_l1_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", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
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 = 0.98;
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();

View File

@@ -0,0 +1,113 @@
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 = 10;
my $array_sql = join(",", ('(random() * 255)::int - 128') x $dim);
sub test_recall
{
my ($min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
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;
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
$correct++;
}
$total++;
}
}
cmp_ok($correct / $total, ">=", $min, $operator);
}
# 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 serial, v intvec($dim));");
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, int(rand(256)) - 128);
}
push(@queries, "[" . join(",", @r) . "]");
}
# Check each index type
my @operators = ("<->", "<#>", "<=>", "<+>");
my @opclasses = ("intvec_l2_ops", "intvec_ip_ops", "intvec_cosine_ops", "intvec_l1_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
# Add index
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v $opclass);");
# Use concurrent inserts
$node->pgbench(
"--no-vacuum --client=10 --transactions=1000",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"040_hnsw_intvec_insert_recall_$opclass" => "INSERT INTO tst (v) VALUES (ARRAY[$array_sql]);"
}
);
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;
));
push(@expected, $res);
}
# Test approximate results
my $min = 0.98;
test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
$node->safe_psql("postgres", "TRUNCATE tst;");
}
done_testing();

View File

@@ -0,0 +1,101 @@
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 = 10;
my $array_sql = join(",", ('(random() * 255)::int - 128') x $dim);
sub test_recall
{
my ($min, $ef_search, $test_name) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$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 = $ef_search;
SELECT i FROM tst ORDER BY v <-> '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
$correct++;
}
$total++;
}
}
cmp_ok($correct / $total, ">=", $min, $test_name);
}
# 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 intvec($dim));");
$node->safe_psql("postgres", "ALTER TABLE tst SET (autovacuum_enabled = false);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
);
# Add index
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v intvec_l2_ops) WITH (m = 4, ef_construction = 8);");
# Delete data
$node->safe_psql("postgres", "DELETE FROM tst WHERE i > 2500;");
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, int(rand(256)) - 128);
}
push(@queries, "[" . join(",", @r) . "]");
}
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
SELECT i FROM tst ORDER BY v <-> '$_' LIMIT $limit;
));
push(@expected, $res);
}
test_recall(0.18, $limit, "before vacuum");
test_recall(0.84, 100, "before vacuum");
# TODO Test concurrent inserts with vacuum
$node->safe_psql("postgres", "VACUUM tst;");
test_recall(0.84, $limit, "after vacuum");
done_testing();

View File

@@ -0,0 +1,58 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
# 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 (v intvec(3));");
sub insert_vectors
{
for my $i (1 .. 20)
{
$node->safe_psql("postgres", "INSERT INTO tst VALUES ('[1,1,1]');");
}
}
sub test_duplicates
{
my $res = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = 1;
SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '[1,1,1]') t;
));
is($res, 10);
}
# Test duplicates with build
insert_vectors();
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v intvec_l2_ops);");
test_duplicates();
# Reset
$node->safe_psql("postgres", "TRUNCATE tst;");
# Test duplicates with inserts
insert_vectors();
test_duplicates();
# Test fallback path for inserts
$node->pgbench(
"--no-vacuum --client=5 --transactions=100",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"042_hnsw_intvec_duplicates" => "INSERT INTO tst VALUES ('[1,1,1]');"
}
);
done_testing();