Compare commits

..

4 Commits

Author SHA1 Message Date
Andrew Kane
d977caa47d Fixed Windows build [skip ci] 2023-10-05 01:31:51 -07:00
Andrew Kane
cc87960109 Fixed more builds 2023-10-05 01:01:45 -07:00
Andrew Kane
c6d1d8bc2c Fixed CI 2023-10-05 00:54:00 -07:00
Andrew Kane
4914511cf6 Added tinyint type 2023-10-05 00:42:52 -07:00
27 changed files with 934 additions and 800 deletions

View File

@@ -8,8 +8,6 @@ jobs:
fail-fast: false fail-fast: false
matrix: matrix:
include: include:
- postgres: 17
os: ubuntu-22.04
- postgres: 16 - postgres: 16
os: ubuntu-22.04 os: ubuntu-22.04
- postgres: 15 - postgres: 15
@@ -23,7 +21,7 @@ jobs:
- postgres: 11 - postgres: 11
os: ubuntu-20.04 os: ubuntu-20.04
steps: steps:
- uses: actions/checkout@v4 - uses: actions/checkout@v3
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1
with: with:
postgres-version: ${{ matrix.postgres }} postgres-version: ${{ matrix.postgres }}
@@ -45,7 +43,7 @@ jobs:
runs-on: macos-latest runs-on: macos-latest
if: ${{ !startsWith(github.ref_name, 'windows') }} if: ${{ !startsWith(github.ref_name, 'windows') }}
steps: steps:
- uses: actions/checkout@v4 - uses: actions/checkout@v3
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1
with: with:
postgres-version: 14 postgres-version: 14
@@ -67,7 +65,7 @@ jobs:
runs-on: windows-latest runs-on: windows-latest
if: ${{ !startsWith(github.ref_name, 'mac') }} if: ${{ !startsWith(github.ref_name, 'mac') }}
steps: steps:
- uses: actions/checkout@v4 - uses: actions/checkout@v3
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1
with: with:
postgres-version: 14 postgres-version: 14

View File

@@ -1,11 +1,7 @@
## 0.6.0 (unreleased) ## 0.5.1 (unreleased)
- Added support for inline filtering with HNSW - Added check for MVCC-compliant snapshot for HNSW index scans
- Improved performance of index scans for IVFFlat after updates and deletes
## 0.5.1 (2023-10-10)
- Improved performance of HNSW index builds
- Added check for MVCC-compliant snapshot for index scans
## 0.5.0 (2023-08-28) ## 0.5.0 (2023-08-28)

View File

@@ -2,7 +2,7 @@
"name": "vector", "name": "vector",
"abstract": "Open-source vector similarity search for Postgres", "abstract": "Open-source vector similarity search for Postgres",
"description": "Supports L2 distance, inner product, and cosine distance", "description": "Supports L2 distance, inner product, and cosine distance",
"version": "0.5.1", "version": "0.5.0",
"maintainer": [ "maintainer": [
"Andrew Kane <andrew@ankane.org>" "Andrew Kane <andrew@ankane.org>"
], ],
@@ -20,7 +20,7 @@
"vector": { "vector": {
"file": "sql/vector.sql", "file": "sql/vector.sql",
"docfile": "README.md", "docfile": "README.md",
"version": "0.5.1", "version": "0.5.0",
"abstract": "Open-source vector similarity search for Postgres" "abstract": "Open-source vector similarity search for Postgres"
} }
}, },

View File

