Compare commits

..

13 Commits

Author SHA1 Message Date
Andrew Kane
70785ff0fe Merge branch 'master' into hqann2 2024-10-10 19:06:26 -07:00
Andrew Kane
a73fe9e9e1 Use same TupleDesc [skip ci] 2024-10-10 02:01:23 -07:00
Andrew Kane
2ffbb2500c Merge branch 'master' into hqann2 2024-10-09 23:59:11 -07:00
Andrew Kane
46209740ce Merge branch 'master' into hqann2 2024-10-09 23:26:57 -07:00
Andrew Kane
d650d23da8 Merge branch 'master' into hqann2 2024-10-09 22:06:29 -07:00
Andrew Kane
ad755d19ca Merge branch 'master' into hqann2 2024-10-09 21:52:45 -07:00
Andrew Kane
f39b515ce7 Merge branch 'master' into hqann2 2024-10-09 21:50:53 -07:00
Andrew Kane
ab88612c83 Always use tuple for duplicates [skip ci] 2024-10-09 21:30:10 -07:00
Andrew Kane
356f175359 Merge branch 'master' into hqann2 2024-10-09 21:25:30 -07:00
Andrew Kane
a49a2b9480 Improved test [skip ci] 2024-10-09 19:54:06 -07:00
Andrew Kane
3b61384367 Added HnswGetElementTupleSize method [skip ci] 2024-10-09 19:37:09 -07:00
Andrew Kane
67f9a3e61c Init collation as well [skip ci] 2024-10-09 19:30:04 -07:00
Andrew Kane
3ccfab8f92 Added support for inline filtering with HNSW 2024-10-09 19:02:40 -07:00
44 changed files with 658 additions and 2668 deletions

View File

