mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-22 12:07:34 +08:00
Compare commits
23 Commits
svector
...
half-index
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
47852e8d15 | ||
|
|
422667f6c6 | ||
|
|
4d6739a7af | ||
|
|
ff744214d0 | ||
|
|
7ca9298163 | ||
|
|
ff3bffd9a8 | ||
|
|
014753eb9c | ||
|
|
6763661d3d | ||
|
|
d287921d15 | ||
|
|
5b12ae8225 | ||
|
|
4549e8aeb1 | ||
|
|
3263b350f5 | ||
|
|
dfee5d4045 | ||
|
|
69a2ce0d43 | ||
|
|
c5e8c46b80 | ||
|
|
94f7304ccd | ||
|
|
d078db3d25 | ||
|
|
fbb904ae2f | ||
|
|
3cf6f62900 | ||
|
|
2a69e22ca4 | ||
|
|
84e073888c | ||
|
|
81a62d55d1 | ||
|
|
3f3463bde5 |
2
.github/workflows/build.yml
vendored
2
.github/workflows/build.yml
vendored
@@ -20,8 +20,6 @@ jobs:
|
||||
os: ubuntu-20.04
|
||||
- postgres: 12
|
||||
os: ubuntu-20.04
|
||||
- postgres: 11
|
||||
os: ubuntu-20.04
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: ankane/setup-postgres@v1
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
## 0.6.0 (unreleased)
|
||||
## 0.5.2 (unreleased)
|
||||
|
||||
- Added support for sparse vectors
|
||||
- Added support for on-disk parallel index builds for HNSW
|
||||
|
||||
## 0.5.1 (2023-10-10)
|
||||
|
||||
|
||||
4
Makefile
4
Makefile
@@ -3,8 +3,8 @@ EXTVERSION = 0.5.1
|
||||
|
||||
MODULE_big = vector
|
||||
DATA = $(wildcard sql/*--*.sql)
|
||||
OBJS = src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/svector.o src/vector.o
|
||||
HEADERS = src/svector.h src/vector.h
|
||||
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
|
||||
HEADERS = src/half.h src/vector.h
|
||||
|
||||
TESTS = $(wildcard test/sql/*.sql)
|
||||
REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS))
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
EXTENSION = vector
|
||||
EXTVERSION = 0.5.1
|
||||
|
||||
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\svector.obj src\vector.obj
|
||||
HEADERS = src\svector.h src\vector.h
|
||||
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
|
||||
HEADERS = src\half.h src\vector.h
|
||||
|
||||
REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged
|
||||
REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION)
|
||||
|
||||
36
README.md
36
README.md
@@ -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
|
||||
|
||||
[](https://github.com/pgvector/pgvector/actions)
|
||||
[](https://github.com/pgvector/pgvector/actions)
|
||||
|
||||
## Installation
|
||||
|
||||
@@ -269,6 +269,8 @@ Specify HNSW parameters
|
||||
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
|
||||
|
||||
Specify the size of the dynamic candidate list for search (40 by default)
|
||||
@@ -369,26 +371,6 @@ To speed up queries with an IVFFlat index, increase the number of inverted lists
|
||||
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 1000);
|
||||
```
|
||||
|
||||
## Sparse Vectors
|
||||
|
||||
Create a sparse vector column with 10 dimensions
|
||||
|
||||
```sql
|
||||
CREATE TABLE items (id bigserial PRIMARY KEY, embedding svector(10));
|
||||
```
|
||||
|
||||
Insert vectors
|
||||
|
||||
```sql
|
||||
INSERT INTO items (embedding) VALUES ('(0,1),(1,2),(2,3)|10|'), ('(0,4),(1,5),(4,6)|10|');
|
||||
```
|
||||
|
||||
Get the nearest neighbors by L2 distance
|
||||
|
||||
```sql
|
||||
SELECT * FROM items ORDER BY embedding <-> '(0,3),(1,1),(2,2)|10|' LIMIT 5;
|
||||
```
|
||||
|
||||
## Languages
|
||||
|
||||
Use pgvector from any language with a Postgres client. You can even generate and store vectors in one language and query them in another.
|
||||
@@ -397,7 +379,7 @@ Language | Libraries / Examples
|
||||
--- | ---
|
||||
C | [pgvector-c](https://github.com/pgvector/pgvector-c)
|
||||
C++ | [pgvector-cpp](https://github.com/pgvector/pgvector-cpp)
|
||||
C# | [pgvector-dotnet](https://github.com/pgvector/pgvector-dotnet)
|
||||
C#, F#, Visual Basic | [pgvector-dotnet](https://github.com/pgvector/pgvector-dotnet)
|
||||
Crystal | [pgvector-crystal](https://github.com/pgvector/pgvector-crystal)
|
||||
Dart | [pgvector-dart](https://github.com/pgvector/pgvector-dart)
|
||||
Elixir | [pgvector-elixir](https://github.com/pgvector/pgvector-elixir)
|
||||
@@ -406,8 +388,10 @@ Haskell | [pgvector-haskell](https://github.com/pgvector/pgvector-haskell)
|
||||
Java, Kotlin, Groovy, 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)
|
||||
Lisp | [pgvector-lisp](https://github.com/pgvector/pgvector-lisp)
|
||||
Lua | [pgvector-lua](https://github.com/pgvector/pgvector-lua)
|
||||
Nim | [pgvector-nim](https://github.com/pgvector/pgvector-nim)
|
||||
OCaml | [pgvector-ocaml](https://github.com/pgvector/pgvector-ocaml)
|
||||
Perl | [pgvector-perl](https://github.com/pgvector/pgvector-perl)
|
||||
PHP | [pgvector-php](https://github.com/pgvector/pgvector-php)
|
||||
Python | [pgvector-python](https://github.com/pgvector/pgvector-python)
|
||||
@@ -480,6 +464,14 @@ and query with:
|
||||
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, you’ll 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
|
||||
|
||||
#### Why isn’t a query using an index?
|
||||
|
||||
@@ -1,79 +1,82 @@
|
||||
-- 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 svector;
|
||||
CREATE TYPE half;
|
||||
|
||||
CREATE FUNCTION svector_in(cstring, oid, integer) RETURNS svector
|
||||
CREATE FUNCTION half_in(cstring, oid, integer) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_out(svector) RETURNS cstring
|
||||
CREATE FUNCTION half_out(half) RETURNS cstring
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_typmod_in(cstring[]) RETURNS integer
|
||||
CREATE FUNCTION half_recv(internal, oid, integer) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_recv(internal, oid, integer) RETURNS svector
|
||||
CREATE FUNCTION half_send(half) RETURNS bytea
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_send(svector) RETURNS bytea
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE TYPE svector (
|
||||
INPUT = svector_in,
|
||||
OUTPUT = svector_out,
|
||||
TYPMOD_IN = svector_typmod_in,
|
||||
RECEIVE = svector_recv,
|
||||
SEND = svector_send,
|
||||
STORAGE = external
|
||||
CREATE TYPE half (
|
||||
INPUT = half_in,
|
||||
OUTPUT = half_out,
|
||||
RECEIVE = half_recv,
|
||||
SEND = half_send,
|
||||
INTERNALLENGTH = 2,
|
||||
PASSEDBYVALUE,
|
||||
ALIGNMENT = int2
|
||||
);
|
||||
|
||||
CREATE FUNCTION l2_distance(svector, svector) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'svector_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
CREATE FUNCTION l2_distance(half[], half[]) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'half_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION inner_product(svector, svector) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'svector_inner_product' 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(svector, svector) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'svector_cosine_distance' 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 jaccard_distance(svector, svector) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'svector_jaccard_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 svector_l2_squared_distance(svector, svector) RETURNS float8
|
||||
CREATE FUNCTION half_l2_squared_distance(half[], half[]) RETURNS float8
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_negative_inner_product(svector, svector) RETURNS float8
|
||||
CREATE FUNCTION half_negative_inner_product(half[], half[]) RETURNS float8
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector(svector, integer, boolean) RETURNS svector
|
||||
CREATE FUNCTION float4_to_half(real, integer, boolean) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION vector_to_svector(vector, integer, boolean) RETURNS svector
|
||||
CREATE FUNCTION integer_to_half(integer, integer, boolean) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_to_vector(svector, integer, boolean) RETURNS vector
|
||||
CREATE FUNCTION numeric_to_half(numeric, integer, boolean) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE CAST (svector AS svector)
|
||||
WITH FUNCTION svector(svector, integer, boolean) AS IMPLICIT;
|
||||
CREATE CAST (real AS half)
|
||||
WITH FUNCTION float4_to_half(real, integer, boolean) AS IMPLICIT;
|
||||
|
||||
CREATE CAST (svector AS vector)
|
||||
WITH FUNCTION svector_to_vector(svector, integer, boolean) AS IMPLICIT;
|
||||
CREATE CAST (integer AS half)
|
||||
WITH FUNCTION integer_to_half(integer, integer, boolean) AS IMPLICIT;
|
||||
|
||||
CREATE CAST (vector AS svector)
|
||||
WITH FUNCTION vector_to_svector(vector, integer, boolean) AS IMPLICIT;
|
||||
CREATE CAST (numeric AS half)
|
||||
WITH FUNCTION numeric_to_half(numeric, integer, boolean) AS IMPLICIT;
|
||||
|
||||
CREATE OPERATOR <-> (
|
||||
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = l2_distance,
|
||||
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = l2_distance,
|
||||
COMMUTATOR = '<->'
|
||||
);
|
||||
|
||||
CREATE OPERATOR <#> (
|
||||
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = svector_negative_inner_product,
|
||||
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = half_negative_inner_product,
|
||||
COMMUTATOR = '<#>'
|
||||
);
|
||||
|
||||
CREATE OPERATOR <=> (
|
||||
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = cosine_distance,
|
||||
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = cosine_distance,
|
||||
COMMUTATOR = '<=>'
|
||||
);
|
||||
|
||||
CREATE OPERATOR CLASS half_l2_ops
|
||||
FOR TYPE half[] USING hnsw AS
|
||||
OPERATOR 1 <-> (half[], half[]) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 half_l2_squared_distance(half[], half[]);
|
||||
|
||||
@@ -291,91 +291,96 @@ CREATE OPERATOR CLASS vector_cosine_ops
|
||||
FUNCTION 1 vector_negative_inner_product(vector, vector),
|
||||
FUNCTION 2 vector_norm(vector);
|
||||
|
||||
--- svector type
|
||||
-- half type
|
||||
|
||||
CREATE TYPE svector;
|
||||
CREATE TYPE half;
|
||||
|
||||
CREATE FUNCTION svector_in(cstring, oid, integer) RETURNS svector
|
||||
CREATE FUNCTION half_in(cstring, oid, integer) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_out(svector) RETURNS cstring
|
||||
CREATE FUNCTION half_out(half) RETURNS cstring
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_typmod_in(cstring[]) RETURNS integer
|
||||
CREATE FUNCTION half_recv(internal, oid, integer) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_recv(internal, oid, integer) RETURNS svector
|
||||
CREATE FUNCTION half_send(half) RETURNS bytea
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_send(svector) RETURNS bytea
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE TYPE svector (
|
||||
INPUT = svector_in,
|
||||
OUTPUT = svector_out,
|
||||
TYPMOD_IN = svector_typmod_in,
|
||||
RECEIVE = svector_recv,
|
||||
SEND = svector_send,
|
||||
STORAGE = external
|
||||
CREATE TYPE half (
|
||||
INPUT = half_in,
|
||||
OUTPUT = half_out,
|
||||
RECEIVE = half_recv,
|
||||
SEND = half_send,
|
||||
INTERNALLENGTH = 2,
|
||||
PASSEDBYVALUE,
|
||||
ALIGNMENT = int2
|
||||
);
|
||||
|
||||
-- svector functions
|
||||
-- half functions
|
||||
|
||||
CREATE FUNCTION l2_distance(svector, svector) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'svector_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
CREATE FUNCTION l2_distance(half[], half[]) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'half_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION inner_product(svector, svector) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'svector_inner_product' 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(svector, svector) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'svector_cosine_distance' 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 jaccard_distance(svector, svector) RETURNS float8
|
||||
AS 'MODULE_PATHNAME', 'svector_jaccard_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;
|
||||
|
||||
-- svector private functions
|
||||
-- half private functions
|
||||
|
||||
CREATE FUNCTION svector_l2_squared_distance(svector, svector) RETURNS float8
|
||||
CREATE FUNCTION half_l2_squared_distance(half[], half[]) RETURNS float8
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_negative_inner_product(svector, svector) RETURNS float8
|
||||
CREATE FUNCTION half_negative_inner_product(half[], half[]) RETURNS float8
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
-- svector cast functions
|
||||
-- half cast functions
|
||||
|
||||
CREATE FUNCTION svector(svector, integer, boolean) RETURNS svector
|
||||
CREATE FUNCTION float4_to_half(real, integer, boolean) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION vector_to_svector(vector, integer, boolean) RETURNS svector
|
||||
CREATE FUNCTION integer_to_half(integer, integer, boolean) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION svector_to_vector(svector, integer, boolean) RETURNS vector
|
||||
CREATE FUNCTION numeric_to_half(numeric, integer, boolean) RETURNS half
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
-- svector casts
|
||||
-- half casts
|
||||
|
||||
CREATE CAST (svector AS svector)
|
||||
WITH FUNCTION svector(svector, integer, boolean) AS IMPLICIT;
|
||||
CREATE CAST (real AS half)
|
||||
WITH FUNCTION float4_to_half(real, integer, boolean) AS IMPLICIT;
|
||||
|
||||
CREATE CAST (svector AS vector)
|
||||
WITH FUNCTION svector_to_vector(svector, integer, boolean) AS IMPLICIT;
|
||||
CREATE CAST (integer AS half)
|
||||
WITH FUNCTION integer_to_half(integer, integer, boolean) AS IMPLICIT;
|
||||
|
||||
CREATE CAST (vector AS svector)
|
||||
WITH FUNCTION vector_to_svector(vector, integer, boolean) AS IMPLICIT;
|
||||
CREATE CAST (numeric AS half)
|
||||
WITH FUNCTION numeric_to_half(numeric, integer, boolean) AS IMPLICIT;
|
||||
|
||||
-- svector operators
|
||||
-- half operators
|
||||
|
||||
CREATE OPERATOR <-> (
|
||||
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = l2_distance,
|
||||
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = l2_distance,
|
||||
COMMUTATOR = '<->'
|
||||
);
|
||||
|
||||
CREATE OPERATOR <#> (
|
||||
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = svector_negative_inner_product,
|
||||
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = half_negative_inner_product,
|
||||
COMMUTATOR = '<#>'
|
||||
);
|
||||
|
||||
CREATE OPERATOR <=> (
|
||||
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = cosine_distance,
|
||||
LEFTARG = half[], RIGHTARG = half[], PROCEDURE = cosine_distance,
|
||||
COMMUTATOR = '<=>'
|
||||
);
|
||||
|
||||
-- half opclasses
|
||||
|
||||
CREATE OPERATOR CLASS half_l2_ops
|
||||
FOR TYPE half[] USING hnsw AS
|
||||
OPERATOR 1 <-> (half[], half[]) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 half_l2_squared_distance(half[], half[]);
|
||||
|
||||
599
src/half.c
Normal file
599
src/half.c
Normal file
@@ -0,0 +1,599 @@
|
||||
#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;
|
||||
}
|
||||
|
||||
/*
|
||||
* Append a half to a StringInfo buffer
|
||||
*/
|
||||
static half
|
||||
pq_getmsghalf(StringInfo msg)
|
||||
{
|
||||
union
|
||||
{
|
||||
half h;
|
||||
uint16 i;
|
||||
} swap;
|
||||
|
||||
/* TODO Likely use float4 for clients */
|
||||
swap.i = pq_getmsgint(msg, 2);
|
||||
return swap.h;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get a half from a message buffer
|
||||
*/
|
||||
static void
|
||||
pq_sendhalf(StringInfo buf, half h)
|
||||
{
|
||||
union
|
||||
{
|
||||
half h;
|
||||
uint16 i;
|
||||
} swap;
|
||||
|
||||
/* TODO Likely use float4 for clients */
|
||||
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
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert a float4 to a half
|
||||
*/
|
||||
static half
|
||||
Float4ToHalf(float num)
|
||||
{
|
||||
half result = Float4ToHalfUnchecked(num);
|
||||
|
||||
/* TODO Perform checks without HalfToFloat4 */
|
||||
if (unlikely(isinf(HalfToFloat4(result))) && !isinf(num))
|
||||
float_overflow_error();
|
||||
if (unlikely(HalfToFloat4(result) == 0.0f) && 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
|
||||
*/
|
||||
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 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);
|
||||
}
|
||||
|
||||
/*
|
||||
* 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);
|
||||
}
|
||||
28
src/half.h
Normal file
28
src/half.h
Normal file
@@ -0,0 +1,28 @@
|
||||
#ifndef HALF_H
|
||||
#define HALF_H
|
||||
|
||||
#define __STDC_WANT_IEC_60559_TYPES_EXT__
|
||||
|
||||
#include <float.h>
|
||||
|
||||
/* _Float16 and __fp16 are not supported on x86_64 with GCC 11 */
|
||||
#if defined(__is_identifier)
|
||||
#if __is_identifier(_Float16)
|
||||
#define FLT16_SUPPORT
|
||||
#endif
|
||||
#elif defined(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
|
||||
@@ -14,6 +14,7 @@
|
||||
#endif
|
||||
|
||||
int hnsw_ef_search;
|
||||
bool hnsw_enable_parallel_build;
|
||||
static relopt_kind hnsw_relopt_kind;
|
||||
|
||||
/*
|
||||
@@ -39,6 +40,11 @@ HnswInit(void)
|
||||
DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for 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);
|
||||
|
||||
/* 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);
|
||||
}
|
||||
|
||||
/*
|
||||
|
||||
62
src/hnsw.h
62
src/hnsw.h
@@ -4,6 +4,7 @@
|
||||
#include "postgres.h"
|
||||
|
||||
#include "access/generic_xlog.h"
|
||||
#include "access/parallel.h"
|
||||
#include "access/reloptions.h"
|
||||
#include "nodes/execnodes.h"
|
||||
#include "port.h" /* for random() */
|
||||
@@ -14,6 +15,10 @@
|
||||
#error "Requires PostgreSQL 11+"
|
||||
#endif
|
||||
|
||||
#if PG_VERSION_NUM < 120000
|
||||
#include "access/relscan.h"
|
||||
#endif
|
||||
|
||||
#define HNSW_MAX_DIM 2000
|
||||
|
||||
/* Support functions */
|
||||
@@ -59,7 +64,7 @@
|
||||
|
||||
#define HNSW_MAX_SIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - sizeof(ItemIdData))
|
||||
|
||||
#define HNSW_ELEMENT_TUPLE_SIZE(_dim) MAXALIGN(offsetof(HnswElementTupleData, vec) + VECTOR_SIZE(_dim))
|
||||
#define HNSW_ELEMENT_TUPLE_SIZE(size) MAXALIGN(offsetof(HnswElementTupleData, data) + (size))
|
||||
#define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData))
|
||||
|
||||
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
|
||||
@@ -90,6 +95,7 @@
|
||||
|
||||
/* Variables */
|
||||
extern int hnsw_ef_search;
|
||||
extern bool hnsw_enable_parallel_build;
|
||||
|
||||
typedef struct HnswNeighborArray HnswNeighborArray;
|
||||
|
||||
@@ -103,7 +109,7 @@ typedef struct HnswElementData
|
||||
OffsetNumber offno;
|
||||
OffsetNumber neighborOffno;
|
||||
BlockNumber neighborPage;
|
||||
Vector *vec;
|
||||
Datum value;
|
||||
} HnswElementData;
|
||||
|
||||
typedef HnswElementData * HnswElement;
|
||||
@@ -136,6 +142,49 @@ typedef struct HnswOptions
|
||||
int efConstruction; /* size of dynamic candidate list */
|
||||
} 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
|
||||
{
|
||||
/* Info */
|
||||
@@ -163,12 +212,16 @@ typedef struct HnswBuildState
|
||||
HnswElement entryPoint;
|
||||
double ml;
|
||||
int maxLevel;
|
||||
double maxInMemoryElements;
|
||||
long memoryLeft;
|
||||
bool flushed;
|
||||
Vector *normvec;
|
||||
|
||||
/* Memory */
|
||||
MemoryContext tmpCtx;
|
||||
|
||||
/* Parallel builds */
|
||||
HnswLeader *hnswleader;
|
||||
HnswShared *hnswshared;
|
||||
} HnswBuildState;
|
||||
|
||||
typedef struct HnswMetaPageData
|
||||
@@ -204,7 +257,7 @@ typedef struct HnswElementTupleData
|
||||
ItemPointerData heaptids[HNSW_HEAPTIDS];
|
||||
ItemPointerData neighbortid;
|
||||
uint16 unused2;
|
||||
Vector vec;
|
||||
Vector data;
|
||||
} HnswElementTupleData;
|
||||
|
||||
typedef HnswElementTupleData * HnswElementTuple;
|
||||
@@ -289,6 +342,7 @@ void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation i
|
||||
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 HnswLoadNeighbors(HnswElement element, Relation index, int m);
|
||||
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
|
||||
|
||||
/* Index access methods */
|
||||
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);
|
||||
|
||||
471
src/hnswbuild.c
471
src/hnswbuild.c
@@ -2,12 +2,16 @@
|
||||
|
||||
#include <math.h>
|
||||
|
||||
#include "access/parallel.h"
|
||||
#include "access/xact.h"
|
||||
#include "catalog/index.h"
|
||||
#include "hnsw.h"
|
||||
#include "miscadmin.h"
|
||||
#include "lib/pairingheap.h"
|
||||
#include "nodes/pg_list.h"
|
||||
#include "storage/bufmgr.h"
|
||||
#include "tcop/tcopprot.h"
|
||||
#include "utils/datum.h"
|
||||
#include "utils/memutils.h"
|
||||
|
||||
#if PG_VERSION_NUM >= 140000
|
||||
@@ -35,6 +39,23 @@
|
||||
#define UpdateProgress(index, val) ((void)val)
|
||||
#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
|
||||
*/
|
||||
@@ -105,8 +126,7 @@ CreateElementPages(HnswBuildState * buildstate)
|
||||
{
|
||||
Relation index = buildstate->index;
|
||||
ForkNumber forkNum = buildstate->forkNum;
|
||||
int dimensions = buildstate->dimensions;
|
||||
Size etupSize;
|
||||
Size etupAllocSize;
|
||||
Size maxSize;
|
||||
HnswElementTuple etup;
|
||||
HnswNeighborTuple ntup;
|
||||
@@ -117,11 +137,11 @@ CreateElementPages(HnswBuildState * buildstate)
|
||||
ListCell *lc;
|
||||
|
||||
/* Calculate sizes */
|
||||
etupAllocSize = BLCKSZ;
|
||||
maxSize = HNSW_MAX_SIZE;
|
||||
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
|
||||
|
||||
/* Allocate once */
|
||||
etup = palloc0(etupSize);
|
||||
etup = palloc0(etupAllocSize);
|
||||
ntup = palloc0(BLCKSZ);
|
||||
|
||||
/* Prepare first page */
|
||||
@@ -133,15 +153,24 @@ CreateElementPages(HnswBuildState * buildstate)
|
||||
foreach(lc, buildstate->elements)
|
||||
{
|
||||
HnswElement element = lfirst(lc);
|
||||
Size etupSize;
|
||||
Size ntupSize;
|
||||
Size combinedSize;
|
||||
|
||||
HnswSetElementTuple(etup, element);
|
||||
/* Zero memory for each element */
|
||||
MemSet(etup, 0, etupAllocSize);
|
||||
|
||||
/* Calculate sizes */
|
||||
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(DatumGetPointer(element->value)));
|
||||
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
|
||||
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 */
|
||||
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
|
||||
HnswBuildAppendPage(index, &buf, &page, &state, forkNum);
|
||||
@@ -264,13 +293,14 @@ FlushPages(HnswBuildState * buildstate)
|
||||
* Insert tuple
|
||||
*/
|
||||
static bool
|
||||
InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup)
|
||||
InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup, MemoryContext outerCtx)
|
||||
{
|
||||
FmgrInfo *procinfo = buildstate->procinfo;
|
||||
Oid collation = buildstate->collation;
|
||||
HnswElement entryPoint = buildstate->entryPoint;
|
||||
int efConstruction = buildstate->efConstruction;
|
||||
int m = buildstate->m;
|
||||
MemoryContext oldCtx;
|
||||
|
||||
/* Detoast once for all calls */
|
||||
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
|
||||
@@ -283,7 +313,9 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
|
||||
}
|
||||
|
||||
/* Copy value to element so accessible outside of memory context */
|
||||
memcpy(element->vec, DatumGetVector(value), VECTOR_SIZE(buildstate->dimensions));
|
||||
oldCtx = MemoryContextSwitchTo(outerCtx);
|
||||
element->value = datumCopy(value, false, -1);
|
||||
MemoryContextSwitchTo(oldCtx);
|
||||
|
||||
/* Insert element in graph */
|
||||
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
|
||||
@@ -313,6 +345,21 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
|
||||
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
|
||||
*/
|
||||
@@ -334,7 +381,7 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
|
||||
if (isnull[0])
|
||||
return;
|
||||
|
||||
if (buildstate->indtuples >= buildstate->maxInMemoryElements)
|
||||
if (buildstate->memoryLeft <= 0)
|
||||
{
|
||||
if (!buildstate->flushed)
|
||||
{
|
||||
@@ -349,7 +396,18 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
|
||||
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
|
||||
|
||||
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 */
|
||||
MemoryContextSwitchTo(oldCtx);
|
||||
@@ -360,13 +418,12 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
|
||||
|
||||
/* Allocate necessary memory outside of memory context */
|
||||
element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel);
|
||||
element->vec = palloc(VECTOR_SIZE(buildstate->dimensions));
|
||||
|
||||
/* Use memory context since detoast can allocate */
|
||||
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
|
||||
|
||||
/* Insert tuple */
|
||||
inserted = InsertTuple(index, values, element, buildstate, &dup);
|
||||
inserted = InsertTuple(index, values, element, buildstate, &dup, oldCtx);
|
||||
|
||||
/* Reset memory context */
|
||||
MemoryContextSwitchTo(oldCtx);
|
||||
@@ -374,31 +431,21 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
|
||||
|
||||
/* Add outside memory context */
|
||||
if (dup != NULL)
|
||||
{
|
||||
HnswAddHeapTid(dup, tid);
|
||||
buildstate->memoryLeft -= sizeof(ItemPointerData);
|
||||
}
|
||||
|
||||
/* Add to buildstate or free */
|
||||
if (inserted)
|
||||
{
|
||||
buildstate->elements = lappend(buildstate->elements, element);
|
||||
buildstate->memoryLeft -= HnswElementMemory(element, buildstate->m);
|
||||
}
|
||||
else
|
||||
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
|
||||
*/
|
||||
@@ -415,8 +462,8 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
|
||||
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
|
||||
|
||||
/* Require column to have dimensions to be indexed */
|
||||
if (buildstate->dimensions < 0)
|
||||
elog(ERROR, "column does not have dimensions");
|
||||
// if (buildstate->dimensions < 0)
|
||||
// elog(ERROR, "column does not have dimensions");
|
||||
|
||||
if (buildstate->dimensions > HNSW_MAX_DIM)
|
||||
elog(ERROR, "column cannot have more than %d dimensions for hnsw index", HNSW_MAX_DIM);
|
||||
@@ -436,7 +483,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
|
||||
buildstate->entryPoint = NULL;
|
||||
buildstate->ml = HnswGetMl(buildstate->m);
|
||||
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
|
||||
buildstate->maxInMemoryElements = HnswGetMaxInMemoryElements(buildstate->m, buildstate->ml, buildstate->dimensions);
|
||||
buildstate->memoryLeft = maintenance_work_mem * 1024L;
|
||||
buildstate->flushed = false;
|
||||
|
||||
/* Reuse for each tuple */
|
||||
@@ -445,6 +492,9 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
|
||||
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
|
||||
"Hnsw build temporary context",
|
||||
ALLOCSET_DEFAULT_SIZES);
|
||||
|
||||
buildstate->hnswleader = NULL;
|
||||
buildstate->hnswshared = NULL;
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -457,21 +507,374 @@ FreeBuildState(HnswBuildState * buildstate)
|
||||
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
|
||||
*/
|
||||
static void
|
||||
BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
|
||||
{
|
||||
int parallel_workers = 0;
|
||||
|
||||
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
|
||||
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
|
||||
true, true, BuildCallback, (void *) buildstate, NULL);
|
||||
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
|
||||
true, true, BuildCallback, (void *) buildstate, NULL);
|
||||
#else
|
||||
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
|
||||
true, BuildCallback, (void *) buildstate, NULL);
|
||||
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
|
||||
true, BuildCallback, (void *) buildstate, NULL);
|
||||
#endif
|
||||
}
|
||||
|
||||
/* End parallel build */
|
||||
if (buildstate->hnswleader)
|
||||
HnswEndParallel(buildstate->hnswleader);
|
||||
}
|
||||
|
||||
/*
|
||||
|
||||
@@ -123,7 +123,6 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
|
||||
Size minCombinedSize;
|
||||
HnswElementTuple etup;
|
||||
BlockNumber currentPage = insertPage;
|
||||
int dimensions = e->vec->dim;
|
||||
HnswNeighborTuple ntup;
|
||||
Buffer nbuf;
|
||||
Page npage;
|
||||
@@ -132,7 +131,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
|
||||
BlockNumber newInsertPage = InvalidBlockNumber;
|
||||
|
||||
/* Calculate sizes */
|
||||
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
|
||||
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(DatumGetPointer(e->value)));
|
||||
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
|
||||
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
|
||||
maxSize = HNSW_MAX_SIZE;
|
||||
@@ -405,8 +404,9 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
|
||||
Buffer buf;
|
||||
Page page;
|
||||
GenericXLogState *state;
|
||||
Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(dup->vec->dim);
|
||||
ItemId itemid;
|
||||
HnswElementTuple etup;
|
||||
Size etupSize;
|
||||
int i;
|
||||
|
||||
/* Read page */
|
||||
@@ -416,7 +416,9 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
|
||||
page = GenericXLogRegisterBuffer(state, buf, 0);
|
||||
|
||||
/* Find space */
|
||||
etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, dup->offno));
|
||||
itemid = PageGetItemId(page, dup->offno);
|
||||
etup = (HnswElementTuple) PageGetItem(page, itemid);
|
||||
etupSize = ItemIdGetLength(itemid);
|
||||
for (i = 0; i < HNSW_HEAPTIDS; i++)
|
||||
{
|
||||
if (!ItemPointerIsValid(&etup->heaptids[i]))
|
||||
@@ -515,7 +517,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
|
||||
|
||||
/* Create an element */
|
||||
element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m));
|
||||
element->vec = DatumGetVector(value);
|
||||
element->value = value;
|
||||
|
||||
/* Prevent concurrent inserts when likely updating entry point */
|
||||
if (entryPoint == NULL || element->level > entryPoint->level)
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
#include "hnsw.h"
|
||||
#include "storage/bufmgr.h"
|
||||
#include "utils/datum.h"
|
||||
#include "vector.h"
|
||||
|
||||
/*
|
||||
@@ -176,6 +177,8 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
|
||||
|
||||
HnswInitNeighbors(element, m);
|
||||
|
||||
element->value = PointerGetDatum(NULL);
|
||||
|
||||
return element;
|
||||
}
|
||||
|
||||
@@ -187,7 +190,8 @@ HnswFreeElement(HnswElement element)
|
||||
{
|
||||
HnswFreeNeighbors(element);
|
||||
list_free_deep(element->heaptids);
|
||||
pfree(element->vec);
|
||||
if (DatumGetPointer(element->value))
|
||||
pfree(DatumGetPointer(element->value));
|
||||
pfree(element);
|
||||
}
|
||||
|
||||
@@ -214,7 +218,7 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
|
||||
element->blkno = blkno;
|
||||
element->offno = offno;
|
||||
element->neighbors = NULL;
|
||||
element->vec = NULL;
|
||||
element->value = PointerGetDatum(NULL);
|
||||
return element;
|
||||
}
|
||||
|
||||
@@ -324,7 +328,7 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
|
||||
else
|
||||
ItemPointerSetInvalid(&etup->heaptids[i]);
|
||||
}
|
||||
memcpy(&etup->vec, element->vec, VECTOR_SIZE(element->vec->dim));
|
||||
memcpy(&etup->data, DatumGetPointer(element->value), VARSIZE_ANY(DatumGetPointer(element->value)));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -446,10 +450,7 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
|
||||
}
|
||||
|
||||
if (loadVec)
|
||||
{
|
||||
element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
|
||||
memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
|
||||
}
|
||||
element->value = datumCopy(PointerGetDatum(&etup->data), false, -1);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -476,7 +477,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
|
||||
|
||||
/* Calculate distance */
|
||||
if (distance != NULL)
|
||||
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->vec)));
|
||||
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->data)));
|
||||
|
||||
UnlockReleaseBuffer(buf);
|
||||
}
|
||||
@@ -487,7 +488,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
|
||||
static float
|
||||
GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
|
||||
{
|
||||
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, PointerGetDatum(hc->element->vec)));
|
||||
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, hc->element->value));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -750,7 +751,7 @@ HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid co
|
||||
}
|
||||
}
|
||||
|
||||
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(a->vec), PointerGetDatum(b->vec)));
|
||||
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, a->value, b->value));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -877,7 +878,7 @@ HnswFindDuplicate(HnswElement e)
|
||||
HnswCandidate *neighbor = &neighbors->items[i];
|
||||
|
||||
/* Exit early since ordered by distance */
|
||||
if (vector_cmp_internal(e->vec, neighbor->element->vec) != 0)
|
||||
if (!datumIsEqual(e->value, neighbor->element->value, false, -1))
|
||||
break;
|
||||
|
||||
/* Check for space */
|
||||
@@ -930,13 +931,13 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
|
||||
/* Load elements on insert */
|
||||
if (index != NULL)
|
||||
{
|
||||
Datum q = PointerGetDatum(hc->element->vec);
|
||||
Datum q = hc->element->value;
|
||||
|
||||
for (int i = 0; i < currentNeighbors->length; i++)
|
||||
{
|
||||
HnswCandidate *hc3 = ¤tNeighbors->items[i];
|
||||
|
||||
if (hc3->element->vec == NULL)
|
||||
if (DatumGetPointer(hc3->element->value) == NULL)
|
||||
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
|
||||
else
|
||||
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
|
||||
@@ -1017,7 +1018,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
|
||||
List *w;
|
||||
int level = element->level;
|
||||
int entryLevel;
|
||||
Datum q = PointerGetDatum(element->vec);
|
||||
Datum q = element->value;
|
||||
HnswElement skipElement = existing ? element : NULL;
|
||||
|
||||
/* No neighbors if no entry point */
|
||||
|
||||
@@ -62,7 +62,8 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
|
||||
/* Iterate over nodes */
|
||||
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
|
||||
{
|
||||
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
|
||||
ItemId itemid = PageGetItemId(page, offno);
|
||||
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, itemid);
|
||||
int idx = 0;
|
||||
bool itemUpdated = false;
|
||||
|
||||
@@ -93,7 +94,7 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
|
||||
|
||||
if (itemUpdated)
|
||||
{
|
||||
Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
|
||||
Size etupSize = ItemIdGetLength(itemid);
|
||||
|
||||
/* Mark rest as invalid */
|
||||
for (int i = idx; i < HNSW_HEAPTIDS; i++)
|
||||
@@ -477,7 +478,8 @@ MarkDeleted(HnswVacuumState * vacuumstate)
|
||||
/* Update element and neighbors together */
|
||||
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
|
||||
{
|
||||
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
|
||||
ItemId itemid = PageGetItemId(page, offno);
|
||||
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, itemid);
|
||||
HnswNeighborTuple ntup;
|
||||
Size etupSize;
|
||||
Size ntupSize;
|
||||
@@ -505,7 +507,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
|
||||
continue;
|
||||
|
||||
/* Calculate sizes */
|
||||
etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
|
||||
etupSize = ItemIdGetLength(itemid);
|
||||
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m);
|
||||
|
||||
/* Get neighbor page */
|
||||
@@ -528,7 +530,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
|
||||
|
||||
/* Overwrite element */
|
||||
etup->deleted = 1;
|
||||
MemSet(&etup->vec.x, 0, etup->vec.dim * sizeof(float));
|
||||
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
|
||||
|
||||
/* Overwrite neighbors */
|
||||
for (int i = 0; i < ntup->count; i++)
|
||||
|
||||
705
src/svector.c
705
src/svector.c
@@ -1,705 +0,0 @@
|
||||
#include "postgres.h"
|
||||
|
||||
#include <math.h>
|
||||
|
||||
#include "fmgr.h"
|
||||
#include "libpq/pqformat.h"
|
||||
#include "svector.h"
|
||||
#include "utils/array.h"
|
||||
#include "vector.h"
|
||||
|
||||
#if PG_VERSION_NUM >= 120000
|
||||
#include "common/shortest_dec.h"
|
||||
#include "utils/float.h"
|
||||
#else
|
||||
#include <float.h>
|
||||
#include "utils/builtins.h"
|
||||
#endif
|
||||
|
||||
/*
|
||||
* Ensure same dimensions
|
||||
*/
|
||||
static inline void
|
||||
CheckDims(SVector * a, SVector * b)
|
||||
{
|
||||
if (a->dim != b->dim)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("different svector dimensions %d and %d", a->dim, b->dim)));
|
||||
}
|
||||
|
||||
/*
|
||||
* Ensure expected dimensions
|
||||
*/
|
||||
static inline void
|
||||
CheckExpectedDim(int32 typmod, int dim)
|
||||
{
|
||||
if (typmod != -1 && typmod != dim)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("expected %d dimensions, not %d", typmod, dim)));
|
||||
}
|
||||
|
||||
/*
|
||||
* Ensure valid dimensions
|
||||
*/
|
||||
static inline void
|
||||
CheckDim(int dim)
|
||||
{
|
||||
if (dim < 1)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("svector must have at least 1 dimension")));
|
||||
|
||||
if (dim > SVECTOR_MAX_DIM)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
|
||||
errmsg("svector cannot have more than %d dimensions", SVECTOR_MAX_DIM)));
|
||||
}
|
||||
|
||||
/*
|
||||
* Ensure valid nnz
|
||||
*/
|
||||
static inline void
|
||||
CheckNnz(int nnz, int dim)
|
||||
{
|
||||
if (nnz < 0)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("svector must have at least one element")));
|
||||
|
||||
if (nnz > dim)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
|
||||
errmsg("svector cannot have more elements than dimensions")));
|
||||
}
|
||||
|
||||
/*
|
||||
* Ensure valid index
|
||||
*/
|
||||
static inline void
|
||||
CheckIndex(int32 *indices, int i, int dim)
|
||||
{
|
||||
int32 index = indices[i];
|
||||
|
||||
if (index < 0)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("index must not be negative")));
|
||||
|
||||
if (index >= dim)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("index must be less than dimensions")));
|
||||
|
||||
if (i > 0)
|
||||
{
|
||||
if (index < indices[i - 1])
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("indexes must be in ascending order")));
|
||||
|
||||
if (index == indices[i - 1])
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("indexes must not contain duplicates")));
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
* Ensure finite element
|
||||
*/
|
||||
static inline void
|
||||
CheckElement(float value)
|
||||
{
|
||||
if (isnan(value))
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("NaN not allowed in svector")));
|
||||
|
||||
if (isinf(value))
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("infinite value not allowed in svector")));
|
||||
}
|
||||
|
||||
/*
|
||||
* Allocate and initialize a new sparse vector
|
||||
*/
|
||||
SVector *
|
||||
InitSVector(int dim, int nnz)
|
||||
{
|
||||
SVector *result;
|
||||
int size;
|
||||
|
||||
size = SVECTOR_SIZE(nnz);
|
||||
result = (SVector *) palloc0(size);
|
||||
SET_VARSIZE(result, size);
|
||||
result->dim = dim;
|
||||
result->nnz = nnz;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert textual representation to internal representation
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_in);
|
||||
Datum
|
||||
svector_in(PG_FUNCTION_ARGS)
|
||||
{
|
||||
char *str = PG_GETARG_CSTRING(0);
|
||||
int32 typmod = PG_GETARG_INT32(2);
|
||||
int dim;
|
||||
char *pt;
|
||||
SVector *result;
|
||||
float *rvalues;
|
||||
char *lit = pstrdup(str);
|
||||
int n;
|
||||
int32 *indices;
|
||||
float *values;
|
||||
int index;
|
||||
float value;
|
||||
int maxNnz;
|
||||
int nnz = 0;
|
||||
|
||||
/* TODO Improve code and checks after deciding on format */
|
||||
|
||||
maxNnz = 1;
|
||||
pt = str;
|
||||
while (*pt != '\0')
|
||||
{
|
||||
if (*pt == ',')
|
||||
maxNnz++;
|
||||
|
||||
pt++;
|
||||
}
|
||||
maxNnz /= 2;
|
||||
|
||||
indices = palloc(maxNnz * sizeof(int32));
|
||||
values = palloc(maxNnz * sizeof(float));
|
||||
|
||||
while (sscanf(str, "(%d,%f)%n", &index, &value, &n) == 2)
|
||||
{
|
||||
/* TODO Better error */
|
||||
if (nnz == maxNnz)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("ran out of buffer: \"%s\"", lit)));
|
||||
|
||||
/* TODO Decide whether to store zero values */
|
||||
indices[nnz] = index;
|
||||
values[nnz] = value;
|
||||
nnz++;
|
||||
|
||||
str += n;
|
||||
|
||||
if (*str == ',')
|
||||
str++;
|
||||
else if (*str == '|')
|
||||
break;
|
||||
else
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed svector literal: \"%s\"", lit)));
|
||||
}
|
||||
|
||||
if (sscanf(str, "|%d|%n", &dim, &n) != 1)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed svector literal: \"%s\"", lit)));
|
||||
|
||||
str += n;
|
||||
|
||||
if (*str != '\0')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed svector literal: \"%s\"", lit),
|
||||
errdetail("Junk after closing pipe.")));
|
||||
|
||||
pfree(lit);
|
||||
|
||||
CheckDim(dim);
|
||||
CheckExpectedDim(typmod, dim);
|
||||
|
||||
result = InitSVector(dim, nnz);
|
||||
rvalues = SVECTOR_VALUES(result);
|
||||
for (int i = 0; i < nnz; i++)
|
||||
{
|
||||
result->indices[i] = indices[i];
|
||||
rvalues[i] = values[i];
|
||||
|
||||
CheckIndex(result->indices, i, dim);
|
||||
CheckElement(rvalues[i]);
|
||||
}
|
||||
|
||||
PG_RETURN_POINTER(result);
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert internal representation to textual representation
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_out);
|
||||
Datum
|
||||
svector_out(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *svector = PG_GETARG_SVECTOR_P(0);
|
||||
float *values = SVECTOR_VALUES(svector);
|
||||
char *buf;
|
||||
char *ptr;
|
||||
int n;
|
||||
|
||||
/* TODO Improve code after deciding on format */
|
||||
|
||||
#if PG_VERSION_NUM < 120000
|
||||
int ndig = FLT_DIG + extra_float_digits;
|
||||
|
||||
if (ndig < 1)
|
||||
ndig = 1;
|
||||
|
||||
#define FLOAT_SHORTEST_DECIMAL_LEN (ndig + 10)
|
||||
#endif
|
||||
|
||||
/* TODO Move */
|
||||
#define APPEND_CHAR(ptr, ch) (*(ptr)++ = (ch))
|
||||
|
||||
/* TODO Improve */
|
||||
buf = (char *) palloc((FLOAT_SHORTEST_DECIMAL_LEN + 20) * svector->nnz + 20);
|
||||
ptr = buf;
|
||||
|
||||
for (int i = 0; i < svector->nnz; i++)
|
||||
{
|
||||
if (i > 0)
|
||||
APPEND_CHAR(ptr, ',');
|
||||
|
||||
n = sprintf(ptr, "(%d,", svector->indices[i]);
|
||||
ptr += n;
|
||||
|
||||
#if PG_VERSION_NUM >= 120000
|
||||
n = float_to_shortest_decimal_bufn(values[i], ptr);
|
||||
#else
|
||||
n = sprintf(ptr, "%.*g", ndig, values[i]);
|
||||
#endif
|
||||
ptr += n;
|
||||
|
||||
APPEND_CHAR(ptr, ')');
|
||||
}
|
||||
|
||||
n = sprintf(ptr, "|%d|", svector->dim);
|
||||
ptr += n;
|
||||
|
||||
APPEND_CHAR(ptr, '\0');
|
||||
|
||||
PG_FREE_IF_COPY(svector, 0);
|
||||
PG_RETURN_CSTRING(buf);
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert type modifier
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_typmod_in);
|
||||
Datum
|
||||
svector_typmod_in(PG_FUNCTION_ARGS)
|
||||
{
|
||||
ArrayType *ta = PG_GETARG_ARRAYTYPE_P(0);
|
||||
int32 *tl;
|
||||
int n;
|
||||
|
||||
tl = ArrayGetIntegerTypmods(ta, &n);
|
||||
|
||||
if (n != 1)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
|
||||
errmsg("invalid type modifier")));
|
||||
|
||||
if (*tl < 1)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
|
||||
errmsg("dimensions for type svector must be at least 1")));
|
||||
|
||||
if (*tl > SVECTOR_MAX_DIM)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
|
||||
errmsg("dimensions for type svector cannot exceed %d", SVECTOR_MAX_DIM)));
|
||||
|
||||
PG_RETURN_INT32(*tl);
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert external binary representation to internal representation
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_recv);
|
||||
Datum
|
||||
svector_recv(PG_FUNCTION_ARGS)
|
||||
{
|
||||
StringInfo buf = (StringInfo) PG_GETARG_POINTER(0);
|
||||
int32 typmod = PG_GETARG_INT32(2);
|
||||
SVector *result;
|
||||
int32 dim;
|
||||
int32 nnz;
|
||||
int32 unused;
|
||||
float *values;
|
||||
|
||||
dim = pq_getmsgint(buf, sizeof(int32));
|
||||
nnz = pq_getmsgint(buf, sizeof(int32));
|
||||
unused = pq_getmsgint(buf, sizeof(int32));
|
||||
|
||||
CheckDim(dim);
|
||||
CheckNnz(nnz, dim);
|
||||
CheckExpectedDim(typmod, dim);
|
||||
|
||||
if (unused != 0)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("expected unused to be 0, not %d", unused)));
|
||||
|
||||
result = InitSVector(dim, nnz);
|
||||
values = SVECTOR_VALUES(result);
|
||||
|
||||
for (int i = 0; i < nnz; i++)
|
||||
{
|
||||
result->indices[i] = pq_getmsgint(buf, sizeof(int32));
|
||||
CheckIndex(result->indices, i, dim);
|
||||
}
|
||||
|
||||
for (int i = 0; i < nnz; i++)
|
||||
{
|
||||
values[i] = pq_getmsgfloat4(buf);
|
||||
CheckElement(values[i]);
|
||||
}
|
||||
|
||||
PG_RETURN_POINTER(result);
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert internal representation to the external binary representation
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_send);
|
||||
Datum
|
||||
svector_send(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *svec = PG_GETARG_SVECTOR_P(0);
|
||||
float *values = SVECTOR_VALUES(svec);
|
||||
StringInfoData buf;
|
||||
|
||||
pq_begintypsend(&buf);
|
||||
pq_sendint(&buf, svec->dim, sizeof(int32));
|
||||
pq_sendint(&buf, svec->nnz, sizeof(int32));
|
||||
pq_sendint(&buf, svec->unused, sizeof(int32));
|
||||
for (int i = 0; i < svec->nnz; i++)
|
||||
pq_sendint(&buf, svec->indices[i], sizeof(int32));
|
||||
for (int i = 0; i < svec->nnz; i++)
|
||||
pq_sendfloat4(&buf, values[i]);
|
||||
|
||||
PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert sparse vector to sparse vector
|
||||
* This is needed to check the type modifier
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector);
|
||||
Datum
|
||||
svector(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *svec = PG_GETARG_SVECTOR_P(0);
|
||||
int32 typmod = PG_GETARG_INT32(1);
|
||||
|
||||
CheckExpectedDim(typmod, svec->dim);
|
||||
|
||||
PG_RETURN_POINTER(svec);
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert dense vector to sparse vector
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_to_svector);
|
||||
Datum
|
||||
vector_to_svector(PG_FUNCTION_ARGS)
|
||||
{
|
||||
Vector *vec = PG_GETARG_VECTOR_P(0);
|
||||
int32 typmod = PG_GETARG_INT32(1);
|
||||
SVector *result;
|
||||
int dim = vec->dim;
|
||||
int nnz = 0;
|
||||
float *values;
|
||||
int j = 0;
|
||||
|
||||
CheckDim(dim);
|
||||
CheckExpectedDim(typmod, dim);
|
||||
|
||||
for (int i = 0; i < dim; i++)
|
||||
{
|
||||
if (vec->x[i] != 0)
|
||||
nnz++;
|
||||
}
|
||||
|
||||
result = InitSVector(dim, nnz);
|
||||
values = SVECTOR_VALUES(result);
|
||||
for (int i = 0; i < dim; i++)
|
||||
{
|
||||
if (vec->x[i] != 0)
|
||||
{
|
||||
/* Safety check */
|
||||
if (j == nnz)
|
||||
elog(ERROR, "safety check failed");
|
||||
|
||||
result->indices[j] = i;
|
||||
values[j] = vec->x[i];
|
||||
j++;
|
||||
}
|
||||
}
|
||||
|
||||
PG_RETURN_POINTER(result);
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the L2 squared distance between sparse vectors
|
||||
*/
|
||||
static double
|
||||
l2_distance_squared_internal(SVector * a, SVector * b)
|
||||
{
|
||||
float *ax = SVECTOR_VALUES(a);
|
||||
float *bx = SVECTOR_VALUES(b);
|
||||
double distance = 0.0;
|
||||
int bpos = 0;
|
||||
|
||||
for (int i = 0; i < a->nnz; i++)
|
||||
{
|
||||
int ai = a->indices[i];
|
||||
int bi = -1;
|
||||
|
||||
for (int j = bpos; j < b->nnz; j++)
|
||||
{
|
||||
bi = b->indices[j];
|
||||
|
||||
if (ai == bi)
|
||||
{
|
||||
double diff = ax[i] - bx[j];
|
||||
|
||||
distance += diff * diff;
|
||||
}
|
||||
else if (ai > bi)
|
||||
distance += bx[j] * bx[j];
|
||||
|
||||
/* Update start for next iteration */
|
||||
if (ai >= bi)
|
||||
bpos = j + 1;
|
||||
|
||||
/* Found or passed it */
|
||||
if (bi >= ai)
|
||||
break;
|
||||
}
|
||||
|
||||
if (ai != bi)
|
||||
distance += ax[i] * ax[i];
|
||||
}
|
||||
|
||||
for (int j = bpos; j < b->nnz; j++)
|
||||
distance += bx[j] * bx[j];
|
||||
|
||||
return distance;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the L2 distance between sparse vectors
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_l2_distance);
|
||||
Datum
|
||||
svector_l2_distance(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *a = PG_GETARG_SVECTOR_P(0);
|
||||
SVector *b = PG_GETARG_SVECTOR_P(1);
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(sqrt(l2_distance_squared_internal(a, b)));
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the L2 squared distance between sparse vectors
|
||||
* This saves a sqrt calculation
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_l2_squared_distance);
|
||||
Datum
|
||||
svector_l2_squared_distance(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *a = PG_GETARG_SVECTOR_P(0);
|
||||
SVector *b = PG_GETARG_SVECTOR_P(1);
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(l2_distance_squared_internal(a, b));
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the inner product of two sparse vectors
|
||||
*/
|
||||
static double
|
||||
inner_product_internal(SVector * a, SVector * b)
|
||||
{
|
||||
float *ax = SVECTOR_VALUES(a);
|
||||
float *bx = SVECTOR_VALUES(b);
|
||||
double distance = 0.0;
|
||||
int bpos = 0;
|
||||
|
||||
for (int i = 0; i < a->nnz; i++)
|
||||
{
|
||||
int ai = a->indices[i];
|
||||
|
||||
for (int j = bpos; j < b->nnz; j++)
|
||||
{
|
||||
int bi = b->indices[j];
|
||||
|
||||
/* Only update when the same index */
|
||||
if (ai == bi)
|
||||
distance += ax[i] * bx[j];
|
||||
|
||||
/* Update start for next iteration */
|
||||
if (ai >= bi)
|
||||
bpos = j + 1;
|
||||
|
||||
/* Found or passed it */
|
||||
if (bi >= ai)
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return distance;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the inner product of two sparse vectors
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_inner_product);
|
||||
Datum
|
||||
svector_inner_product(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *a = PG_GETARG_SVECTOR_P(0);
|
||||
SVector *b = PG_GETARG_SVECTOR_P(1);
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(inner_product_internal(a, b));
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the negative inner product of two sparse vectors
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_negative_inner_product);
|
||||
Datum
|
||||
svector_negative_inner_product(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *a = PG_GETARG_SVECTOR_P(0);
|
||||
SVector *b = PG_GETARG_SVECTOR_P(1);
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(-inner_product_internal(a, b));
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the cosine distance between two sparse vectors
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_cosine_distance);
|
||||
Datum
|
||||
svector_cosine_distance(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *a = PG_GETARG_SVECTOR_P(0);
|
||||
SVector *b = PG_GETARG_SVECTOR_P(1);
|
||||
float *ax = SVECTOR_VALUES(a);
|
||||
float *bx = SVECTOR_VALUES(b);
|
||||
float norma = 0.0;
|
||||
float normb = 0.0;
|
||||
double similarity;
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
similarity = inner_product_internal(a, b);
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < a->nnz; i++)
|
||||
norma += ax[i] * ax[i];
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < b->nnz; i++)
|
||||
normb += bx[i] * bx[i];
|
||||
|
||||
/* Use sqrt(a * b) over sqrt(a) * sqrt(b) */
|
||||
similarity /= 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.0;
|
||||
else if (similarity < -1)
|
||||
similarity = -1.0;
|
||||
|
||||
PG_RETURN_FLOAT8(1.0 - similarity);
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the weighted Jaccard distance between two sparse vectors
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_jaccard_distance);
|
||||
Datum
|
||||
svector_jaccard_distance(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *a = PG_GETARG_SVECTOR_P(0);
|
||||
SVector *b = PG_GETARG_SVECTOR_P(1);
|
||||
float *ax = SVECTOR_VALUES(a);
|
||||
float *bx = SVECTOR_VALUES(b);
|
||||
double num = 0.0;
|
||||
double denom = 0.0;
|
||||
int bpos = 0;
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
/*
|
||||
* Weighted Jaccard distance is not defined for vectors with negative
|
||||
* values. Could check and return NaN if minimal impact on performance.
|
||||
*/
|
||||
|
||||
for (int i = 0; i < a->nnz; i++)
|
||||
{
|
||||
int ai = a->indices[i];
|
||||
int bi = -1;
|
||||
|
||||
for (int j = bpos; j < b->nnz; j++)
|
||||
{
|
||||
bi = b->indices[j];
|
||||
|
||||
if (ai == bi)
|
||||
{
|
||||
num += ax[i] < bx[j] ? ax[i] : bx[j];
|
||||
denom += ax[i] > bx[j] ? ax[i] : bx[j];
|
||||
}
|
||||
else if (ai > bi)
|
||||
denom += bx[j];
|
||||
|
||||
/* Update start for next iteration */
|
||||
if (ai >= bi)
|
||||
bpos = j + 1;
|
||||
|
||||
/* Found or passed it */
|
||||
if (bi >= ai)
|
||||
break;
|
||||
}
|
||||
|
||||
if (ai != bi)
|
||||
denom += ax[i];
|
||||
}
|
||||
|
||||
for (int j = bpos; j < b->nnz; j++)
|
||||
denom += bx[j];
|
||||
|
||||
if (denom > 0)
|
||||
PG_RETURN_FLOAT8(1.0 - (num / denom));
|
||||
else
|
||||
PG_RETURN_FLOAT8(NAN);
|
||||
}
|
||||
@@ -1,23 +0,0 @@
|
||||
#ifndef SVECTOR_H
|
||||
#define SVECTOR_H
|
||||
|
||||
#define SVECTOR_MAX_DIM 100000
|
||||
|
||||
#define SVECTOR_SIZE(_nnz) (offsetof(SVector, indices) + (_nnz) * sizeof(int32) + (_nnz * sizeof(float)))
|
||||
#define SVECTOR_VALUES(x) ((float *) (((char *) (x)) + offsetof(SVector, indices) + (x)->nnz * sizeof(int32)))
|
||||
#define DatumGetSVector(x) ((SVector *) PG_DETOAST_DATUM(x))
|
||||
#define PG_GETARG_SVECTOR_P(x) DatumGetSVector(PG_GETARG_DATUM(x))
|
||||
#define PG_RETURN_SVECTOR_P(x) PG_RETURN_POINTER(x)
|
||||
|
||||
typedef struct SVector
|
||||
{
|
||||
int32 vl_len_; /* varlena header (do not touch directly!) */
|
||||
int32 dim; /* number of dimensions */
|
||||
int32 nnz;
|
||||
int32 unused;
|
||||
int32 indices[FLEXIBLE_ARRAY_MEMBER];
|
||||
} SVector;
|
||||
|
||||
SVector *InitSVector(int dim, int nnz);
|
||||
|
||||
#endif
|
||||
24
src/vector.c
24
src/vector.c
@@ -9,7 +9,6 @@
|
||||
#include "lib/stringinfo.h"
|
||||
#include "libpq/pqformat.h"
|
||||
#include "port.h" /* for strtof() */
|
||||
#include "svector.h"
|
||||
#include "utils/array.h"
|
||||
#include "utils/builtins.h"
|
||||
#include "utils/lsyscache.h"
|
||||
@@ -1152,26 +1151,3 @@ vector_avg(PG_FUNCTION_ARGS)
|
||||
|
||||
PG_RETURN_POINTER(result);
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert sparse vector to dense vector
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_to_vector);
|
||||
Datum
|
||||
svector_to_vector(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SVector *svec = PG_GETARG_SVECTOR_P(0);
|
||||
int32 typmod = PG_GETARG_INT32(1);
|
||||
Vector *result;
|
||||
int dim = svec->dim;
|
||||
float *values = SVECTOR_VALUES(svec);
|
||||
|
||||
CheckDim(dim);
|
||||
CheckExpectedDim(typmod, dim);
|
||||
|
||||
result = InitVector(dim);
|
||||
for (int i = 0; i < svec->nnz; i++)
|
||||
result->x[svec->indices[i]] = values[i];
|
||||
|
||||
PG_RETURN_POINTER(result);
|
||||
}
|
||||
|
||||
@@ -1,15 +1,15 @@
|
||||
CREATE TABLE t (val vector(3));
|
||||
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
|
||||
CREATE TABLE t2 (val vector(3));
|
||||
CREATE TABLE t (val vector(3), val2 half[]);
|
||||
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);
|
||||
CREATE TABLE t2 (val vector(3), val2 half[]);
|
||||
\copy t TO 'results/data.bin' WITH (FORMAT binary)
|
||||
\copy t2 FROM 'results/data.bin' WITH (FORMAT binary)
|
||||
SELECT * FROM t2 ORDER BY val;
|
||||
val
|
||||
---------
|
||||
[0,0,0]
|
||||
[1,1,1]
|
||||
[1,2,3]
|
||||
|
||||
val | val2
|
||||
---------+---------
|
||||
[0,0,0] | {0,0,0}
|
||||
[1,1,1] | {1,1,1}
|
||||
[1,2,3] | {1,2,3}
|
||||
|
|
||||
(4 rows)
|
||||
|
||||
DROP TABLE t;
|
||||
|
||||
@@ -138,21 +138,21 @@ SELECT cosine_distance('[3e38]'::vector, '[3e38]');
|
||||
NaN
|
||||
(1 row)
|
||||
|
||||
SELECT l1_distance('[0,0]', '[3,4]');
|
||||
SELECT l1_distance('[0,0]'::vector, '[3,4]');
|
||||
l1_distance
|
||||
-------------
|
||||
7
|
||||
(1 row)
|
||||
|
||||
SELECT l1_distance('[0,0]', '[0,1]');
|
||||
SELECT l1_distance('[0,0]'::vector, '[0,1]');
|
||||
l1_distance
|
||||
-------------
|
||||
1
|
||||
(1 row)
|
||||
|
||||
SELECT l1_distance('[1,2]', '[3]');
|
||||
SELECT l1_distance('[1,2]'::vector, '[3]');
|
||||
ERROR: different vector dimensions 2 and 1
|
||||
SELECT l1_distance('[3e38]', '[-3e38]');
|
||||
SELECT l1_distance('[3e38]'::vector, '[-3e38]');
|
||||
l1_distance
|
||||
-------------
|
||||
Infinity
|
||||
|
||||
182
test/expected/half.out
Normal file
182
test/expected/half.out
Normal file
@@ -0,0 +1,182 @@
|
||||
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 '65505'::integer::half;
|
||||
half
|
||||
-------
|
||||
65504
|
||||
(1 row)
|
||||
|
||||
SELECT 'NaN'::real::half;
|
||||
half
|
||||
------
|
||||
NaN
|
||||
(1 row)
|
||||
|
||||
SELECT 'Infinity'::real::half;
|
||||
half
|
||||
----------
|
||||
Infinity
|
||||
(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
|
||||
@@ -1,140 +0,0 @@
|
||||
SELECT '(0,1.5),(2,3.5)|5|'::svector;
|
||||
svector
|
||||
--------------------
|
||||
(0,1.5),(2,3.5)|5|
|
||||
(1 row)
|
||||
|
||||
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector;
|
||||
vector
|
||||
-----------------
|
||||
[1.5,0,3.5,0,0]
|
||||
(1 row)
|
||||
|
||||
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector(5);
|
||||
vector
|
||||
-----------------
|
||||
[1.5,0,3.5,0,0]
|
||||
(1 row)
|
||||
|
||||
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector(4);
|
||||
ERROR: expected 4 dimensions, not 5
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::svector;
|
||||
svector
|
||||
--------------------
|
||||
(1,1.5),(3,3.5)|5|
|
||||
(1 row)
|
||||
|
||||
SELECT '(0,0),(1,1),(2,0)|3|'::svector;
|
||||
svector
|
||||
----------------------
|
||||
(0,0),(1,1),(2,0)|3|
|
||||
(1 row)
|
||||
|
||||
SELECT '|5|'::svector;
|
||||
svector
|
||||
---------
|
||||
|5|
|
||||
(1 row)
|
||||
|
||||
SELECT '|-1|'::svector;
|
||||
ERROR: svector must have at least 1 dimension
|
||||
LINE 1: SELECT '|-1|'::svector;
|
||||
^
|
||||
SELECT '|100001|'::svector;
|
||||
ERROR: svector cannot have more than 100000 dimensions
|
||||
LINE 1: SELECT '|100001|'::svector;
|
||||
^
|
||||
SELECT '|16001|'::svector::vector;
|
||||
ERROR: vector cannot have more than 16000 dimensions
|
||||
SELECT '(-1,1)|1|'::svector;
|
||||
ERROR: index must not be negative
|
||||
LINE 1: SELECT '(-1,1)|1|'::svector;
|
||||
^
|
||||
SELECT '(1,1)|1|'::svector;
|
||||
ERROR: index must be less than dimensions
|
||||
LINE 1: SELECT '(1,1)|1|'::svector;
|
||||
^
|
||||
SELECT '|1|'::svector(2);
|
||||
ERROR: expected 2 dimensions, not 1
|
||||
SELECT l2_distance('|2|'::svector, '(0,3),(1,4)|2|');
|
||||
l2_distance
|
||||
-------------
|
||||
5
|
||||
(1 row)
|
||||
|
||||
SELECT l2_distance('|2|'::svector, '(1,1)|2|');
|
||||
l2_distance
|
||||
-------------
|
||||
1
|
||||
(1 row)
|
||||
|
||||
SELECT '|2|'::svector <-> '(0,3),(1,4)|2|';
|
||||
?column?
|
||||
----------
|
||||
5
|
||||
(1 row)
|
||||
|
||||
SELECT inner_product('(0,1),(1,2)|2|'::svector, '(0,2),(1,4)|2|');
|
||||
inner_product
|
||||
---------------
|
||||
10
|
||||
(1 row)
|
||||
|
||||
SELECT svector_negative_inner_product('(0,1),(1,2)|2|', '(0,2),(1,4)|2|');
|
||||
svector_negative_inner_product
|
||||
--------------------------------
|
||||
-10
|
||||
(1 row)
|
||||
|
||||
SELECT cosine_distance('(0,1),(1,2)|2|'::svector, '(0,2),(1,4)|2|');
|
||||
cosine_distance
|
||||
-----------------
|
||||
0
|
||||
(1 row)
|
||||
|
||||
SELECT cosine_distance('(0,1),(1,2)|2|'::svector, '|2|');
|
||||
cosine_distance
|
||||
-----------------
|
||||
NaN
|
||||
(1 row)
|
||||
|
||||
SELECT cosine_distance('(0,1),(1,1)|2|'::svector, '(0,-1),(1,-1)|2|');
|
||||
cosine_distance
|
||||
-----------------
|
||||
2
|
||||
(1 row)
|
||||
|
||||
SELECT cosine_distance('(0,1)|2|'::svector, '(1,2)|2|');
|
||||
cosine_distance
|
||||
-----------------
|
||||
1
|
||||
(1 row)
|
||||
|
||||
SELECT cosine_distance('|1|'::svector, '|1|');
|
||||
cosine_distance
|
||||
-----------------
|
||||
NaN
|
||||
(1 row)
|
||||
|
||||
SELECT cosine_distance('(0,1)|2|'::svector, '(0,1)|3|');
|
||||
ERROR: different svector dimensions 2 and 3
|
||||
SELECT jaccard_distance('(0,1)|2|', '(0,1)|2|');
|
||||
jaccard_distance
|
||||
------------------
|
||||
0
|
||||
(1 row)
|
||||
|
||||
SELECT jaccard_distance('(0,1)|2|', '(1,1)|2|');
|
||||
jaccard_distance
|
||||
------------------
|
||||
1
|
||||
(1 row)
|
||||
|
||||
SELECT jaccard_distance('|1|', '|1|');
|
||||
jaccard_distance
|
||||
------------------
|
||||
NaN
|
||||
(1 row)
|
||||
|
||||
SELECT jaccard_distance('(0,1)|2|', '(0,1)|3|');
|
||||
ERROR: different svector dimensions 2 and 3
|
||||
@@ -1,7 +1,7 @@
|
||||
CREATE TABLE t (val vector(3));
|
||||
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
|
||||
CREATE TABLE t (val vector(3), val2 half[]);
|
||||
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);
|
||||
|
||||
CREATE TABLE t2 (val vector(3));
|
||||
CREATE TABLE t2 (val vector(3), val2 half[]);
|
||||
|
||||
\copy t TO 'results/data.bin' WITH (FORMAT binary)
|
||||
\copy t2 FROM 'results/data.bin' WITH (FORMAT binary)
|
||||
|
||||
@@ -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('[3e38]'::vector, '[3e38]');
|
||||
|
||||
SELECT l1_distance('[0,0]', '[3,4]');
|
||||
SELECT l1_distance('[0,0]', '[0,1]');
|
||||
SELECT l1_distance('[1,2]', '[3]');
|
||||
SELECT l1_distance('[3e38]', '[-3e38]');
|
||||
SELECT l1_distance('[0,0]'::vector, '[3,4]');
|
||||
SELECT l1_distance('[0,0]'::vector, '[0,1]');
|
||||
SELECT l1_distance('[1,2]'::vector, '[3]');
|
||||
SELECT l1_distance('[3e38]'::vector, '[-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]', NULL]) v;
|
||||
|
||||
42
test/sql/half.sql
Normal file
42
test/sql/half.sql
Normal file
@@ -0,0 +1,42 @@
|
||||
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 '65505'::integer::half;
|
||||
SELECT 'NaN'::real::half;
|
||||
SELECT 'Infinity'::real::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}');
|
||||
@@ -1,36 +0,0 @@
|
||||
SELECT '(0,1.5),(2,3.5)|5|'::svector;
|
||||
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector;
|
||||
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector(5);
|
||||
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector(4);
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::svector;
|
||||
|
||||
SELECT '(0,0),(1,1),(2,0)|3|'::svector;
|
||||
|
||||
SELECT '|5|'::svector;
|
||||
SELECT '|-1|'::svector;
|
||||
SELECT '|100001|'::svector;
|
||||
SELECT '|16001|'::svector::vector;
|
||||
|
||||
SELECT '(-1,1)|1|'::svector;
|
||||
SELECT '(1,1)|1|'::svector;
|
||||
|
||||
SELECT '|1|'::svector(2);
|
||||
|
||||
SELECT l2_distance('|2|'::svector, '(0,3),(1,4)|2|');
|
||||
SELECT l2_distance('|2|'::svector, '(1,1)|2|');
|
||||
SELECT '|2|'::svector <-> '(0,3),(1,4)|2|';
|
||||
|
||||
SELECT inner_product('(0,1),(1,2)|2|'::svector, '(0,2),(1,4)|2|');
|
||||
SELECT svector_negative_inner_product('(0,1),(1,2)|2|', '(0,2),(1,4)|2|');
|
||||
|
||||
SELECT cosine_distance('(0,1),(1,2)|2|'::svector, '(0,2),(1,4)|2|');
|
||||
SELECT cosine_distance('(0,1),(1,2)|2|'::svector, '|2|');
|
||||
SELECT cosine_distance('(0,1),(1,1)|2|'::svector, '(0,-1),(1,-1)|2|');
|
||||
SELECT cosine_distance('(0,1)|2|'::svector, '(1,2)|2|');
|
||||
SELECT cosine_distance('|1|'::svector, '|1|');
|
||||
SELECT cosine_distance('(0,1)|2|'::svector, '(0,1)|3|');
|
||||
|
||||
SELECT jaccard_distance('(0,1)|2|', '(0,1)|2|');
|
||||
SELECT jaccard_distance('(0,1)|2|', '(1,1)|2|');
|
||||
SELECT jaccard_distance('|1|', '|1|');
|
||||
SELECT jaccard_distance('(0,1)|2|', '(0,1)|3|');
|
||||
@@ -83,11 +83,32 @@ for my $i (0 .. $#operators)
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
# Add index
|
||||
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v $opclass);");
|
||||
# Build index serially
|
||||
$node->safe_psql("postgres", qq(
|
||||
SET max_parallel_maintenance_workers = 0;
|
||||
CREATE INDEX idx ON tst USING hnsw (v $opclass);
|
||||
));
|
||||
|
||||
# Test approximate results
|
||||
my $min = $operator eq "<#>" ? 0.80 : 0.99;
|
||||
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();
|
||||
|
||||
93
test/t/019_hnsw_half.pl
Normal file
93
test/t/019_hnsw_half.pl
Normal file
@@ -0,0 +1,93 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
my $node;
|
||||
my @queries = ();
|
||||
my @expected;
|
||||
my $limit = 20;
|
||||
|
||||
sub test_recall
|
||||
{
|
||||
my ($min, $operator) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v $operator '$queries[0]' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan/);
|
||||
|
||||
for my $i (0 .. $#queries)
|
||||
{
|
||||
my $actual = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
|
||||
));
|
||||
my @actual_ids = split("\n", $actual);
|
||||
my %actual_set = map { $_ => 1 } @actual_ids;
|
||||
|
||||
my @expected_ids = split("\n", $expected[$i]);
|
||||
|
||||
foreach (@expected_ids)
|
||||
{
|
||||
if (exists($actual_set{$_}))
|
||||
{
|
||||
$correct++;
|
||||
}
|
||||
$total++;
|
||||
}
|
||||
}
|
||||
|
||||
cmp_ok($correct / $total, ">=", $min, $operator);
|
||||
}
|
||||
|
||||
# Initialize node
|
||||
$node = get_new_node('node');
|
||||
$node->init;
|
||||
$node->start;
|
||||
|
||||
# Create table
|
||||
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
|
||||
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v half[3]);");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()]::numeric[]::half[] FROM generate_series(1, 10000) i;"
|
||||
);
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my $r1 = rand();
|
||||
my $r2 = rand();
|
||||
my $r3 = rand();
|
||||
push(@queries, "{$r1,$r2,$r3}");
|
||||
}
|
||||
|
||||
# Check each index type
|
||||
my @operators = ("<->");
|
||||
my @opclasses = ("half_l2_ops");
|
||||
|
||||
for my $i (0 .. $#operators)
|
||||
{
|
||||
my $operator = $operators[$i];
|
||||
my $opclass = $opclasses[$i];
|
||||
|
||||
# Get exact results
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
# Add index
|
||||
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v $opclass);");
|
||||
|
||||
my $min = $operator eq "<#>" ? 0.80 : 0.99;
|
||||
test_recall($min, $operator);
|
||||
}
|
||||
|
||||
done_testing();
|
||||
Reference in New Issue
Block a user