@@ -1,10 +1,10 @@
EXTENSION = vector EXTENSION = vector
EXTVERSION = 0.5.1 EXTVERSION = 0.5.0
MODULE_big = vector MODULE_big = vector
DATA = $(wildcard sql/*--*.sql) DATA = $(wildcard sql/*--*.sql)
OBJS = src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/vector.o OBJS = src/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/tinyint.o src/vector.o
HEADERS = src/vector.h HEADERS = src/tinyint.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

@@ -1,8 +1,8 @@
EXTENSION = vector EXTENSION = vector
EXTVERSION = 0.5.1 EXTVERSION = 0.5.0
OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj OBJS = src\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\tinyint.obj src\vector.obj
HEADERS = src\vector.h HEADERS = src\tinyint.h src\vector.h
REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged
REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION) REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION)

108
README.md
View File

@@ -18,7 +18,7 @@ Compile and install the extension (supports Postgres 11+)
```sh ```sh
cd /tmp cd /tmp
git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git git clone --branch v0.5.0 https://github.com/pgvector/pgvector.git
cd pgvector cd pgvector
make make
make install # may need sudo make install # may need sudo
@@ -26,7 +26,7 @@ make install # may need sudo
See the [installation notes](#installation-notes) if you run into issues See the [installation notes](#installation-notes) if you run into issues
You can also install it with [Docker](#docker), [Homebrew](#homebrew), [PGXN](#pgxn), [APT](#apt), [Yum](#yum), or [conda-forge](#conda-forge), and it comes preinstalled with [Postgres.app](#postgresapp) and many [hosted providers](#hosted-postgres). There are also instructions for [GitHub Actions](https://github.com/pgvector/setup-pgvector). You can also install it with [Docker](#docker), [Homebrew](#homebrew), [PGXN](#pgxn), [APT](#apt), [Yum](#yum), or [conda-forge](#conda-forge), and it comes preinstalled with [Postgres.app](#postgresapp) and many [hosted providers](#hosted-postgres)
## Getting Started ## Getting Started
@@ -215,23 +215,6 @@ SELECT ...
COMMIT; COMMIT;
``` ```
### Indexing Progress
Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+
```sql
SELECT phase, round(100.0 * tuples_done / nullif(tuples_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;
```
The phases for IVFFlat are:
1. `initializing`
2. `performing k-means`
3. `assigning tuples`
4. `loading tuples`
Note: `%` is only populated during the `loading tuples` phase
## HNSW ## HNSW
An HNSW index creates a multilayer graph. It has slower build times and uses more memory than IVFFlat, but has better query performance (in terms of speed-recall tradeoff). Theres no training step like IVFFlat, so the index can be created without any data in the table. An HNSW index creates a multilayer graph. It has slower build times and uses more memory than IVFFlat, but has better query performance (in terms of speed-recall tradeoff). Theres no training step like IVFFlat, so the index can be created without any data in the table.
@@ -288,18 +271,22 @@ SELECT ...
COMMIT; COMMIT;
``` ```
### Indexing Progress ## Indexing Progress
Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+ Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+
```sql ```sql
SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index; SELECT phase, tuples_done, tuples_total FROM pg_stat_progress_create_index;
``` ```
The phases for HNSW are: The phases are:
1. `initializing` 1. `initializing`
2. `loading tuples` 2. `performing k-means` - IVFFlat only
3. `assigning tuples` - IVFFlat only
4. `loading tuples`
Note: `tuples_done` and `tuples_total` are only populated during the `loading tuples` phase
## Filtering ## Filtering
@@ -315,12 +302,6 @@ 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.6.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
@@ -336,15 +317,13 @@ CREATE TABLE items (embedding vector(3), category_id int) PARTITION BY LIST(cate
## Hybrid Search ## Hybrid Search
Use together with Postgres [full-text search](https://www.postgresql.org/docs/current/textsearch-intro.html) for hybrid search. Use together with Postgres [full-text search](https://www.postgresql.org/docs/current/textsearch-intro.html) for hybrid search ([Python example](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py)).
```sql ```sql
SELECT id, content FROM items, plainto_tsquery('hello search') query SELECT id, content FROM items, plainto_tsquery('hello search') query
WHERE textsearch @@ query ORDER BY ts_rank_cd(textsearch, query) DESC LIMIT 5; WHERE textsearch @@ query ORDER BY ts_rank_cd(textsearch, query) DESC LIMIT 5;
``` ```
You can use [Reciprocal Rank Fusion](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) or a [cross-encoder](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) to combine results.
## Performance ## Performance
Use `EXPLAIN ANALYZE` to debug performance. Use `EXPLAIN ANALYZE` to debug performance.
@@ -381,7 +360,6 @@ Use pgvector from any language with a Postgres client. You can even generate and
Language | Libraries / Examples Language | Libraries / Examples
--- | --- --- | ---
C | [pgvector-c](https://github.com/pgvector/pgvector-c)
C++ | [pgvector-cpp](https://github.com/pgvector/pgvector-cpp) C++ | [pgvector-cpp](https://github.com/pgvector/pgvector-cpp)
C# | [pgvector-dotnet](https://github.com/pgvector/pgvector-dotnet) C# | [pgvector-dotnet](https://github.com/pgvector/pgvector-dotnet)
Crystal | [pgvector-crystal](https://github.com/pgvector/pgvector-crystal) Crystal | [pgvector-crystal](https://github.com/pgvector/pgvector-crystal)
@@ -389,11 +367,10 @@ Dart | [pgvector-dart](https://github.com/pgvector/pgvector-dart)
Elixir | [pgvector-elixir](https://github.com/pgvector/pgvector-elixir) Elixir | [pgvector-elixir](https://github.com/pgvector/pgvector-elixir)
Go | [pgvector-go](https://github.com/pgvector/pgvector-go) Go | [pgvector-go](https://github.com/pgvector/pgvector-go)
Haskell | [pgvector-haskell](https://github.com/pgvector/pgvector-haskell) Haskell | [pgvector-haskell](https://github.com/pgvector/pgvector-haskell)
Java, Kotlin, Groovy, Scala | [pgvector-java](https://github.com/pgvector/pgvector-java) Java, Scala | [pgvector-java](https://github.com/pgvector/pgvector-java)
JavaScript, TypeScript | [pgvector-node](https://github.com/pgvector/pgvector-node)
Julia | [pgvector-julia](https://github.com/pgvector/pgvector-julia) Julia | [pgvector-julia](https://github.com/pgvector/pgvector-julia)
Lua | [pgvector-lua](https://github.com/pgvector/pgvector-lua) Lua | [pgvector-lua](https://github.com/pgvector/pgvector-lua)
Nim | [pgvector-nim](https://github.com/pgvector/pgvector-nim) Node.js | [pgvector-node](https://github.com/pgvector/pgvector-node)
Perl | [pgvector-perl](https://github.com/pgvector/pgvector-perl) Perl | [pgvector-perl](https://github.com/pgvector/pgvector-perl)
PHP | [pgvector-php](https://github.com/pgvector/pgvector-php) PHP | [pgvector-php](https://github.com/pgvector/pgvector-php)
Python | [pgvector-python](https://github.com/pgvector/pgvector-python) Python | [pgvector-python](https://github.com/pgvector/pgvector-python)
@@ -401,7 +378,6 @@ R | [pgvector-r](https://github.com/pgvector/pgvector-r)
Ruby | [pgvector-ruby](https://github.com/pgvector/pgvector-ruby), [Neighbor](https://github.com/ankane/neighbor) Ruby | [pgvector-ruby](https://github.com/pgvector/pgvector-ruby), [Neighbor](https://github.com/ankane/neighbor)
Rust | [pgvector-rust](https://github.com/pgvector/pgvector-rust) Rust | [pgvector-rust](https://github.com/pgvector/pgvector-rust)
Swift | [pgvector-swift](https://github.com/pgvector/pgvector-swift) Swift | [pgvector-swift](https://github.com/pgvector/pgvector-swift)
Zig | [pgvector-zig](https://github.com/pgvector/pgvector-zig)
## Frequently Asked Questions ## Frequently Asked Questions
@@ -417,55 +393,6 @@ Yes, pgvector uses the write-ahead log (WAL), which allows for replication and p
Youll need to use [dimensionality reduction](https://en.wikipedia.org/wiki/Dimensionality_reduction) at the moment. Youll need to use [dimensionality reduction](https://en.wikipedia.org/wiki/Dimensionality_reduction) at the moment.
#### Can I store vectors with different dimensions in the same column?
You can use `vector` as the type (instead of `vector(3)`).
```sql
CREATE TABLE embeddings (model_id bigint, item_id bigint, embedding vector, PRIMARY KEY (model_id, item_id));
```
However, you can only create indexes on rows with the same number of dimensions (using [expression](https://www.postgresql.org/docs/current/indexes-expressional.html) and [partial](https://www.postgresql.org/docs/current/indexes-partial.html) indexing):
```sql
CREATE INDEX ON embeddings USING hnsw ((embedding::vector(3)) vector_l2_ops) WHERE (model_id = 123);
```
and query with:
```sql
SELECT * FROM embeddings WHERE model_id = 123 ORDER BY embedding::vector(3) <-> '[3,1,2]' LIMIT 5;
```
#### Can I store vectors with more precision?
You can use the `double precision[]` or `numeric[]` type to store vectors with more precision.
```sql
CREATE TABLE items (id bigserial PRIMARY KEY, embedding double precision[]);
-- use {} instead of [] for Postgres arrays
INSERT INTO items (embedding) VALUES ('{1,2,3}'), ('{4,5,6}');
```
Optionally, add a [check constraint](https://www.postgresql.org/docs/current/ddl-constraints.html) to ensure data can be converted to the `vector` type and has the expected dimensions.
```sql
ALTER TABLE items ADD CHECK (vector_dims(embedding::vector) = 3);
```
Use [expression indexing](https://www.postgresql.org/docs/current/indexes-expressional.html) to index (at a lower precision):
```sql
CREATE INDEX ON items USING hnsw ((embedding::vector(3)) vector_l2_ops);
```
and query with:
```sql
SELECT * FROM items ORDER BY embedding::vector(3) <-> '[3,1,2]' LIMIT 5;
```
## Troubleshooting ## Troubleshooting
#### Why isnt a query using an index? #### Why isnt a query using an index?
@@ -479,8 +406,6 @@ SELECT ...
COMMIT; COMMIT;
``` ```
Also, if the table is small, a table scan may be faster.
#### Why isnt a query using a parallel table scan? #### Why isnt a query using a parallel table scan?
The planner doesnt consider [out-of-line storage](https://www.postgresql.org/docs/current/storage-toast.html) in cost estimates, which can make a serial scan look cheaper. You can reduce the cost of a parallel scan for a query with: The planner doesnt consider [out-of-line storage](https://www.postgresql.org/docs/current/storage-toast.html) in cost estimates, which can make a serial scan look cheaper. You can reduce the cost of a parallel scan for a query with:
@@ -584,7 +509,7 @@ Then use `nmake` to build:
```cmd ```cmd
set "PGROOT=C:\Program Files\PostgreSQL\15" set "PGROOT=C:\Program Files\PostgreSQL\15"
git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git git clone --branch v0.5.0 https://github.com/pgvector/pgvector.git
cd pgvector cd pgvector
nmake /F Makefile.win nmake /F Makefile.win
nmake /F Makefile.win install nmake /F Makefile.win install
@@ -605,7 +530,7 @@ This adds pgvector to the [Postgres image](https://hub.docker.com/_/postgres) (r
You can also build the image manually: You can also build the image manually:
```sh ```sh
git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git git clone --branch v0.5.0 https://github.com/pgvector/pgvector.git
cd pgvector cd pgvector
docker build --build-arg PG_MAJOR=15 -t myuser/pgvector . docker build --build-arg PG_MAJOR=15 -t myuser/pgvector .
``` ```
@@ -670,7 +595,7 @@ pgvector is available on [these providers](https://github.com/pgvector/pgvector/
## Upgrading ## Upgrading
[Install](#installation) the latest version (use the same method as the original installation). Then in each database you want to upgrade, run: Install the latest version. Then in each database you want to upgrade, run:
```sql ```sql
ALTER EXTENSION vector UPDATE; ALTER EXTENSION vector UPDATE;
@@ -718,7 +643,6 @@ 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,2 +1,67 @@
-- 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.5.1'" to load this file. \quit \echo Use "ALTER EXTENSION vector UPDATE TO '0.5.1'" to load this file. \quit
-- tinyint
CREATE TYPE tinyint;
CREATE FUNCTION tinyint_in(cstring, oid, integer) RETURNS tinyint
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION tinyint_out(tinyint) RETURNS cstring
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION tinyint_recv(internal, oid, integer) RETURNS tinyint
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION tinyint_send(tinyint) RETURNS bytea
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE TYPE tinyint (
INPUT = tinyint_in,
OUTPUT = tinyint_out,
RECEIVE = tinyint_recv,
SEND = tinyint_send,
INTERNALLENGTH = 1,
PASSEDBYVALUE,
ALIGNMENT = char
);
CREATE FUNCTION integer_to_tinyint(integer, integer, boolean) RETURNS tinyint
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION numeric_to_tinyint(numeric, integer, boolean) RETURNS tinyint
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE CAST (integer AS tinyint)
WITH FUNCTION integer_to_tinyint(integer, integer, boolean) AS IMPLICIT;
CREATE CAST (numeric AS tinyint)
WITH FUNCTION numeric_to_tinyint(numeric, integer, boolean) AS IMPLICIT;
CREATE FUNCTION l2_distance(tinyint[], tinyint[]) RETURNS float8
AS 'MODULE_PATHNAME', 'tinyint_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION inner_product(tinyint[], tinyint[]) RETURNS float8
AS 'MODULE_PATHNAME', 'tinyint_inner_product' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION cosine_distance(tinyint[], tinyint[]) RETURNS float8
AS 'MODULE_PATHNAME', 'tinyint_cosine_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION tinyint_negative_inner_product(tinyint[], tinyint[]) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR <-> (
LEFTARG = tinyint[], RIGHTARG = tinyint[], PROCEDURE = l2_distance,
COMMUTATOR = '<->'
);
CREATE OPERATOR <#> (
LEFTARG = tinyint[], RIGHTARG = tinyint[], PROCEDURE = tinyint_negative_inner_product,
COMMUTATOR = '<#>'
);
CREATE OPERATOR <=> (
LEFTARG = tinyint[], RIGHTARG = tinyint[], PROCEDURE = cosine_distance,
COMMUTATOR = '<=>'
);

View File

@@ -1,18 +0,0 @@
-- complain if script is sourced in psql, rather than via CREATE EXTENSION
\echo Use "ALTER EXTENSION vector UPDATE TO '0.6.0'" to load this file. \quit
CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8
AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION hnsw_attribute_distance(bigint, bigint) RETURNS float8
AS 'MODULE_PATHNAME', 'hnsw_int8_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 3 hnsw_attribute_distance(integer, integer);
CREATE OPERATOR CLASS vector_bigint_ops
DEFAULT FOR TYPE bigint USING hnsw AS
OPERATOR 2 = (bigint, bigint),
FUNCTION 3 hnsw_attribute_distance(bigint, bigint);

View File

@@ -291,20 +291,67 @@ CREATE OPERATOR CLASS vector_cosine_ops
FUNCTION 1 vector_negative_inner_product(vector, vector), FUNCTION 1 vector_negative_inner_product(vector, vector),
FUNCTION 2 vector_norm(vector); FUNCTION 2 vector_norm(vector);
-- hnsw attributes -- tinyint
CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8 CREATE TYPE tinyint;
AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION hnsw_attribute_distance(bigint, bigint) RETURNS float8 CREATE FUNCTION tinyint_in(cstring, oid, integer) RETURNS tinyint
AS 'MODULE_PATHNAME', 'hnsw_int8_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR CLASS vector_integer_ops CREATE FUNCTION tinyint_out(tinyint) RETURNS cstring
DEFAULT FOR TYPE integer USING hnsw AS AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
OPERATOR 2 = (integer, integer),
FUNCTION 3 hnsw_attribute_distance(integer, integer);
CREATE OPERATOR CLASS vector_bigint_ops CREATE FUNCTION tinyint_recv(internal, oid, integer) RETURNS tinyint
DEFAULT FOR TYPE bigint USING hnsw AS AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
OPERATOR 2 = (bigint, bigint),
FUNCTION 3 hnsw_attribute_distance(bigint, bigint); CREATE FUNCTION tinyint_send(tinyint) RETURNS bytea
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE TYPE tinyint (
INPUT = tinyint_in,
OUTPUT = tinyint_out,
RECEIVE = tinyint_recv,
SEND = tinyint_send,
INTERNALLENGTH = 1,
PASSEDBYVALUE,
ALIGNMENT = char
);
CREATE FUNCTION integer_to_tinyint(integer, integer, boolean) RETURNS tinyint
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION numeric_to_tinyint(numeric, integer, boolean) RETURNS tinyint
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE CAST (integer AS tinyint)
WITH FUNCTION integer_to_tinyint(integer, integer, boolean) AS IMPLICIT;
CREATE CAST (numeric AS tinyint)
WITH FUNCTION numeric_to_tinyint(numeric, integer, boolean) AS IMPLICIT;
CREATE FUNCTION l2_distance(tinyint[], tinyint[]) RETURNS float8
AS 'MODULE_PATHNAME', 'tinyint_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION inner_product(tinyint[], tinyint[]) RETURNS float8
AS 'MODULE_PATHNAME', 'tinyint_inner_product' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION cosine_distance(tinyint[], tinyint[]) RETURNS float8
AS 'MODULE_PATHNAME', 'tinyint_cosine_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION tinyint_negative_inner_product(tinyint[], tinyint[]) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR <-> (
LEFTARG = tinyint[], RIGHTARG = tinyint[], PROCEDURE = l2_distance,
COMMUTATOR = '<->'
);
CREATE OPERATOR <#> (
LEFTARG = tinyint[], RIGHTARG = tinyint[], PROCEDURE = tinyint_negative_inner_product,
COMMUTATOR = '<#>'
);
CREATE OPERATOR <=> (
LEFTARG = tinyint[], RIGHTARG = tinyint[], PROCEDURE = cosine_distance,
COMMUTATOR = '<=>'
);

View File

@@ -167,7 +167,7 @@ hnswhandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine); IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0; amroutine->amstrategies = 0;
amroutine->amsupport = 3; amroutine->amsupport = 2;
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
amroutine->amoptsprocnum = 0; amroutine->amoptsprocnum = 0;
#endif #endif
@@ -175,7 +175,7 @@ hnswhandler(PG_FUNCTION_ARGS)
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 = true; amroutine->amcanmulticol = false;
amroutine->amoptionalkey = true; amroutine->amoptionalkey = true;
amroutine->amsearcharray = false; amroutine->amsearcharray = false;
amroutine->amsearchnulls = false; amroutine->amsearchnulls = false;
@@ -222,31 +222,3 @@ 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);
}
/*
* Get the distance between two int8 attributes
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(hnsw_int8_attribute_distance);
Datum
hnsw_int8_attribute_distance(PG_FUNCTION_ARGS)
{
int64 a = PG_GETARG_INT64(0);
int64 b = PG_GETARG_INT64(1);
double distance = ((double) a) - ((double) b);
PG_RETURN_FLOAT8(distance);
}

View File

@@ -19,7 +19,6 @@
/* Support functions */ /* Support functions */
#define HNSW_DISTANCE_PROC 1 #define HNSW_DISTANCE_PROC 1
#define HNSW_NORM_PROC 2 #define HNSW_NORM_PROC 2
#define HNSW_ATTRIBUTE_DISTANCE_PROC 3
#define HNSW_VERSION 1 #define HNSW_VERSION 1
#define HNSW_MAGIC_NUMBER 0xA953A953 #define HNSW_MAGIC_NUMBER 0xA953A953
@@ -60,7 +59,7 @@
#define HNSW_MAX_SIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - sizeof(ItemIdData)) #define HNSW_MAX_SIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - sizeof(ItemIdData))
#define HNSW_ELEMENT_TUPLE_SIZE(size) MAXALIGN(offsetof(HnswElementTupleData, data) + (size)) #define HNSW_ELEMENT_TUPLE_SIZE(_dim) MAXALIGN(offsetof(HnswElementTupleData, vec) + VECTOR_SIZE(_dim))
#define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData)) #define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData))
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page)) #define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
@@ -104,8 +103,7 @@ typedef struct HnswElementData
OffsetNumber offno; OffsetNumber offno;
OffsetNumber neighborOffno; OffsetNumber neighborOffno;
BlockNumber neighborPage; BlockNumber neighborPage;
Datum value; Vector *vec;
IndexTuple itup;
} HnswElementData; } HnswElementData;
typedef HnswElementData * HnswElement; typedef HnswElementData * HnswElement;
@@ -114,13 +112,11 @@ typedef struct HnswCandidate
{ {
HnswElement element; HnswElement element;
float distance; float distance;
bool closer;
} HnswCandidate; } HnswCandidate;
typedef struct HnswNeighborArray typedef struct HnswNeighborArray
{ {
int length; int length;
bool closerSet;
HnswCandidate *items; HnswCandidate *items;
} HnswNeighborArray; } HnswNeighborArray;
@@ -156,18 +152,17 @@ typedef struct HnswBuildState
double reltuples; double reltuples;
/* Support functions */ /* Support functions */
FmgrInfo **procinfos; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
Oid *collations; Oid collation;
/* Variables */ /* Variables */
List *elements; List *elements;
HnswElement entryPoint; HnswElement entryPoint;
double ml; double ml;
int maxLevel; int maxLevel;
long memoryLeft; double maxInMemoryElements;
bool flushed; bool flushed;
bool useIndexTuple;
Vector *normvec; Vector *normvec;
/* Memory */ /* Memory */
@@ -207,7 +202,7 @@ typedef struct HnswElementTupleData
ItemPointerData heaptids[HNSW_HEAPTIDS]; ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid; ItemPointerData neighbortid;
uint16 unused2; uint16 unused2;
Vector data; Vector vec;
} HnswElementTupleData; } HnswElementTupleData;
typedef HnswElementTupleData * HnswElementTuple; typedef HnswElementTupleData * HnswElementTuple;
@@ -229,9 +224,9 @@ typedef struct HnswScanOpaqueData
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Support functions */ /* Support functions */
FmgrInfo **procinfos; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
Oid *collations; Oid collation;
} HnswScanOpaqueData; } HnswScanOpaqueData;
typedef HnswScanOpaqueData * HnswScanOpaque; typedef HnswScanOpaqueData * HnswScanOpaque;
@@ -249,8 +244,8 @@ typedef struct HnswVacuumState
int efConstruction; int efConstruction;
/* Support functions */ /* Support functions */
FmgrInfo **procinfos; FmgrInfo *procinfo;
Oid *collations; Oid collation;
/* Variables */ /* Variables */
HTAB *deleted; HTAB *deleted;
@@ -272,27 +267,26 @@ Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page); void HnswInitPage(Buffer buf, Page page);
void HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state); void HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
void HnswInit(void); void HnswInit(void);
List *HnswSearchLayer(Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef, int lc, Relation index, FmgrInfo **procinfos, Oid *collations, int m, bool loadVec, HnswElement skipElement, bool inMemory); List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement);
HnswElement HnswGetEntryPoint(Relation index); HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint); void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel); HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
void HnswFreeElement(HnswElement element); void HnswFreeElement(HnswElement element);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno); HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo **procinfos, Oid *collations, int m, int efConstruction, bool existing, bool inMemory); void HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswElement HnswFindDuplicate(HnswElement e, Relation index); HnswElement HnswFindDuplicate(HnswElement e);
HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation rel, FmgrInfo **procinfos, Oid *collations, bool loadVec, bool inMemory); HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum); void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum);
void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m); void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
void HnswInitNeighbors(HnswElement element, int m); void HnswInitNeighbors(HnswElement element, int m);
bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel); bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel);
void HnswUpdateNeighborPages(Relation index, FmgrInfo **procinfos, Oid *collations, HnswElement e, int m, bool checkExisting); void HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index); void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, float *distance, Datum *q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfos, Oid *collations, bool loadVec); void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswSetElementTuple(HnswElementTuple etup, HnswElement element, bool useIndexTuple); void HnswSetElementTuple(HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo **procinfos, Oid *collations, bool inMemory); void HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index, int m); void HnswLoadNeighbors(HnswElement element, Relation index, int m);
void HnswElementSetData(HnswElement element, Relation index, Datum value, Datum *values, bool *isnull);
/* Index access methods */ /* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo); IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);
@@ -309,6 +303,5 @@ IndexScanDesc hnswbeginscan(Relation index, int nkeys, int norderbys);
void hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int norderbys); void hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int norderbys);
bool hnswgettuple(IndexScanDesc scan, ScanDirection dir); bool hnswgettuple(IndexScanDesc scan, ScanDirection dir);
void hnswendscan(IndexScanDesc scan); void hnswendscan(IndexScanDesc scan);
FmgrInfo **HnswInitProcinfos(Relation index);
#endif #endif

View File

@@ -8,7 +8,6 @@
#include "lib/pairingheap.h" #include "lib/pairingheap.h"
#include "nodes/pg_list.h" #include "nodes/pg_list.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
@@ -106,8 +105,8 @@ CreateElementPages(HnswBuildState * buildstate)
{ {
Relation index = buildstate->index; Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
bool useIndexTuple = buildstate->useIndexTuple; int dimensions = buildstate->dimensions;
Size etupAllocSize; Size etupSize;
Size maxSize; Size maxSize;
HnswElementTuple etup; HnswElementTuple etup;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
@@ -118,11 +117,11 @@ CreateElementPages(HnswBuildState * buildstate)
ListCell *lc; ListCell *lc;
/* Calculate sizes */ /* Calculate sizes */
etupAllocSize = BLCKSZ;
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
/* Allocate once */ /* Allocate once */
etup = palloc0(etupAllocSize); etup = palloc0(etupSize);
ntup = palloc0(BLCKSZ); ntup = palloc0(BLCKSZ);
/* Prepare first page */ /* Prepare first page */
@@ -134,24 +133,15 @@ CreateElementPages(HnswBuildState * buildstate)
foreach(lc, buildstate->elements) foreach(lc, buildstate->elements)
{ {
HnswElement element = lfirst(lc); HnswElement element = lfirst(lc);
Size etupSize;
Size ntupSize; Size ntupSize;
Size combinedSize; Size combinedSize;
/* Zero memory for each element */ HnswSetElementTuple(etup, element);
MemSet(etup, 0, etupAllocSize);
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(useIndexTuple ? IndexTupleSize(element->itup) : VARSIZE_ANY(DatumGetPointer(element->value)));
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);
/* Initial size check */
if (etupSize > etupAllocSize)
elog(ERROR, "index tuple too large");
HnswSetElementTuple(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))
HnswBuildAppendPage(index, &buf, &page, &state, forkNum); HnswBuildAppendPage(index, &buf, &page, &state, forkNum);
@@ -274,15 +264,13 @@ FlushPages(HnswBuildState * buildstate)
* Insert tuple * Insert tuple
*/ */
static bool static bool
InsertTuple(Relation index, Datum *values, bool *isnull, HnswElement element, HnswBuildState * buildstate, HnswElement * dup, MemoryContext outerCtx) InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup)
{ {
FmgrInfo **procinfos = buildstate->procinfos; FmgrInfo *procinfo = buildstate->procinfo;
Oid *collations = buildstate->collations; Oid collation = buildstate->collation;
HnswElement entryPoint = buildstate->entryPoint; HnswElement entryPoint = buildstate->entryPoint;
int efConstruction = buildstate->efConstruction; int efConstruction = buildstate->efConstruction;
int m = buildstate->m; int m = buildstate->m;
bool inMemory = true;
MemoryContext oldCtx;
/* 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]));
@@ -290,20 +278,18 @@ InsertTuple(Relation index, Datum *values, bool *isnull, HnswElement element, Hn
/* Normalize if needed */ /* Normalize if needed */
if (buildstate->normprocinfo != NULL) if (buildstate->normprocinfo != NULL)
{ {
if (!HnswNormValue(buildstate->normprocinfo, collations[0], &value, buildstate->normvec)) if (!HnswNormValue(buildstate->normprocinfo, collation, &value, buildstate->normvec))
return false; return false;
} }
/* Copy value to element so accessible outside of memory context */ /* Copy value to element so accessible outside of memory context */
oldCtx = MemoryContextSwitchTo(outerCtx); memcpy(element->vec, DatumGetVector(value), VECTOR_SIZE(buildstate->dimensions));
HnswElementSetData(element, index, value, values, isnull);
MemoryContextSwitchTo(oldCtx);
/* Insert element in graph */ /* Insert element in graph */
HnswInsertElement(element, entryPoint, index, procinfos, collations, m, efConstruction, false, inMemory); HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */ /* Look for duplicate */
*dup = HnswFindDuplicate(element, index); *dup = HnswFindDuplicate(element);
/* Update neighbors if needed */ /* Update neighbors if needed */
if (*dup == NULL) if (*dup == NULL)
@@ -314,7 +300,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, HnswElement element, Hn
HnswNeighborArray *neighbors = &element->neighbors[lc]; HnswNeighborArray *neighbors = &element->neighbors[lc];
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, index, procinfos, collations, inMemory); HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation);
} }
} }
@@ -327,21 +313,6 @@ InsertTuple(Relation index, Datum *values, bool *isnull, HnswElement element, Hn
return *dup == NULL; return *dup == NULL;
} }
/*
* Get the memory used by an element
*/
static long
HnswElementMemory(HnswElement e, int m)
{
long elementSize = sizeof(HnswElementData);
elementSize += sizeof(HnswNeighborArray) * (e->level + 1);
elementSize += sizeof(HnswCandidate) * (m * (e->level + 2));
elementSize += sizeof(ItemPointerData);
elementSize += IndexTupleSize(e->itup);
return elementSize;
}
/* /*
* Callback for table_index_build_scan * Callback for table_index_build_scan
*/ */
@@ -363,7 +334,7 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
if (isnull[0]) if (isnull[0])
return; return;
if (buildstate->memoryLeft <= 0) if (buildstate->indtuples >= buildstate->maxInMemoryElements)
{ {
if (!buildstate->flushed) if (!buildstate->flushed)
{ {
@@ -389,12 +360,13 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
/* Allocate necessary memory outside of memory context */ /* Allocate necessary memory outside of memory context */
element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel); element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel);
element->vec = palloc(VECTOR_SIZE(buildstate->dimensions));
/* Use memory context since detoast can allocate */ /* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Insert tuple */ /* Insert tuple */
inserted = InsertTuple(index, values, isnull, element, buildstate, &dup, oldCtx); inserted = InsertTuple(index, values, element, buildstate, &dup);
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -402,21 +374,31 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
/* Add outside memory context */ /* Add outside memory context */
if (dup != NULL) if (dup != NULL)
{
HnswAddHeapTid(dup, tid); HnswAddHeapTid(dup, tid);
buildstate->memoryLeft -= sizeof(ItemPointerData);
}
/* Add to buildstate or free */ /* Add to buildstate or free */
if (inserted) if (inserted)
{
buildstate->elements = lappend(buildstate->elements, element); buildstate->elements = lappend(buildstate->elements, element);
buildstate->memoryLeft -= HnswElementMemory(element, buildstate->m);
}
else else
HnswFreeElement(element); HnswFreeElement(element);
} }
/*
* Get the max number of elements that fit into maintenance_work_mem
*/
static double
HnswGetMaxInMemoryElements(int m, double ml, int dimensions)
{
Size elementSize = sizeof(HnswElementData);
double avgLevel = -log(0.5) * ml;
elementSize += sizeof(HnswNeighborArray) * (avgLevel + 1);
elementSize += sizeof(HnswCandidate) * (m * (avgLevel + 2));
elementSize += sizeof(ItemPointerData);
elementSize += VECTOR_SIZE(dimensions);
return (maintenance_work_mem * 1024L) / elementSize;
}
/* /*
* Initialize the build state * Initialize the build state
*/ */
@@ -432,19 +414,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->efConstruction = HnswGetEfConstruction(index); buildstate->efConstruction = HnswGetEfConstruction(index);
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod; buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
/* 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)
elog(ERROR, "column does not have dimensions"); elog(ERROR, "column does not have dimensions");
@@ -459,17 +428,16 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->indtuples = 0; buildstate->indtuples = 0;
/* Get support functions */ /* Get support functions */
buildstate->procinfos = HnswInitProcinfos(index); buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->collations = index->rd_indcollation; buildstate->collation = index->rd_indcollation[0];
buildstate->elements = NIL; buildstate->elements = NIL;
buildstate->entryPoint = NULL; buildstate->entryPoint = NULL;
buildstate->ml = HnswGetMl(buildstate->m); buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m); buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->memoryLeft = maintenance_work_mem * 1024L; buildstate->maxInMemoryElements = HnswGetMaxInMemoryElements(buildstate->m, buildstate->ml, buildstate->dimensions);
buildstate->flushed = false; buildstate->flushed = false;
buildstate->useIndexTuple = IndexRelationGetNumberOfAttributes(index) > 1;
/* Reuse for each tuple */ /* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions); buildstate->normvec = InitVector(buildstate->dimensions);
@@ -485,7 +453,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
static void static void
FreeBuildState(HnswBuildState * buildstate) FreeBuildState(HnswBuildState * buildstate)
{ {
pfree(buildstate->procinfos);
pfree(buildstate->normvec); pfree(buildstate->normvec);
MemoryContextDelete(buildstate->tmpCtx); MemoryContextDelete(buildstate->tmpCtx);
} }

View File

@@ -123,16 +123,16 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
Size minCombinedSize; Size minCombinedSize;
HnswElementTuple etup; HnswElementTuple etup;
BlockNumber currentPage = insertPage; BlockNumber currentPage = insertPage;
int dimensions = e->vec->dim;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
Buffer nbuf; Buffer nbuf;
Page npage; Page npage;
OffsetNumber freeOffno = InvalidOffsetNumber; OffsetNumber freeOffno = InvalidOffsetNumber;
OffsetNumber freeNeighborOffno = InvalidOffsetNumber; OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
bool useIndexTuple = IndexRelationGetNumberOfAttributes(index) > 1;
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(useIndexTuple ? IndexTupleSize(e->itup) : VARSIZE_ANY(DatumGetPointer(e->value))); etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
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;
@@ -140,7 +140,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
/* Prepare element tuple */ /* Prepare element tuple */
etup = palloc0(etupSize); etup = palloc0(etupSize);
HnswSetElementTuple(etup, e, useIndexTuple); HnswSetElementTuple(etup, e);
/* Prepare neighbor tuple */ /* Prepare neighbor tuple */
ntup = palloc0(ntupSize); ntup = palloc0(ntupSize);
@@ -302,7 +302,7 @@ ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
* Update neighbors * Update neighbors
*/ */
void void
HnswUpdateNeighborPages(Relation index, FmgrInfo **procinfos, Oid *collations, HnswElement e, int m, bool checkExisting) HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting)
{ {
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
@@ -334,7 +334,7 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo **procinfos, Oid *collations, H
*/ */
/* Select neighbors */ /* Select neighbors */
HnswUpdateConnection(e, hc, lm, lc, &idx, index, procinfos, collations, false); HnswUpdateConnection(e, hc, lm, lc, &idx, index, procinfo, collation);
/* New element was not selected as a neighbor */ /* New element was not selected as a neighbor */
if (idx == -1) if (idx == -1)
@@ -405,9 +405,8 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
ItemId itemid; Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(dup->vec->dim);
HnswElementTuple etup; HnswElementTuple etup;
Size etupSize;
int i; int i;
/* Read page */ /* Read page */
@@ -417,9 +416,7 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
page = GenericXLogRegisterBuffer(state, buf, 0); page = GenericXLogRegisterBuffer(state, buf, 0);
/* Find space */ /* Find space */
itemid = PageGetItemId(page, dup->offno); etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, dup->offno));
etup = (HnswElementTuple) PageGetItem(page, itemid);
etupSize = ItemIdGetLength(itemid);
for (i = 0; i < HNSW_HEAPTIDS; i++) for (i = 0; i < HNSW_HEAPTIDS; i++)
{ {
if (!ItemPointerIsValid(&etup->heaptids[i])) if (!ItemPointerIsValid(&etup->heaptids[i]))
@@ -452,7 +449,7 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
* Write changes to disk * Write changes to disk
*/ */
static void static void
WriteElement(Relation index, FmgrInfo **procinfos, Oid *collations, HnswElement element, int m, int efConstruction, HnswElement dup, HnswElement entryPoint) WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement dup, HnswElement entryPoint)
{ {
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
@@ -471,7 +468,7 @@ WriteElement(Relation index, FmgrInfo **procinfos, Oid *collations, HnswElement
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM); HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborPages(index, procinfos, collations, element, m, false); HnswUpdateNeighborPages(index, procinfo, collation, element, m, false);
/* Update metapage if needed */ /* Update metapage if needed */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -490,8 +487,8 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
HnswElement element; HnswElement element;
int m; int m;
int efConstruction = HnswGetEfConstruction(index); int efConstruction = HnswGetEfConstruction(index);
FmgrInfo **procinfos = HnswInitProcinfos(index); FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid *collations = index->rd_indcollation; Oid collation = index->rd_indcollation[0];
HnswElement dup; HnswElement dup;
LOCKMODE lockmode = ShareLock; LOCKMODE lockmode = ShareLock;
@@ -502,7 +499,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
if (normprocinfo != NULL) if (normprocinfo != NULL)
{ {
if (!HnswNormValue(normprocinfo, collations[0], &value, NULL)) if (!HnswNormValue(normprocinfo, collation, &value, NULL))
return false; return false;
} }
@@ -518,7 +515,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
/* Create an element */ /* Create an element */
element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m)); element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m));
HnswElementSetData(element, index, value, values, isnull); element->vec = DatumGetVector(value);
/* 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)
@@ -535,13 +532,13 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
} }
/* Insert element in graph */ /* Insert element in graph */
HnswInsertElement(element, entryPoint, index, procinfos, collations, m, efConstruction, false, false); HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */ /* Look for duplicate */
dup = HnswFindDuplicate(element, index); dup = HnswFindDuplicate(element);
/* Write to disk */ /* Write to disk */
WriteElement(index, procinfos, collations, element, m, efConstruction, dup, entryPoint); WriteElement(index, procinfo, collation, element, m, efConstruction, dup, entryPoint);
/* Release lock */ /* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode); UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);

View File

@@ -15,13 +15,12 @@ GetScanItems(IndexScanDesc scan, Datum q)
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation; Relation index = scan->indexRelation;
FmgrInfo **procinfos = so->procinfos; FmgrInfo *procinfo = so->procinfo;
Oid *collations = so->collations; Oid collation = so->collation;
List *ep; List *ep;
List *w; List *w;
int m; int m;
HnswElement entryPoint; HnswElement entryPoint;
ScanKeyData *keyData = scan->keyData;
/* Get m and entry point */ /* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
@@ -29,15 +28,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
if (entryPoint == NULL) if (entryPoint == NULL)
return NIL; return NIL;
ep = list_make1(HnswEntryCandidate(entryPoint, q, NULL, keyData, index, procinfos, collations, false, false)); ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, false));
for (int lc = entryPoint->level; lc >= 1; lc--) for (int lc = entryPoint->level; lc >= 1; lc--)
{ {
w = HnswSearchLayer(q, NULL, keyData, ep, 1, lc, index, procinfos, collations, m, false, NULL, false); w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, m, false, NULL);
ep = w; ep = w;
} }
return HnswSearchLayer(q, NULL, keyData, ep, hnsw_ef_search, 0, index, procinfos, collations, m, false, NULL, false); return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
} }
/* /*
@@ -84,7 +83,7 @@ GetScanValue(IndexScanDesc scan)
/* Fine if normalization fails */ /* Fine if normalization fails */
if (so->normprocinfo != NULL) if (so->normprocinfo != NULL)
HnswNormValue(so->normprocinfo, so->collations[0], &value, NULL); HnswNormValue(so->normprocinfo, so->collation, &value, NULL);
} }
return value; return value;
@@ -108,9 +107,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
/* Set support functions */ /* Set support functions */
so->procinfos = HnswInitProcinfos(index); so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
so->collations = index->rd_indcollation; so->collation = index->rd_indcollation[0];
scan->opaque = so; scan->opaque = so;
@@ -207,9 +206,6 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
scan->xs_ctup.t_self = *heaptid; scan->xs_ctup.t_self = *heaptid;
#endif #endif
/* TODO Check during scan */
scan->xs_recheck = scan->numberOfKeys > 0;
scan->xs_recheckorderby = false; scan->xs_recheckorderby = false;
return true; return true;
} }
@@ -226,7 +222,6 @@ hnswendscan(IndexScanDesc scan)
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
pfree(so->procinfos);
MemoryContextDelete(so->tmpCtx); MemoryContextDelete(so->tmpCtx);
pfree(so); pfree(so);