@@ -1,7 +1,7 @@
## 0.8.0 (unreleased) ## 0.8.0 (unreleased)
- Added support for inline filtering with HNSW
- 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/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 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
HEADERS = src/halfvec.h src/intvec.h src/sparsevec.h src/vector.h HEADERS = src/halfvec.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\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 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
HEADERS = src\halfvec.h src\intvec.h src\sparsevec.h src\vector.h HEADERS = src\halfvec.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)

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`, `intvec_l2_ops` for `intvec`, and `sparsevec_l2_ops` for `sparsevec` (and similar with the other distance functions) Note: Use `halfvec_l2_ops` for `halfvec` and `sparsevec_l2_ops` for `sparsevec` (and similar with the other distance functions)
Inner product Inner product
@@ -259,7 +259,6 @@ 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)
@@ -440,6 +439,12 @@ Create an index on one [or more](https://www.postgresql.org/docs/current/indexes
CREATE INDEX ON items (category_id); CREATE INDEX ON items (category_id);
``` ```
Or a composite HNSW index for approximate search (added in 0.8.0)
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops, category_id);
```
Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search
```sql ```sql
@@ -478,16 +483,6 @@ 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))
@@ -951,30 +946,6 @@ 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.
@@ -1224,6 +1195,7 @@ Thanks to:
- [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf) - [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf)
- [Concept Decompositions for Large Sparse Text Data using Clustering](https://www.cs.utexas.edu/users/inderjit/public_papers/concept_mlj.pdf) - [Concept Decompositions for Large Sparse Text Data using Clustering](https://www.cs.utexas.edu/users/inderjit/public_papers/concept_mlj.pdf)
- [Efficient and Robust Approximate Nearest Neighbor Search using Hierarchical Navigable Small World Graphs](https://arxiv.org/ftp/arxiv/papers/1603/1603.09320.pdf) - [Efficient and Robust Approximate Nearest Neighbor Search using Hierarchical Navigable Small World Graphs](https://arxiv.org/ftp/arxiv/papers/1603/1603.09320.pdf)
- [HQANN: Efficient and Robust Similarity Search for Hybrid Queries with Structured and Unstructured Constraints](https://arxiv.org/pdf/2207.07940.pdf)
## History ## History

View File

@@ -1,188 +1,6 @@
-- 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;
@@ -206,3 +24,11 @@ CREATE CAST (double precision[] AS sparsevec)
CREATE CAST (numeric[] AS sparsevec) CREATE CAST (numeric[] AS sparsevec)
WITH FUNCTION array_to_sparsevec(numeric[], integer, boolean) AS ASSIGNMENT; WITH FUNCTION array_to_sparsevec(numeric[], integer, boolean) AS ASSIGNMENT;
CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8
AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING hnsw AS
OPERATOR 2 = (integer, integer),
FUNCTION 4 hnsw_attribute_distance(integer, integer);

View File

@@ -272,9 +272,6 @@ 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;
@@ -650,199 +647,6 @@ 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
@@ -1112,3 +916,13 @@ CREATE OPERATOR CLASS sparsevec_l1_ops
OPERATOR 1 <+> (sparsevec, sparsevec) FOR ORDER BY float_ops, OPERATOR 1 <+> (sparsevec, sparsevec) FOR ORDER BY float_ops,
FUNCTION 1 l1_distance(sparsevec, sparsevec), FUNCTION 1 l1_distance(sparsevec, sparsevec),
FUNCTION 3 hnsw_sparsevec_support(internal); FUNCTION 3 hnsw_sparsevec_support(internal);
-- hnsw attributes
CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8
AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING hnsw AS
OPERATOR 2 = (integer, integer),
FUNCTION 4 hnsw_attribute_distance(integer, integer);

View File

@@ -20,13 +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},
{"relaxed_order", HNSW_ITERATIVE_SEARCH_RELAXED, false}, {"on", HNSW_ITERATIVE_SEARCH_RELAXED, false},
{"strict_order", HNSW_ITERATIVE_SEARCH_STRICT, false}, {"strict", HNSW_ITERATIVE_SEARCH_STRICT, false},
{NULL, 0, false} {NULL, 0, false}
}; };
int hnsw_ef_search; int hnsw_ef_search;
int hnsw_max_search_tuples; int hnsw_iterative_search_max_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;
@@ -78,13 +78,13 @@ 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);
DefineCustomEnumVariable("hnsw.iterative_search", "Sets the iterative search mode", DefineCustomEnumVariable("hnsw.iterative_search", "Sets iterative search",
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);
/* 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.iterative_search_max_tuples", "Sets the max number of candidates to visit for iterative search",
"-1 means no limit", &hnsw_max_search_tuples, "-1 means all", &hnsw_iterative_search_max_tuples,
-1, -1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL); -1, -1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("hnsw"); MarkGUCPrefixReserved("hnsw");
@@ -245,13 +245,13 @@ hnswhandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine); IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0; amroutine->amstrategies = 0;
amroutine->amsupport = 3; amroutine->amsupport = 4;
amroutine->amoptsprocnum = 0; amroutine->amoptsprocnum = 0;
amroutine->amcanorder = false; amroutine->amcanorder = false;
amroutine->amcanorderbyop = true; amroutine->amcanorderbyop = true;
amroutine->amcanbackward = false; /* can change direction mid-scan */ amroutine->amcanbackward = false; /* can change direction mid-scan */
amroutine->amcanunique = false; amroutine->amcanunique = false;
amroutine->amcanmulticol = false; amroutine->amcanmulticol = true;
amroutine->amoptionalkey = true; amroutine->amoptionalkey = true;
amroutine->amsearcharray = false; amroutine->amsearcharray = false;
amroutine->amsearchnulls = false; amroutine->amsearchnulls = false;
@@ -303,3 +303,17 @@ hnswhandler(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(amroutine); PG_RETURN_POINTER(amroutine);
} }
/*
* Get the distance between two int4 attributes
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(hnsw_int4_attribute_distance);
Datum
hnsw_int4_attribute_distance(PG_FUNCTION_ARGS)
{
int32 a = PG_GETARG_INT32(0);
int32 b = PG_GETARG_INT32(1);
double distance = ((double) a) - ((double) b);
PG_RETURN_FLOAT8(distance);
}

View File

@@ -19,6 +19,7 @@
#define HNSW_DISTANCE_PROC 1 #define HNSW_DISTANCE_PROC 1
#define HNSW_NORM_PROC 2 #define HNSW_NORM_PROC 2
#define HNSW_TYPE_INFO_PROC 3 #define HNSW_TYPE_INFO_PROC 3
#define HNSW_ATTRIBUTE_DISTANCE_PROC 4
#define HNSW_VERSION 1 #define HNSW_VERSION 1
#define HNSW_MAGIC_NUMBER 0xA953A953 #define HNSW_MAGIC_NUMBER 0xA953A953
@@ -107,18 +108,20 @@
#define HnswPtrPointer(hp) (hp).ptr #define HnswPtrPointer(hp) (hp).ptr
#define HnswPtrOffset(hp) relptr_offset((hp).relptr) #define HnswPtrOffset(hp) relptr_offset((hp).relptr)
#define HnswUseIndexTuple(index) (IndexRelationGetNumberOfAttributes(index) > 1)
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_iterative_search; extern int hnsw_iterative_search;
extern int hnsw_max_search_tuples; extern int hnsw_iterative_search_max_tuples;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef enum HnswIterativeSearchMode typedef enum HnswIterativeSearchType
{ {
HNSW_ITERATIVE_SEARCH_OFF, HNSW_ITERATIVE_SEARCH_OFF,
HNSW_ITERATIVE_SEARCH_RELAXED, HNSW_ITERATIVE_SEARCH_RELAXED,
HNSW_ITERATIVE_SEARCH_STRICT HNSW_ITERATIVE_SEARCH_STRICT
} HnswIterativeSearchMode; } HnswIterativeSearchType;
typedef struct HnswElementData HnswElementData; typedef struct HnswElementData HnswElementData;
typedef struct HnswNeighborArray HnswNeighborArray; typedef struct HnswNeighborArray HnswNeighborArray;
@@ -133,6 +136,7 @@ HnswPtrDeclare(HnswElementData, HnswElementRelptr, HnswElementPtr);
HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr); HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr);
HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr); HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr);
HnswPtrDeclare(char, DatumRelptr, DatumPtr); HnswPtrDeclare(char, DatumRelptr, DatumPtr);
HnswPtrDeclare(IndexTupleData, IndexTupleRelptr, IndexTuplePtr);
struct HnswElementData struct HnswElementData
{ {
@@ -149,6 +153,7 @@ struct HnswElementData
OffsetNumber neighborOffno; OffsetNumber neighborOffno;
BlockNumber neighborPage; BlockNumber neighborPage;
DatumPtr value; DatumPtr value;
IndexTuplePtr itup;
LWLock lock; LWLock lock;
}; };
@@ -174,6 +179,7 @@ typedef struct HnswSearchCandidate
pairingheap_node w_node; pairingheap_node w_node;
HnswElementPtr element; HnswElementPtr element;
double distance; double distance;
bool matches;
} HnswSearchCandidate; } HnswSearchCandidate;
/* HNSW index options */ /* HNSW index options */
@@ -252,14 +258,16 @@ typedef struct HnswTypeInfo
typedef struct HnswSupport typedef struct HnswSupport
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo[2];
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
Oid collation; Oid *collation;
} HnswSupport; } HnswSupport;
typedef struct HnswQuery typedef struct HnswQuery
{ {
Datum value; Datum value;
IndexTuple itup;
ScanKeyData *keyData;
} HnswQuery; } HnswQuery;
typedef struct HnswBuildState typedef struct HnswBuildState
@@ -288,6 +296,8 @@ typedef struct HnswBuildState
HnswGraph *graph; HnswGraph *graph;
double ml; double ml;
int maxLevel; int maxLevel;
bool useIndexTuple;
TupleDesc tupdesc;
/* Memory */ /* Memory */
MemoryContext graphCtx; MemoryContext graphCtx;
@@ -415,30 +425,32 @@ bool HnswCheckNorm(HnswSupport * support, Datum value);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page); void HnswInitPage(Buffer buf, Page page);
void HnswInit(void); void HnswInit(void);
List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples); List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
HnswElement HnswGetEntryPoint(Relation index); HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint); void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
void *HnswAlloc(HnswAllocator * allocator, Size size); void *HnswAlloc(HnswAllocator * allocator, Size size);
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc); HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno); HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing); void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, bool inMemory);
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, HnswQuery * q, Relation rel, HnswSupport * support, bool loadVec); HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, HnswQuery * q, Relation rel, HnswSupport * support, bool loadVec, bool inMemory);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building); void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m); void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator); HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator);
void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc); void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
bool HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPointer heaptid, bool building); bool HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc);
void HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building); void HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec); void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index);
void HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance); void HnswLoadElement(HnswElement element, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance);
bool HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support); void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple);
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support); void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support);
bool HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc);
bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc); bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc);
void HnswInitLockTranche(void); void HnswInitLockTranche(void);
const HnswTypeInfo *HnswGetTypeInfo(Relation index); const HnswTypeInfo *HnswGetTypeInfo(Relation index);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc); PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
Size HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple);
bool HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc);
/* Index access methods */ /* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo); IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);

View File

@@ -148,6 +148,7 @@ CreateGraphPages(HnswBuildState * buildstate)
Page page; Page page;
HnswElementPtr iter = buildstate->graph->head; HnswElementPtr iter = buildstate->graph->head;
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
bool useIndexTuple = buildstate->useIndexTuple;
/* Calculate sizes */ /* Calculate sizes */
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
@@ -167,7 +168,6 @@ CreateGraphPages(HnswBuildState * buildstate)
Size etupSize; Size etupSize;
Size ntupSize; Size ntupSize;
Size combinedSize; Size combinedSize;
Pointer valuePtr = HnswPtrAccess(base, element->value);
/* Update iterator */ /* Update iterator */
iter = element->next; iter = element->next;
@@ -176,7 +176,7 @@ CreateGraphPages(HnswBuildState * buildstate)
MemSet(etup, 0, HNSW_TUPLE_ALLOC_SIZE); MemSet(etup, 0, HNSW_TUPLE_ALLOC_SIZE);
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(valuePtr)); etupSize = HnswGetElementTupleSize(base, element, useIndexTuple);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
@@ -186,7 +186,7 @@ CreateGraphPages(HnswBuildState * buildstate)
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED), (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("index tuple too large"))); errmsg("index tuple too large")));
HnswSetElementTuple(base, etup, element); HnswSetElementTuple(base, etup, element, useIndexTuple);
/* Keep element and neighbors on the same page if possible */ /* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize)) if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
@@ -327,19 +327,18 @@ AddDuplicateInMemory(HnswElement element, HnswElement dup)
* Find duplicate element * Find duplicate element
*/ */
static bool static bool
FindDuplicateInMemory(char *base, HnswElement element) FindDuplicateInMemory(char *base, HnswElement element, bool useIndexTuple, TupleDesc tupdesc)
{ {
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0); HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element); IndexTuple itup = HnswPtrAccess(base, element->itup);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
HnswCandidate *neighbor = &neighbors->items[i]; HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element); HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */ /* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1)) if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
return false; return false;
/* Check for space */ /* Check for space */
@@ -366,7 +365,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
* Update neighbors * Update neighbors
*/ */
static void static void
UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m) UpdateNeighborsInMemory(char *base, Relation index, HnswSupport * support, HnswElement e, int m)
{ {
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
@@ -388,7 +387,7 @@ UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
Assert(neighborElement); Assert(neighborElement);
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE); LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, NULL, support); HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, index, support);
LWLockRelease(&neighborElement->lock); LWLockRelease(&neighborElement->lock);
} }
} }
@@ -404,14 +403,14 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
/* Look for duplicate */ /* Look for duplicate */
if (FindDuplicateInMemory(base, element)) if (FindDuplicateInMemory(base, element, buildstate->useIndexTuple, buildstate->tupdesc))
return; return;
/* Add element */ /* Add element */
AddElementInMemory(base, graph, element); AddElementInMemory(base, graph, element);
/* Update neighbors */ /* Update neighbors */
UpdateNeighborsInMemory(base, support, element, m); UpdateNeighborsInMemory(base, buildstate->index, support, element, m);
/* Update entry point if needed (already have lock) */ /* Update entry point if needed (already have lock) */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -424,6 +423,7 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
static void static void
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element) InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
{ {
Relation index = buildstate->index;
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
HnswSupport *support = &buildstate->support; HnswSupport *support = &buildstate->support;
HnswElement entryPoint; HnswElement entryPoint;
@@ -457,7 +457,7 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
} }
/* Find neighbors for element */ /* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, NULL, support, m, efConstruction, false); HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, true);
/* Update graph in memory */ /* Update graph in memory */
UpdateGraphInMemory(support, element, m, efConstruction, entryPoint, buildstate); UpdateGraphInMemory(support, element, m, efConstruction, entryPoint, buildstate);
@@ -476,18 +476,20 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
HnswElement element; HnswElement element;
HnswAllocator *allocator = &buildstate->allocator; HnswAllocator *allocator = &buildstate->allocator;
HnswSupport *support = &buildstate->support; HnswSupport *support = &buildstate->support;
Size valueSize;
Pointer valuePtr;
LWLock *flushLock = &graph->flushLock; LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
Datum value; TupleDesc tupdesc = buildstate->tupdesc;
IndexTuple itup;
Size itupSize;
IndexTuple itupShared;
bool unused;
/* Form index value */ /* Form index tuple */
if (!HnswFormIndexValue(&value, values, isnull, buildstate->typeInfo, support)) if (!HnswFormIndexTuple(&itup, values, isnull, buildstate->typeInfo, support, tupdesc))
return false; return false;
/* Get datum size */ /* Get tuple size */
valueSize = VARSIZE_ANY(DatumGetPointer(value)); itupSize = IndexTupleSize(itup);
/* Ensure graph not flushed when inserting */ /* Ensure graph not flushed when inserting */
LWLockAcquire(flushLock, LW_SHARED); LWLockAcquire(flushLock, LW_SHARED);
@@ -497,7 +499,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
{ {
LWLockRelease(flushLock); LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, support, value, heaptid, true); return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc);
} }
/* /*
@@ -529,12 +531,12 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
LWLockRelease(flushLock); LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, support, value, heaptid, true); return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc);
} }
/* Ok, we can proceed to allocate the element */ /* Ok, we can proceed to allocate the element */
element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator); element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
valuePtr = HnswAlloc(allocator, valueSize); itupShared = HnswAlloc(allocator, itupSize);
/* /*
* We have now allocated the space needed for the element, so we don't * We have now allocated the space needed for the element, so we don't
@@ -543,9 +545,10 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
*/ */
LWLockRelease(&graph->allocatorLock); LWLockRelease(&graph->allocatorLock);
/* Copy the datum */ /* Copy the tuple */
memcpy(valuePtr, DatumGetPointer(value), valueSize); memcpy(itupShared, itup, itupSize);
HnswPtrStore(base, element->value, valuePtr); HnswPtrStore(base, element->itup, itupShared);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itupShared, 1, tupdesc, &unused)));
/* Create a lock for the element */ /* Create a lock for the element */
LWLockInitialize(&element->lock, hnsw_lock_tranche_id); LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
@@ -672,6 +675,19 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED), (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
errmsg("type not supported for hnsw index"))); errmsg("type not supported for hnsw index")));
/* TODO See if needed */
if (IndexRelationGetNumberOfKeyAttributes(index) > 2)
elog(ERROR, "index cannot have more than two columns");
if (!OidIsValid(index_getprocid(index, 1, HNSW_DISTANCE_PROC)))
elog(ERROR, "first column must be a vector");
for (int i = 1; i < IndexRelationGetNumberOfKeyAttributes(index); i++)
{
if (!OidIsValid(index_getprocid(index, i + 1, HNSW_ATTRIBUTE_DISTANCE_PROC)))
elog(ERROR, "column %d cannot be a vector", i + 1);
}
/* Require column to have dimensions to be indexed */ /* Require column to have dimensions to be indexed */
if (buildstate->dimensions < 0) if (buildstate->dimensions < 0)
ereport(ERROR, ereport(ERROR,
@@ -698,6 +714,8 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->graph = &buildstate->graphData; buildstate->graph = &buildstate->graphData;
buildstate->ml = HnswGetMl(buildstate->m); buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m); buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->useIndexTuple = HnswUseIndexTuple(index);
buildstate->tupdesc = RelationGetDescr(index);
buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext, buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw build graph context", "Hnsw build graph context",

View File

@@ -156,9 +156,10 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
uint8 tupleVersion; uint8 tupleVersion;
char *base = NULL; char *base = NULL;
bool useIndexTuple = HnswUseIndexTuple(index);
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(HnswPtrAccess(base, e->value))); etupSize = HnswGetElementTupleSize(base, e, useIndexTuple);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
@@ -166,7 +167,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
/* Prepare element tuple */ /* Prepare element tuple */
etup = palloc0(etupSize); etup = palloc0(etupSize);
HnswSetElementTuple(base, etup, e); HnswSetElementTuple(base, etup, e, useIndexTuple);
/* Prepare neighbor tuple */ /* Prepare neighbor tuple */
ntup = palloc0(ntupSize); ntup = palloc0(ntupSize);
@@ -383,8 +384,9 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, HnswQuery * q, int *idx, Re
HnswCandidate *hc = &neighbors->items[i]; HnswCandidate *hc = &neighbors->items[i];
HnswElement element = HnswPtrAccess(base, hc->element); HnswElement element = HnswPtrAccess(base, hc->element);
double distance; double distance;
bool matches;
HnswLoadElement(element, &distance, q, index, support, true, NULL); HnswLoadElement(element, &distance, &matches, q, index, support, true, NULL);
hc->distance = distance; hc->distance = distance;
/* Prune element if being deleted */ /* Prune element if being deleted */
@@ -428,6 +430,8 @@ GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int
HnswQuery q; HnswQuery q;
q.value = HnswGetValue(base, element); q.value = HnswGetValue(base, element);
q.itup = HnswPtrAccess(base, element->itup);
q.keyData = NULL;
LoadElementsForInsert(neighbors, &q, &idx, index, support); LoadElementsForInsert(neighbors, &q, &idx, index, support);
@@ -633,21 +637,30 @@ AddDuplicateOnDisk(Relation index, HnswElement element, HnswElement dup, bool bu
* Find duplicate element * Find duplicate element
*/ */
static bool static bool
FindDuplicateOnDisk(Relation index, HnswElement element, bool building) FindDuplicateOnDisk(Relation index, HnswElement element, bool building, TupleDesc tupdesc)
{ {
char *base = NULL; char *base = NULL;
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0); HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element); Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
HnswCandidate *neighbor = &neighbors->items[i]; HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element); HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
if (HnswUseIndexTuple(index))
{
/* Exit early since ordered by distance */ /* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1)) if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
return false; return false;
}
else
{
/* Exit early since ordered by distance */
if (!datumIsEqual(value, HnswGetValue(base, neighborElement), false, -1))
return false;
}
if (AddDuplicateOnDisk(index, element, neighborElement, building)) if (AddDuplicateOnDisk(index, element, neighborElement, building))
return true; return true;
@@ -660,12 +673,12 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
* Update graph on disk * Update graph on disk
*/ */
static void static void
UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building) UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building, TupleDesc tupdesc)
{ {
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Look for duplicate */ /* Look for duplicate */
if (FindDuplicateOnDisk(index, element, building)) if (FindDuplicateOnDisk(index, element, building, tupdesc))
return; return;
/* Add element */ /* Add element */
@@ -687,7 +700,7 @@ UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, in
* Insert a tuple into the index * Insert a tuple into the index
*/ */
bool bool
HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPointer heaptid, bool building) HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc)
{ {
HnswElement entryPoint; HnswElement entryPoint;
HnswElement element; HnswElement element;
@@ -695,6 +708,7 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPo
int efConstruction = HnswGetEfConstruction(index); int efConstruction = HnswGetEfConstruction(index);
LOCKMODE lockmode = ShareLock; LOCKMODE lockmode = ShareLock;
char *base = NULL; char *base = NULL;
bool unused;
/* /*
* Get a shared lock. This allows vacuum to ensure no in-flight inserts * Get a shared lock. This allows vacuum to ensure no in-flight inserts
@@ -708,7 +722,8 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPo
/* Create an element */ /* Create an element */
element = HnswInitElement(base, heaptid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL); element = HnswInitElement(base, heaptid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
HnswPtrStore(base, element->value, DatumGetPointer(value)); HnswPtrStore(base, element->itup, itup);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
/* Prevent concurrent inserts when likely updating entry point */ /* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -725,10 +740,10 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPo
} }
/* Find neighbors for element */ /* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false); HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, false);
/* Update graph on disk */ /* Update graph on disk */
UpdateGraphOnDisk(index, support, element, m, efConstruction, entryPoint, building); UpdateGraphOnDisk(index, support, element, m, efConstruction, entryPoint, building, tupdesc);
/* Release lock */ /* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode); UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
@@ -742,17 +757,18 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPo
static void static void
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid) HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid)
{ {
Datum value; IndexTuple itup;
const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index); const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index);
TupleDesc tupdesc = RelationGetDescr(index);
HnswSupport support; HnswSupport support;
HnswInitSupport(&support, index); HnswInitSupport(&support, index);
/* Form index value */ /* Form index tuple */
if (!HnswFormIndexValue(&value, values, isnull, typeInfo, &support)) if (!HnswFormIndexTuple(&itup, values, isnull, typeInfo, &support, tupdesc))
return; return;
HnswInsertTupleOnDisk(index, &support, value, heaptid, false); HnswInsertTupleOnDisk(index, &support, itup, heaptid, false, tupdesc);
} }
/* /*

View File

@@ -22,26 +22,30 @@ GetScanItems(IndexScanDesc scan, Datum value)
int m; int m;
HnswElement entryPoint; HnswElement entryPoint;
char *base = NULL; char *base = NULL;
bool inMemory = false;
HnswQuery *q = &so->q; HnswQuery *q = &so->q;
q->value = value;
q->itup = NULL;
q->keyData = scan->keyData;
/* Get m and entry point */ /* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
q->value = value;
so->m = m; so->m = m;
if (entryPoint == NULL) if (entryPoint == NULL)
return NIL; return NIL;
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, false)); ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, false, inMemory));
for (int lc = entryPoint->level; lc >= 1; lc--) for (int lc = entryPoint->level; lc >= 1; lc--)
{ {
w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, NULL, NULL, true, NULL); w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, inMemory, NULL, NULL, true, NULL);
ep = w; ep = w;
} }
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, &so->v, hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF ? &so->discarded : NULL, true, &so->tuples); return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, inMemory, &so->v, hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF ? &so->discarded : NULL, true, &so->tuples);
} }
/* /*
@@ -72,7 +76,7 @@ ResumeScanItems(IndexScanDesc scan)
ep = lappend(ep, sc); ep = lappend(ep, sc);
} }
return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples); return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, false, &so->v, &so->discarded, false, &so->tuples);
} }
/* /*
@@ -96,23 +100,12 @@ GetScanValue(IndexScanDesc scan)
/* Normalize if needed */ /* Normalize if needed */
if (so->support.normprocinfo != NULL) if (so->support.normprocinfo != NULL)
value = HnswNormValue(so->typeInfo, so->support.collation, value); value = HnswNormValue(so->typeInfo, so->support.collation[0], value);
} }
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
*/ */
@@ -129,14 +122,9 @@ 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",
0, 8 * 1024, 512 * 1024); ALLOCSET_DEFAULT_SIZES);
/* Set support functions */ /* Set support functions */
HnswInitSupport(&so->support, index); HnswInitSupport(&so->support, index);
@@ -220,7 +208,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->first = false; so->first = false;
#if defined(HNSW_MEMORY) #if defined(HNSW_MEMORY)
ShowMemoryUsage(so); elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
#endif #endif
} }
@@ -240,8 +228,8 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
if (so->discarded == NULL) if (so->discarded == NULL)
break; break;
/* Reached max number of tuples */ /* Reached max number of additional tuples */
if (hnsw_max_search_tuples != -1 && so->tuples >= hnsw_max_search_tuples) if (hnsw_iterative_search_max_tuples != -1 && so->tuples >= hnsw_iterative_search_max_tuples)
{ {
if (pairingheap_is_empty(so->discarded)) if (pairingheap_is_empty(so->discarded))
break; break;
@@ -282,7 +270,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)
ShowMemoryUsage(so); elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
#endif #endif
} }
@@ -294,7 +282,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
element = HnswPtrAccess(base, sc->element); element = HnswPtrAccess(base, sc->element);
/* Move to next element if no valid heap TIDs */ /* Move to next element if no valid heap TIDs */
if (element->heaptidsLength == 0) if (!sc->matches || element->heaptidsLength == 0)
{ {
so->w = list_delete_last(so->w); so->w = list_delete_last(so->w);

View File

@@ -1,6 +1,5 @@
#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"
@@ -147,20 +146,45 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
void void
HnswInitSupport(HnswSupport * support, Relation index) HnswInitSupport(HnswSupport * support, Relation index)
{ {
support->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); support->procinfo[0] = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
support->collation = index->rd_indcollation[0];
if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
support->procinfo[1] = index_getprocinfo(index, 2, HNSW_ATTRIBUTE_DISTANCE_PROC);
support->collation = index->rd_indcollation;
support->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); support->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
} }
/*
* Get element tuple size
*/
Size
HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple)
{
Size size;
if (useIndexTuple)
{
IndexTuple itup = HnswPtrAccess(base, element->itup);
size = IndexTupleSize(itup);
}
else
{
Pointer valuePtr = HnswPtrAccess(base, element->value);
size = VARSIZE_ANY(valuePtr);
}
return HNSW_ELEMENT_TUPLE_SIZE(size);
}
/* /*
* Normalize value * Normalize value
*/ */
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);
} }
@@ -170,7 +194,38 @@ HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value)
bool bool
HnswCheckNorm(HnswSupport * support, Datum value) HnswCheckNorm(HnswSupport * support, Datum value)
{ {
return DatumGetFloat8(FunctionCall1Coll(support->normprocinfo, support->collation, value)) > 0; return DatumGetFloat8(FunctionCall1Coll(support->normprocinfo, support->collation[0], value)) > 0;
}
/*
* Check if index tuples are equal
*/
bool
HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc)
{
for (int i = 0; i < tupdesc->natts; i++)
{
bool nullA;
bool nullB;
Datum datumA = index_getattr(a, i + 1, tupdesc, &nullA);
Datum datumB = index_getattr(b, i + 1, tupdesc, &nullB);
if (nullA || nullB)
{
if (nullA != nullB)
return false;
}
else
{
Form_pg_attribute att = TupleDescAttr(tupdesc, i);
if (!datumIsEqual(datumA, datumB, att->attbyval, att->attlen))
return false;
}
}
return true;
} }
/* /*
@@ -261,6 +316,7 @@ HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel,
HnswInitNeighbors(base, element, m, allocator); HnswInitNeighbors(base, element, m, allocator);
HnswPtrStore(base, element->value, (Pointer) NULL); HnswPtrStore(base, element->value, (Pointer) NULL);
HnswPtrStore(base, element->itup, (IndexTuple) NULL);
return element; return element;
} }
@@ -287,6 +343,7 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
element->offno = offno; element->offno = offno;
HnswPtrStore(base, element->neighbors, (HnswNeighborArrayPtr *) NULL); HnswPtrStore(base, element->neighbors, (HnswNeighborArrayPtr *) NULL);
HnswPtrStore(base, element->value, (Pointer) NULL); HnswPtrStore(base, element->value, (Pointer) NULL);
HnswPtrStore(base, element->itup, (IndexTuple) NULL);
return element; return element;
} }
@@ -399,11 +456,13 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
} }
/* /*
* Form index value * Form index tuple
*/ */
bool bool
HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support) HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc)
{ {
Datum newValues[2];
/* Detoast once for all calls */ /* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -417,10 +476,14 @@ HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo *
if (!HnswCheckNorm(support, value)) if (!HnswCheckNorm(support, value))
return false; return false;
value = HnswNormValue(typeInfo, support->collation, value); value = HnswNormValue(typeInfo, support->collation[0], value);
} }
*out = value; newValues[0] = value;
for (int i = 1; i < tupdesc->natts; i++)
newValues[i] = values[i];
*out = index_form_tuple(tupdesc, newValues, isnull);
return true; return true;
} }
@@ -429,10 +492,8 @@ HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo *
* Set element tuple, except for neighbor info * Set element tuple, except for neighbor info
*/ */
void void
HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element) HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple)
{ {
Pointer valuePtr = HnswPtrAccess(base, element->value);
etup->type = HNSW_ELEMENT_TUPLE_TYPE; etup->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level; etup->level = element->level;
etup->deleted = 0; etup->deleted = 0;
@@ -444,8 +505,20 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
else else
ItemPointerSetInvalid(&etup->heaptids[i]); ItemPointerSetInvalid(&etup->heaptids[i]);
} }
if (useIndexTuple)
{
IndexTuple itup = HnswPtrAccess(base, element->itup);
memcpy(&etup->data, itup, IndexTupleSize(itup));
}
else
{
Pointer valuePtr = HnswPtrAccess(base, element->value);
memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr)); memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr));
} }
}
/* /*
* Set neighbor tuple * Set neighbor tuple
@@ -486,7 +559,7 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
* Load an element from a tuple * Load an element from a tuple
*/ */
void void
HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec) HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index)
{ {
element->level = etup->level; element->level = etup->level;
element->deleted = etup->deleted; element->deleted = etup->deleted;
@@ -510,26 +583,128 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
if (loadVec) if (loadVec)
{ {
char *base = NULL; char *base = NULL;
if (HnswUseIndexTuple(index))
{
IndexTuple itup = CopyIndexTuple((IndexTuple) &etup->data);
TupleDesc tupdesc = RelationGetDescr(index);
bool unused;
HnswPtrStore(base, element->itup, itup);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
}
else
{
Datum value = datumCopy(PointerGetDatum(&etup->data), false, -1); Datum value = datumCopy(PointerGetDatum(&etup->data), false, -1);
HnswPtrStore(base, element->value, DatumGetPointer(value)); HnswPtrStore(base, element->value, DatumGetPointer(value));
} }
} }
}
/*
* Get the attribute distance
*/
static inline double
AttributeDistance(double e)
{
/* TODO Better bias */
/* must be >> max(w * g) + 1 / log10(2) */
double bias = 4.32;
return e > 0 ? bias - 1.0 / log10(e + 1) : 0;
}
/* /*
* Calculate the distance between values * Calculate the distance between values
*/ */
static inline double static double
HnswGetDistance(Datum a, Datum b, HnswSupport * support) HnswGetDistance(IndexTuple itup, Datum vec, HnswQuery * q, Relation index, HnswSupport * support, bool *matches)
{ {
return DatumGetFloat8(FunctionCall2Coll(support->procinfo, support->collation, a, b)); double g;
if (DatumGetPointer(q->value) == NULL)
g = 0;
else
g = DatumGetFloat8(FunctionCall2Coll(support->procinfo[0], support->collation[0], q->value, vec));
Assert(PointerIsValid(matches));
*matches = true;
if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
{
double w = 0.25;
double e = 0.0;
TupleDesc tupdesc = RelationGetDescr(index);
if (q->keyData)
{
/* TODO need to pass length of key data */
int keyCount = 1;
for (int i = 0; i < keyCount; i++)
{
ScanKey key = &q->keyData[i];
bool isnull;
Datum value = index_getattr(itup, key->sk_attno, tupdesc, &isnull);
bool attnull = key->sk_flags & SK_ISNULL;
if (isnull || attnull)
{
if (isnull != attnull)
{
e += 1000;
*matches = false;
}
}
else if (!DatumGetBool(FunctionCall2Coll(&key->sk_func, key->sk_collation, value, key->sk_argument)))
{
double ei = fabs(DatumGetFloat8(FunctionCall2Coll(support->procinfo[key->sk_attno - 1], support->collation[key->sk_attno - 1], value, key->sk_argument)));
if (ei > 0)
e += ei;
else
/* Distance is zero for inequality */
e += 1000;
*matches = false;
}
}
return w * g + AttributeDistance(e);
}
else if (q->itup)
{
int keyCount = IndexRelationGetNumberOfKeyAttributes(index) - 1;
for (int i = 0; i < keyCount; i++)
{
bool isnull;
bool attnull;
Datum value = index_getattr(itup, i + 2, tupdesc, &isnull);
Datum value2 = index_getattr(q->itup, i + 2, tupdesc, &attnull);
if (isnull || attnull)
{
if (isnull != attnull)
e += 1000;
}
else
e += fabs(DatumGetFloat8(FunctionCall2Coll(support->procinfo[i + 1], support->collation[i + 1], value, value2)));
}
return w * g + AttributeDistance(e);
}
}
return g;
} }
/* /*
* Load an element and optionally get its distance from q * Load an element and optionally get its distance from q
*/ */
static void static void
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element) HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -547,17 +722,23 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Hns
/* Calculate distance */ /* Calculate distance */
if (distance != NULL) if (distance != NULL)
{ {
if (DatumGetPointer(q->value) == NULL) IndexTuple itup = NULL;
*distance = 0; Datum value;
if (HnswUseIndexTuple(index))
{
TupleDesc tupdesc = RelationGetDescr(index);
bool unused;
itup = (IndexTuple) &etup->data;
value = index_getattr(itup, 1, tupdesc, &unused);
}
else else
{ {
*distance = HnswGetDistance(q->value, PointerGetDatum(&etup->data), support); value = PointerGetDatum(&etup->data);
/* Needed for intvec cosine distance */
/* TODO Improve */
if (isnan(*distance))
*distance = DBL_MAX;
} }
*distance = HnswGetDistance(itup, value, q, index, support, matches);
} }
/* Load element */ /* Load element */
@@ -566,7 +747,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Hns
if (*element == NULL) if (*element == NULL)
*element = HnswInitElementFromBlock(blkno, offno); *element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(*element, etup, true, loadVec); HnswLoadElementFromTuple(*element, etup, true, loadVec, index);
} }
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
@@ -576,50 +757,37 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Hns
* Load an element and optionally get its distance from q * Load an element and optionally get its distance from q
*/ */
void void
HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance) HnswLoadElement(HnswElement element, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
{ {
HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, support, loadVec, maxDistance, &element); HnswLoadElementImpl(element->blkno, element->offno, distance, matches, q, index, support, loadVec, maxDistance, &element);
} }
/* /*
* Get the distance for an element * Get the distance for an element
*/ */
static double static double
GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport * support) GetElementDistance(char *base, HnswElement element, bool *matches, HnswQuery * q, Relation index, HnswSupport * support)
{ {
Datum value = HnswGetValue(base, element); Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup);
return HnswGetDistance(q->value, value, support); return HnswGetDistance(itup, value, q, index, support, matches);
}
/*
* 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, bool inMemory)
{ {
bool inMemory = index == NULL; HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
double distance;
HnswPtrStore(base, sc->element, entryPoint);
if (inMemory) if (inMemory)
distance = GetElementDistance(base, entryPoint, q, support); sc->distance = GetElementDistance(base, entryPoint, &sc->matches, q, index, support);
else else
HnswLoadElement(entryPoint, &distance, q, index, support, loadVec, NULL); HnswLoadElement(entryPoint, &sc->distance, &sc->matches, q, index, support, loadVec, NULL);
return sc;
return HnswInitSearchCandidate(base, entryPoint, distance);
} }
/* /*
@@ -822,7 +990,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
* Algorithm 2 from paper * Algorithm 2 from paper
*/ */
List * List *
HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples) HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples)
{ {
List *w = NIL; List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -835,7 +1003,8 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited)); HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
int unvisitedLength; int unvisitedLength;
bool inMemory = index == NULL; uint64 additional = 0;
uint64 maxAdditional = q->keyData && lc == 0 ? 10000 : 0;
if (v == NULL) if (v == NULL)
{ {
@@ -875,6 +1044,10 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
pairingheap_add(C, &sc->c_node); pairingheap_add(C, &sc->c_node);
pairingheap_add(W, &sc->w_node); pairingheap_add(W, &sc->w_node);
/* Do not count elements that do not match filter towards ef */
if (!sc->matches && ++additional <= maxAdditional)
continue;
/* /*
* Do not count elements being deleted towards ef when vacuuming. It * Do not count elements being deleted towards ef when vacuuming. It
* would be ideal to do this for inserts as well, but this could * would be ideal to do this for inserts as well, but this could
@@ -908,6 +1081,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
HnswElement eElement; HnswElement eElement;
HnswSearchCandidate *e; HnswSearchCandidate *e;
double eDistance; double eDistance;
bool eMatches;
bool alwaysAdd = wlen < ef; bool alwaysAdd = wlen < ef;
f = HnswGetSearchCandidate(w_node, pairingheap_first(W)); f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
@@ -915,7 +1089,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
if (inMemory) if (inMemory)
{ {
eElement = unvisited[i].element; eElement = unvisited[i].element;
eDistance = GetElementDistance(base, eElement, q, support); eDistance = GetElementDistance(base, eElement, &eMatches, q, index, support);
} }
else else
{ {
@@ -925,7 +1099,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
/* Avoid any allocations if not adding */ /* Avoid any allocations if not adding */
eElement = NULL; eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement); HnswLoadElementImpl(blkno, offno, &eDistance, &eMatches, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
if (eElement == NULL) if (eElement == NULL)
continue; continue;
@@ -936,7 +1110,9 @@ 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 = HnswInitSearchCandidate(base, eElement, eDistance); e = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
pairingheap_add(*discarded, &e->w_node); pairingheap_add(*discarded, &e->w_node);
} }
@@ -948,7 +1124,10 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
continue; continue;
/* Create a new candidate */ /* Create a new candidate */
e = HnswInitSearchCandidate(base, eElement, eDistance); e = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
e->matches = eMatches;
pairingheap_add(C, &e->c_node); pairingheap_add(C, &e->c_node);
pairingheap_add(W, &e->w_node); pairingheap_add(W, &e->w_node);
@@ -959,6 +1138,10 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
*/ */
if (CountElement(skipElement, eElement)) if (CountElement(skipElement, eElement))
{ {
/* Do not count elements that do not match filter towards ef */
if (!e->matches && ++additional <= maxAdditional)
continue;
wlen++; wlen++;
/* No need to decrement wlen */ /* No need to decrement wlen */
@@ -1036,18 +1219,24 @@ CompareCandidateDistancesOffset(const ListCell *a, const ListCell *b)
* Check if an element is closer to q than any element from R * Check if an element is closer to q than any element from R
*/ */
static bool static bool
CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support) CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, HnswSupport * support)
{ {
HnswElement eElement = HnswPtrAccess(base, e->element); HnswElement eElement = HnswPtrAccess(base, e->element);
Datum eValue = HnswGetValue(base, eElement); HnswQuery q;
ListCell *lc2; ListCell *lc2;
q.value = HnswGetValue(base, eElement);
q.itup = HnswPtrAccess(base, eElement->itup);
q.keyData = NULL;
foreach(lc2, r) foreach(lc2, r)
{ {
HnswCandidate *ri = lfirst(lc2); HnswCandidate *ri = lfirst(lc2);
HnswElement riElement = HnswPtrAccess(base, ri->element); HnswElement riElement = HnswPtrAccess(base, ri->element);
Datum riValue = HnswGetValue(base, riElement); Datum riValue = HnswGetValue(base, riElement);
float distance = HnswGetDistance(eValue, riValue, support); IndexTuple ritup = HnswPtrAccess(base, riElement->itup);
bool matches;
float distance = HnswGetDistance(ritup, riValue, &q, index, support, &matches);
if (distance <= e->distance) if (distance <= e->distance)
return false; return false;
@@ -1060,7 +1249,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support
* Algorithm 4 from paper * Algorithm 4 from paper
*/ */
static List * static List *
SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closerSet, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates) SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * support, bool *closerSet, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
{ {
List *r = NIL; List *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
@@ -1094,7 +1283,7 @@ SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closer
/* Use previous state of r and wd to skip work when possible */ /* Use previous state of r and wd to skip work when possible */
if (mustCalculate) if (mustCalculate)
e->closer = CheckElementCloser(base, e, r, support); e->closer = CheckElementCloser(base, e, r, index, support);
else if (list_length(added) > 0) else if (list_length(added) > 0)
{ {
/* Keep Valgrind happy for in-memory, parallel builds */ /* Keep Valgrind happy for in-memory, parallel builds */
@@ -1107,8 +1296,7 @@ SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closer
*/ */
if (e->closer) if (e->closer)
{ {
e->closer = CheckElementCloser(base, e, added, support); e->closer = CheckElementCloser(base, e, added, index, support);
if (!e->closer) if (!e->closer)
removedAny = true; removedAny = true;
} }
@@ -1120,7 +1308,7 @@ SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closer
*/ */
if (removedAny) if (removedAny)
{ {
e->closer = CheckElementCloser(base, e, r, support); e->closer = CheckElementCloser(base, e, r, index, support);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -1128,7 +1316,7 @@ SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closer
} }
else if (e == newCandidate) else if (e == newCandidate)
{ {
e->closer = CheckElementCloser(base, e, r, support); e->closer = CheckElementCloser(base, e, r, index, support);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -1205,7 +1393,7 @@ HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newE
c = lappend(c, &neighbors->items[i]); c = lappend(c, &neighbors->items[i]);
c = lappend(c, &newHc); c = lappend(c, &newHc);
SelectNeighbors(base, c, lm, support, &neighbors->closerSet, &newHc, &pruned, true); SelectNeighbors(base, c, lm, index, support, &neighbors->closerSet, &newHc, &pruned, true);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -1276,17 +1464,19 @@ PrecomputeHash(char *base, HnswElement element)
* Algorithm 1 from paper * Algorithm 1 from paper
*/ */
void void
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing) HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, bool inMemory)
{ {
List *ep; List *ep;
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
HnswQuery q; HnswQuery q;
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
bool inMemory = index == NULL;
q.value = HnswGetValue(base, element); q.value = HnswGetValue(base, element);
q.itup = HnswPtrAccess(base, element->itup);
q.keyData = NULL;
/* Precompute hash */ /* Precompute hash */
if (inMemory) if (inMemory)
@@ -1297,13 +1487,13 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
return; return;
/* Get entry point and level */ /* Get entry point and level */
ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true)); ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true, inMemory));
entryLevel = entryPoint->level; entryLevel = entryPoint->level;
/* 1st phase: greedy search to insert level */ /* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--) for (int lc = entryLevel; lc >= level + 1; lc--)
{ {
w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL); w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL);
ep = w; ep = w;
} }
@@ -1322,7 +1512,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
List *lw = NIL; List *lw = NIL;
ListCell *lc2; ListCell *lc2;
w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL); w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL);
/* Convert search candidates to candidates */ /* Convert search candidates to candidates */
foreach(lc2, w) foreach(lc2, w)
@@ -1346,7 +1536,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
* sortCandidates to true for in-memory builds to enable closer * sortCandidates to true for in-memory builds to enable closer
* caching, but there does not seem to be a difference in performance. * caching, but there does not seem to be a difference in performance.
*/ */
neighbors = SelectNeighbors(base, lw, lm, support, &HnswGetNeighbors(base, element, lc)->closerSet, NULL, NULL, false); neighbors = SelectNeighbors(base, lw, lm, index, support, &HnswGetNeighbors(base, element, lc)->closerSet, NULL, NULL, false);
AddConnections(base, element, neighbors, lc); AddConnections(base, element, neighbors, lc);
@@ -1404,19 +1594,6 @@ 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)

View File

@@ -204,7 +204,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
element->heaptidsLength = 0; element->heaptidsLength = 0;
/* Find neighbors for element, skipping itself */ /* Find neighbors for element, skipping itself */
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true); HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true, false);
/* Zero memory for each element */ /* Zero memory for each element */
MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE); MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);
@@ -256,7 +256,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
LockPage(index, HNSW_UPDATE_LOCK, ShareLock); LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
/* Load element */ /* Load element */
HnswLoadElement(highestPoint, NULL, NULL, index, support, true, NULL); HnswLoadElement(highestPoint, NULL, NULL, NULL, index, support, true, NULL);
/* Repair if needed */ /* Repair if needed */
if (NeedsUpdated(vacuumstate, highestPoint)) if (NeedsUpdated(vacuumstate, highestPoint))
@@ -294,7 +294,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
* is outdated, this can remove connections at higher levels in * is outdated, this can remove connections at higher levels in
* the graph until they are repaired, but this should be fine. * the graph until they are repaired, but this should be fine.
*/ */
HnswLoadElement(entryPoint, NULL, NULL, index, support, true, NULL); HnswLoadElement(entryPoint, NULL, NULL, NULL, index, support, true, NULL);
if (NeedsUpdated(vacuumstate, entryPoint)) if (NeedsUpdated(vacuumstate, entryPoint))
{ {
@@ -370,7 +370,7 @@ RepairGraph(HnswVacuumState * vacuumstate)
/* Create an element */ /* Create an element */
element = HnswInitElementFromBlock(blkno, offno); element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(element, etup, false, true); HnswLoadElementFromTuple(element, etup, false, true, index);
elements = lappend(elements, element); elements = lappend(elements, element);
} }
@@ -440,6 +440,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
BlockNumber insertPage = InvalidBlockNumber; BlockNumber insertPage = InvalidBlockNumber;
Relation index = vacuumstate->index; Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas; BufferAccessStrategy bas = vacuumstate->bas;
bool useIndexTuple = HnswUseIndexTuple(index);
/* /*
* Wait for index scans to complete. Scans before this point may contain * Wait for index scans to complete. Scans before this point may contain
@@ -521,6 +522,13 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Overwrite element */ /* Overwrite element */
etup->deleted = 1; etup->deleted = 1;
if (useIndexTuple)
{
IndexTuple itup = (IndexTuple) &etup->data;
MemSet(itup, 0, IndexTupleSize(itup));
}
else
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data)); MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
/* Overwrite neighbors */ /* Overwrite neighbors */

