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
32 changed files with 824 additions and 1973 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
@@ -20,8 +18,10 @@ jobs:
os: ubuntu-20.04 os: ubuntu-20.04
- postgres: 12 - postgres: 12
os: ubuntu-20.04 os: ubuntu-20.04
- postgres: 11
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 }}
@@ -43,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
@@ -65,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.5.2 (unreleased) ## 0.5.1 (unreleased)
- Added support for on-disk parallel index builds for 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/half.o src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/vector.o 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/half.h 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\half.obj src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj 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\half.h 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)

117
README.md
View File

@@ -10,7 +10,7 @@ Store your vectors with the rest of your data. Supports:
Plus [ACID](https://en.wikipedia.org/wiki/ACID) compliance, point-in-time recovery, JOINs, and all of the other [great features](https://www.postgresql.org/about/) of Postgres Plus [ACID](https://en.wikipedia.org/wiki/ACID) compliance, point-in-time recovery, JOINs, and all of the other [great features](https://www.postgresql.org/about/) of Postgres
[![Build Status](https://github.com/pgvector/pgvector/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector/actions) [![Build Status](https://github.com/pgvector/pgvector/workflows/build/badge.svg?branch=master)](https://github.com/pgvector/pgvector/actions)
## Installation ## Installation
@@ -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.
@@ -269,8 +252,6 @@ Specify HNSW parameters
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64); CREATE INDEX ON items USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64);
``` ```
A higher value of `ef_construction` provides better recall at the cost of index build time / insert speed.
### Query Options ### Query Options
Specify the size of the dynamic candidate list for search (40 by default) Specify the size of the dynamic candidate list for search (40 by default)
@@ -290,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
@@ -332,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.
@@ -377,21 +360,17 @@ 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#, F#, Visual Basic | [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)
Dart | [pgvector-dart](https://github.com/pgvector/pgvector-dart) 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)
Lisp | [pgvector-lisp](https://github.com/pgvector/pgvector-lisp)
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)
OCaml | [pgvector-ocaml](https://github.com/pgvector/pgvector-ocaml)
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)
@@ -399,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
@@ -415,63 +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;
```
#### Do indexes need to fit into memory?
No, but like other index types, youll likely see better performance if they do. You can get the size of an index with:
```sql
SELECT pg_size_pretty(pg_relation_size('index_name'));
```
## Troubleshooting ## Troubleshooting
#### Why isnt a query using an index? #### Why isnt a query using an index?
@@ -485,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:
@@ -590,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
@@ -611,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 .
``` ```
@@ -676,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;

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,101 +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 TYPE half;
CREATE FUNCTION half_in(cstring, oid, integer) RETURNS half
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_out(half) RETURNS cstring
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_recv(internal, oid, integer) RETURNS half
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_send(half) RETURNS bytea
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE TYPE half (
INPUT = half_in,
OUTPUT = half_out,
RECEIVE = half_recv,
SEND = half_send,
INTERNALLENGTH = 2,
PASSEDBYVALUE,
ALIGNMENT = int2
);
CREATE FUNCTION l2_distance(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME', 'half_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION inner_product(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME', 'half_inner_product' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION cosine_distance(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME', 'half_cosine_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION l1_distance(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME', 'half_l1_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_l2_squared_distance(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_negative_inner_product(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION float4_to_half(real, integer, boolean) RETURNS half
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_to_float4(half, integer, boolean) RETURNS real
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION float8_to_half(float8, integer, boolean) RETURNS half
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_to_float8(half, integer, boolean) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION integer_to_half(integer, integer, boolean) RETURNS half
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION numeric_to_half(numeric, integer, boolean) RETURNS half
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_to_numeric(half, integer, boolean) RETURNS numeric
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE CAST (real AS half)
WITH FUNCTION float4_to_half(real, integer, boolean) AS IMPLICIT;
CREATE CAST (half AS real)
WITH FUNCTION half_to_float4(half, integer, boolean) AS IMPLICIT;
CREATE CAST (float8 AS half)
WITH FUNCTION float8_to_half(float8, integer, boolean) AS IMPLICIT;
CREATE CAST (half AS float8)
WITH FUNCTION half_to_float8(half, integer, boolean) AS IMPLICIT;
CREATE CAST (integer AS half)
WITH FUNCTION integer_to_half(integer, integer, boolean) AS IMPLICIT;
CREATE CAST (numeric AS half)
WITH FUNCTION numeric_to_half(numeric, integer, boolean) AS IMPLICIT;
CREATE CAST (half AS numeric)
WITH FUNCTION half_to_numeric(half, integer, boolean) AS IMPLICIT;
CREATE OPERATOR <-> (
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = l2_distance,
COMMUTATOR = '<->'
);
CREATE OPERATOR <#> (
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = half_negative_inner_product,
COMMUTATOR = '<#>'
);
CREATE OPERATOR <=> (
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = cosine_distance,
COMMUTATOR = '<=>'
);

View File

@@ -291,113 +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);
-- half type -- tinyint
CREATE TYPE half; CREATE TYPE tinyint;
CREATE FUNCTION half_in(cstring, oid, integer) RETURNS half CREATE FUNCTION tinyint_in(cstring, oid, integer) RETURNS tinyint
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_out(half) RETURNS cstring CREATE FUNCTION tinyint_out(tinyint) RETURNS cstring
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_recv(internal, oid, integer) RETURNS half CREATE FUNCTION tinyint_recv(internal, oid, integer) RETURNS tinyint
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_send(half) RETURNS bytea CREATE FUNCTION tinyint_send(tinyint) RETURNS bytea
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE TYPE half ( CREATE TYPE tinyint (
INPUT = half_in, INPUT = tinyint_in,
OUTPUT = half_out, OUTPUT = tinyint_out,
RECEIVE = half_recv, RECEIVE = tinyint_recv,
SEND = half_send, SEND = tinyint_send,
INTERNALLENGTH = 2, INTERNALLENGTH = 1,
PASSEDBYVALUE, PASSEDBYVALUE,
ALIGNMENT = int2 ALIGNMENT = char
); );
-- half functions CREATE FUNCTION integer_to_tinyint(integer, integer, boolean) RETURNS tinyint
CREATE FUNCTION l2_distance(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME', 'half_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION inner_product(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME', 'half_inner_product' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION cosine_distance(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME', 'half_cosine_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION l1_distance(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME', 'half_l1_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- half private functions
CREATE FUNCTION half_l2_squared_distance(half[], half[]) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_negative_inner_product(half[], half[]) RETURNS float8 CREATE FUNCTION numeric_to_tinyint(numeric, integer, boolean) RETURNS tinyint
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE; AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- half cast functions CREATE CAST (integer AS tinyint)
WITH FUNCTION integer_to_tinyint(integer, integer, boolean) AS IMPLICIT;
CREATE FUNCTION float4_to_half(real, integer, boolean) RETURNS half 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; AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_to_float4(half, integer, boolean) RETURNS real
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION float8_to_half(float8, integer, boolean) RETURNS half
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_to_float8(half, integer, boolean) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION integer_to_half(integer, integer, boolean) RETURNS half
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION numeric_to_half(numeric, integer, boolean) RETURNS half
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION half_to_numeric(half, integer, boolean) RETURNS numeric
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- half casts
CREATE CAST (real AS half)
WITH FUNCTION float4_to_half(real, integer, boolean) AS IMPLICIT;
CREATE CAST (half AS real)
WITH FUNCTION half_to_float4(half, integer, boolean) AS IMPLICIT;
CREATE CAST (float8 AS half)
WITH FUNCTION float8_to_half(float8, integer, boolean) AS IMPLICIT;
CREATE CAST (half AS float8)
WITH FUNCTION half_to_float8(half, integer, boolean) AS IMPLICIT;
CREATE CAST (integer AS half)
WITH FUNCTION integer_to_half(integer, integer, boolean) AS IMPLICIT;
CREATE CAST (numeric AS half)
WITH FUNCTION numeric_to_half(numeric, integer, boolean) AS IMPLICIT;
CREATE CAST (half AS numeric)
WITH FUNCTION half_to_numeric(half, integer, boolean) AS IMPLICIT;
-- half operators
CREATE OPERATOR <-> ( CREATE OPERATOR <-> (
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = l2_distance, LEFTARG = tinyint[], RIGHTARG = tinyint[], PROCEDURE = l2_distance,
COMMUTATOR = '<->' COMMUTATOR = '<->'
); );
CREATE OPERATOR <#> ( CREATE OPERATOR <#> (
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = half_negative_inner_product, LEFTARG = tinyint[], RIGHTARG = tinyint[], PROCEDURE = tinyint_negative_inner_product,
COMMUTATOR = '<#>' COMMUTATOR = '<#>'
); );
CREATE OPERATOR <=> ( CREATE OPERATOR <=> (
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = cosine_distance, LEFTARG = tinyint[], RIGHTARG = tinyint[], PROCEDURE = cosine_distance,
COMMUTATOR = '<=>' COMMUTATOR = '<=>'
); );

View File

@@ -1,694 +0,0 @@
#include "postgres.h"
#include <math.h>
#include "common/shortest_dec.h"
#include "fmgr.h"
#include "half.h"
#include "lib/stringinfo.h"
#include "libpq/pqformat.h"
#include "utils/array.h"
#include "utils/builtins.h"
#include "utils/float.h"
#include "utils/numeric.h"
#if PG_VERSION_NUM < 120003
static pg_noinline void
float_overflow_error(void)
{
ereport(ERROR,
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("value out of range: overflow")));
}
static pg_noinline void
float_underflow_error(void)
{
ereport(ERROR,
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("value out of range: underflow")));
}
#endif
/*
* Check if array is a vector
*/
static void
CheckArrayIsVector(ArrayType *array)
{
if (ARR_NDIM(array) > 1)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("array must be 1-D")));
if (ARR_HASNULL(array) && array_contains_nulls(array))
ereport(ERROR,
(errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
errmsg("array must not contain nulls")));
}
/*
* Check if dimensions are the same
*/
static int
CheckDims(ArrayType *a, ArrayType *b)
{
int dima;
int dimb;
CheckArrayIsVector(a);
CheckArrayIsVector(b);
dima = ARR_DIMS(a)[0];
dimb = ARR_DIMS(b)[0];
if (dima != dimb)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("different dimensions %d and %d", dima, dimb)));
return dima;
}
/*
* Return the datum representation for a half
*/
static inline Datum
HalfGetDatum(half X)
{
union
{
half value;
int16 retval;
} myunion;
myunion.value = X;
return Int16GetDatum(myunion.retval);
}
/*
* Return the half value of a datum
*/
static inline half
DatumGetHalf(Datum X)
{
union
{
int16 value;
half retval;
} myunion;
myunion.value = DatumGetInt16(X);
return myunion.retval;
}
/*
* Get a half from a message buffer
*/
static half
pq_getmsghalf(StringInfo msg)
{
union
{
half h;
uint16 i;
} swap;
swap.i = pq_getmsgint(msg, 2);
return swap.h;
}
/*
* Append a half to a StringInfo buffer
*/
static void
pq_sendhalf(StringInfo buf, half h)
{
union
{
half h;
uint16 i;
} swap;
swap.h = h;
pq_sendint16(buf, swap.i);
}
/*
* Convert a half to a float4
*/
static float
HalfToFloat4(half num)
{
#ifdef FLT16_SUPPORT
return (float) num;
#else
/* TODO Improve performance */
/* TODO Check endianness */
uint16 bin = *((uint16 *) &num);
uint32 exponent = (bin & 0x7C00) >> 10;
uint32 mantissa = bin & 0x03FF;
/* Sign */
uint32 result = (bin & 0x8000) << 16;
if (exponent == 31)
{
if (mantissa == 0)
{
/* Infinite */
result |= 0x7F800000;
}
else
{
/* NaN */
result |= 0x7FC00000;
result |= mantissa << 13;
}
}
else if (exponent == 0)
{
/* Subnormal */
if (mantissa != 0)
{
exponent = -14;
for (int i = 0; i < 10; i++)
{
mantissa <<= 1;
exponent -= 1;
if ((mantissa >> 10) % 2 == 1)
{
mantissa &= 0x03ff;
break;
}
}
result |= (exponent + 127) << 23;
result |= mantissa << 13;
}
}
else
{
/* Normal */
result |= (exponent - 15 + 127) << 23;
result |= mantissa << 13;
}
return *((float *) &result);
#endif
}
/*
* Convert a float4 to a half
*/
static half
Float4ToHalfUnchecked(float num)
{
#ifdef FLT16_SUPPORT
return (_Float16) num;
#else
/* TODO Improve performance */
/* TODO Check endianness */
uint32 bin = *((uint32 *) &num);
int exponent = (bin & 0x7F800000) >> 23;
int mantissa = bin & 0x007FFFFF;
/* Sign */
uint16 result = (bin & 0x80000000) >> 16;
if (isinf(num))
{
/* Infinite */
result |= 0x7C00;
}
else if (isnan(num))
{
/* NaN */
result |= 0x7E00;
result |= mantissa >> 13;
}
else if (exponent > 98)
{
int m;
int gr;
int s;
exponent -= 127;
s = mantissa & 0x00000FFF;
/* Subnormal */
if (exponent < -14)
{
int diff = -exponent - 14;
mantissa >>= diff;
mantissa += 1 << (23 - diff);
s |= mantissa & 0x00000FFF;
}
m = mantissa >> 13;
/* Round */
gr = (mantissa >> 12) % 4;
if (gr == 3 || (gr == 1 && s != 0))
m += 1;
if (m == 1024)
{
m = 0;
exponent += 1;
}
if (exponent > 15)
{
/* Infinite */
result |= 0x7C00;
}
else
{
if (exponent >= -14)
result |= (exponent + 15) << 10;
result |= m;
}
}
return *((half *) & result);
#endif
}
/*
* Check if half is infinite
*/
static inline bool
HalfIsInf(half num)
{
#ifdef FLT16_SUPPORT
return isinf(num);
#else
return (num << 1) == 0xF800;
#endif
}
/*
* Check if half is zero
*/
static inline bool
HalfIsZero(half num)
{
#ifdef FLT16_SUPPORT
return num == 0;
#else
return (num << 1) == 0x0000;
#endif
}
/*
* Convert a float4 to a half
*/
static half
Float4ToHalf(float num)
{
half result = Float4ToHalfUnchecked(num);
if (unlikely(HalfIsInf(result)) && !isinf(num))
float_overflow_error();
if (unlikely(HalfIsZero(result)) && num != 0.0)
float_underflow_error();
return result;
}
/*
* Convert a float8 to a half
*/
static half
Float8ToHalf(double num)
{
/* TODO Convert directly for greater accuracy */
half result = Float4ToHalfUnchecked((float) num);
if (unlikely(HalfIsInf(result)) && !isinf(num))
float_overflow_error();
if (unlikely(HalfIsZero(result)) && num != 0.0)
float_underflow_error();
return result;
}
/*
* Convert textual representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_in);
Datum
half_in(PG_FUNCTION_ARGS)
{
char *num = PG_GETARG_CSTRING(0);
char *orig_num;
float val;
char *endptr;
orig_num = num;
/* Skip leading whitespace */
while (*num != '\0' && isspace((unsigned char) *num))
num++;
/*
* Check for an empty-string input to begin with, to avoid the vagaries of
* strtof() on different platforms.
*/
if (*num == '\0')
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type %s: \"%s\"",
"half", orig_num)));
val = strtof(num, &endptr);
if (val < -HALF_MAX || val > HALF_MAX)
ereport(ERROR,
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("\"%s\" is out of range for type %s",
orig_num, "half")));
/* Skip trailing whitespace */
while (*endptr != '\0' && isspace((unsigned char) *endptr))
endptr++;
/* If there is any junk left at the end of the string, bail out */
if (*endptr != '\0')
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("invalid input syntax for type %s: \"%s\"",
"half", orig_num)));
PG_RETURN_HALF(Float4ToHalf(val));
}
/*
* Convert internal representation to textual representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_out);
Datum
half_out(PG_FUNCTION_ARGS)
{
float num = HalfToFloat4(PG_GETARG_HALF(0));
char *ascii = (char *) palloc(32);
int ndig = FLT_DIG + extra_float_digits;
if (extra_float_digits > 0)
{
float_to_shortest_decimal_buf(num, ascii);
PG_RETURN_CSTRING(ascii);
}
(void) pg_strfromd(ascii, 32, ndig, num);
PG_RETURN_CSTRING(ascii);
}
/*
* Convert external binary representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_recv);
Datum
half_recv(PG_FUNCTION_ARGS)
{
StringInfo buf = (StringInfo) PG_GETARG_POINTER(0);
PG_RETURN_HALF(pq_getmsghalf(buf));
}
/*
* Convert internal representation to the external binary representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_send);
Datum
half_send(PG_FUNCTION_ARGS)
{
half arg1 = PG_GETARG_HALF(0);
StringInfoData buf;
pq_begintypsend(&buf);
pq_sendhalf(&buf, arg1);
PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
}
/*
* Convert integer to half
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(integer_to_half);
Datum
integer_to_half(PG_FUNCTION_ARGS)
{
int32 i = PG_GETARG_INT32(0);
/* TODO Figure out correct error */
float f = (float) i;
half h = Float4ToHalf(f);
PG_RETURN_HALF(h);
}
/*
* Convert numeric to half
*
* TODO Improve error message if out of range
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(numeric_to_half);
Datum
numeric_to_half(PG_FUNCTION_ARGS)
{
Numeric num = PG_GETARG_NUMERIC(0);
float f = DatumGetFloat4(DirectFunctionCall1(numeric_float4, NumericGetDatum(num)));
half h = Float4ToHalf(f);
PG_RETURN_HALF(h);
}
/*
* Convert half to numeric
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_to_numeric);
Datum
half_to_numeric(PG_FUNCTION_ARGS)
{
half h = PG_GETARG_HALF(0);
float f = HalfToFloat4(h);
Numeric num = DatumGetNumeric(DirectFunctionCall1(float4_numeric, Float4GetDatum(f)));
PG_RETURN_NUMERIC(num);
}
/*
* Convert float4 to half
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(float4_to_half);
Datum
float4_to_half(PG_FUNCTION_ARGS)
{
float f = PG_GETARG_FLOAT4(0);
half h = Float4ToHalf(f);
PG_RETURN_HALF(h);
}
/*
* Convert half to float4
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_to_float4);
Datum
half_to_float4(PG_FUNCTION_ARGS)
{
half h = PG_GETARG_HALF(0);
float f = HalfToFloat4(h);
PG_RETURN_FLOAT4(f);
}
/*
* Convert float8 to half
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(float8_to_half);
Datum
float8_to_half(PG_FUNCTION_ARGS)
{
float8 d = PG_GETARG_FLOAT8(0);
half h = Float8ToHalf(d);
PG_RETURN_HALF(h);
}
/*
* Convert half to float8
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_to_float8);
Datum
half_to_float8(PG_FUNCTION_ARGS)
{
half h = PG_GETARG_HALF(0);
float f = HalfToFloat4(h);
PG_RETURN_FLOAT8((double) f);
}
/*
* Get the L2 distance between half arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_l2_distance);
Datum
half_l2_distance(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
half *ax = (half *) ARR_DATA_PTR(a);
half *bx = (half *) ARR_DATA_PTR(b);
float distance = 0.0;
int dim = CheckDims(a, b);
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
{
float diff = HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]);
distance += diff * diff;
}
PG_RETURN_FLOAT8(sqrt((double) distance));
}
/*
* Get the L2 squared distance between half arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_l2_squared_distance);
Datum
half_l2_squared_distance(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
half *ax = (half *) ARR_DATA_PTR(a);
half *bx = (half *) ARR_DATA_PTR(b);
float distance = 0.0;
int dim = CheckDims(a, b);
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
{
float diff = HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]);
distance += diff * diff;
}
PG_RETURN_FLOAT8((double) distance);
}
/*
* Get the inner product of two half arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_inner_product);
Datum
half_inner_product(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
half *ax = (half *) ARR_DATA_PTR(a);
half *bx = (half *) ARR_DATA_PTR(b);
float distance = 0.0;
int dim = CheckDims(a, b);
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);
PG_RETURN_FLOAT8((double) distance);
}
/*
* Get the negative inner product of two half arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_negative_inner_product);
Datum
half_negative_inner_product(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
half *ax = (half *) ARR_DATA_PTR(a);
half *bx = (half *) ARR_DATA_PTR(b);
float distance = 0.0;
int dim = CheckDims(a, b);
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);
PG_RETURN_FLOAT8((double) distance * -1);
}
/*
* Get the cosine distance between two half arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_cosine_distance);
Datum
half_cosine_distance(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
half *ax = (half *) ARR_DATA_PTR(a);
half *bx = (half *) ARR_DATA_PTR(b);
float distance = 0.0;
float norma = 0.0;
float normb = 0.0;
double similarity;
int dim = CheckDims(a, b);
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
{
float axi = HalfToFloat4(ax[i]);
float bxi = HalfToFloat4(bx[i]);
distance += axi * bxi;
norma += axi * axi;
normb += bxi * bxi;
}
/* Use sqrt(a * b) over sqrt(a) * sqrt(b) */
similarity = (double) distance / sqrt((double) norma * (double) 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);
}
/*
* Get the L1 distance between two half arrays
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(half_l1_distance);
Datum
half_l1_distance(PG_FUNCTION_ARGS)
{
ArrayType *a = PG_GETARG_ARRAYTYPE_P(0);
ArrayType *b = PG_GETARG_ARRAYTYPE_P(1);
half *ax = (half *) ARR_DATA_PTR(a);
half *bx = (half *) ARR_DATA_PTR(b);
float distance = 0.0;
int dim = CheckDims(a, b);
/* Auto-vectorized */
for (int i = 0; i < dim; i++)
distance += fabsf(HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]));
PG_RETURN_FLOAT8((double) distance);
}

View File

@@ -1,23 +0,0 @@
#ifndef HALF_H
#define HALF_H
#define __STDC_WANT_IEC_60559_TYPES_EXT__
#include <float.h>
#ifdef __FLT16_MAX__
#define FLT16_SUPPORT
#endif
#ifdef FLT16_SUPPORT
#define half _Float16
#define HALF_MAX FLT16_MAX
#else
#define half uint16
#define HALF_MAX 65504
#endif
#define PG_GETARG_HALF(n) DatumGetHalf(PG_GETARG_DATUM(n))
#define PG_RETURN_HALF(x) return HalfGetDatum(x)
#endif

View File

@@ -14,7 +14,6 @@
#endif #endif
int hnsw_ef_search; int hnsw_ef_search;
bool hnsw_enable_parallel_build;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
/* /*
@@ -40,11 +39,6 @@ HnswInit(void)
DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search", DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search",
"Valid range is 1..1000.", &hnsw_ef_search, "Valid range is 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
/* Behind a variable for now since can be slower than building in memory */
DefineCustomBoolVariable("hnsw.enable_parallel_build", "Enables or disables building indexes in parallel",
NULL, &hnsw_enable_parallel_build,
false, PGC_USERSET, 0, NULL, NULL, NULL);
} }
/* /*

View File

@@ -4,7 +4,6 @@
#include "postgres.h" #include "postgres.h"
#include "access/generic_xlog.h" #include "access/generic_xlog.h"
#include "access/parallel.h"
#include "access/reloptions.h" #include "access/reloptions.h"
#include "nodes/execnodes.h" #include "nodes/execnodes.h"
#include "port.h" /* for random() */ #include "port.h" /* for random() */
@@ -15,10 +14,6 @@
#error "Requires PostgreSQL 11+" #error "Requires PostgreSQL 11+"
#endif #endif
#if PG_VERSION_NUM < 120000
#include "access/relscan.h"
#endif
#define HNSW_MAX_DIM 2000 #define HNSW_MAX_DIM 2000
/* Support functions */ /* Support functions */
@@ -64,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))
@@ -95,7 +90,6 @@
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern bool hnsw_enable_parallel_build;
typedef struct HnswNeighborArray HnswNeighborArray; typedef struct HnswNeighborArray HnswNeighborArray;
@@ -109,7 +103,7 @@ typedef struct HnswElementData
OffsetNumber offno; OffsetNumber offno;
OffsetNumber neighborOffno; OffsetNumber neighborOffno;
BlockNumber neighborPage; BlockNumber neighborPage;
Datum value; Vector *vec;
} HnswElementData; } HnswElementData;
typedef HnswElementData * HnswElement; typedef HnswElementData * HnswElement;
@@ -118,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;
@@ -142,49 +134,6 @@ typedef struct HnswOptions
int efConstruction; /* size of dynamic candidate list */ int efConstruction; /* size of dynamic candidate list */
} HnswOptions; } HnswOptions;
typedef struct HnswSpool
{
Relation heap;
Relation index;
} HnswSpool;
typedef struct HnswShared
{
/* Immutable state */
Oid heaprelid;
Oid indexrelid;
bool isconcurrent;
int scantuplesortstates;
/* Worker progress */
ConditionVariable workersdonecv;
/* Mutex for mutable state */
slock_t mutex;
/* Mutable state */
int nparticipantsdone;
double reltuples;
double indtuples;
#if PG_VERSION_NUM < 120000
ParallelHeapScanDescData heapdesc; /* must come last */
#endif
} HnswShared;
#if PG_VERSION_NUM >= 120000
#define ParallelTableScanFromHnswShared(shared) \
(ParallelTableScanDesc) ((char *) (shared) + BUFFERALIGN(sizeof(HnswShared)))
#endif
typedef struct HnswLeader
{
ParallelContext *pcxt;
int nparticipanttuplesorts;
HnswShared *hnswshared;
Snapshot snapshot;
} HnswLeader;
typedef struct HnswBuildState typedef struct HnswBuildState
{ {
/* Info */ /* Info */
@@ -212,16 +161,12 @@ typedef struct HnswBuildState
HnswElement entryPoint; HnswElement entryPoint;
double ml; double ml;
int maxLevel; int maxLevel;
long memoryLeft; double maxInMemoryElements;
bool flushed; bool flushed;
Vector *normvec; Vector *normvec;
/* Memory */ /* Memory */
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Parallel builds */
HnswLeader *hnswleader;
HnswShared *hnswshared;
} HnswBuildState; } HnswBuildState;
typedef struct HnswMetaPageData typedef struct HnswMetaPageData
@@ -257,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;
@@ -342,7 +287,6 @@ void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation i
void HnswSetElementTuple(HnswElementTuple etup, HnswElement element); void HnswSetElementTuple(HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation); 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);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
/* Index access methods */ /* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo); IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);

View File

@@ -2,16 +2,12 @@
#include <math.h> #include <math.h>
#include "access/parallel.h"
#include "access/xact.h"
#include "catalog/index.h" #include "catalog/index.h"
#include "hnsw.h" #include "hnsw.h"
#include "miscadmin.h" #include "miscadmin.h"
#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 "tcop/tcopprot.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
@@ -39,23 +35,6 @@
#define UpdateProgress(index, val) ((void)val) #define UpdateProgress(index, val) ((void)val)
#endif #endif
#if PG_VERSION_NUM >= 140000
#include "utils/backend_status.h"
#include "utils/wait_event.h"
#endif
#if PG_VERSION_NUM >= 120000
#include "access/table.h"
#include "optimizer/optimizer.h"
#else
#include "access/heapam.h"
#include "optimizer/planner.h"
#include "pgstat.h"
#endif
#define PARALLEL_KEY_HNSW_SHARED UINT64CONST(0xA000000000000001)
#define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000002)
/* /*
* Create the metapage * Create the metapage
*/ */
@@ -126,7 +105,8 @@ CreateElementPages(HnswBuildState * buildstate)
{ {
Relation index = buildstate->index; Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
Size etupAllocSize; int dimensions = buildstate->dimensions;
Size etupSize;
Size maxSize; Size maxSize;
HnswElementTuple etup; HnswElementTuple etup;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
@@ -137,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 */
@@ -153,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(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);
/* 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);
@@ -293,14 +264,13 @@ FlushPages(HnswBuildState * buildstate)
* Insert tuple * Insert tuple
*/ */
static bool static bool
InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup, MemoryContext outerCtx) InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup)
{ {
FmgrInfo *procinfo = buildstate->procinfo; FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation; 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;
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]));
@@ -313,9 +283,7 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
} }
/* 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));
element->value = datumCopy(value, false, -1);
MemoryContextSwitchTo(oldCtx);
/* Insert element in graph */ /* Insert element in graph */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false); HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
@@ -345,21 +313,6 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
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 += VARSIZE_ANY(DatumGetPointer(e->value));
return elementSize;
}
/* /*
* Callback for table_index_build_scan * Callback for table_index_build_scan
*/ */
@@ -381,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)
{ {
@@ -396,18 +349,7 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
if (HnswInsertTuple(buildstate->index, values, isnull, tid, buildstate->heap)) if (HnswInsertTuple(buildstate->index, values, isnull, tid, buildstate->heap))
{ UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
if (buildstate->hnswshared)
{
HnswShared *hnswshared = buildstate->hnswshared;
SpinLockAcquire(&hnswshared->mutex);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++hnswshared->indtuples);
SpinLockRelease(&hnswshared->mutex);
}
else
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
}
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -418,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, element, buildstate, &dup, oldCtx); inserted = InsertTuple(index, values, element, buildstate, &dup);
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -431,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
*/ */
@@ -483,7 +436,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
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;
/* Reuse for each tuple */ /* Reuse for each tuple */
@@ -492,9 +445,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw build temporary context", "Hnsw build temporary context",
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
buildstate->hnswleader = NULL;
buildstate->hnswshared = NULL;
} }
/* /*
@@ -507,374 +457,21 @@ FreeBuildState(HnswBuildState * buildstate)
MemoryContextDelete(buildstate->tmpCtx); MemoryContextDelete(buildstate->tmpCtx);
} }
/*
* Within leader, wait for end of heap scan
*/
static double
ParallelHeapScan(HnswBuildState * buildstate)
{
HnswShared *hnswshared = buildstate->hnswleader->hnswshared;
int nparticipanttuplesorts;
double reltuples;
nparticipanttuplesorts = buildstate->hnswleader->nparticipanttuplesorts;
for (;;)
{
SpinLockAcquire(&hnswshared->mutex);
if (hnswshared->nparticipantsdone == nparticipanttuplesorts)
{
buildstate->indtuples = hnswshared->indtuples;
reltuples = hnswshared->reltuples;
SpinLockRelease(&hnswshared->mutex);
break;
}
SpinLockRelease(&hnswshared->mutex);
ConditionVariableSleep(&hnswshared->workersdonecv,
WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN);
}
ConditionVariableCancelSleep();
return reltuples;
}
/*
* Perform a worker's portion of a parallel insert
*/
static void
HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool progress)
{
HnswBuildState buildstate;
#if PG_VERSION_NUM >= 120000
TableScanDesc scan;
#else
HeapScanDesc scan;
#endif
double reltuples;
IndexInfo *indexInfo;
/* Join parallel scan */
indexInfo = BuildIndexInfo(hnswspool->index);
indexInfo->ii_Concurrent = hnswshared->isconcurrent;
InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM);
/* TODO Support in-memory builds */
buildstate.memoryLeft = 0;
buildstate.flushed = true;
buildstate.hnswshared = hnswshared;
#if PG_VERSION_NUM >= 120000
scan = table_beginscan_parallel(hnswspool->heap,
ParallelTableScanFromHnswShared(hnswshared));
reltuples = table_index_build_scan(hnswspool->heap, hnswspool->index, indexInfo,
true, progress, BuildCallback,
(void *) &buildstate, scan);
#else
scan = heap_beginscan_parallel(hnswspool->heap, &hnswshared->heapdesc);
reltuples = IndexBuildHeapScan(hnswspool->heap, hnswspool->index, indexInfo,
true, BuildCallback,
(void *) &buildstate, scan);
#endif
/* Record statistics */
SpinLockAcquire(&hnswshared->mutex);
hnswshared->nparticipantsdone++;
hnswshared->reltuples += reltuples;
SpinLockRelease(&hnswshared->mutex);
/* Log statistics */
if (progress)
ereport(DEBUG1, (errmsg("leader processed " INT64_FORMAT " tuples", (int64) reltuples)));
else
ereport(DEBUG1, (errmsg("worker processed " INT64_FORMAT " tuples", (int64) reltuples)));
/* Notify leader */
ConditionVariableSignal(&hnswshared->workersdonecv);
FreeBuildState(&buildstate);
}
/*
* Perform work within a launched parallel process
*/
void
HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
{
char *sharedquery;
HnswSpool *hnswspool;
HnswShared *hnswshared;
Relation heapRel;
Relation indexRel;
LOCKMODE heapLockmode;
LOCKMODE indexLockmode;
/* Set debug_query_string for individual workers first */
sharedquery = shm_toc_lookup(toc, PARALLEL_KEY_QUERY_TEXT, true);
debug_query_string = sharedquery;
/* Report the query string from leader */
pgstat_report_activity(STATE_RUNNING, debug_query_string);
/* Look up shared state */
hnswshared = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_SHARED, false);
/* Open relations using lock modes known to be obtained by index.c */
if (!hnswshared->isconcurrent)
{
heapLockmode = ShareLock;
indexLockmode = AccessExclusiveLock;
}
else
{
heapLockmode = ShareUpdateExclusiveLock;
indexLockmode = RowExclusiveLock;
}
/* Open relations within worker */
#if PG_VERSION_NUM >= 120000
heapRel = table_open(hnswshared->heaprelid, heapLockmode);
#else
heapRel = heap_open(hnswshared->heaprelid, heapLockmode);
#endif
indexRel = index_open(hnswshared->indexrelid, indexLockmode);
/* Initialize worker's own spool */
hnswspool = (HnswSpool *) palloc0(sizeof(HnswSpool));
hnswspool->heap = heapRel;
hnswspool->index = indexRel;
/* Perform inserts */
HnswParallelScanAndInsert(hnswspool, hnswshared, false);
/* Close relations within worker */
index_close(indexRel, indexLockmode);
#if PG_VERSION_NUM >= 120000
table_close(heapRel, heapLockmode);
#else
heap_close(heapRel, heapLockmode);
#endif
}
/*
* End parallel build
*/
static void
HnswEndParallel(HnswLeader * hnswleader)
{
/* Shutdown worker processes */
WaitForParallelWorkersToFinish(hnswleader->pcxt);
/* Free last reference to MVCC snapshot, if one was used */
if (IsMVCCSnapshot(hnswleader->snapshot))
UnregisterSnapshot(hnswleader->snapshot);
DestroyParallelContext(hnswleader->pcxt);
ExitParallelMode();
}
/*
* Return size of shared memory required for parallel index build
*/
static Size
ParallelEstimateShared(Relation heap, Snapshot snapshot)
{
#if PG_VERSION_NUM >= 120000
return add_size(BUFFERALIGN(sizeof(HnswShared)), table_parallelscan_estimate(heap, snapshot));
#else
if (!IsMVCCSnapshot(snapshot))
{
Assert(snapshot == SnapshotAny);
return sizeof(HnswShared);
}
return add_size(offsetof(HnswShared, heapdesc) +
offsetof(ParallelHeapScanDescData, phs_snapshot_data),
EstimateSnapshotSpace(snapshot));
#endif
}
/*
* Within leader, participate as a parallel worker
*/
static void
HnswLeaderParticipateAsWorker(HnswBuildState * buildstate)
{
HnswLeader *hnswleader = buildstate->hnswleader;
HnswSpool *leaderworker;
/* Allocate memory and initialize private spool */
leaderworker = (HnswSpool *) palloc0(sizeof(HnswSpool));
leaderworker->heap = buildstate->heap;
leaderworker->index = buildstate->index;
/* Perform work common to all participants */
HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, true);
}
/*
* Begin parallel build
*/
static void
HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
{
ParallelContext *pcxt;
int scantuplesortstates;
Snapshot snapshot;
Size esthnswshared;
HnswShared *hnswshared;
HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader));
bool leaderparticipates = true;
int querylen;
#ifdef DISABLE_LEADER_PARTICIPATION
leaderparticipates = false;
#endif
/* Enter parallel mode and create context */
EnterParallelMode();
Assert(request > 0);
#if PG_VERSION_NUM >= 120000
pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request);
#else
pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request, true);
#endif
scantuplesortstates = leaderparticipates ? request + 1 : request;
/* Get snapshot for table scan */
if (!isconcurrent)
snapshot = SnapshotAny;
else
snapshot = RegisterSnapshot(GetTransactionSnapshot());
/* Estimate size of workspaces */
esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot);
shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared);
shm_toc_estimate_keys(&pcxt->estimator, 1);
/* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */
if (debug_query_string)
{
querylen = strlen(debug_query_string);
shm_toc_estimate_chunk(&pcxt->estimator, querylen + 1);
shm_toc_estimate_keys(&pcxt->estimator, 1);
}
else
querylen = 0; /* keep compiler quiet */
/* Everyone's had a chance to ask for space, so now create the DSM */
InitializeParallelDSM(pcxt);
/* If no DSM segment was available, back out (do serial build) */
if (pcxt->seg == NULL)
{
if (IsMVCCSnapshot(snapshot))
UnregisterSnapshot(snapshot);
DestroyParallelContext(pcxt);
ExitParallelMode();
return;
}
/* Store shared build state, for which we reserved space */
hnswshared = (HnswShared *) shm_toc_allocate(pcxt->toc, esthnswshared);
/* Initialize immutable state */
hnswshared->heaprelid = RelationGetRelid(buildstate->heap);
hnswshared->indexrelid = RelationGetRelid(buildstate->index);
hnswshared->isconcurrent = isconcurrent;
hnswshared->scantuplesortstates = scantuplesortstates;
ConditionVariableInit(&hnswshared->workersdonecv);
SpinLockInit(&hnswshared->mutex);
/* Initialize mutable state */
hnswshared->nparticipantsdone = 0;
hnswshared->reltuples = 0;
hnswshared->indtuples = 0;
#if PG_VERSION_NUM >= 120000
table_parallelscan_initialize(buildstate->heap,
ParallelTableScanFromHnswShared(hnswshared),
snapshot);
#else
heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot);
#endif
shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_SHARED, hnswshared);
/* Store query string for workers */
if (debug_query_string)
{
char *sharedquery;
sharedquery = (char *) shm_toc_allocate(pcxt->toc, querylen + 1);
memcpy(sharedquery, debug_query_string, querylen + 1);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_QUERY_TEXT, sharedquery);
}
/* Launch workers, saving status for leader/caller */
LaunchParallelWorkers(pcxt);
hnswleader->pcxt = pcxt;
hnswleader->nparticipanttuplesorts = pcxt->nworkers_launched;
if (leaderparticipates)
hnswleader->nparticipanttuplesorts++;
hnswleader->hnswshared = hnswshared;
hnswleader->snapshot = snapshot;
/* If no workers were successfully launched, back out (do serial build) */
if (pcxt->nworkers_launched == 0)
{
HnswEndParallel(hnswleader);
return;
}
/* Log participants */
ereport(DEBUG1, (errmsg("using %d parallel workers", pcxt->nworkers_launched)));
/* Save leader state now that it's clear build will be parallel */
buildstate->hnswleader = hnswleader;
/* Join heap scan ourselves */
if (leaderparticipates)
HnswLeaderParticipateAsWorker(buildstate);
/* Wait for all launched workers */
WaitForParallelWorkersToAttach(pcxt);
}
/* /*
* Build graph * Build graph
*/ */
static void static void
BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum) BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
{ {
int parallel_workers = 0;
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD); UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD);
/* Calculate parallel workers */
if (hnsw_enable_parallel_build)
parallel_workers = plan_create_index_workers(RelationGetRelid(buildstate->heap), RelationGetRelid(buildstate->index));
/* Attempt to launch parallel worker scan when required */
if (parallel_workers > 0)
{
/* TODO Support in-memory builds */
FlushPages(buildstate);
HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
}
/* Add tuples to sort */
if (buildstate->hnswleader)
buildstate->reltuples = ParallelHeapScan(buildstate);
else
{
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo, buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL); true, true, BuildCallback, (void *) buildstate, NULL);
#else #else
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo, buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL); true, BuildCallback, (void *) buildstate, NULL);
#endif #endif
}
/* End parallel build */
if (buildstate->hnswleader)
HnswEndParallel(buildstate->hnswleader);
} }
/* /*

View File

@@ -123,6 +123,7 @@ 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;
@@ -131,7 +132,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(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;
@@ -404,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 */
@@ -416,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]))
@@ -517,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));
element->value = value; 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)

View File

@@ -4,7 +4,6 @@
#include "hnsw.h" #include "hnsw.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h"
#include "vector.h" #include "vector.h"
/* /*
@@ -140,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;
} }
} }
@@ -177,8 +175,6 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
HnswInitNeighbors(element, m); HnswInitNeighbors(element, m);
element->value = PointerGetDatum(NULL);
return element; return element;
} }
@@ -190,8 +186,7 @@ HnswFreeElement(HnswElement element)
{ {
HnswFreeNeighbors(element); HnswFreeNeighbors(element);
list_free_deep(element->heaptids); list_free_deep(element->heaptids);
if (DatumGetPointer(element->value)) pfree(element->vec);
pfree(DatumGetPointer(element->value));
pfree(element); pfree(element);
} }
@@ -218,7 +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;
return element; return element;
} }
@@ -328,7 +323,7 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
else else
ItemPointerSetInvalid(&etup->heaptids[i]); ItemPointerSetInvalid(&etup->heaptids[i]);
} }
memcpy(&etup->data, DatumGetPointer(element->value), VARSIZE_ANY(DatumGetPointer(element->value))); memcpy(&etup->vec, element->vec, VECTOR_SIZE(element->vec->dim));
} }
/* /*
@@ -450,7 +445,10 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
} }
if (loadVec) if (loadVec)
element->value = datumCopy(PointerGetDatum(&etup->data), false, -1); {
element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
}
} }
/* /*
@@ -477,7 +475,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
/* Calculate distance */ /* Calculate distance */
if (distance != NULL) if (distance != NULL)
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->data))); *distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->vec)));
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -488,7 +486,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
static float static float
GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation) GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
{ {
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, hc->element->value)); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, PointerGetDatum(hc->element->vec)));
} }
/* /*
@@ -694,34 +692,6 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
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
*/ */
@@ -751,7 +721,7 @@ HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid co
} }
} }
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, a->value, b->value)); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(a->vec), PointerGetDatum(b->vec)));
} }
/* /*
@@ -778,77 +748,33 @@ CheckElementCloser(HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid c
* Algorithm 4 from paper * Algorithm 4 from paper
*/ */
static List * static List *
SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, 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, procinfo, collation);
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, procinfo, collation);
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, procinfo, collation);
if (e->closer)
added = lappend(added, e);
}
}
}
else if (e == newCandidate)
{
e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
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);
@@ -878,7 +804,7 @@ HnswFindDuplicate(HnswElement e)
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 */
@@ -902,6 +828,28 @@ 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
*/ */
@@ -931,13 +879,13 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
/* Load elements on insert */ /* Load elements on insert */
if (index != NULL) if (index != NULL)
{ {
Datum q = hc->element->value; Datum q = PointerGetDatum(hc->element->vec);
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, index, procinfo, collation, true); HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
else else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation); hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
@@ -955,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, procinfo, collation, hc->element, &hc2, &pruned, true); SelectNeighbors(c, m, lc, procinfo, collation, &pruned);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -1018,7 +967,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
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);
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
/* No neighbors if no entry point */ /* No neighbors if no entry point */
@@ -1059,12 +1008,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
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, procinfo, collation, 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++)
@@ -478,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;
@@ -507,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 */
@@ -530,7 +528,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Overwrite element */ /* Overwrite element */
etup->deleted = 1; etup->deleted = 1;
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data)); MemSet(&etup->vec.x, 0, etup->vec.dim * sizeof(float));
/* Overwrite neighbors */ /* Overwrite neighbors */
for (int i = 0; i < ntup->count; i++) for (int i = 0; i < ntup->count; i++)

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

@@ -1,15 +1,15 @@
CREATE TABLE t (val vector(3), val2 half[]); CREATE TABLE t (val vector(3));
INSERT INTO t (val, val2) VALUES ('[0,0,0]', '{0,0,0}'), ('[1,2,3]', '{1,2,3}'), ('[1,1,1]', '{1,1,1}'), (NULL, NULL); INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE TABLE t2 (val vector(3), val2 half[]); CREATE TABLE t2 (val vector(3));
\copy t TO 'results/data.bin' WITH (FORMAT binary) \copy t TO 'results/data.bin' WITH (FORMAT binary)
\copy t2 FROM 'results/data.bin' WITH (FORMAT binary) \copy t2 FROM 'results/data.bin' WITH (FORMAT binary)
SELECT * FROM t2 ORDER BY val; SELECT * FROM t2 ORDER BY val;
val | val2 val
---------+--------- ---------
[0,0,0] | {0,0,0} [0,0,0]
[1,1,1] | {1,1,1} [1,1,1]
[1,2,3] | {1,2,3} [1,2,3]
|
(4 rows) (4 rows)
DROP TABLE t; DROP TABLE t;

View File

@@ -138,21 +138,21 @@ SELECT cosine_distance('[3e38]'::vector, '[3e38]');
NaN NaN
(1 row) (1 row)
SELECT l1_distance('[0,0]'::vector, '[3,4]'); SELECT l1_distance('[0,0]', '[3,4]');
l1_distance l1_distance
------------- -------------
7 7
(1 row) (1 row)
SELECT l1_distance('[0,0]'::vector, '[0,1]'); SELECT l1_distance('[0,0]', '[0,1]');
l1_distance l1_distance
------------- -------------
1 1
(1 row) (1 row)
SELECT l1_distance('[1,2]'::vector, '[3]'); SELECT l1_distance('[1,2]', '[3]');
ERROR: different vector dimensions 2 and 1 ERROR: different vector dimensions 2 and 1
SELECT l1_distance('[3e38]'::vector, '[-3e38]'); SELECT l1_distance('[3e38]', '[-3e38]');
l1_distance l1_distance
------------- -------------
Infinity Infinity

View File

@@ -1,226 +0,0 @@
SELECT '1.5'::half;
half
------
1.5
(1 row)
SELECT '65504'::half;
half
-------
65504
(1 row)
SELECT '65505'::half;
ERROR: "65505" is out of range for type half
LINE 1: SELECT '65505'::half;
^
SELECT '-65504'::half;
half
--------
-65504
(1 row)
SELECT '-65505'::half;
ERROR: "-65505" is out of range for type half
LINE 1: SELECT '-65505'::half;
^
SELECT ''::half;
ERROR: invalid input syntax for type half: ""
LINE 1: SELECT ''::half;
^
SELECT ' '::half;
ERROR: invalid input syntax for type half: " "
LINE 1: SELECT ' '::half;
^
SELECT '-'::half;
ERROR: invalid input syntax for type half: "-"
LINE 1: SELECT '-'::half;
^
SELECT ' 1.5'::half;
half
------
1.5
(1 row)
SELECT '1.5 '::half;
half
------
1.5
(1 row)
SELECT '1.5a'::half;
ERROR: invalid input syntax for type half: "1.5a"
LINE 1: SELECT '1.5a'::half;
^
SELECT '{1,2,3}'::half[];
half
---------
{1,2,3}
(1 row)
SELECT '{1,2,3}'::half[]::real[];
float4
---------
{1,2,3}
(1 row)
SELECT '65505'::integer::half;
half
-------
65504
(1 row)
SELECT 'NaN'::real::half;
half
------
NaN
(1 row)
SELECT 'Infinity'::real::half;
half
----------
Infinity
(1 row)
SELECT '1e-38'::real::half;
ERROR: value out of range: underflow
SELECT '1.5'::half::real;
float4
--------
1.5
(1 row)
SELECT '1.5'::real::half;
half
------
1.5
(1 row)
SELECT '1.5'::half::double precision;
float8
--------
1.5
(1 row)
SELECT '1.5'::double precision::half;
half
------
1.5
(1 row)
SELECT '1.5'::half::numeric;
numeric
---------
1.5
(1 row)
SELECT '1.5'::numeric::half;
half
------
1.5
(1 row)
SELECT l2_distance('{0,0}'::half[], '{3,4}'::half[]);
l2_distance
-------------
5
(1 row)
SELECT l2_distance('{0,0}'::half[], '{0,1}'::half[]);
l2_distance
-------------
1
(1 row)
SELECT l2_distance('{1,2}'::half[], '{3}'::half[]);
ERROR: different dimensions 2 and 1
SELECT '{0,0}'::half[] <-> '{3,4}'::half[];
?column?
----------
5
(1 row)
SELECT inner_product('{1,2}'::half[], '{3,4}'::half[]);
inner_product
---------------
11
(1 row)
SELECT inner_product('{1,2}'::half[], '{3}'::half[]);
ERROR: different dimensions 2 and 1
SELECT inner_product('{65504}'::half[], '{65504}'::half[]);
inner_product
---------------
4290774016
(1 row)
SELECT '{1,2}'::half[] <#> '{3,4}'::half[];
?column?
----------
-11
(1 row)
SELECT cosine_distance('{1,2}'::half[], '{2,4}'::half[]);
cosine_distance
-----------------
0
(1 row)
SELECT cosine_distance('{1,2}'::half[], '{0,0}'::half[]);
cosine_distance
-----------------
NaN
(1 row)
SELECT cosine_distance('{1,1}'::half[], '{1,1}'::half[]);
cosine_distance
-----------------
0
(1 row)
SELECT cosine_distance('{1,0}'::half[], '{0,2}'::half[]);
cosine_distance
-----------------
1
(1 row)
SELECT cosine_distance('{1,1}'::half[], '{-1,-1}'::half[]);
cosine_distance
-----------------
2
(1 row)
SELECT cosine_distance('{1,2}'::half[], '{3}'::half[]);
ERROR: different dimensions 2 and 1
SELECT cosine_distance('{1,1}'::half[], '{1.1,1.1}'::half[]);
cosine_distance
-----------------
0
(1 row)
SELECT cosine_distance('{1,1}'::half[], '{-1.1,-1.1}'::half[]);
cosine_distance
-----------------
2
(1 row)
SELECT '{1,2}'::half[] <=> '{2,4}'::half[];
?column?
----------
0
(1 row)
SELECT l1_distance('{0,0}'::half[], '{3,4}');
l1_distance
-------------
7
(1 row)
SELECT l1_distance('{0,0}'::half[], '{0,1}');
l1_distance
-------------
1
(1 row)
SELECT l1_distance('{1,2}'::half[], '{3}');
ERROR: different dimensions 2 and 1

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

@@ -1,7 +1,7 @@
CREATE TABLE t (val vector(3), val2 half[]); CREATE TABLE t (val vector(3));
INSERT INTO t (val, val2) VALUES ('[0,0,0]', '{0,0,0}'), ('[1,2,3]', '{1,2,3}'), ('[1,1,1]', '{1,1,1}'), (NULL, NULL); INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE TABLE t2 (val vector(3), val2 half[]); CREATE TABLE t2 (val vector(3));
\copy t TO 'results/data.bin' WITH (FORMAT binary) \copy t TO 'results/data.bin' WITH (FORMAT binary)
\copy t2 FROM 'results/data.bin' WITH (FORMAT binary) \copy t2 FROM 'results/data.bin' WITH (FORMAT binary)

View File

@@ -32,10 +32,10 @@ SELECT cosine_distance('[1,1]'::vector, '[1.1,1.1]');
SELECT cosine_distance('[1,1]'::vector, '[-1.1,-1.1]'); SELECT cosine_distance('[1,1]'::vector, '[-1.1,-1.1]');
SELECT cosine_distance('[3e38]'::vector, '[3e38]'); SELECT cosine_distance('[3e38]'::vector, '[3e38]');
SELECT l1_distance('[0,0]'::vector, '[3,4]'); SELECT l1_distance('[0,0]', '[3,4]');
SELECT l1_distance('[0,0]'::vector, '[0,1]'); SELECT l1_distance('[0,0]', '[0,1]');
SELECT l1_distance('[1,2]'::vector, '[3]'); SELECT l1_distance('[1,2]', '[3]');
SELECT l1_distance('[3e38]'::vector, '[-3e38]'); SELECT l1_distance('[3e38]', '[-3e38]');
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v; SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v; SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v;

View File

@@ -1,53 +0,0 @@
SELECT '1.5'::half;
SELECT '65504'::half;
SELECT '65505'::half;
SELECT '-65504'::half;
SELECT '-65505'::half;
SELECT ''::half;
SELECT ' '::half;
SELECT '-'::half;
SELECT ' 1.5'::half;
SELECT '1.5 '::half;
SELECT '1.5a'::half;
SELECT '{1,2,3}'::half[];
SELECT '{1,2,3}'::half[]::real[];
SELECT '65505'::integer::half;
SELECT 'NaN'::real::half;
SELECT 'Infinity'::real::half;
SELECT '1e-38'::real::half;
SELECT '1.5'::half::real;
SELECT '1.5'::real::half;
SELECT '1.5'::half::double precision;
SELECT '1.5'::double precision::half;
SELECT '1.5'::half::numeric;
SELECT '1.5'::numeric::half;
SELECT l2_distance('{0,0}'::half[], '{3,4}'::half[]);
SELECT l2_distance('{0,0}'::half[], '{0,1}'::half[]);
SELECT l2_distance('{1,2}'::half[], '{3}'::half[]);
SELECT '{0,0}'::half[] <-> '{3,4}'::half[];
SELECT inner_product('{1,2}'::half[], '{3,4}'::half[]);
SELECT inner_product('{1,2}'::half[], '{3}'::half[]);
SELECT inner_product('{65504}'::half[], '{65504}'::half[]);
SELECT '{1,2}'::half[] <#> '{3,4}'::half[];
SELECT cosine_distance('{1,2}'::half[], '{2,4}'::half[]);
SELECT cosine_distance('{1,2}'::half[], '{0,0}'::half[]);
SELECT cosine_distance('{1,1}'::half[], '{1,1}'::half[]);
SELECT cosine_distance('{1,0}'::half[], '{0,2}'::half[]);
SELECT cosine_distance('{1,1}'::half[], '{-1,-1}'::half[]);
SELECT cosine_distance('{1,2}'::half[], '{3}'::half[]);
SELECT cosine_distance('{1,1}'::half[], '{1.1,1.1}'::half[]);
SELECT cosine_distance('{1,1}'::half[], '{-1.1,-1.1}'::half[]);
SELECT '{1,2}'::half[] <=> '{2,4}'::half[];
SELECT l1_distance('{0,0}'::half[], '{3,4}');
SELECT l1_distance('{0,0}'::half[], '{0,1}');
SELECT l1_distance('{1,2}'::half[], '{3}');

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

@@ -83,32 +83,11 @@ for my $i (0 .. $#operators)
push(@expected, $res); push(@expected, $res);
} }
# Build index serially # Add index
$node->safe_psql("postgres", qq( $node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v $opclass);");
SET max_parallel_maintenance_workers = 0;
CREATE INDEX idx ON tst USING hnsw (v $opclass);
));
# Test approximate results
my $min = $operator eq "<#>" ? 0.80 : 0.99; my $min = $operator eq "<#>" ? 0.80 : 0.99;
test_recall($min, $operator); test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
# Build index in parallel
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET client_min_messages = DEBUG;
SET min_parallel_table_scan_size = 1;
SET hnsw.enable_parallel_build = on;
CREATE INDEX idx ON tst USING hnsw (v $opclass);
));
is($ret, 0, $stderr);
like($stderr, qr/using \d+ parallel workers/);
# Test approximate results
test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
} }
done_testing(); 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