View File

@@ -4,13 +4,8 @@
#include "hnsw.h" #include "hnsw.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h"
#include "vector.h" #include "vector.h"
#if PG_VERSION_NUM < 130000
#define TYPSTORAGE_PLAIN 'p'
#endif
/* /*
* Get the max number of connections in an upper layer for each element in the index * Get the max number of connections in an upper layer for each element in the index
*/ */
@@ -51,22 +46,6 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
return index_getprocinfo(index, 1, procnum); return index_getprocinfo(index, 1, procnum);
} }
/*
* Init procs
*/
FmgrInfo **
HnswInitProcinfos(Relation index)
{
int keyAttributes = IndexRelationGetNumberOfKeyAttributes(index);
FmgrInfo **procinfos = palloc(keyAttributes * sizeof(FmgrInfo *));
procinfos[0] = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
for (int i = 1; i < keyAttributes; i++)
procinfos[i] = index_getprocinfo(index, i + 1, HNSW_ATTRIBUTE_DISTANCE_PROC);
return procinfos;
}
/* /*
* Divide by the norm * Divide by the norm
* *
@@ -160,7 +139,6 @@ HnswInitNeighbors(HnswElement element, int m)
a = &element->neighbors[lc]; a = &element->neighbors[lc];
a->length = 0; a->length = 0;
a->items = palloc(sizeof(HnswCandidate) * lm); a->items = palloc(sizeof(HnswCandidate) * lm);
a->closerSet = false;
} }
} }
@@ -194,7 +172,6 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
element->level = level; element->level = level;
element->deleted = 0; element->deleted = 0;
element->itup = NULL;
HnswInitNeighbors(element, m); HnswInitNeighbors(element, m);
@@ -209,8 +186,7 @@ HnswFreeElement(HnswElement element)
{ {
HnswFreeNeighbors(element); HnswFreeNeighbors(element);
list_free_deep(element->heaptids); list_free_deep(element->heaptids);
if (element->itup) pfree(element->vec);
pfree(element->itup);
pfree(element); pfree(element);
} }
@@ -237,8 +213,7 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
element->blkno = blkno; element->blkno = blkno;
element->offno = offno; element->offno = offno;
element->neighbors = NULL; element->neighbors = NULL;
element->value = PointerGetDatum(NULL); element->vec = NULL;
element->itup = NULL;
return element; return element;
} }
@@ -336,7 +311,7 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
* Set element tuple, except for neighbor info * Set element tuple, except for neighbor info
*/ */
void void
HnswSetElementTuple(HnswElementTuple etup, HnswElement element, bool useIndexTuple) HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
{ {
etup->type = HNSW_ELEMENT_TUPLE_TYPE; etup->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level; etup->level = element->level;
@@ -348,11 +323,7 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element, bool useIndexTup
else else
ItemPointerSetInvalid(&etup->heaptids[i]); ItemPointerSetInvalid(&etup->heaptids[i]);
} }
memcpy(&etup->vec, element->vec, VECTOR_SIZE(element->vec->dim));
if (useIndexTuple)
memcpy(&etup->data, element->itup, IndexTupleSize(element->itup));
else
memcpy(&etup->data, DatumGetPointer(element->value), VARSIZE_ANY(DatumGetPointer(element->value)));
} }
/* /*
@@ -453,7 +424,7 @@ HnswLoadNeighbors(HnswElement element, Relation index, 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, Relation index) HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec)
{ {
element->level = etup->level; element->level = etup->level;
element->deleted = etup->deleted; element->deleted = etup->deleted;
@@ -475,149 +446,16 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
if (loadVec) if (loadVec)
{ {
if (IndexRelationGetNumberOfAttributes(index) > 1) element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
{ memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
TupleDesc tupdesc = RelationGetDescr(index);
bool unused;
element->itup = CopyIndexTuple((IndexTuple) &etup->data);
element->value = index_getattr(element->itup, 1, tupdesc, &unused);
}
else
{
Vector *vec = palloc(VARSIZE_ANY(&etup->data));
memcpy(vec, &etup->data, VARSIZE_ANY(&etup->data));
element->value = PointerGetDatum(vec);
}
} }
} }
/*
* Get the tuple descriptor
*/
static TupleDesc
HnswTupleDesc(Relation index)
{
TupleDesc tupdesc = CreateTupleDescCopyConstr(RelationGetDescr(index));
/* Prevent compression */
TupleDescAttr(tupdesc, 0)->attstorage = TYPSTORAGE_PLAIN;
return tupdesc;
}
/*
* Set element data
*/
void
HnswElementSetData(HnswElement element, Relation index, Datum value, Datum *values, bool *isnull)
{
/* TODO Create once per index build */
TupleDesc tupdesc = HnswTupleDesc(index);
bool unused;
Datum tmp;
tmp = values[0];
values[0] = value;
element->itup = index_form_tuple(tupdesc, values, isnull);
values[0] = tmp;
element->value = index_getattr(element->itup, 1, tupdesc, &unused);
FreeTupleDesc(tupdesc);
}
/*
* 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;
}
/*
* Get the distance
*/
static double
GetDistance(IndexTuple itup, Datum vec, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfos, Oid *collations)
{
double g = DatumGetFloat8(FunctionCall2Coll(procinfos[0], collations[0], q, vec));
if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
{
double w = 0.25;
double e = 0.0;
TupleDesc tupdesc = RelationGetDescr(index);
if (keyData)
{
/* TODO need to pass length of key data */
int keyCount = 1;
for (int i = 0; i < keyCount; i++)
{
ScanKey key = &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;
}
else if (!DatumGetBool(FunctionCall2Coll(&key->sk_func, key->sk_collation, value, key->sk_argument)))
{
double ei = fabs(DatumGetFloat8(FunctionCall2Coll(procinfos[key->sk_attno - 1], collations[key->sk_attno - 1], value, key->sk_argument)));
if (ei > 0)
e += ei;
else
/* Distance is zero for inequality */
e += 1000;
}
}
return w * g + AttributeDistance(e);
}
else if (qtup)
{
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(qtup, i + 2, tupdesc, &attnull);
if (isnull || attnull)
{
if (isnull != attnull)
e += 1000;
}
else
e += fabs(DatumGetFloat8(FunctionCall2Coll(procinfos[i + 1], collations[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
*/ */
void void
HnswLoadElement(HnswElement element, float *distance, Datum *q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfos, Oid *collations, bool loadVec) HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -633,27 +471,11 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, IndexTuple qtup,
Assert(HnswIsElementTuple(etup)); Assert(HnswIsElementTuple(etup));
/* Load element */ /* Load element */
HnswLoadElementFromTuple(element, etup, true, loadVec, index); HnswLoadElementFromTuple(element, etup, true, loadVec);
/* Calculate distance */ /* Calculate distance */
if (distance != NULL) if (distance != NULL)
{ *distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->vec)));
IndexTuple itup = NULL;
Datum value;
if (IndexRelationGetNumberOfAttributes(index) > 1)
{
TupleDesc tupdesc = RelationGetDescr(index);
bool unused;
itup = (IndexTuple) &etup->data;
value = index_getattr(itup, 1, tupdesc, &unused);
}
else
value = PointerGetDatum(&etup->data);
*distance = GetDistance(itup, value, *q, qtup, keyData, index, procinfos, collations);
}
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -662,24 +484,24 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, IndexTuple qtup,
* Get the distance for a candidate * Get the distance for a candidate
*/ */
static float static float
GetCandidateDistance(HnswCandidate * hc, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfos, Oid *collations) GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
{ {
return GetDistance(hc->element->itup, hc->element->value, q, qtup, keyData, index, procinfos, collations); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, PointerGetDatum(hc->element->vec)));
} }
/* /*
* Create a candidate for the entry point * Create a candidate for the entry point
*/ */
HnswCandidate * HnswCandidate *
HnswEntryCandidate(HnswElement entryPoint, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfos, Oid *collations, bool loadVec, bool inMemory) HnswEntryCandidate(HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{ {
HnswCandidate *hc = palloc(sizeof(HnswCandidate)); HnswCandidate *hc = palloc(sizeof(HnswCandidate));
hc->element = entryPoint; hc->element = entryPoint;
if (inMemory) if (index == NULL)
hc->distance = GetCandidateDistance(hc, q, qtup, keyData, index, procinfos, collations); hc->distance = GetCandidateDistance(hc, q, procinfo, collation);
else else
HnswLoadElement(hc->element, &hc->distance, &q, qtup, keyData, index, procinfos, collations, loadVec); HnswLoadElement(hc->element, &hc->distance, &q, index, procinfo, collation, loadVec);
return hc; return hc;
} }
@@ -729,9 +551,9 @@ CreatePairingHeapNode(HnswCandidate * c)
* Add to visited * Add to visited
*/ */
static inline void static inline void
AddToVisited(HTAB *v, HnswCandidate * hc, bool inMemory, bool *found) AddToVisited(HTAB *v, HnswCandidate * hc, Relation index, bool *found)
{ {
if (inMemory) if (index == NULL)
hash_search(v, &hc->element, HASH_ENTER, found); hash_search(v, &hc->element, HASH_ENTER, found);
else else
{ {
@@ -746,7 +568,7 @@ AddToVisited(HTAB *v, HnswCandidate * hc, bool inMemory, bool *found)
* Algorithm 2 from paper * Algorithm 2 from paper
*/ */
List * List *
HnswSearchLayer(Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef, int lc, Relation index, FmgrInfo **procinfos, Oid *collations, int m, bool loadVec, HnswElement skipElement, bool inMemory) HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
{ {
ListCell *lc2; ListCell *lc2;
@@ -758,7 +580,7 @@ HnswSearchLayer(Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef
HTAB *v; HTAB *v;
/* Create hash table */ /* Create hash table */
if (inMemory) if (index == NULL)
{ {
hash_ctl.keysize = sizeof(HnswElement *); hash_ctl.keysize = sizeof(HnswElement *);
hash_ctl.entrysize = sizeof(HnswElement *); hash_ctl.entrysize = sizeof(HnswElement *);
@@ -777,7 +599,7 @@ HnswSearchLayer(Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef
{ {
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2); HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
AddToVisited(v, hc, inMemory, NULL); AddToVisited(v, hc, index, NULL);
pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node)); pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node));
pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node)); pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node));
@@ -811,7 +633,7 @@ HnswSearchLayer(Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef
HnswCandidate *e = &neighborhood->items[i]; HnswCandidate *e = &neighborhood->items[i];
bool visited; bool visited;
AddToVisited(v, e, inMemory, &visited); AddToVisited(v, e, index, &visited);
if (!visited) if (!visited)
{ {
@@ -819,10 +641,10 @@ HnswSearchLayer(Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef
f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner; f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
if (inMemory) if (index == NULL)
eDistance = GetCandidateDistance(e, q, qtup, keyData, index, procinfos, collations); eDistance = GetCandidateDistance(e, q, procinfo, collation);
else else
HnswLoadElement(e->element, &eDistance, &q, qtup, keyData, index, procinfos, collations, loadVec); HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, inserting);
Assert(!e->element->deleted); Assert(!e->element->deleted);
@@ -870,39 +692,11 @@ HnswSearchLayer(Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef
return w; return w;
} }
/*
* Compare candidate distances
*/
static int
#if PG_VERSION_NUM >= 130000
CompareCandidateDistances(const ListCell *a, const ListCell *b)
#else
CompareCandidateDistances(const void *a, const void *b)
#endif
{
HnswCandidate *hca = lfirst((ListCell *) a);
HnswCandidate *hcb = lfirst((ListCell *) b);
if (hca->distance < hcb->distance)
return 1;
if (hca->distance > hcb->distance)
return -1;
if (hca->element < hcb->element)
return 1;
if (hca->element > hcb->element)
return -1;
return 0;
}
/* /*
* Calculate the distance between elements * Calculate the distance between elements
*/ */
static float static float
HnswGetCachedDistance(HnswElement a, HnswElement b, int lc, Relation index, FmgrInfo **procinfos, Oid *collations) HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid collation)
{ {
/* Look for cached distance */ /* Look for cached distance */
if (a->neighbors != NULL) if (a->neighbors != NULL)
@@ -927,21 +721,21 @@ HnswGetCachedDistance(HnswElement a, HnswElement b, int lc, Relation index, Fmgr
} }
} }
return GetDistance(a->itup, a->value, b->value, b->itup, NULL, index, procinfos, collations); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(a->vec), PointerGetDatum(b->vec)));
} }
/* /*
* 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(HnswCandidate * e, List *r, int lc, Relation index, FmgrInfo **procinfos, Oid *collations) CheckElementCloser(HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid collation)
{ {
ListCell *lc2; ListCell *lc2;
foreach(lc2, r) foreach(lc2, r)
{ {
HnswCandidate *ri = lfirst(lc2); HnswCandidate *ri = lfirst(lc2);
float distance = HnswGetCachedDistance(e->element, ri->element, lc, index, procinfos, collations); float distance = HnswGetDistance(e->element, ri->element, lc, procinfo, collation);
if (distance <= e->distance) if (distance <= e->distance)
return false; return false;
@@ -954,77 +748,33 @@ CheckElementCloser(HnswCandidate * e, List *r, int lc, Relation index, FmgrInfo
* Algorithm 4 from paper * Algorithm 4 from paper
*/ */
static List * static List *
SelectNeighbors(List *c, int m, int lc, Relation index, FmgrInfo **procinfos, Oid *collations, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates) SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswCandidate * *pruned)
{ {
List *r = NIL; List *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
pairingheap *wd; pairingheap *wd;
bool mustCalculate = !e2->neighbors[lc].closerSet;
List *added = NIL;
bool removedAny = false;
if (list_length(w) <= m) if (list_length(w) <= m)
return w; return w;
wd = pairingheap_allocate(CompareNearestCandidates, NULL); wd = pairingheap_allocate(CompareNearestCandidates, NULL);
/* Ensure order of candidates is deterministic for closer caching */
if (sortCandidates)
list_sort(w, CompareCandidateDistances);
while (list_length(w) > 0 && list_length(r) < m) while (list_length(w) > 0 && list_length(r) < m)
{ {
/* Assumes w is already ordered desc */ /* Assumes w is already ordered desc */
HnswCandidate *e = llast(w); HnswCandidate *e = llast(w);
bool closer;
w = list_delete_last(w); w = list_delete_last(w);
/* Use previous state of r and wd to skip work when possible */ closer = CheckElementCloser(e, r, lc, procinfo, collation);
if (mustCalculate)
e->closer = CheckElementCloser(e, r, lc, index, procinfos, collations);
else if (list_length(added) > 0)
{
/*
* If the current candidate was closer, we only need to compare it
* with the other candidates that we have added.
*/
if (e->closer)
{
e->closer = CheckElementCloser(e, added, lc, index, procinfos, collations);
if (!e->closer) if (closer)
removedAny = true;
}
else
{
/*
* If we have removed any candidates from closer, a candidate
* that was not closer earlier might now be.
*/
if (removedAny)
{
e->closer = CheckElementCloser(e, r, lc, index, procinfos, collations);
if (e->closer)
added = lappend(added, e);
}
}
}
else if (e == newCandidate)
{
e->closer = CheckElementCloser(e, r, lc, index, procinfos, collations);
if (e->closer)
added = lappend(added, e);
}
if (e->closer)
r = lappend(r, e); r = lappend(r, e);
else else
pairingheap_add(wd, &(CreatePairingHeapNode(e)->ph_node)); pairingheap_add(wd, &(CreatePairingHeapNode(e)->ph_node));
} }
/* Cached value can only be used in future if sorted deterministically */
e2->neighbors[lc].closerSet = sortCandidates;
/* Keep pruned connections */ /* Keep pruned connections */
while (!pairingheap_is_empty(wd) && list_length(r) < m) while (!pairingheap_is_empty(wd) && list_length(r) < m)
r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner); r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner);
@@ -1045,20 +795,16 @@ SelectNeighbors(List *c, int m, int lc, Relation index, FmgrInfo **procinfos, Oi
* Find duplicate element * Find duplicate element
*/ */
HnswElement HnswElement
HnswFindDuplicate(HnswElement e, Relation index) HnswFindDuplicate(HnswElement e)
{ {
HnswNeighborArray *neighbors = &e->neighbors[0]; HnswNeighborArray *neighbors = &e->neighbors[0];
/* TODO Implement */
if (IndexRelationGetNumberOfAttributes(index) > 1)
return NULL;
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];
/* Exit early since ordered by distance */ /* Exit early since ordered by distance */
if (!datumIsEqual(e->value, neighbor->element->value, false, -1)) if (vector_cmp_internal(e->vec, neighbor->element->vec) != 0)
break; break;
/* Check for space */ /* Check for space */
@@ -1082,11 +828,33 @@ AddConnections(HnswElement element, List *neighbors, int m, int lc)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2)); a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
} }
/*
* Compare candidate distances
*/
static int
#if PG_VERSION_NUM >= 130000
CompareCandidateDistances(const ListCell *a, const ListCell *b)
#else
CompareCandidateDistances(const void *a, const void *b)
#endif
{
HnswCandidate *hca = lfirst((ListCell *) a);
HnswCandidate *hcb = lfirst((ListCell *) b);
if (hca->distance < hcb->distance)
return 1;
if (hca->distance > hcb->distance)
return -1;
return 0;
}
/* /*
* Update connections * Update connections
*/ */
void void
HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo **procinfos, Oid *collations, bool inMemory) HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
{ {
HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc]; HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc];
@@ -1109,20 +877,18 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
HnswCandidate *pruned = NULL; HnswCandidate *pruned = NULL;
/* Load elements on insert */ /* Load elements on insert */
if (!inMemory) if (index != NULL)
{ {
Datum q = hc->element->value; Datum q = PointerGetDatum(hc->element->vec);
IndexTuple qtup = hc->element->itup;
ScanKeyData *keyData = NULL;
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
{ {
HnswCandidate *hc3 = &currentNeighbors->items[i]; HnswCandidate *hc3 = &currentNeighbors->items[i];
if (DatumGetPointer(hc3->element->value) == NULL) if (hc3->element->vec == NULL)
HnswLoadElement(hc3->element, &hc3->distance, &q, qtup, keyData, index, procinfos, collations, true); HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
else else
hc3->distance = GetCandidateDistance(hc3, q, qtup, keyData, index, procinfos, collations); hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
/* Prune element if being deleted */ /* Prune element if being deleted */
if (list_length(hc3->element->heaptids) == 0) if (list_length(hc3->element->heaptids) == 0)
@@ -1137,12 +903,13 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
{ {
List *c = NIL; List *c = NIL;
/* Add candidates */ /* Add and sort candidates */
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
c = lappend(c, &currentNeighbors->items[i]); c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2); c = lappend(c, &hc2);
list_sort(c, CompareCandidateDistances);
SelectNeighbors(c, m, lc, index, procinfos, collations, hc->element, &hc2, &pruned, true); SelectNeighbors(c, m, lc, procinfo, collation, &pruned);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -1194,15 +961,13 @@ RemoveElements(List *w, HnswElement skipElement)
* Algorithm 1 from paper * Algorithm 1 from paper
*/ */
void void
HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo **procinfos, Oid *collations, int m, int efConstruction, bool existing, bool inMemory) HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing)
{ {
List *ep; List *ep;
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
Datum q = element->value; Datum q = PointerGetDatum(element->vec);
IndexTuple qtup = element->itup;
ScanKeyData *keyData = NULL;
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
/* No neighbors if no entry point */ /* No neighbors if no entry point */
@@ -1210,13 +975,13 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
return; return;
/* Get entry point and level */ /* Get entry point and level */
ep = list_make1(HnswEntryCandidate(entryPoint, q, qtup, keyData, index, procinfos, collations, true, inMemory)); ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, true));
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(q, qtup, keyData, ep, 1, lc, index, procinfos, collations, m, true, skipElement, inMemory); w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, m, true, skipElement);
ep = w; ep = w;
} }
@@ -1234,21 +999,16 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
List *neighbors; List *neighbors;
List *lw; List *lw;
w = HnswSearchLayer(q, qtup, keyData, ep, efConstruction, lc, index, procinfos, collations, m, true, skipElement, inMemory); w = HnswSearchLayer(q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement);
/* Elements being deleted or skipped can help with search */ /* Elements being deleted or skipped can help with search */
/* but should be removed before selecting neighbors */ /* but should be removed before selecting neighbors */
if (!inMemory) if (index != NULL)
lw = RemoveElements(w, skipElement); lw = RemoveElements(w, skipElement);
else else
lw = w; lw = w;
/* neighbors = SelectNeighbors(lw, lm, lc, procinfo, collation, NULL);
* Candidates are sorted, but not deterministically. Could set
* sortCandidates to true for in-memory builds to enable closer
* caching, but there does not seem to be a difference in performance.
*/
neighbors = SelectNeighbors(lw, lm, lc, index, procinfos, collations, element, NULL, NULL, false);
AddConnections(element, neighbors, lm, lc); AddConnections(element, neighbors, lm, lc);