View File

@@ -1,754 +0,0 @@
#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));
}

View File

@@ -1,23 +0,0 @@
#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

@@ -138,7 +138,7 @@ SampleRows(IvfflatBuildState * buildstate)
* Add tuple to sort * Add tuple to sort
*/ */
static void static void
AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState * buildstate) AddTupleToSort(Relation index, ItemPointer tid, Datum *values, bool *isnull, IvfflatBuildState * buildstate)
{ {
double distance; double distance;
double minDistance = DBL_MAX; double minDistance = DBL_MAX;
@@ -184,6 +184,11 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState
slot->tts_isnull[1] = false; slot->tts_isnull[1] = false;
slot->tts_values[2] = value; slot->tts_values[2] = value;
slot->tts_isnull[2] = false; slot->tts_isnull[2] = false;
for (int i = 1; i < buildstate->tupdesc->natts; i++)
{
slot->tts_values[2 + i] = values[i];
slot->tts_isnull[2 + i] = isnull[i];
}
ExecStoreVirtualTuple(slot); ExecStoreVirtualTuple(slot);
/* /*
@@ -215,7 +220,7 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Add tuple to sort */ /* Add tuple to sort */
AddTupleToSort(index, tid, values, buildstate); AddTupleToSort(index, tid, values, isnull, buildstate);
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -226,19 +231,20 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
* Get index tuple from sort state * Get index tuple from sort state
*/ */
static inline void static inline void
GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, IndexTuple *itup, int *list) GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, Datum *values, bool *isnull, IndexTuple *itup, int *list)
{ {
if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL)) if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
{ {
Datum value; bool unused;
bool isnull;
*list = DatumGetInt32(slot_getattr(slot, 1, &isnull)); *list = DatumGetInt32(slot_getattr(slot, 1, &unused));
value = slot_getattr(slot, 3, &isnull);
for (int i = 0; i < tupdesc->natts; i++)
values[i] = slot_getattr(slot, 3 + i, &isnull[i]);
/* Form the index tuple */ /* Form the index tuple */
*itup = index_form_tuple(tupdesc, &value, &isnull); *itup = index_form_tuple(tupdesc, values, isnull);
(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &isnull))); (*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &unused)));
} }
else else
*list = -1; *list = -1;
@@ -256,12 +262,14 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple); TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple);
TupleDesc tupdesc = buildstate->tupdesc; TupleDesc tupdesc = buildstate->tupdesc;
Datum *values = palloc(tupdesc->natts * sizeof(Datum));
bool *isnull = palloc(tupdesc->natts * sizeof(bool));
pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD); pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD);
pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_TOTAL, buildstate->indtuples); pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_TOTAL, buildstate->indtuples);
GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list); GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list);
for (int i = 0; i < buildstate->centers->length; i++) for (int i = 0; i < buildstate->centers->length; i++)
{ {
@@ -297,7 +305,7 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_DONE, ++inserted); pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_DONE, ++inserted);
GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list); GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list);
} }
insertPage = BufferGetBlockNumber(buf); insertPage = BufferGetBlockNumber(buf);
@@ -307,6 +315,9 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
/* Set the start and insert pages */ /* Set the start and insert pages */
IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum); IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum);
} }
pfree(values);
pfree(isnull);
} }
/* /*
@@ -357,10 +368,11 @@ 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->sortdesc = CreateTemplateTupleDesc(3); buildstate->sortdesc = CreateTemplateTupleDesc(2 + buildstate->tupdesc->natts);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 3, "vector", buildstate->tupdesc->attrs[0].atttypid, -1, 0); for (int i = 0; i < buildstate->tupdesc->natts; i++)
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) (3 + i), NULL, buildstate->tupdesc->attrs[i].atttypid, -1, 0);
buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual); buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual);

View File

@@ -18,12 +18,12 @@
int ivfflat_probes; int ivfflat_probes;
int ivfflat_iterative_search; int ivfflat_iterative_search;
int ivfflat_max_probes; int ivfflat_iterative_search_max_probes;
static relopt_kind ivfflat_relopt_kind; static relopt_kind ivfflat_relopt_kind;
static const struct config_enum_entry ivfflat_iterative_search_options[] = { static const struct config_enum_entry ivfflat_iterative_search_options[] = {
{"off", IVFFLAT_ITERATIVE_SEARCH_OFF, false}, {"off", IVFFLAT_ITERATIVE_SEARCH_OFF, false},
{"relaxed_order", IVFFLAT_ITERATIVE_SEARCH_RELAXED, false}, {"on", IVFFLAT_ITERATIVE_SEARCH_RELAXED, false},
{NULL, 0, false} {NULL, 0, false}
}; };
@@ -41,14 +41,13 @@ 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", DefineCustomEnumVariable("ivfflat.iterative_search", "Sets whether to use iterative search",
NULL, &ivfflat_iterative_search, NULL, &ivfflat_iterative_search,
IVFFLAT_ITERATIVE_SEARCH_OFF, ivfflat_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL); 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.iterative_search_max_probes", "Sets the max number of probes for iterative search",
DefineCustomIntVariable("ivfflat.max_probes", "Sets the max number of probes for iterative search", "Zero sets to the number of lists", &ivfflat_iterative_search_max_probes,
"-1 means no limit", &ivfflat_max_probes, 0, 0, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
-1, -1, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("ivfflat"); MarkGUCPrefixReserved("ivfflat");
} }

View File

@@ -81,13 +81,13 @@
/* Variables */ /* Variables */
extern int ivfflat_probes; extern int ivfflat_probes;
extern int ivfflat_iterative_search; extern int ivfflat_iterative_search;
extern int ivfflat_max_probes; extern int ivfflat_iterative_search_max_probes;
typedef enum IvfflatIterativeSearchMode typedef enum IvfflatIterativeSearchType
{ {
IVFFLAT_ITERATIVE_SEARCH_OFF, IVFFLAT_ITERATIVE_SEARCH_OFF,
IVFFLAT_ITERATIVE_SEARCH_RELAXED IVFFLAT_ITERATIVE_SEARCH_RELAXED
} IvfflatIterativeSearchMode; } IvfflatIterativeSearchType;
typedef struct VectorArrayData typedef struct VectorArrayData
{ {
@@ -260,7 +260,6 @@ typedef struct IvfflatScanOpaqueData
int dimensions; int dimensions;
bool first; bool first;
Datum value; Datum value;
MemoryContext tmpCtx;
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
@@ -279,7 +278,7 @@ typedef struct IvfflatScanOpaqueData
pairingheap *listQueue; pairingheap *listQueue;
BlockNumber *listPages; BlockNumber *listPages;
int listIndex; int listIndex;
IvfflatScanList *lists; IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */
} IvfflatScanOpaqueData; } IvfflatScanOpaqueData;
typedef IvfflatScanOpaqueData * IvfflatScanOpaque; typedef IvfflatScanOpaqueData * IvfflatScanOpaque;