View File

@@ -62,8 +62,7 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
/* Iterate over nodes */ /* Iterate over nodes */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{ {
ItemId itemid = PageGetItemId(page, offno); HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, itemid);
int idx = 0; int idx = 0;
bool itemUpdated = false; bool itemUpdated = false;
@@ -94,7 +93,7 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
if (itemUpdated) if (itemUpdated)
{ {
Size etupSize = ItemIdGetLength(itemid); Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
/* Mark rest as invalid */ /* Mark rest as invalid */
for (int i = idx; i < HNSW_HEAPTIDS; i++) for (int i = idx; i < HNSW_HEAPTIDS; i++)
@@ -195,8 +194,8 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
GenericXLogState *state; GenericXLogState *state;
int m = vacuumstate->m; int m = vacuumstate->m;
int efConstruction = vacuumstate->efConstruction; int efConstruction = vacuumstate->efConstruction;
FmgrInfo **procinfos = vacuumstate->procinfos; FmgrInfo *procinfo = vacuumstate->procinfo;
Oid *collations = vacuumstate->collations; Oid collation = vacuumstate->collation;
BufferAccessStrategy bas = vacuumstate->bas; BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup; HnswNeighborTuple ntup = vacuumstate->ntup;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m); Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
@@ -210,7 +209,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
element->heaptids = NIL; element->heaptids = NIL;
/* Add element to graph, skipping itself */ /* Add element to graph, skipping itself */
HnswInsertElement(element, entryPoint, index, procinfos, collations, m, efConstruction, true, false); HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, true);
/* Update neighbor tuple */ /* Update neighbor tuple */
/* Do this before getting page to minimize locking */ /* Do this before getting page to minimize locking */
@@ -231,7 +230,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborPages(index, procinfos, collations, element, m, true); HnswUpdateNeighborPages(index, procinfo, collation, element, m, true);
} }
/* /*
@@ -258,7 +257,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
LockPage(index, HNSW_UPDATE_LOCK, ShareLock); LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
/* Load element */ /* Load element */
HnswLoadElement(highestPoint, NULL, NULL, NULL, NULL, index, vacuumstate->procinfos, vacuumstate->collations, true); HnswLoadElement(highestPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
/* Repair if needed */ /* Repair if needed */
if (NeedsUpdated(vacuumstate, highestPoint)) if (NeedsUpdated(vacuumstate, highestPoint))
@@ -296,7 +295,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, NULL, NULL, index, vacuumstate->procinfos, vacuumstate->collations, true); HnswLoadElement(entryPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
if (NeedsUpdated(vacuumstate, entryPoint)) if (NeedsUpdated(vacuumstate, entryPoint))
{ {
@@ -372,7 +371,7 @@ RepairGraph(HnswVacuumState * vacuumstate)
/* Create an element */ /* Create an element */
element = HnswInitElementFromBlock(blkno, offno); element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(element, etup, false, true, index); HnswLoadElementFromTuple(element, etup, false, true);
elements = lappend(elements, element); elements = lappend(elements, element);
} }
@@ -442,7 +441,6 @@ 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 = IndexRelationGetNumberOfAttributes(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
@@ -479,8 +477,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Update element and neighbors together */ /* Update element and neighbors together */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{ {
ItemId itemid = PageGetItemId(page, offno); HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, itemid);
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
Size etupSize; Size etupSize;
Size ntupSize; Size ntupSize;
@@ -508,7 +505,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
continue; continue;
/* Calculate sizes */ /* Calculate sizes */
etupSize = ItemIdGetLength(itemid); etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m);
/* Get neighbor page */ /* Get neighbor page */
@@ -531,18 +528,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Overwrite element */ /* Overwrite element */
etup->deleted = 1; etup->deleted = 1;
if (useIndexTuple) MemSet(&etup->vec.x, 0, etup->vec.dim * sizeof(float));
{
IndexTuple itup = (IndexTuple) &etup->data;
MemSet(itup, 0, IndexTupleSize(itup));
}
else
{
Vector *vec = (Vector *) (&etup->data);
MemSet(vec, 0, VARSIZE_ANY(vec));
}
/* Overwrite neighbors */ /* Overwrite neighbors */
for (int i = 0; i < ntup->count; i++) for (int i = 0; i < ntup->count; i++)
@@ -598,8 +584,8 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
vacuumstate->callback_state = callback_state; vacuumstate->callback_state = callback_state;
vacuumstate->efConstruction = HnswGetEfConstruction(index); vacuumstate->efConstruction = HnswGetEfConstruction(index);
vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD); vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD);
vacuumstate->procinfos = HnswInitProcinfos(index); vacuumstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
vacuumstate->collations = index->rd_indcollation; vacuumstate->collation = index->rd_indcollation[0];
vacuumstate->ntup = palloc0(BLCKSZ); vacuumstate->ntup = palloc0(BLCKSZ);
vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw vacuum temporary context", "Hnsw vacuum temporary context",
@@ -623,7 +609,6 @@ FreeVacuumState(HnswVacuumState * vacuumstate)
{ {
hash_destroy(vacuumstate->deleted); hash_destroy(vacuumstate->deleted);
FreeAccessStrategy(vacuumstate->bas); FreeAccessStrategy(vacuumstate->bas);
pfree(vacuumstate->procinfos);
pfree(vacuumstate->ntup); pfree(vacuumstate->ntup);
MemoryContextDelete(vacuumstate->tmpCtx); MemoryContextDelete(vacuumstate->tmpCtx);
} }

View File

@@ -246,6 +246,8 @@ typedef struct IvfflatScanOpaqueData
int probes; int probes;
int dimensions; int dimensions;
bool first; bool first;
Buffer buf;
ItemPointerData heaptid;
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;

View File

@@ -143,6 +143,10 @@ GetScanItems(IndexScanDesc scan, Datum value)
bool isnull; bool isnull;
ItemId itemid = PageGetItemId(page, offno); ItemId itemid = PageGetItemId(page, offno);
/* Skip dead tuples */
if (scan->ignore_killed_tuples && ItemIdIsDead(itemid))
continue;
itup = (IndexTuple) PageGetItem(page, itemid); itup = (IndexTuple) PageGetItem(page, itemid);
datum = index_getattr(itup, 1, tupdesc, &isnull); datum = index_getattr(itup, 1, tupdesc, &isnull);
@@ -157,6 +161,8 @@ GetScanItems(IndexScanDesc scan, Datum value)
slot->tts_isnull[0] = false; slot->tts_isnull[0] = false;
slot->tts_values[1] = PointerGetDatum(&itup->t_tid); slot->tts_values[1] = PointerGetDatum(&itup->t_tid);
slot->tts_isnull[1] = false; slot->tts_isnull[1] = false;
slot->tts_values[2] = Int32GetDatum((int) searchPage);
slot->tts_isnull[2] = false;
ExecStoreVirtualTuple(slot); ExecStoreVirtualTuple(slot);
tuplesort_puttupleslot(so->sortstate, slot); tuplesort_puttupleslot(so->sortstate, slot);
@@ -181,6 +187,55 @@ GetScanItems(IndexScanDesc scan, Datum value)
tuplesort_performsort(so->sortstate); tuplesort_performsort(so->sortstate);
} }
/*
* Mark prior tuple as dead
*/
static void
MarkPriorTupleDead(IndexScanDesc scan)
{
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
Buffer buf = so->buf;
Page page;
OffsetNumber maxoffno;
/* Safety check */
if (!BufferIsValid(so->buf) || !ItemPointerIsValid(&so->heaptid))
return;
/* Only a shared locked is needed for ItemIdMarkDead */
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
maxoffno = PageGetMaxOffsetNumber(page);
for (OffsetNumber offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
ItemId itemid = PageGetItemId(page, offno);
IndexTuple itup = (IndexTuple) PageGetItem(page, itemid);
/*
* Find tuple. Since buffer has been pinned, tuple cannot have been
* vacuumed (and heap TID reused).
*/
if (ItemPointerEquals(&itup->t_tid, &so->heaptid))
{
/*
* Make sure tuple has not already been marked dead to avoid extra
* WAL if wal_log_hints or data checksums enabled
*/
if (!ItemIdIsDead(itemid))
{
ItemIdMarkDead(itemid);
MarkBufferDirtyHint(buf, true);
}
break;
}
}
/* Unlock buffer */
LockBuffer(buf, BUFFER_LOCK_UNLOCK);
}
/* /*
* Prepare for an index scan * Prepare for an index scan
*/ */
@@ -206,7 +261,9 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
probes = lists; probes = lists;
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList)); so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
so->buf = InvalidBuffer;
so->first = true; so->first = true;
ItemPointerSetInvalid(&so->heaptid);
so->probes = probes; so->probes = probes;
so->dimensions = dimensions; so->dimensions = dimensions;
@@ -217,12 +274,13 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
/* Create tuple description for sorting */ /* Create tuple description for sorting */
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
so->tupdesc = CreateTemplateTupleDesc(2); so->tupdesc = CreateTemplateTupleDesc(3);
#else #else
so->tupdesc = CreateTemplateTupleDesc(2, false); so->tupdesc = CreateTemplateTupleDesc(3, false);
#endif #endif
TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0);
/* Prep sort */ /* Prep sort */
so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false); so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
@@ -254,6 +312,7 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
#endif #endif
so->first = true; so->first = true;
ItemPointerSetInvalid(&so->heaptid);
pairingheap_reset(so->listQueue); pairingheap_reset(so->listQueue);
if (keys && scan->numberOfKeys > 0) if (keys && scan->numberOfKeys > 0)
@@ -288,11 +347,6 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
if (scan->orderByData == NULL) if (scan->orderByData == NULL)
elog(ERROR, "cannot scan ivfflat index without order"); elog(ERROR, "cannot scan ivfflat index without order");
/* Requires MVCC-compliant snapshot as not able to pin during sorting */
/* https://www.postgresql.org/docs/current/index-locking.html */
if (!IsMVCCSnapshot(scan->xs_snapshot))
elog(ERROR, "non-MVCC snapshots are not supported with ivfflat");
if (scan->orderByData->sk_flags & SK_ISNULL) if (scan->orderByData->sk_flags & SK_ISNULL)
value = PointerGetDatum(InitVector(so->dimensions)); value = PointerGetDatum(InitVector(so->dimensions));
else else
@@ -316,10 +370,17 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
if (value != scan->orderByData->sk_argument) if (value != scan->orderByData->sk_argument)
pfree(DatumGetPointer(value)); pfree(DatumGetPointer(value));
} }
else
{
/* Mark prior tuple as dead */
if (scan->kill_prior_tuple)
MarkPriorTupleDead(scan);
}
if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL)) if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
{ {
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull)); ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;
@@ -327,6 +388,21 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
scan->xs_ctup.t_self = *heaptid; scan->xs_ctup.t_self = *heaptid;
#endif #endif
/* Keep track of info needed to mark tuple as dead */
so->heaptid = *heaptid;
/* Unpin buffer */
if (BufferIsValid(so->buf))
ReleaseBuffer(so->buf);
/*
* An index scan must maintain a pin on the index page holding the
* item last returned by amgettuple
*
* https://www.postgresql.org/docs/current/index-locking.html
*/
so->buf = ReadBuffer(scan->indexRelation, indexblkno);
scan->xs_recheckorderby = false; scan->xs_recheckorderby = false;
return true; return true;
} }
@@ -342,6 +418,10 @@ ivfflatendscan(IndexScanDesc scan)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
/* Release pin */
if (BufferIsValid(so->buf))
ReleaseBuffer(so->buf);
pairingheap_free(so->listQueue); pairingheap_free(so->listQueue);
tuplesort_end(so->sortstate); tuplesort_end(so->sortstate);