View File

@@ -78,6 +78,8 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
BlockNumber insertPage = InvalidBlockNumber; BlockNumber insertPage = InvalidBlockNumber;
ListInfo listInfo; ListInfo listInfo;
BlockNumber originalInsertPage; BlockNumber originalInsertPage;
TupleDesc tupdesc = RelationGetDescr(index);
Datum *newValues = palloc(tupdesc->natts * sizeof(Datum));
/* Detoast once for all calls */ /* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -102,8 +104,12 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
Assert(BlockNumberIsValid(insertPage)); Assert(BlockNumberIsValid(insertPage));
originalInsertPage = insertPage; originalInsertPage = insertPage;
newValues[0] = value;
for (int i = 1; i < tupdesc->natts; i++)
newValues[i] = values[i];
/* Form tuple */ /* Form tuple */
itup = index_form_tuple(RelationGetDescr(index), &value, isnull); itup = index_form_tuple(tupdesc, newValues, isnull);
itup->t_tid = *heap_tid; itup->t_tid = *heap_tid;
/* Get tuple size */ /* Get tuple size */

View File

@@ -10,7 +10,10 @@
#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)
@@ -218,13 +221,7 @@ 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;
@@ -256,35 +253,26 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
int dimensions; int dimensions;
int probes = ivfflat_probes; int probes = ivfflat_probes;
int maxProbes; 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 (probes > lists)
probes = lists;
if (ivfflat_iterative_search != IVFFLAT_ITERATIVE_SEARCH_OFF) if (ivfflat_iterative_search != IVFFLAT_ITERATIVE_SEARCH_OFF)
{ {
maxProbes = ivfflat_max_probes; if (ivfflat_iterative_search_max_probes == 0)
if (maxProbes < 0)
maxProbes = lists; maxProbes = lists;
else if (maxProbes < probes) else
{ maxProbes = Min(ivfflat_iterative_search_max_probes, lists);
/* TODO Show notice */
maxProbes = probes;
}
} }
else else
maxProbes = probes; maxProbes = probes;
if (probes > lists) so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + maxProbes * sizeof(IvfflatScanList));
probes = lists;
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;
@@ -296,12 +284,6 @@ 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);
@@ -324,9 +306,6 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->listQueue = pairingheap_allocate(CompareLists, scan); so->listQueue = pairingheap_allocate(CompareLists, scan);
so->listPages = palloc(maxProbes * sizeof(BlockNumber)); so->listPages = palloc(maxProbes * sizeof(BlockNumber));
so->listIndex = 0; so->listIndex = 0;
so->lists = palloc(maxProbes * sizeof(IvfflatScanList));
MemoryContextSwitchTo(oldCtx);
scan->opaque = so; scan->opaque = so;
@@ -389,6 +368,8 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
IvfflatBench("GetScanItems", GetScanItems(scan, value)); IvfflatBench("GetScanItems", GetScanItems(scan, value));
so->first = false; so->first = false;
so->value = value; so->value = value;
/* TODO clean up if we allocated a new value */
} }
while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL)) while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
@@ -415,10 +396,13 @@ ivfflatendscan(IndexScanDesc scan)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
/* Free any temporary files */ pairingheap_free(so->listQueue);
pfree(so->listPages);
tuplesort_end(so->sortstate); tuplesort_end(so->sortstate);
FreeAccessStrategy(so->bas);
FreeTupleDesc(so->tupdesc);
MemoryContextDelete(so->tmpCtx); /* TODO Free vslot and mslot without freeing TupleDesc */
pfree(so); pfree(so);
scan->opaque = NULL; scan->opaque = NULL;