294
src/tinyint.c Normal file
View File

@@ -0,0 +1,294 @@
#include "postgres.h"
#include <math.h>
#include <stdint.h>
#include "fmgr.h"
#include "lib/stringinfo.h"
#include "libpq/pqformat.h"
#include "tinyint.h"
#include "utils/array.h"
#include "utils/builtins.h"
#include "utils/numeric.h"
/*
* Check if array is a vector
*/
static bool
ArrayIsVector(ArrayType *a)
{
return ARR_NDIM(a) == 1 && !array_contains_nulls(a);
}
/*
* Check if dimensions are the same
*/
static int
CheckDims(ArrayType *a, ArrayType *b)
{
int dima;
int dimb;
if (!ArrayIsVector(a) || !ArrayIsVector(b))
return 0;
dima = ARR_DIMS(a)[0];
dimb = ARR_DIMS(b)[0];
if (dima != dimb)
return 0;
return dima;
}
/*
* Check range
*/
static void
CheckRange(long i)
{
if (i < INT8_MIN || i > INT8_MAX)
ereport(ERROR,
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("value \"%ld\" is out of range for type tinyint", i)));
}
/*
* Convert textual representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(tinyint_in);
Datum
tinyint_in(PG_FUNCTION_ARGS)
{
char *s = PG_GETARG_CSTRING(0);
const char *ptr = s;
long i;
char *end;
/* skip leading spaces */
while (*ptr != '\0' && isspace((unsigned char) *ptr))
ptr++;
if (*ptr == '\0')
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type tinyint: \"%s\"", s)));
i = strtol(ptr, &end, 10);
ptr = end;
if (i < INT8_MIN || i > INT8_MAX)
ereport(ERROR,
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("value \"%s\" is out of range for type tinyint", s)));
/* allow trailing whitespace, but not other trailing chars */
while (*ptr != '\0' && isspace((unsigned char) *ptr))
ptr++;
if (*ptr != '\0')
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type tinyint: \"%s\"", s)));
PG_RETURN_INT8(i);
}
/*
* Convert internal representation to textual representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(tinyint_out);
Datum
tinyint_out(PG_FUNCTION_ARGS)
{
int8 num = PG_GETARG_INT8(0);
char *result = (char *) palloc(5); /* sign, 3 digits, '\0' */
pg_ltoa((int32) num, result);
PG_RETURN_CSTRING(result);
}
/*
* Convert external binary representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(tinyint_recv);
Datum
tinyint_recv(PG_FUNCTION_ARGS)
{
StringInfo buf = (StringInfo) PG_GETARG_POINTER(0);
PG_RETURN_INT8((int8) pq_getmsgint(buf, sizeof(int8)));
}
/*
* Convert internal representation to the external binary representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(tinyint_send);
Datum
tinyint_send(PG_FUNCTION_ARGS)
{
int8 arg1 = PG_GETARG_INT8(0);
StringInfoData buf;
pq_begintypsend(&buf);
pq_sendint8(&buf, arg1);
PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
}
/*
* Convert integer to tinyint
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(integer_to_tinyint);
Datum
integer_to_tinyint(PG_FUNCTION_ARGS)
{
int32 i = PG_GETARG_INT32(0);
CheckRange(i);
PG_RETURN_INT8(i);
}
/*
* Convert numeric to tinyint
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(numeric_to_tinyint);
Datum
numeric_to_tinyint(PG_FUNCTION_ARGS)
{
Numeric num = PG_GETARG_NUMERIC(0);
int32 i = numeric_int4_opt_error(num, NULL);
CheckRange(i);
PG_RETURN_INT8(i);
}
/*
* Get the L2 distance between tinyint arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(tinyint_l2_distance);
Datum
tinyint_l2_distance(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
int8 *ax = (int8 *) ARR_DATA_PTR(a);
int8 *bx = (int8 *) ARR_DATA_PTR(b);
double distance = 0.0;
int dim = CheckDims(a, b);
/* TODO Decide on error or NULL */
if (!dim)
PG_RETURN_NULL();
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
{
double diff = ax[i] - bx[i];
distance += diff * diff;
}
PG_RETURN_FLOAT8(sqrt(distance));
}
/*
* Get the inner product of two tinyint arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(tinyint_inner_product);
Datum
tinyint_inner_product(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
int8 *ax = (int8 *) ARR_DATA_PTR(a);
int8 *bx = (int8 *) ARR_DATA_PTR(b);
double distance = 0.0;
int dim = CheckDims(a, b);
/* TODO Decide on error or NULL */
if (!dim)
PG_RETURN_NULL();
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
distance += ax[i] * bx[i];
PG_RETURN_FLOAT8(distance);
}
/*
* Get the negative inner product of two tinyint arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(tinyint_negative_inner_product);
Datum
tinyint_negative_inner_product(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
int8 *ax = (int8 *) ARR_DATA_PTR(a);
int8 *bx = (int8 *) ARR_DATA_PTR(b);
double distance = 0.0;
int dim = CheckDims(a, b);
/* TODO Decide on error or NULL */
if (!dim)
PG_RETURN_NULL();
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
distance += ax[i] * bx[i];
PG_RETURN_FLOAT8(distance * -1);
}
/*
* Get the cosine distance between two float2 arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(tinyint_cosine_distance);
Datum
tinyint_cosine_distance(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
int8 *ax = (int8 *) ARR_DATA_PTR(a);
int8 *bx = (int8 *) ARR_DATA_PTR(b);
double distance = 0.0;
double norma = 0.0;
double normb = 0.0;
double similarity;
int dim = CheckDims(a, b);
/* TODO Decide on error or NULL */
if (!dim)
PG_RETURN_NULL();
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
{
float axi = ax[i];
float bxi = bx[i];
distance += axi * bxi;
norma += axi * axi;
normb += bxi * bxi;
}
/* Use sqrt(a * b) over sqrt(a) * sqrt(b) */
similarity = distance / sqrt(norma * normb);
#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);
}

8
src/tinyint.h Normal file
View File

@@ -0,0 +1,8 @@
#ifndef TINYINT_H
#define TINYINT_H
#define DatumGetInt8(X) ((int8) (X))
#define PG_GETARG_INT8(n) DatumGetInt8(PG_GETARG_DATUM(n))
#define PG_RETURN_INT8(x) return Int8GetDatum(x)
#endif

View File

@@ -89,7 +89,7 @@ CheckDim(int dim)
} }
/* /*
* Ensure finite element * Ensure finite elements
*/ */
static inline void static inline void
CheckElement(float value) CheckElement(float value)
@@ -437,18 +437,17 @@ vector_send(PG_FUNCTION_ARGS)
/* /*
* Convert vector to vector * Convert vector to vector
* This is needed to check the type modifier
*/ */
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector); PGDLLEXPORT PG_FUNCTION_INFO_V1(vector);
Datum Datum
vector(PG_FUNCTION_ARGS) vector(PG_FUNCTION_ARGS)
{ {
Vector *vec = PG_GETARG_VECTOR_P(0); Vector *arg = PG_GETARG_VECTOR_P(0);
int32 typmod = PG_GETARG_INT32(1); int32 typmod = PG_GETARG_INT32(1);
CheckExpectedDim(typmod, vec->dim); CheckExpectedDim(typmod, arg->dim);
PG_RETURN_POINTER(vec); PG_RETURN_POINTER(arg);
} }
/* /*
@@ -465,6 +464,7 @@ array_to_vector(PG_FUNCTION_ARGS)
bool typbyval; bool typbyval;
char typalign; char typalign;
Datum *elemsp; Datum *elemsp;
bool *nullsp;
int nelemsp; int nelemsp;
if (ARR_NDIM(array) > 1) if (ARR_NDIM(array) > 1)
@@ -478,7 +478,7 @@ array_to_vector(PG_FUNCTION_ARGS)
errmsg("array must not contain nulls"))); errmsg("array must not contain nulls")));
get_typlenbyvalalign(ARR_ELEMTYPE(array), &typlen, &typbyval, &typalign); get_typlenbyvalalign(ARR_ELEMTYPE(array), &typlen, &typbyval, &typalign);
deconstruct_array(array, ARR_ELEMTYPE(array), typlen, typbyval, typalign, &elemsp, NULL, &nelemsp); deconstruct_array(array, ARR_ELEMTYPE(array), typlen, typbyval, typalign, &elemsp, &nullsp, &nelemsp);
CheckDim(nelemsp); CheckDim(nelemsp);
CheckExpectedDim(typmod, nelemsp); CheckExpectedDim(typmod, nelemsp);
@@ -512,12 +512,6 @@ array_to_vector(PG_FUNCTION_ARGS)
errmsg("unsupported array type"))); 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);
/* Check elements */ /* Check elements */
for (int i = 0; i < result->dim; i++) for (int i = 0; i < result->dim; i++)
CheckElement(result->x[i]); CheckElement(result->x[i]);

View File