View File

@@ -140,34 +140,6 @@ 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,23 +30,6 @@ 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

@@ -1,102 +0,0 @@
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,32 +99,6 @@ 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));
@@ -165,21 +139,4 @@ 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;

View File

@@ -1,328 +0,0 @@
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,44 +81,6 @@ 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));
@@ -147,27 +109,4 @@ 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,14 +38,6 @@ 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,21 +28,6 @@ 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));

View File

@@ -1,58 +0,0 @@
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,23 +57,6 @@ 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));
@@ -98,12 +81,4 @@ 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;

View File

@@ -1,73 +0,0 @@
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,28 +44,6 @@ 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));
@@ -84,16 +62,4 @@ 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

@@ -40,10 +40,6 @@ 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,10 +45,6 @@ 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", "intvec", "sparsevec"); my @types = ("vector", "halfvec", "sparsevec");
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 @inputs = ("[1.23,4.56,7.89]", "[1.23,4.56,7.89]", "{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

@@ -23,7 +23,7 @@ $node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops
my $count = $node->safe_psql("postgres", qq( my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET ivfflat.probes = 10; SET ivfflat.probes = 10;
SET ivfflat.iterative_search = relaxed_order; SET ivfflat.iterative_search = on;
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;
)); ));
is($count, 10); is($count, 10);
@@ -39,8 +39,8 @@ foreach ((30, 50, 70))
$count = $node->safe_psql("postgres", qq( $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET ivfflat.probes = 10; SET ivfflat.probes = 10;
SET ivfflat.iterative_search = relaxed_order; SET ivfflat.iterative_search = on;
SET ivfflat.max_probes = $max_probes; SET ivfflat.iterative_search_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; 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; $sum += $count;

View File

@@ -19,7 +19,7 @@ sub test_recall
my $explain = $node->safe_psql("postgres", qq( my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET ivfflat.probes = $probes; SET ivfflat.probes = $probes;
SET ivfflat.iterative_search = relaxed_order; SET ivfflat.iterative_search = on;
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/);
@@ -29,7 +29,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 ivfflat.probes = $probes; SET ivfflat.probes = $probes;
SET ivfflat.iterative_search = relaxed_order; SET ivfflat.iterative_search = on;
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);
@@ -48,7 +48,7 @@ sub test_recall
$total += $limit; $total += $limit;
} }
cmp_ok($correct / $total, ">=", $min, "$operator $c"); cmp_ok($correct / $total, ">=", $min, $operator);
} }
# Initialize node # Initialize node
@@ -103,7 +103,7 @@ for my $i (0 .. $#operators)
if ($c == 100) if ($c == 100)
{ {
test_recall($c, 1, 0.57, $operator); test_recall($c, 1, 0.58, $operator);
test_recall($c, 10, 0.98, $operator); test_recall($c, 10, 0.98, $operator);
} }
else else

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.iterative_search = relaxed_order; SET hnsw.iterative_search = on;
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;
)); ));
@@ -42,8 +42,8 @@ foreach ((30000, 50000, 70000))
{ {
$count = $node->safe_psql("postgres", qq( $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.iterative_search = relaxed_order; SET hnsw.iterative_search = on;
SET hnsw.max_search_tuples = $max_tuples; SET hnsw.iterative_search_max_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.iterative_search = relaxed_order; SET hnsw.iterative_search = on;
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 = (50, 500); my @cs = (100, 1000);
sub test_recall sub test_recall
{ {
my ($c, $ef_search, $min, $operator, $mode) = @_; my ($c, $ef_search, $min, $operator) = @_;
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.iterative_search = $mode; SET hnsw.iterative_search = on;
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.iterative_search = $mode; SET hnsw.iterative_search = on;
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 $mode $c"); cmp_ok($correct / $total, ">=", $min, $operator);
} }
# 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, 50000) i;" "INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
); );
# Generate queries # Generate queries
@@ -108,8 +108,21 @@ for my $i (0 .. $#operators)
push(@expected, $res); push(@expected, $res);
} }
test_recall($c, 40, 0.99, $operator, "strict_order"); if ($c == 100)
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;");

113
test/t/045_hnsw_hqann.pl Normal file
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 @cs = ();
my @expected;
my $limit = 20;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
my $nc = 1000;
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 WHERE c = $cs[0] ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
is(scalar(@actual_ids), $limit);
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 vector($dim), c int4);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc FROM generate_series(1, 20000) i;"
);
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
push(@queries, "[" . join(",", @r) . "]");
push(@cs, int(rand() * $nc));
}
# Get exact results
@expected = ();
for my $i (0 .. $#queries)
{
my $res = $node->safe_psql("postgres", "SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v <=> '$queries[$i]' LIMIT $limit;");
push(@expected, $res);
}
# Add index
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '256MB';
SET max_parallel_maintenance_workers = 2;
CREATE INDEX ON tst USING hnsw (v vector_cosine_ops, c);
));
# Test recall
test_recall(0.99, '<=>');
# Test vacuum
$node->safe_psql("postgres", "DELETE FROM tst WHERE c > 5;");
$node->safe_psql("postgres", "VACUUM tst;");
# Test columns
my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c);");
like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c, v vector_cosine_ops);");
like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_cosine_ops, c, c);");
like($stderr, qr/index cannot have more than two columns/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_cosine_ops, v vector_cosine_ops);");
like($stderr, qr/column 2 cannot be a vector/);
done_testing();

View File

@@ -1,132 +0,0 @@
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

@@ -1,113 +0,0 @@
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

@@ -1,101 +0,0 @@
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

@@ -1,58 +0,0 @@
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();