@@ -54,85 +54,85 @@ SELECT vector_norm('[3e37,4e37]')::real;
5e+37 5e+37
(1 row) (1 row)
SELECT l2_distance('[0,0]', '[3,4]'); SELECT l2_distance('[0,0]'::vector, '[3,4]');
l2_distance l2_distance
------------- -------------
5 5
(1 row) (1 row)
SELECT l2_distance('[0,0]', '[0,1]'); SELECT l2_distance('[0,0]'::vector, '[0,1]');
l2_distance l2_distance
------------- -------------
1 1
(1 row) (1 row)
SELECT l2_distance('[1,2]', '[3]'); SELECT l2_distance('[1,2]'::vector, '[3]');
ERROR: different vector dimensions 2 and 1 ERROR: different vector dimensions 2 and 1
SELECT l2_distance('[3e38]', '[-3e38]'); SELECT l2_distance('[3e38]'::vector, '[-3e38]');
l2_distance l2_distance
------------- -------------
Infinity Infinity
(1 row) (1 row)
SELECT inner_product('[1,2]', '[3,4]'); SELECT inner_product('[1,2]'::vector, '[3,4]');
inner_product inner_product
--------------- ---------------
11 11
(1 row) (1 row)
SELECT inner_product('[1,2]', '[3]'); SELECT inner_product('[1,2]'::vector, '[3]');
ERROR: different vector dimensions 2 and 1 ERROR: different vector dimensions 2 and 1
SELECT inner_product('[3e38]', '[3e38]'); SELECT inner_product('[3e38]'::vector, '[3e38]');
inner_product inner_product
--------------- ---------------
Infinity Infinity
(1 row) (1 row)
SELECT cosine_distance('[1,2]', '[2,4]'); SELECT cosine_distance('[1,2]'::vector, '[2,4]');
cosine_distance cosine_distance
----------------- -----------------
0 0
(1 row) (1 row)
SELECT cosine_distance('[1,2]', '[0,0]'); SELECT cosine_distance('[1,2]'::vector, '[0,0]');
cosine_distance cosine_distance
----------------- -----------------
NaN NaN
(1 row) (1 row)
SELECT cosine_distance('[1,1]', '[1,1]'); SELECT cosine_distance('[1,1]'::vector, '[1,1]');
cosine_distance cosine_distance
----------------- -----------------
0 0
(1 row) (1 row)
SELECT cosine_distance('[1,0]', '[0,2]'); SELECT cosine_distance('[1,0]'::vector, '[0,2]');
cosine_distance cosine_distance
----------------- -----------------
1 1
(1 row) (1 row)
SELECT cosine_distance('[1,1]', '[-1,-1]'); SELECT cosine_distance('[1,1]'::vector, '[-1,-1]');
cosine_distance cosine_distance
----------------- -----------------
2 2
(1 row) (1 row)
SELECT cosine_distance('[1,2]', '[3]'); SELECT cosine_distance('[1,2]'::vector, '[3]');
ERROR: different vector dimensions 2 and 1 ERROR: different vector dimensions 2 and 1
SELECT cosine_distance('[1,1]', '[1.1,1.1]'); SELECT cosine_distance('[1,1]'::vector, '[1.1,1.1]');
cosine_distance cosine_distance
----------------- -----------------
0 0
(1 row) (1 row)
SELECT cosine_distance('[1,1]', '[-1.1,-1.1]'); SELECT cosine_distance('[1,1]'::vector, '[-1.1,-1.1]');
cosine_distance cosine_distance
----------------- -----------------
2 2
(1 row) (1 row)
SELECT cosine_distance('[3e38]', '[3e38]'); SELECT cosine_distance('[3e38]'::vector, '[3e38]');
cosine_distance cosine_distance
----------------- -----------------
NaN NaN

148
test/expected/tinyint.out Normal file
View File

@@ -0,0 +1,148 @@
SELECT '127'::tinyint;
tinyint
---------
127
(1 row)
SELECT '128'::tinyint;
ERROR: value "128" is out of range for type tinyint
LINE 1: SELECT '128'::tinyint;
^
SELECT '-128'::tinyint;
tinyint
---------
-128
(1 row)
SELECT '-129'::tinyint;
ERROR: value "-129" is out of range for type tinyint
LINE 1: SELECT '-129'::tinyint;
^
SELECT ''::tinyint;
ERROR: invalid input syntax for type tinyint: ""
LINE 1: SELECT ''::tinyint;
^
SELECT ' 1'::tinyint;
tinyint
---------
1
(1 row)
SELECT '1 '::tinyint;
tinyint
---------
1
(1 row)
SELECT '1a'::tinyint;
ERROR: invalid input syntax for type tinyint: "1a"
LINE 1: SELECT '1a'::tinyint;
^
SELECT '{1,2,3}'::tinyint[];
tinyint
---------
{1,2,3}
(1 row)
SELECT '128'::numeric::tinyint;
ERROR: value "128" is out of range for type tinyint
SELECT 'NaN'::numeric::tinyint;
ERROR: cannot convert NaN to integer
SELECT l2_distance('{0,0}'::tinyint[], '{3,4}'::tinyint[]);
l2_distance
-------------
5
(1 row)
SELECT l2_distance('{0,0}'::tinyint[], '{0,1}'::tinyint[]);
l2_distance
-------------
1
(1 row)
SELECT l2_distance('{1,2}'::tinyint[], '{3}'::tinyint[]);
l2_distance
-------------
(1 row)
SELECT l2_distance('{3e38}'::tinyint[], '{-3e38}'::tinyint[]);
ERROR: invalid input syntax for type tinyint: "3e38"
LINE 1: SELECT l2_distance('{3e38}'::tinyint[], '{-3e38}'::tinyint[]...
^
SELECT '{0,0}'::tinyint[] <-> '{3,4}'::tinyint[];
?column?
----------
5
(1 row)
SELECT inner_product('{1,2}'::tinyint[], '{3,4}'::tinyint[]);
inner_product
---------------
11
(1 row)
SELECT inner_product('{1,2}'::tinyint[], '{3}'::tinyint[]);
inner_product
---------------
(1 row)
SELECT inner_product('{127}'::tinyint[], '{127}'::tinyint[]);
inner_product
---------------
16129
(1 row)
SELECT '{1,2}'::tinyint[] <#> '{3,4}'::tinyint[];
?column?
----------
-11
(1 row)
SELECT cosine_distance('{1,2}'::tinyint[], '{2,4}'::tinyint[]);
cosine_distance
-----------------
0
(1 row)
SELECT cosine_distance('{1,2}'::tinyint[], '{0,0}'::tinyint[]);
cosine_distance
-----------------
NaN
(1 row)
SELECT cosine_distance('{1,1}'::tinyint[], '{1,1}'::tinyint[]);
cosine_distance
-----------------
0
(1 row)
SELECT cosine_distance('{1,0}'::tinyint[], '{0,2}'::tinyint[]);
cosine_distance
-----------------
1
(1 row)
SELECT cosine_distance('{1,1}'::tinyint[], '{-1,-1}'::tinyint[]);
cosine_distance
-----------------
2
(1 row)
SELECT cosine_distance('{1,2}'::tinyint[], '{3}'::tinyint[]);
cosine_distance
-----------------
(1 row)
SELECT cosine_distance('{3e38}'::tinyint[], '{3e38}'::tinyint[]);
ERROR: invalid input syntax for type tinyint: "3e38"
LINE 1: SELECT cosine_distance('{3e38}'::tinyint[], '{3e38}'::tinyin...
^
SELECT '{1,2}'::tinyint[] <=> '{2,4}'::tinyint[];
?column?
----------
0
(1 row)

View File

@@ -13,24 +13,24 @@ SELECT vector_norm('[3,4]');
SELECT vector_norm('[0,1]'); SELECT vector_norm('[0,1]');
SELECT vector_norm('[3e37,4e37]')::real; SELECT vector_norm('[3e37,4e37]')::real;
SELECT l2_distance('[0,0]', '[3,4]'); SELECT l2_distance('[0,0]'::vector, '[3,4]');
SELECT l2_distance('[0,0]', '[0,1]'); SELECT l2_distance('[0,0]'::vector, '[0,1]');
SELECT l2_distance('[1,2]', '[3]'); SELECT l2_distance('[1,2]'::vector, '[3]');
SELECT l2_distance('[3e38]', '[-3e38]'); SELECT l2_distance('[3e38]'::vector, '[-3e38]');
SELECT inner_product('[1,2]', '[3,4]'); SELECT inner_product('[1,2]'::vector, '[3,4]');
SELECT inner_product('[1,2]', '[3]'); SELECT inner_product('[1,2]'::vector, '[3]');
SELECT inner_product('[3e38]', '[3e38]'); SELECT inner_product('[3e38]'::vector, '[3e38]');
SELECT cosine_distance('[1,2]', '[2,4]'); SELECT cosine_distance('[1,2]'::vector, '[2,4]');
SELECT cosine_distance('[1,2]', '[0,0]'); SELECT cosine_distance('[1,2]'::vector, '[0,0]');
SELECT cosine_distance('[1,1]', '[1,1]'); SELECT cosine_distance('[1,1]'::vector, '[1,1]');
SELECT cosine_distance('[1,0]', '[0,2]'); SELECT cosine_distance('[1,0]'::vector, '[0,2]');
SELECT cosine_distance('[1,1]', '[-1,-1]'); SELECT cosine_distance('[1,1]'::vector, '[-1,-1]');
SELECT cosine_distance('[1,2]', '[3]'); SELECT cosine_distance('[1,2]'::vector, '[3]');
SELECT cosine_distance('[1,1]', '[1.1,1.1]'); SELECT cosine_distance('[1,1]'::vector, '[1.1,1.1]');
SELECT cosine_distance('[1,1]', '[-1.1,-1.1]'); SELECT cosine_distance('[1,1]'::vector, '[-1.1,-1.1]');
SELECT cosine_distance('[3e38]', '[3e38]'); SELECT cosine_distance('[3e38]'::vector, '[3e38]');
SELECT l1_distance('[0,0]', '[3,4]'); SELECT l1_distance('[0,0]', '[3,4]');
SELECT l1_distance('[0,0]', '[0,1]'); SELECT l1_distance('[0,0]', '[0,1]');

34
test/sql/tinyint.sql Normal file
View File

@@ -0,0 +1,34 @@
SELECT '127'::tinyint;
SELECT '128'::tinyint;
SELECT '-128'::tinyint;
SELECT '-129'::tinyint;
SELECT ''::tinyint;
SELECT ' 1'::tinyint;
SELECT '1 '::tinyint;
SELECT '1a'::tinyint;
SELECT '{1,2,3}'::tinyint[];
SELECT '128'::numeric::tinyint;
SELECT 'NaN'::numeric::tinyint;
SELECT l2_distance('{0,0}'::tinyint[], '{3,4}'::tinyint[]);
SELECT l2_distance('{0,0}'::tinyint[], '{0,1}'::tinyint[]);
SELECT l2_distance('{1,2}'::tinyint[], '{3}'::tinyint[]);
SELECT l2_distance('{3e38}'::tinyint[], '{-3e38}'::tinyint[]);
SELECT '{0,0}'::tinyint[] <-> '{3,4}'::tinyint[];
SELECT inner_product('{1,2}'::tinyint[], '{3,4}'::tinyint[]);
SELECT inner_product('{1,2}'::tinyint[], '{3}'::tinyint[]);
SELECT inner_product('{127}'::tinyint[], '{127}'::tinyint[]);
SELECT '{1,2}'::tinyint[] <#> '{3,4}'::tinyint[];
SELECT cosine_distance('{1,2}'::tinyint[], '{2,4}'::tinyint[]);
SELECT cosine_distance('{1,2}'::tinyint[], '{0,0}'::tinyint[]);
SELECT cosine_distance('{1,1}'::tinyint[], '{1,1}'::tinyint[]);
SELECT cosine_distance('{1,0}'::tinyint[], '{0,2}'::tinyint[]);
SELECT cosine_distance('{1,1}'::tinyint[], '{-1,-1}'::tinyint[]);
SELECT cosine_distance('{1,2}'::tinyint[], '{3}'::tinyint[]);
SELECT cosine_distance('{3e38}'::tinyint[], '{3e38}'::tinyint[]);
SELECT '{1,2}'::tinyint[] <=> '{2,4}'::tinyint[];

View File

@@ -1,107 +0,0 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
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 = 50;
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;
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 = get_new_node('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int8);");
$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", "CREATE INDEX ON tst USING hnsw (v vector_l2_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_l2_ops);");
like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_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_l2_ops, v vector_l2_ops);");
like($stderr, qr/column 2 cannot be a vector/);
done_testing();

View File

@@ -1,4 +1,4 @@
comment = 'vector data type and ivfflat and hnsw access methods' comment = 'vector data type and ivfflat and hnsw access methods'
default_version = '0.5.1' default_version = '0.5.0'
module_pathname = '$libdir/vector' module_pathname = '$libdir/vector'
relocatable = true relocatable = true