mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-22 03:57:34 +08:00
Compare commits
102 Commits
windows-si
...
target-clo
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
665b3dfa1c | ||
|
|
ce597f770d | ||
|
|
1234fffcb0 | ||
|
|
e70a47c154 | ||
|
|
0c9ad67a1c | ||
|
|
8a4845b52e | ||
|
|
e14fdba8b3 | ||
|
|
5abf83c415 | ||
|
|
96b30fd25d | ||
|
|
1c791480ff | ||
|
|
545ae30057 | ||
|
|
20fb2e0797 | ||
|
|
89629abc08 | ||
|
|
db112862a5 | ||
|
|
bb84d69e57 | ||
|
|
7ea1590ea3 | ||
|
|
d98eb0a28d | ||
|
|
72eee91d00 | ||
|
|
f5e6e58359 | ||
|
|
8f93d02c71 | ||
|
|
c5c82bafda | ||
|
|
f627e69c5f | ||
|
|
23944302fe | ||
|
|
9957ba6817 | ||
|
|
aed463dbad | ||
|
|
f64abe3aae | ||
|
|
06d90fdc76 | ||
|
|
cc4b01bd49 | ||
|
|
3621a84ef8 | ||
|
|
b9d5269547 | ||
|
|
fbc96bb488 | ||
|
|
5510ae5b8c | ||
|
|
70eee06e83 | ||
|
|
f2bffff460 | ||
|
|
af7d9b74a9 | ||
|
|
f1a71524f0 | ||
|
|
7710cc8c28 | ||
|
|
fea2eb262e | ||
|
|
1bc6f954f4 | ||
|
|
6fd6b0bd5f | ||
|
|
546006b7ca | ||
|
|
cca0edf458 | ||
|
|
8d9400bae3 | ||
|
|
a4531ca51f | ||
|
|
94ee6b13c3 | ||
|
|
e6a591275a | ||
|
|
7fddd296ee | ||
|
|
c581db9f98 | ||
|
|
626bc053e5 | ||
|
|
66a29dbdf3 | ||
|
|
1c26da6ef5 | ||
|
|
000cc13c29 | ||
|
|
bbfe1e8b61 | ||
|
|
7e8be16e37 | ||
|
|
17c2f9c0b6 | ||
|
|
bd52ed29e0 | ||
|
|
245bac5e8e | ||
|
|
d84fc303ee | ||
|
|
4ff36af67e | ||
|
|
5834b58c5a | ||
|
|
381216a956 | ||
|
|
c3c6af8a84 | ||
|
|
d45f561d75 | ||
|
|
e1647473c9 | ||
|
|
e8bd7cd2f5 | ||
|
|
bed40ec0b5 | ||
|
|
d64295dfd7 | ||
|
|
8178a902ce | ||
|
|
f9f53b64e8 | ||
|
|
fc83cd8d12 | ||
|
|
1e41ed6f15 | ||
|
|
e6ca831f3d | ||
|
|
a824af42fb | ||
|
|
8a29bf4619 | ||
|
|
33369e0744 | ||
|
|
34b09cc062 | ||
|
|
d8f3e18db6 | ||
|
|
e6e7d7c1bc | ||
|
|
bf355792b2 | ||
|
|
12f4a34708 | ||
|
|
f6becf01aa | ||
|
|
dd5b856f68 | ||
|
|
9c96164f2c | ||
|
|
b583803b2c | ||
|
|
a04bf7fce2 | ||
|
|
62b411f94c | ||
|
|
82a17b81f5 | ||
|
|
3b2ca7df7a | ||
|
|
21d125abed | ||
|
|
f42ccd451d | ||
|
|
85345e3f8c | ||
|
|
05ce20990c | ||
|
|
f3477cf28d | ||
|
|
434f3f5e88 | ||
|
|
bd6fdb73eb | ||
|
|
ab382a2053 | ||
|
|
191c8e1cca | ||
|
|
3eba34e5e3 | ||
|
|
862f17c1de | ||
|
|
21bcff6722 | ||
|
|
60b4bb2ad1 | ||
|
|
c27363fbf7 |
@@ -7,6 +7,7 @@
|
||||
- Added `jaccard_distance` function
|
||||
- Added `quantize_binary` function
|
||||
- Added `subvector` function
|
||||
- Added CPU dispatching for distance functions on Linux x86-64
|
||||
- Updated comparison operators to support vectors with different dimensions
|
||||
|
||||
## 0.6.2 (2024-03-18)
|
||||
|
||||
2
Makefile
2
Makefile
@@ -3,7 +3,7 @@ EXTVERSION = 0.6.2
|
||||
|
||||
MODULE_big = vector
|
||||
DATA = $(wildcard sql/*--*.sql)
|
||||
OBJS = src/bitvector.o src/halfvec.o src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/sparsevec.o src/vector.o
|
||||
OBJS = src/bitvector.o src/halfutils.o src/halfvec.o src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/sparsevec.o src/vector.o
|
||||
HEADERS = src/halfvec.h src/sparsevec.h src/vector.h
|
||||
|
||||
TESTS = $(wildcard test/sql/*.sql)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
EXTENSION = vector
|
||||
EXTVERSION = 0.6.2
|
||||
|
||||
OBJS = src\bitvector.obj src\halfvec.obj src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\sparsevec.obj src\vector.obj
|
||||
OBJS = src\bitvector.obj src\halfutils.obj src\halfvec.obj src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\sparsevec.obj src\vector.obj
|
||||
HEADERS = src\halfvec.h src\sparsevec.h src\vector.h
|
||||
|
||||
REGRESS = bit_functions btree cast copy halfvec_functions halfvec_input hnsw_bit_hamming hnsw_bit_jaccard hnsw_halfvec_cosine hnsw_halfvec_ip hnsw_halfvec_l2 hnsw_options hnsw_sparsevec_cosine hnsw_sparsevec_ip hnsw_sparsevec_l2 hnsw_unlogged hnsw_vector_cosine hnsw_vector_ip hnsw_vector_l2 ivfflat_options ivfflat_unlogged ivfflat_vector_cosine ivfflat_vector_ip ivfflat_vector_l2 sparsevec_functions sparsevec_input vector_functions vector_input
|
||||
|
||||
25
README.md
25
README.md
@@ -346,6 +346,7 @@ CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists =
|
||||
Supported types are:
|
||||
|
||||
- `vector` - up to 2,000 dimensions
|
||||
- `halfvec` - up to 4,000 dimensions (unreleased)
|
||||
|
||||
### Query Options
|
||||
|
||||
@@ -527,6 +528,30 @@ SELECT id, content FROM items, plainto_tsquery('hello search') query
|
||||
|
||||
You can use [Reciprocal Rank Fusion](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) or a [cross-encoder](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) to combine results.
|
||||
|
||||
## Subvector Indexing
|
||||
|
||||
*Unreleased*
|
||||
|
||||
Use expression indexing to index subvectors
|
||||
|
||||
```sql
|
||||
CREATE INDEX ON items USING hnsw ((subvector(embedding, 1, 3)::vector(3)) vector_cosine_ops);
|
||||
```
|
||||
|
||||
Get the nearest neighbors by cosine distance
|
||||
|
||||
```sql
|
||||
SELECT * FROM items ORDER BY subvector(embedding, 1, 3)::vector(3) <=> subvector('[1,2,3,4,5]'::vector, 1, 3) LIMIT 5;
|
||||
```
|
||||
|
||||
Re-rank by the full vectors for better recall
|
||||
|
||||
```sql
|
||||
SELECT * FROM (
|
||||
SELECT * FROM items ORDER BY subvector(embedding, 1, 3)::vector(3) <=> subvector('[1,2,3,4,5]'::vector, 1, 3) LIMIT 20
|
||||
) ORDER BY embedding <=> '[1,2,3,4,5]' LIMIT 5;
|
||||
```
|
||||
|
||||
## Performance
|
||||
|
||||
### Tuning
|
||||
|
||||
@@ -86,6 +86,9 @@ CREATE FUNCTION halfvec_l2_squared_distance(halfvec, halfvec) RETURNS float8
|
||||
CREATE FUNCTION halfvec_negative_inner_product(halfvec, halfvec) RETURNS float8
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION halfvec_spherical_distance(halfvec, halfvec) RETURNS float8
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION halfvec(halfvec, integer, boolean) RETURNS halfvec
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
@@ -137,6 +140,27 @@ CREATE OPERATOR <=> (
|
||||
COMMUTATOR = '<=>'
|
||||
);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_l2_ops
|
||||
FOR TYPE halfvec USING ivfflat AS
|
||||
OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 halfvec_l2_squared_distance(halfvec, halfvec),
|
||||
FUNCTION 3 l2_distance(halfvec, halfvec);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_ip_ops
|
||||
FOR TYPE halfvec USING ivfflat AS
|
||||
OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec),
|
||||
FUNCTION 3 halfvec_spherical_distance(halfvec, halfvec),
|
||||
FUNCTION 4 halfvec_norm(halfvec);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_cosine_ops
|
||||
FOR TYPE halfvec USING ivfflat AS
|
||||
OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec),
|
||||
FUNCTION 2 halfvec_norm(halfvec),
|
||||
FUNCTION 3 halfvec_spherical_distance(halfvec, halfvec),
|
||||
FUNCTION 4 halfvec_norm(halfvec);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_l2_ops
|
||||
FOR TYPE halfvec USING hnsw AS
|
||||
OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
|
||||
@@ -381,6 +381,9 @@ CREATE FUNCTION halfvec_l2_squared_distance(halfvec, halfvec) RETURNS float8
|
||||
CREATE FUNCTION halfvec_negative_inner_product(halfvec, halfvec) RETURNS float8
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION halfvec_spherical_distance(halfvec, halfvec) RETURNS float8
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
-- halfvec cast functions
|
||||
|
||||
CREATE FUNCTION halfvec(halfvec, integer, boolean) RETURNS halfvec
|
||||
@@ -440,6 +443,27 @@ CREATE OPERATOR <=> (
|
||||
|
||||
-- halfvec opclasses
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_l2_ops
|
||||
FOR TYPE halfvec USING ivfflat AS
|
||||
OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 halfvec_l2_squared_distance(halfvec, halfvec),
|
||||
FUNCTION 3 l2_distance(halfvec, halfvec);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_ip_ops
|
||||
FOR TYPE halfvec USING ivfflat AS
|
||||
OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec),
|
||||
FUNCTION 3 halfvec_spherical_distance(halfvec, halfvec),
|
||||
FUNCTION 4 halfvec_norm(halfvec);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_cosine_ops
|
||||
FOR TYPE halfvec USING ivfflat AS
|
||||
OPERATOR 1 <=> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec),
|
||||
FUNCTION 2 halfvec_norm(halfvec),
|
||||
FUNCTION 3 halfvec_spherical_distance(halfvec, halfvec),
|
||||
FUNCTION 4 halfvec_norm(halfvec);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_l2_ops
|
||||
FOR TYPE halfvec USING hnsw AS
|
||||
OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
|
||||
156
src/halfutils.c
Normal file
156
src/halfutils.c
Normal file
@@ -0,0 +1,156 @@
|
||||
#include "postgres.h"
|
||||
|
||||
#include "halfutils.h"
|
||||
#include "halfvec.h"
|
||||
|
||||
#ifdef HALFVEC_DISPATCH
|
||||
#include <immintrin.h>
|
||||
|
||||
#if defined(HAVE__GET_CPUID)
|
||||
#include <cpuid.h>
|
||||
#elif defined(HAVE__CPUID)
|
||||
#include <intrin.h>
|
||||
#endif
|
||||
|
||||
#ifdef _MSC_VER
|
||||
#define TARGET_F16C_FMA
|
||||
#else
|
||||
#define TARGET_F16C_FMA __attribute__((target("f16c,fma")))
|
||||
#endif
|
||||
#endif
|
||||
|
||||
float (*HalfvecL2SquaredDistance) (int dim, half * ax, half * bx);
|
||||
float (*HalfvecInnerProduct) (int dim, half * ax, half * bx);
|
||||
|
||||
static float
|
||||
HalfvecL2SquaredDistanceDefault(int dim, half * ax, half * bx)
|
||||
{
|
||||
float distance = 0.0;
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < dim; i++)
|
||||
{
|
||||
float diff = HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]);
|
||||
|
||||
distance += diff * diff;
|
||||
}
|
||||
|
||||
return distance;
|
||||
}
|
||||
|
||||
#ifdef HALFVEC_DISPATCH
|
||||
TARGET_F16C_FMA static float
|
||||
HalfvecL2SquaredDistanceF16cFma(int dim, half * ax, half * bx)
|
||||
{
|
||||
float distance;
|
||||
int i;
|
||||
float s[8];
|
||||
int count = (dim / 8) * 8;
|
||||
__m256 dist = _mm256_setzero_ps();
|
||||
|
||||
for (i = 0; i < count; i += 8)
|
||||
{
|
||||
__m128i axi = _mm_loadu_si128((__m128i *) (ax + i));
|
||||
__m128i bxi = _mm_loadu_si128((__m128i *) (bx + i));
|
||||
__m256 axs = _mm256_cvtph_ps(axi);
|
||||
__m256 bxs = _mm256_cvtph_ps(bxi);
|
||||
__m256 diff = _mm256_sub_ps(axs, bxs);
|
||||
|
||||
dist = _mm256_fmadd_ps(diff, diff, dist);
|
||||
}
|
||||
|
||||
_mm256_storeu_ps(s, dist);
|
||||
|
||||
distance = s[0] + s[1] + s[2] + s[3] + s[4] + s[5] + s[6] + s[7];
|
||||
|
||||
for (; i < dim; i++)
|
||||
{
|
||||
float diff = HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]);
|
||||
|
||||
distance += diff * diff;
|
||||
}
|
||||
|
||||
return distance;
|
||||
}
|
||||
#endif
|
||||
|
||||
static float
|
||||
HalfvecInnerProductDefault(int dim, half * ax, half * bx)
|
||||
{
|
||||
float distance = 0.0;
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < dim; i++)
|
||||
distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);
|
||||
|
||||
return distance;
|
||||
}
|
||||
|
||||
#ifdef HALFVEC_DISPATCH
|
||||
TARGET_F16C_FMA static float
|
||||
HalfvecInnerProductF16cFma(int dim, half * ax, half * bx)
|
||||
{
|
||||
float distance;
|
||||
int i;
|
||||
float s[8];
|
||||
int count = (dim / 8) * 8;
|
||||
__m256 dist = _mm256_setzero_ps();
|
||||
|
||||
for (i = 0; i < count; i += 8)
|
||||
{
|
||||
__m128i axi = _mm_loadu_si128((__m128i *) (ax + i));
|
||||
__m128i bxi = _mm_loadu_si128((__m128i *) (bx + i));
|
||||
__m256 axs = _mm256_cvtph_ps(axi);
|
||||
__m256 bxs = _mm256_cvtph_ps(bxi);
|
||||
|
||||
dist = _mm256_fmadd_ps(axs, bxs, dist);
|
||||
}
|
||||
|
||||
_mm256_storeu_ps(s, dist);
|
||||
|
||||
distance = s[0] + s[1] + s[2] + s[3] + s[4] + s[5] + s[6] + s[7];
|
||||
|
||||
for (; i < dim; i++)
|
||||
distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);
|
||||
|
||||
return distance;
|
||||
}
|
||||
#endif
|
||||
|
||||
#ifdef HALFVEC_DISPATCH
|
||||
#define CPU_FEATURE_FMA (1 << 12)
|
||||
#define CPU_FEATURE_F16C (1 << 29)
|
||||
|
||||
static bool
|
||||
SupportsCpuFeature(unsigned int feature)
|
||||
{
|
||||
unsigned int exx[4] = {0, 0, 0, 0};
|
||||
|
||||
#if defined(HAVE__GET_CPUID)
|
||||
__get_cpuid(1, &exx[0], &exx[1], &exx[2], &exx[3]);
|
||||
#elif defined(HAVE__CPUID)
|
||||
__cpuid(exx, 1);
|
||||
#endif
|
||||
|
||||
return (exx[2] & feature) == feature;
|
||||
}
|
||||
#endif
|
||||
|
||||
void
|
||||
HalfvecInit(void)
|
||||
{
|
||||
/*
|
||||
* Could skip pointer when single function, but no difference in
|
||||
* performance
|
||||
*/
|
||||
HalfvecL2SquaredDistance = HalfvecL2SquaredDistanceDefault;
|
||||
HalfvecInnerProduct = HalfvecInnerProductDefault;
|
||||
|
||||
#ifdef HALFVEC_DISPATCH
|
||||
if (SupportsCpuFeature(CPU_FEATURE_FMA | CPU_FEATURE_F16C))
|
||||
{
|
||||
HalfvecL2SquaredDistance = HalfvecL2SquaredDistanceF16cFma;
|
||||
HalfvecInnerProduct = HalfvecInnerProductF16cFma;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
254
src/halfutils.h
Normal file
254
src/halfutils.h
Normal file
@@ -0,0 +1,254 @@
|
||||
#ifndef HALFUTILS_H
|
||||
#define HALFUTILS_H
|
||||
|
||||
#include <math.h>
|
||||
|
||||
#include "common/shortest_dec.h"
|
||||
#include "halfvec.h"
|
||||
|
||||
#ifdef F16C_SUPPORT
|
||||
#include <immintrin.h>
|
||||
#endif
|
||||
|
||||
extern float (*HalfvecL2SquaredDistance) (int dim, half * ax, half * bx);
|
||||
extern float (*HalfvecInnerProduct) (int dim, half * ax, half * bx);
|
||||
|
||||
void HalfvecInit(void);
|
||||
|
||||
/*
|
||||
* Check if half is NaN
|
||||
*/
|
||||
static inline bool
|
||||
HalfIsNan(half num)
|
||||
{
|
||||
#ifdef FLT16_SUPPORT
|
||||
return isnan(num);
|
||||
#else
|
||||
return (num & 0x7C00) == 0x7C00 && (num & 0x7FFF) != 0x7C00;
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
* Check if half is infinite
|
||||
*/
|
||||
static inline bool
|
||||
HalfIsInf(half num)
|
||||
{
|
||||
#ifdef FLT16_SUPPORT
|
||||
return isinf(num);
|
||||
#else
|
||||
return (num & 0x7FFF) == 0x7C00;
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert a half to a float4
|
||||
*/
|
||||
static inline float
|
||||
HalfToFloat4(half num)
|
||||
{
|
||||
#if defined(F16C_SUPPORT)
|
||||
return _cvtsh_ss(num);
|
||||
#elif defined(FLT16_SUPPORT)
|
||||
return (float) num;
|
||||
#else
|
||||
/* TODO Improve performance */
|
||||
|
||||
/* Assumes same endianness for floats and integers */
|
||||
union
|
||||
{
|
||||
float f;
|
||||
uint32 i;
|
||||
} swapfloat;
|
||||
|
||||
union
|
||||
{
|
||||
half h;
|
||||
uint16 i;
|
||||
} swaphalf;
|
||||
|
||||
uint16 bin;
|
||||
uint32 exponent;
|
||||
uint32 mantissa;
|
||||
uint32 result;
|
||||
|
||||
swaphalf.h = num;
|
||||
bin = swaphalf.i;
|
||||
exponent = (bin & 0x7C00) >> 10;
|
||||
mantissa = bin & 0x03FF;
|
||||
|
||||
/* Sign */
|
||||
result = (bin & 0x8000) << 16;
|
||||
|
||||
if (unlikely(exponent == 31))
|
||||
{
|
||||
if (mantissa == 0)
|
||||
{
|
||||
/* Infinite */
|
||||
result |= 0x7F800000;
|
||||
}
|
||||
else
|
||||
{
|
||||
/* NaN */
|
||||
result |= 0x7FC00000;
|
||||
}
|
||||
}
|
||||
else if (unlikely(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;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
/* Normal */
|
||||
result |= (exponent - 15 + 127) << 23;
|
||||
}
|
||||
|
||||
result |= mantissa << 13;
|
||||
|
||||
swapfloat.i = result;
|
||||
return swapfloat.f;
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert a float4 to a half
|
||||
*/
|
||||
static inline half
|
||||
Float4ToHalfUnchecked(float num)
|
||||
{
|
||||
#if defined(F16C_SUPPORT)
|
||||
return _cvtss_sh(num, 0);
|
||||
#elif defined(FLT16_SUPPORT)
|
||||
return (_Float16) num;
|
||||
#else
|
||||
/* TODO Improve performance */
|
||||
|
||||
/* Assumes same endianness for floats and integers */
|
||||
union
|
||||
{
|
||||
float f;
|
||||
uint32 i;
|
||||
} swapfloat;
|
||||
|
||||
union
|
||||
{
|
||||
half h;
|
||||
uint16 i;
|
||||
} swaphalf;
|
||||
|
||||
uint32 bin;
|
||||
int exponent;
|
||||
int mantissa;
|
||||
uint16 result;
|
||||
|
||||
swapfloat.f = num;
|
||||
bin = swapfloat.i;
|
||||
exponent = (bin & 0x7F800000) >> 23;
|
||||
mantissa = bin & 0x007FFFFF;
|
||||
|
||||
/* Sign */
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
swaphalf.i = result;
|
||||
return swaphalf.h;
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert a float4 to a half
|
||||
*/
|
||||
static inline half
|
||||
Float4ToHalf(float num)
|
||||
{
|
||||
half result = Float4ToHalfUnchecked(num);
|
||||
|
||||
if (unlikely(HalfIsInf(result)) && !isinf(num))
|
||||
{
|
||||
char *buf = palloc(FLOAT_SHORTEST_DECIMAL_LEN);
|
||||
|
||||
float_to_shortest_decimal_buf(num, buf);
|
||||
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("\"%s\" is out of range for type halfvec", buf)));
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
#endif
|
||||
512
src/halfvec.c
512
src/halfvec.c
@@ -6,6 +6,7 @@
|
||||
#include "catalog/pg_type.h"
|
||||
#include "common/shortest_dec.h"
|
||||
#include "fmgr.h"
|
||||
#include "halfutils.h"
|
||||
#include "halfvec.h"
|
||||
#include "lib/stringinfo.h"
|
||||
#include "libpq/pqformat.h"
|
||||
@@ -22,49 +23,6 @@
|
||||
#define TYPALIGN_INT 'i'
|
||||
#endif
|
||||
|
||||
#ifdef F16C_SUPPORT
|
||||
#include <immintrin.h>
|
||||
#endif
|
||||
|
||||
/*
|
||||
* Check if half is NaN
|
||||
*/
|
||||
static inline bool
|
||||
HalfIsNan(half num)
|
||||
{
|
||||
#ifdef FLT16_SUPPORT
|
||||
return isnan(num);
|
||||
#else
|
||||
return (num & 0x7C00) == 0x7C00 && (num & 0x7FFF) != 0x7C00;
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
* Check if half is infinite
|
||||
*/
|
||||
static inline bool
|
||||
HalfIsInf(half num)
|
||||
{
|
||||
#ifdef FLT16_SUPPORT
|
||||
return isinf(num);
|
||||
#else
|
||||
return (num & 0x7FFF) == 0x7C00;
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
* Check if half is zero
|
||||
*/
|
||||
static inline bool
|
||||
HalfIsZero(half num)
|
||||
{
|
||||
#ifdef FLT16_SUPPORT
|
||||
return num == 0;
|
||||
#else
|
||||
return (num & 0x7FFF) == 0x0000;
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
* Get a half from a message buffer
|
||||
*/
|
||||
@@ -97,212 +55,6 @@ pq_sendhalf(StringInfo buf, half h)
|
||||
pq_sendint16(buf, swap.i);
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert a half to a float4
|
||||
*/
|
||||
float
|
||||
HalfToFloat4(half num)
|
||||
{
|
||||
#if defined(F16C_SUPPORT) && !defined(_MSC_VER)
|
||||
/* TODO Use instrinsics for Windows */
|
||||
return _cvtsh_ss(num);
|
||||
#elif defined(FLT16_SUPPORT)
|
||||
return (float) num;
|
||||
#else
|
||||
/* TODO Improve performance */
|
||||
|
||||
/* Assumes same endianness for floats and integers */
|
||||
union
|
||||
{
|
||||
float f;
|
||||
uint32 i;
|
||||
} swapfloat;
|
||||
|
||||
union
|
||||
{
|
||||
half h;
|
||||
uint16 i;
|
||||
} swaphalf;
|
||||
|
||||
uint16 bin;
|
||||
uint32 exponent;
|
||||
uint32 mantissa;
|
||||
uint32 result;
|
||||
|
||||
swaphalf.h = num;
|
||||
bin = swaphalf.i;
|
||||
exponent = (bin & 0x7C00) >> 10;
|
||||
mantissa = bin & 0x03FF;
|
||||
|
||||
/* Sign */
|
||||
result = (bin & 0x8000) << 16;
|
||||
|
||||
if (unlikely(exponent == 31))
|
||||
{
|
||||
if (mantissa == 0)
|
||||
{
|
||||
/* Infinite */
|
||||
result |= 0x7F800000;
|
||||
}
|
||||
else
|
||||
{
|
||||
/* NaN */
|
||||
result |= 0x7FC00000;
|
||||
}
|
||||
}
|
||||
else if (unlikely(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;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
/* Normal */
|
||||
result |= (exponent - 15 + 127) << 23;
|
||||
}
|
||||
|
||||
result |= mantissa << 13;
|
||||
|
||||
swapfloat.i = result;
|
||||
return swapfloat.f;
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert a float4 to a half
|
||||
*/
|
||||
half
|
||||
Float4ToHalfUnchecked(float num)
|
||||
{
|
||||
#if defined(F16C_SUPPORT) && !defined(_MSC_VER)
|
||||
/* TODO Use instrinsics for Windows */
|
||||
return _cvtss_sh(num, 0);
|
||||
#elif defined(FLT16_SUPPORT)
|
||||
return (_Float16) num;
|
||||
#else
|
||||
/* TODO Improve performance */
|
||||
|
||||
/* Assumes same endianness for floats and integers */
|
||||
union
|
||||
{
|
||||
float f;
|
||||
uint32 i;
|
||||
} swapfloat;
|
||||
|
||||
union
|
||||
{
|
||||
half h;
|
||||
uint16 i;
|
||||
} swaphalf;
|
||||
|
||||
uint32 bin;
|
||||
int exponent;
|
||||
int mantissa;
|
||||
uint16 result;
|
||||
|
||||
swapfloat.f = num;
|
||||
bin = swapfloat.i;
|
||||
exponent = (bin & 0x7F800000) >> 23;
|
||||
mantissa = bin & 0x007FFFFF;
|
||||
|
||||
/* Sign */
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
swaphalf.i = result;
|
||||
return swaphalf.h;
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
* Convert a float4 to a half
|
||||
*/
|
||||
half
|
||||
Float4ToHalf(float num)
|
||||
{
|
||||
half result = Float4ToHalfUnchecked(num);
|
||||
|
||||
if (unlikely(HalfIsInf(result)) && !isinf(num))
|
||||
float_overflow_error();
|
||||
if (unlikely(HalfIsZero(result)) && num != 0.0)
|
||||
float_underflow_error();
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/*
|
||||
* Ensure same dimensions
|
||||
*/
|
||||
@@ -394,24 +146,6 @@ halfvec_isspace(char ch)
|
||||
return false;
|
||||
}
|
||||
|
||||
#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
|
||||
|
||||
/*
|
||||
* Convert textual representation to internal representation
|
||||
*/
|
||||
@@ -423,28 +157,32 @@ halfvec_in(PG_FUNCTION_ARGS)
|
||||
int32 typmod = PG_GETARG_INT32(2);
|
||||
half x[HALFVEC_MAX_DIM];
|
||||
int dim = 0;
|
||||
char *pt;
|
||||
char *stringEnd;
|
||||
char *pt = lit;
|
||||
HalfVector *result;
|
||||
char *litcopy = pstrdup(lit);
|
||||
char *str = litcopy;
|
||||
|
||||
while (halfvec_isspace(*str))
|
||||
str++;
|
||||
while (halfvec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*str != '[')
|
||||
if (*pt != '[')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed halfvec literal: \"%s\"", lit),
|
||||
errmsg("invalid input syntax for type halfvec: \"%s\"", lit),
|
||||
errdetail("Vector contents must start with \"[\".")));
|
||||
|
||||
str++;
|
||||
pt = strtok(str, ",");
|
||||
stringEnd = pt;
|
||||
pt++;
|
||||
|
||||
while (pt != NULL && *stringEnd != ']')
|
||||
while (halfvec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*pt == ']')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("halfvec must have at least 1 dimension")));
|
||||
|
||||
for (;;)
|
||||
{
|
||||
float val;
|
||||
char *stringEnd;
|
||||
|
||||
if (dim == HALFVEC_MAX_DIM)
|
||||
ereport(ERROR,
|
||||
@@ -460,8 +198,9 @@ halfvec_in(PG_FUNCTION_ARGS)
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type halfvec: \"%s\"", lit)));
|
||||
|
||||
/* Use strtof like float4in to avoid a double-rounding problem */
|
||||
errno = 0;
|
||||
|
||||
/* Postgres sets LC_NUMERIC to C on startup */
|
||||
val = strtof(pt, &stringEnd);
|
||||
|
||||
if (stringEnd == pt)
|
||||
@@ -471,59 +210,44 @@ halfvec_in(PG_FUNCTION_ARGS)
|
||||
|
||||
x[dim] = Float4ToHalfUnchecked(val);
|
||||
|
||||
if ((errno == ERANGE && (isinf(val) || val == 0)) || (HalfIsInf(x[dim]) && !isinf(val)) || (HalfIsZero(x[dim]) && val != 0))
|
||||
/* Check for range error like float4in */
|
||||
if ((errno == ERANGE && isinf(val)) || (HalfIsInf(x[dim]) && !isinf(val)))
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("\"%s\" is out of range for type halfvec", pt)));
|
||||
errmsg("\"%s\" is out of range for type halfvec", pnstrdup(pt, stringEnd - pt))));
|
||||
|
||||
CheckElement(x[dim]);
|
||||
dim++;
|
||||
|
||||
while (halfvec_isspace(*stringEnd))
|
||||
stringEnd++;
|
||||
pt = stringEnd;
|
||||
|
||||
if (*stringEnd != '\0' && *stringEnd != ']')
|
||||
while (halfvec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*pt == ',')
|
||||
pt++;
|
||||
else if (*pt == ']')
|
||||
{
|
||||
pt++;
|
||||
break;
|
||||
}
|
||||
else
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type halfvec: \"%s\"", lit)));
|
||||
|
||||
pt = strtok(NULL, ",");
|
||||
}
|
||||
|
||||
if (stringEnd == NULL || *stringEnd != ']')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed halfvec literal: \"%s\"", lit),
|
||||
errdetail("Unexpected end of input.")));
|
||||
|
||||
stringEnd++;
|
||||
|
||||
/* Only whitespace is allowed after the closing brace */
|
||||
while (halfvec_isspace(*stringEnd))
|
||||
stringEnd++;
|
||||
while (halfvec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*stringEnd != '\0')
|
||||
if (*pt != '\0')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed halfvec literal: \"%s\"", lit),
|
||||
errmsg("invalid input syntax for type halfvec: \"%s\"", lit),
|
||||
errdetail("Junk after closing right brace.")));
|
||||
|
||||
/* Ensure no consecutive delimiters since strtok skips */
|
||||
for (pt = lit + 1; *pt != '\0'; pt++)
|
||||
{
|
||||
if (pt[-1] == ',' && *pt == ',')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed halfvec literal: \"%s\"", lit)));
|
||||
}
|
||||
|
||||
if (dim < 1)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("halfvec must have at least 1 dimension")));
|
||||
|
||||
pfree(litcopy);
|
||||
|
||||
CheckDim(dim);
|
||||
CheckExpectedDim(typmod, dim);
|
||||
|
||||
result = InitHalfVector(dim);
|
||||
@@ -793,65 +517,11 @@ vector_to_halfvec(PG_FUNCTION_ARGS)
|
||||
result = InitHalfVector(vec->dim);
|
||||
|
||||
for (int i = 0; i < vec->dim; i++)
|
||||
{
|
||||
result->x[i] = Float4ToHalfUnchecked(vec->x[i]);
|
||||
/* TODO Better error for overflow */
|
||||
CheckElement(result->x[i]);
|
||||
}
|
||||
result->x[i] = Float4ToHalf(vec->x[i]);
|
||||
|
||||
PG_RETURN_POINTER(result);
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the L2 squared distance between half vectors
|
||||
*/
|
||||
static double
|
||||
l2_distance_squared_internal(HalfVector * a, HalfVector * b)
|
||||
{
|
||||
half *ax = a->x;
|
||||
half *bx = b->x;
|
||||
float distance = 0.0;
|
||||
|
||||
#if defined(F16C_SUPPORT) && defined(__FMA__)
|
||||
int i;
|
||||
float s[8];
|
||||
int count = (a->dim / 8) * 8;
|
||||
__m256 dist = _mm256_setzero_ps();
|
||||
|
||||
for (i = 0; i < count; i += 8)
|
||||
{
|
||||
__m128i axi = _mm_loadu_si128((__m128i *) (ax + i));
|
||||
__m128i bxi = _mm_loadu_si128((__m128i *) (bx + i));
|
||||
__m256 axs = _mm256_cvtph_ps(axi);
|
||||
__m256 bxs = _mm256_cvtph_ps(bxi);
|
||||
__m256 diff = _mm256_sub_ps(axs, bxs);
|
||||
|
||||
dist = _mm256_fmadd_ps(diff, diff, dist);
|
||||
}
|
||||
|
||||
_mm256_storeu_ps(s, dist);
|
||||
|
||||
distance = s[0] + s[1] + s[2] + s[3] + s[4] + s[5] + s[6] + s[7];
|
||||
|
||||
for (; i < a->dim; i++)
|
||||
{
|
||||
float diff = HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]);
|
||||
|
||||
distance += diff * diff;
|
||||
}
|
||||
#else
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < a->dim; i++)
|
||||
{
|
||||
float diff = HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]);
|
||||
|
||||
distance += diff * diff;
|
||||
}
|
||||
#endif
|
||||
|
||||
return (double) distance;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the L2 distance between half vectors
|
||||
*/
|
||||
@@ -864,7 +534,7 @@ halfvec_l2_distance(PG_FUNCTION_ARGS)
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(sqrt(l2_distance_squared_internal(a, b)));
|
||||
PG_RETURN_FLOAT8(sqrt((double) HalfvecL2SquaredDistance(a->dim, a->x, b->x)));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -879,48 +549,7 @@ halfvec_l2_squared_distance(PG_FUNCTION_ARGS)
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(l2_distance_squared_internal(a, b));
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the inner product of two half vectors
|
||||
*/
|
||||
static double
|
||||
inner_product_internal(HalfVector * a, HalfVector * b)
|
||||
{
|
||||
half *ax = a->x;
|
||||
half *bx = b->x;
|
||||
float distance = 0.0;
|
||||
|
||||
#if defined(F16C_SUPPORT) && defined(__FMA__)
|
||||
int i;
|
||||
float s[8];
|
||||
int count = (a->dim / 8) * 8;
|
||||
__m256 dist = _mm256_setzero_ps();
|
||||
|
||||
for (i = 0; i < count; i += 8)
|
||||
{
|
||||
__m128i axi = _mm_loadu_si128((__m128i *) (ax + i));
|
||||
__m128i bxi = _mm_loadu_si128((__m128i *) (bx + i));
|
||||
__m256 axs = _mm256_cvtph_ps(axi);
|
||||
__m256 bxs = _mm256_cvtph_ps(bxi);
|
||||
|
||||
dist = _mm256_fmadd_ps(axs, bxs, dist);
|
||||
}
|
||||
|
||||
_mm256_storeu_ps(s, dist);
|
||||
|
||||
distance = s[0] + s[1] + s[2] + s[3] + s[4] + s[5] + s[6] + s[7];
|
||||
|
||||
for (; i < a->dim; i++)
|
||||
distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);
|
||||
#else
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < a->dim; i++)
|
||||
distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);
|
||||
#endif
|
||||
|
||||
return (double) distance;
|
||||
PG_RETURN_FLOAT8((double) HalfvecL2SquaredDistance(a->dim, a->x, b->x));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -935,7 +564,7 @@ halfvec_inner_product(PG_FUNCTION_ARGS)
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(inner_product_internal(a, b));
|
||||
PG_RETURN_FLOAT8((double) HalfvecInnerProduct(a->dim, a->x, b->x));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -950,7 +579,7 @@ halfvec_negative_inner_product(PG_FUNCTION_ARGS)
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(-inner_product_internal(a, b));
|
||||
PG_RETURN_FLOAT8((double) -HalfvecInnerProduct(a->dim, a->x, b->x));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -1000,6 +629,32 @@ halfvec_cosine_distance(PG_FUNCTION_ARGS)
|
||||
PG_RETURN_FLOAT8(1 - similarity);
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the distance for spherical k-means
|
||||
* Currently uses angular distance since needs to satisfy triangle inequality
|
||||
* Assumes inputs are unit vectors (skips norm)
|
||||
*/
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_spherical_distance);
|
||||
Datum
|
||||
halfvec_spherical_distance(PG_FUNCTION_ARGS)
|
||||
{
|
||||
HalfVector *a = PG_GETARG_HALFVEC_P(0);
|
||||
HalfVector *b = PG_GETARG_HALFVEC_P(1);
|
||||
double distance;
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
distance = (double) HalfvecInnerProduct(a->dim, a->x, b->x);
|
||||
|
||||
/* Prevent NaN with acos with loss of precision */
|
||||
if (distance > 1)
|
||||
distance = 1;
|
||||
else if (distance < -1)
|
||||
distance = -1;
|
||||
|
||||
PG_RETURN_FLOAT8(acos(distance) / M_PI);
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the L1 distance between two half vectors
|
||||
*/
|
||||
@@ -1093,3 +748,30 @@ halfvec_subvector(PG_FUNCTION_ARGS)
|
||||
|
||||
PG_RETURN_POINTER(result);
|
||||
}
|
||||
|
||||
/*
|
||||
* Internal helper to compare half vectors
|
||||
*/
|
||||
int
|
||||
halfvec_cmp_internal(HalfVector * a, HalfVector * b)
|
||||
{
|
||||
int dim = Min(a->dim, b->dim);
|
||||
|
||||
/* Check values before dimensions to be consistent with Postgres arrays */
|
||||
for (int i = 0; i < dim; i++)
|
||||
{
|
||||
if (HalfToFloat4(a->x[i]) < HalfToFloat4(b->x[i]))
|
||||
return -1;
|
||||
|
||||
if (HalfToFloat4(a->x[i]) > HalfToFloat4(b->x[i]))
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (a->dim < b->dim)
|
||||
return -1;
|
||||
|
||||
if (a->dim > b->dim)
|
||||
return 1;
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
@@ -7,10 +7,14 @@
|
||||
|
||||
#include "vector.h"
|
||||
|
||||
#if defined(__x86_64__) || defined(_M_AMD64)
|
||||
#define HALFVEC_DISPATCH
|
||||
#endif
|
||||
|
||||
/* F16C has better performance than _Float16 (on x86-64) */
|
||||
#if defined(__F16C__) || defined(_MSC_VER)
|
||||
#if defined(__F16C__)
|
||||
#define F16C_SUPPORT
|
||||
#elif defined(__FLT16_MAX__)
|
||||
#elif defined(__FLT16_MAX__) && !defined(HALFVEC_DISPATCH)
|
||||
#define FLT16_SUPPORT
|
||||
#endif
|
||||
|
||||
@@ -38,8 +42,6 @@ typedef struct HalfVector
|
||||
} HalfVector;
|
||||
|
||||
HalfVector *InitHalfVector(int dim);
|
||||
float HalfToFloat4(half num);
|
||||
half Float4ToHalf(float num);
|
||||
half Float4ToHalfUnchecked(float num);
|
||||
int halfvec_cmp_internal(HalfVector * a, HalfVector * b);
|
||||
|
||||
#endif
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
#include "access/generic_xlog.h"
|
||||
#include "catalog/pg_type.h"
|
||||
#include "catalog/pg_type_d.h"
|
||||
#include "halfutils.h"
|
||||
#include "halfvec.h"
|
||||
#include "hnsw.h"
|
||||
#include "lib/pairingheap.h"
|
||||
@@ -222,10 +223,7 @@ HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, HnswType type)
|
||||
HalfVector *result = InitHalfVector(v->dim);
|
||||
|
||||
for (int i = 0; i < v->dim; i++)
|
||||
{
|
||||
/* TODO Fix */
|
||||
result->x[i] = Float4ToHalfUnchecked(HalfToFloat4(v->x[i]) / norm);
|
||||
}
|
||||
|
||||
*value = PointerGetDatum(result);
|
||||
}
|
||||
|
||||
@@ -10,12 +10,14 @@
|
||||
#include "catalog/pg_operator_d.h"
|
||||
#include "catalog/pg_type_d.h"
|
||||
#include "commands/progress.h"
|
||||
#include "halfvec.h"
|
||||
#include "ivfflat.h"
|
||||
#include "miscadmin.h"
|
||||
#include "optimizer/optimizer.h"
|
||||
#include "storage/bufmgr.h"
|
||||
#include "tcop/tcopprot.h"
|
||||
#include "utils/memutils.h"
|
||||
#include "vector.h"
|
||||
|
||||
#if PG_VERSION_NUM >= 140000
|
||||
#include "utils/backend_progress.h"
|
||||
@@ -63,7 +65,7 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
|
||||
|
||||
if (samples->length < targsamples)
|
||||
{
|
||||
VectorArraySet(samples, samples->length, DatumGetVector(value));
|
||||
VectorArraySet(samples, samples->length, DatumGetPointer(value));
|
||||
samples->length++;
|
||||
}
|
||||
else
|
||||
@@ -80,7 +82,7 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
|
||||
#endif
|
||||
|
||||
Assert(k >= 0 && k < targsamples);
|
||||
VectorArraySet(samples, k, DatumGetVector(value));
|
||||
VectorArraySet(samples, k, DatumGetPointer(value));
|
||||
}
|
||||
|
||||
buildstate->rowstoskip -= 1;
|
||||
@@ -318,7 +320,26 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
|
||||
static int
|
||||
GetMaxDimensions(IvfflatType type)
|
||||
{
|
||||
return IVFFLAT_MAX_DIM;
|
||||
int maxDimensions = IVFFLAT_MAX_DIM;
|
||||
|
||||
if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
maxDimensions *= 2;
|
||||
|
||||
return maxDimensions;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get item size
|
||||
*/
|
||||
static Size
|
||||
GetItemSize(IvfflatType type, int dimensions)
|
||||
{
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
return VECTOR_SIZE(dimensions);
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
return HALFVEC_SIZE(dimensions);
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -367,7 +388,7 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
|
||||
|
||||
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual);
|
||||
|
||||
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions);
|
||||
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, GetItemSize(buildstate->type, buildstate->dimensions));
|
||||
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
|
||||
|
||||
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
|
||||
@@ -422,7 +443,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
|
||||
|
||||
/* Sample rows */
|
||||
/* TODO Ensure within maintenance_work_mem */
|
||||
buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions);
|
||||
buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions, buildstate->centers->itemsize);
|
||||
if (buildstate->heap != NULL)
|
||||
{
|
||||
SampleRows(buildstate);
|
||||
@@ -437,7 +458,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
|
||||
}
|
||||
|
||||
/* Calculate centers */
|
||||
IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers));
|
||||
IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers, buildstate->type));
|
||||
|
||||
/* Free samples before we allocate more memory */
|
||||
VectorArrayFree(buildstate->samples);
|
||||
@@ -482,7 +503,7 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
|
||||
Size listSize;
|
||||
IvfflatList list;
|
||||
|
||||
listSize = MAXALIGN(IVFFLAT_LIST_SIZE(dimensions));
|
||||
listSize = MAXALIGN(IVFFLAT_LIST_SIZE(centers->itemsize));
|
||||
list = palloc0(listSize);
|
||||
|
||||
buf = IvfflatNewBuffer(index, forkNum);
|
||||
@@ -495,7 +516,7 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
|
||||
/* Load list */
|
||||
list->startPage = InvalidBlockNumber;
|
||||
list->insertPage = InvalidBlockNumber;
|
||||
memcpy(&list->center, VectorArrayGet(centers, i), VECTOR_SIZE(dimensions));
|
||||
memcpy(&list->center, VectorArrayGet(centers, i), centers->itemsize);
|
||||
|
||||
/* Ensure free space */
|
||||
if (PageGetFreeSpace(page) < listSize)
|
||||
@@ -600,7 +621,7 @@ ParallelHeapScan(IvfflatBuildState * buildstate)
|
||||
* Perform a worker's portion of a parallel sort
|
||||
*/
|
||||
static void
|
||||
IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, Sharedsort *sharedsort, Vector * ivfcenters, int sortmem, bool progress)
|
||||
IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, Sharedsort *sharedsort, char *ivfcenters, int sortmem, bool progress)
|
||||
{
|
||||
SortCoordinate coordinate;
|
||||
IvfflatBuildState buildstate;
|
||||
@@ -624,7 +645,7 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
|
||||
indexInfo = BuildIndexInfo(ivfspool->index);
|
||||
indexInfo->ii_Concurrent = ivfshared->isconcurrent;
|
||||
InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo);
|
||||
memcpy(buildstate.centers->items, ivfcenters, VECTOR_SIZE(buildstate.centers->dim) * buildstate.centers->maxlen);
|
||||
memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen);
|
||||
buildstate.centers->length = buildstate.centers->maxlen;
|
||||
ivfspool->sortstate = tuplesort_begin_heap(buildstate.tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, sortmem, coordinate, false);
|
||||
buildstate.sortstate = ivfspool->sortstate;
|
||||
@@ -672,7 +693,7 @@ IvfflatParallelBuildMain(dsm_segment *seg, shm_toc *toc)
|
||||
IvfflatSpool *ivfspool;
|
||||
IvfflatShared *ivfshared;
|
||||
Sharedsort *sharedsort;
|
||||
Vector *ivfcenters;
|
||||
char *ivfcenters;
|
||||
Relation heapRel;
|
||||
Relation indexRel;
|
||||
LOCKMODE heapLockmode;
|
||||
@@ -786,7 +807,7 @@ IvfflatBeginParallel(IvfflatBuildState * buildstate, bool isconcurrent, int requ
|
||||
Size estcenters;
|
||||
IvfflatShared *ivfshared;
|
||||
Sharedsort *sharedsort;
|
||||
Vector *ivfcenters;
|
||||
char *ivfcenters;
|
||||
IvfflatLeader *ivfleader = (IvfflatLeader *) palloc0(sizeof(IvfflatLeader));
|
||||
bool leaderparticipates = true;
|
||||
int querylen;
|
||||
@@ -813,7 +834,7 @@ IvfflatBeginParallel(IvfflatBuildState * buildstate, bool isconcurrent, int requ
|
||||
shm_toc_estimate_chunk(&pcxt->estimator, estivfshared);
|
||||
estsort = tuplesort_estimate_shared(scantuplesortstates);
|
||||
shm_toc_estimate_chunk(&pcxt->estimator, estsort);
|
||||
estcenters = VECTOR_SIZE(buildstate->dimensions) * buildstate->lists;
|
||||
estcenters = buildstate->centers->itemsize * buildstate->centers->maxlen;
|
||||
shm_toc_estimate_chunk(&pcxt->estimator, estcenters);
|
||||
shm_toc_estimate_keys(&pcxt->estimator, 3);
|
||||
|
||||
@@ -865,7 +886,7 @@ IvfflatBeginParallel(IvfflatBuildState * buildstate, bool isconcurrent, int requ
|
||||
tuplesort_initialize_shared(sharedsort, scantuplesortstates,
|
||||
pcxt->seg);
|
||||
|
||||
ivfcenters = (Vector *) shm_toc_allocate(pcxt->toc, estcenters);
|
||||
ivfcenters = shm_toc_allocate(pcxt->toc, estcenters);
|
||||
memcpy(ivfcenters, buildstate->centers->items, estcenters);
|
||||
|
||||
shm_toc_insert(pcxt->toc, PARALLEL_KEY_IVFFLAT_SHARED, ivfshared);
|
||||
|
||||
@@ -45,7 +45,8 @@
|
||||
|
||||
typedef enum IvfflatType
|
||||
{
|
||||
IVFFLAT_TYPE_VECTOR
|
||||
IVFFLAT_TYPE_VECTOR,
|
||||
IVFFLAT_TYPE_HALFVEC
|
||||
} IvfflatType;
|
||||
|
||||
/* Build phases */
|
||||
@@ -54,7 +55,7 @@ typedef enum IvfflatType
|
||||
#define PROGRESS_IVFFLAT_PHASE_ASSIGN 3
|
||||
#define PROGRESS_IVFFLAT_PHASE_LOAD 4
|
||||
|
||||
#define IVFFLAT_LIST_SIZE(_dim) (offsetof(IvfflatListData, center) + VECTOR_SIZE(_dim))
|
||||
#define IVFFLAT_LIST_SIZE(size) (offsetof(IvfflatListData, center) + size)
|
||||
|
||||
#define IvfflatPageGetOpaque(page) ((IvfflatPageOpaque) PageGetSpecialPointer(page))
|
||||
#define IvfflatPageGetMeta(page) ((IvfflatMetaPageData *) PageGetContents(page))
|
||||
@@ -90,7 +91,8 @@ typedef struct VectorArrayData
|
||||
int length;
|
||||
int maxlen;
|
||||
int dim;
|
||||
Vector *items;
|
||||
Size itemsize;
|
||||
char *items;
|
||||
} VectorArrayData;
|
||||
|
||||
typedef VectorArrayData * VectorArray;
|
||||
@@ -149,7 +151,7 @@ typedef struct IvfflatLeader
|
||||
IvfflatShared *ivfshared;
|
||||
Sharedsort *sharedsort;
|
||||
Snapshot snapshot;
|
||||
Vector *ivfcenters;
|
||||
char *ivfcenters;
|
||||
} IvfflatLeader;
|
||||
|
||||
typedef struct IvfflatBuildState
|
||||
@@ -261,16 +263,15 @@ typedef struct IvfflatScanOpaqueData
|
||||
|
||||
typedef IvfflatScanOpaqueData * IvfflatScanOpaque;
|
||||
|
||||
#define VECTOR_ARRAY_SIZE(_length, _dim) (sizeof(VectorArrayData) + (_length) * VECTOR_SIZE(_dim))
|
||||
#define VECTOR_ARRAY_OFFSET(_arr, _offset) ((char*) (_arr)->items + (_offset) * VECTOR_SIZE((_arr)->dim))
|
||||
#define VectorArrayGet(_arr, _offset) ((Vector *) VECTOR_ARRAY_OFFSET(_arr, _offset))
|
||||
#define VectorArraySet(_arr, _offset, _val) memcpy(VECTOR_ARRAY_OFFSET(_arr, _offset), _val, VECTOR_SIZE((_arr)->dim))
|
||||
#define VECTOR_ARRAY_SIZE(_length, _size) (sizeof(VectorArrayData) + (_length) * _size)
|
||||
#define VECTOR_ARRAY_OFFSET(_arr, _offset) ((char*) (_arr)->items + (_offset) * (_arr)->itemsize)
|
||||
#define VectorArrayGet(_arr, _offset) VECTOR_ARRAY_OFFSET(_arr, _offset)
|
||||
#define VectorArraySet(_arr, _offset, _val) memcpy(VECTOR_ARRAY_OFFSET(_arr, _offset), _val, (_arr)->itemsize)
|
||||
|
||||
/* Methods */
|
||||
VectorArray VectorArrayInit(int maxlen, int dimensions);
|
||||
VectorArray VectorArrayInit(int maxlen, int dimensions, Size itemsize);
|
||||
void VectorArrayFree(VectorArray arr);
|
||||
void PrintVectorArray(char *msg, VectorArray arr);
|
||||
void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);
|
||||
void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatType type);
|
||||
FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum);
|
||||
IvfflatType IvfflatGetType(Relation index);
|
||||
bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, IvfflatType type);
|
||||
|
||||
400
src/ivfkmeans.c
400
src/ivfkmeans.c
@@ -3,12 +3,13 @@
|
||||
#include <float.h>
|
||||
#include <math.h>
|
||||
|
||||
#include "halfutils.h"
|
||||
#include "halfvec.h"
|
||||
#include "ivfflat.h"
|
||||
#include "miscadmin.h"
|
||||
|
||||
#ifdef IVFFLAT_MEMORY
|
||||
#include "utils/datum.h"
|
||||
#include "utils/memutils.h"
|
||||
#endif
|
||||
#include "vector.h"
|
||||
|
||||
/*
|
||||
* Initialize with kmeans++
|
||||
@@ -46,12 +47,12 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
|
||||
|
||||
for (j = 0; j < numSamples; j++)
|
||||
{
|
||||
Vector *vec = VectorArrayGet(samples, j);
|
||||
Datum vec = PointerGetDatum(VectorArrayGet(samples, j));
|
||||
double distance;
|
||||
|
||||
/* Only need to compute distance for new center */
|
||||
/* TODO Use triangle inequality to reduce distance calculations */
|
||||
distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, i))));
|
||||
distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, vec, PointerGetDatum(VectorArrayGet(centers, i))));
|
||||
|
||||
/* Set lower bound */
|
||||
lowerBound[j * numCenters + i] = distance;
|
||||
@@ -89,15 +90,29 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
|
||||
* Apply norm to vector
|
||||
*/
|
||||
static inline void
|
||||
ApplyNorm(FmgrInfo *normprocinfo, Oid collation, Vector * vec)
|
||||
ApplyNorm(FmgrInfo *normprocinfo, Oid collation, Datum value, IvfflatType type)
|
||||
{
|
||||
double norm = DatumGetFloat8(FunctionCall1Coll(normprocinfo, collation, PointerGetDatum(vec)));
|
||||
double norm = DatumGetFloat8(FunctionCall1Coll(normprocinfo, collation, value));
|
||||
|
||||
/* TODO Handle zero norm */
|
||||
if (norm > 0)
|
||||
{
|
||||
for (int i = 0; i < vec->dim; i++)
|
||||
vec->x[i] /= norm;
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
{
|
||||
Vector *vec = DatumGetVector(value);
|
||||
|
||||
for (int i = 0; i < vec->dim; i++)
|
||||
vec->x[i] /= norm;
|
||||
}
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
{
|
||||
HalfVector *vec = DatumGetHalfVector(value);
|
||||
|
||||
for (int i = 0; i < vec->dim; i++)
|
||||
vec->x[i] = Float4ToHalfUnchecked(HalfToFloat4(vec->x[i]) / norm);
|
||||
}
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -110,11 +125,20 @@ CompareVectors(const void *a, const void *b)
|
||||
return vector_cmp_internal((Vector *) a, (Vector *) b);
|
||||
}
|
||||
|
||||
/*
|
||||
* Compare half vectors
|
||||
*/
|
||||
static int
|
||||
CompareHalfVectors(const void *a, const void *b)
|
||||
{
|
||||
return halfvec_cmp_internal((HalfVector *) a, (HalfVector *) b);
|
||||
}
|
||||
|
||||
/*
|
||||
* Quick approach if we have little data
|
||||
*/
|
||||
static void
|
||||
QuickCenters(Relation index, VectorArray samples, VectorArray centers)
|
||||
QuickCenters(Relation index, VectorArray samples, VectorArray centers, IvfflatType type)
|
||||
{
|
||||
int dimensions = centers->dim;
|
||||
Oid collation = index->rd_indcollation[0];
|
||||
@@ -123,14 +147,20 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
|
||||
/* Copy existing vectors while avoiding duplicates */
|
||||
if (samples->length > 0)
|
||||
{
|
||||
qsort(samples->items, samples->length, VECTOR_SIZE(samples->dim), CompareVectors);
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
qsort(samples->items, samples->length, samples->itemsize, CompareVectors);
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
qsort(samples->items, samples->length, samples->itemsize, CompareHalfVectors);
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
|
||||
for (int i = 0; i < samples->length; i++)
|
||||
{
|
||||
Vector *vec = VectorArrayGet(samples, i);
|
||||
Datum vec = PointerGetDatum(VectorArrayGet(samples, i));
|
||||
|
||||
if (i == 0 || CompareVectors(vec, VectorArrayGet(samples, i - 1)) != 0)
|
||||
if (i == 0 || !datumIsEqual(vec, PointerGetDatum(VectorArrayGet(samples, i - 1)), false, -1))
|
||||
{
|
||||
VectorArraySet(centers, centers->length, vec);
|
||||
VectorArraySet(centers, centers->length, DatumGetPointer(vec));
|
||||
centers->length++;
|
||||
}
|
||||
}
|
||||
@@ -139,17 +169,34 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
|
||||
/* Fill remaining with random data */
|
||||
while (centers->length < centers->maxlen)
|
||||
{
|
||||
Vector *vec = VectorArrayGet(centers, centers->length);
|
||||
Datum center = PointerGetDatum(VectorArrayGet(centers, centers->length));
|
||||
|
||||
SET_VARSIZE(vec, VECTOR_SIZE(dimensions));
|
||||
vec->dim = dimensions;
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
{
|
||||
Vector *vec = DatumGetVector(center);
|
||||
|
||||
for (int j = 0; j < dimensions; j++)
|
||||
vec->x[j] = RandomDouble();
|
||||
SET_VARSIZE(vec, VECTOR_SIZE(dimensions));
|
||||
vec->dim = dimensions;
|
||||
|
||||
for (int j = 0; j < dimensions; j++)
|
||||
vec->x[j] = RandomDouble();
|
||||
}
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
{
|
||||
HalfVector *vec = DatumGetHalfVector(center);
|
||||
|
||||
SET_VARSIZE(vec, HALFVEC_SIZE(dimensions));
|
||||
vec->dim = dimensions;
|
||||
|
||||
for (int j = 0; j < dimensions; j++)
|
||||
vec->x[j] = Float4ToHalfUnchecked((float) RandomDouble());
|
||||
}
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
|
||||
/* Normalize if needed (only needed for random centers) */
|
||||
if (normprocinfo != NULL)
|
||||
ApplyNorm(normprocinfo, collation, vec);
|
||||
ApplyNorm(normprocinfo, collation, center, type);
|
||||
|
||||
centers->length++;
|
||||
}
|
||||
@@ -160,18 +207,120 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
|
||||
* Show memory usage
|
||||
*/
|
||||
static void
|
||||
ShowMemoryUsage(Size estimatedSize)
|
||||
ShowMemoryUsage(MemoryContext context, Size estimatedSize)
|
||||
{
|
||||
#if PG_VERSION_NUM >= 130000
|
||||
elog(INFO, "total memory: %zu MB",
|
||||
MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
|
||||
MemoryContextMemAllocated(context, true) / (1024 * 1024));
|
||||
#else
|
||||
MemoryContextStats(CurrentMemoryContext);
|
||||
MemoryContextStats(context);
|
||||
#endif
|
||||
elog(INFO, "estimated memory: %zu MB", estimatedSize / (1024 * 1024));
|
||||
}
|
||||
#endif
|
||||
|
||||
/*
|
||||
* Compute new centers
|
||||
*/
|
||||
static void
|
||||
ComputeNewCenters(VectorArray samples, VectorArray aggCenters, VectorArray newCenters, int *centerCounts, int *closestCenters, FmgrInfo *normprocinfo, Oid collation, IvfflatType type)
|
||||
{
|
||||
int dimensions = aggCenters->dim;
|
||||
int numCenters = aggCenters->maxlen;
|
||||
int numSamples = samples->length;
|
||||
|
||||
/* Reset sum and count */
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
Vector *vec = (Vector *) VectorArrayGet(aggCenters, j);
|
||||
|
||||
for (int k = 0; k < dimensions; k++)
|
||||
vec->x[k] = 0.0;
|
||||
|
||||
centerCounts[j] = 0;
|
||||
}
|
||||
|
||||
/* Increment sum of closest center */
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
{
|
||||
for (int j = 0; j < numSamples; j++)
|
||||
{
|
||||
Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, closestCenters[j]);
|
||||
Vector *vec = (Vector *) VectorArrayGet(samples, j);
|
||||
|
||||
for (int k = 0; k < dimensions; k++)
|
||||
aggCenter->x[k] += vec->x[k];
|
||||
}
|
||||
}
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
{
|
||||
for (int j = 0; j < numSamples; j++)
|
||||
{
|
||||
Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, closestCenters[j]);
|
||||
HalfVector *vec = (HalfVector *) VectorArrayGet(samples, j);
|
||||
|
||||
for (int k = 0; k < dimensions; k++)
|
||||
aggCenter->x[k] += HalfToFloat4(vec->x[k]);
|
||||
}
|
||||
}
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
|
||||
/* Increment count of closest center */
|
||||
for (int j = 0; j < numSamples; j++)
|
||||
centerCounts[closestCenters[j]] += 1;
|
||||
|
||||
/* Divide sum by count */
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
Vector *vec = (Vector *) VectorArrayGet(aggCenters, j);
|
||||
|
||||
if (centerCounts[j] > 0)
|
||||
{
|
||||
/* Double avoids overflow, but requires more memory */
|
||||
/* TODO Update bounds */
|
||||
for (int k = 0; k < dimensions; k++)
|
||||
{
|
||||
if (isinf(vec->x[k]))
|
||||
vec->x[k] = vec->x[k] > 0 ? FLT_MAX : -FLT_MAX;
|
||||
}
|
||||
|
||||
for (int k = 0; k < dimensions; k++)
|
||||
vec->x[k] /= centerCounts[j];
|
||||
}
|
||||
else
|
||||
{
|
||||
/* TODO Handle empty centers properly */
|
||||
for (int k = 0; k < dimensions; k++)
|
||||
vec->x[k] = RandomDouble();
|
||||
}
|
||||
}
|
||||
|
||||
/* Set new centers if different from agg centers */
|
||||
if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
{
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, j);
|
||||
HalfVector *newCenter = (HalfVector *) VectorArrayGet(newCenters, j);
|
||||
|
||||
for (int k = 0; k < dimensions; k++)
|
||||
newCenter->x[k] = Float4ToHalfUnchecked(aggCenter->x[k]);
|
||||
}
|
||||
}
|
||||
|
||||
/* Normalize if needed */
|
||||
if (normprocinfo != NULL)
|
||||
{
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
Datum newCenter = PointerGetDatum(VectorArrayGet(newCenters, j));
|
||||
|
||||
ApplyNorm(normprocinfo, collation, newCenter, type);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
* Use Elkan for performance. This requires distance function to satisfy triangle inequality.
|
||||
*
|
||||
@@ -181,19 +330,16 @@ ShowMemoryUsage(Size estimatedSize)
|
||||
* https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf
|
||||
*/
|
||||
static void
|
||||
ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatType type)
|
||||
{
|
||||
FmgrInfo *procinfo;
|
||||
FmgrInfo *normprocinfo;
|
||||
Oid collation;
|
||||
Vector *vec;
|
||||
Vector *newCenter;
|
||||
int64 j;
|
||||
int64 k;
|
||||
int dimensions = centers->dim;
|
||||
int numCenters = centers->maxlen;
|
||||
int numSamples = samples->length;
|
||||
VectorArray newCenters;
|
||||
VectorArray aggCenters;
|
||||
int *centerCounts;
|
||||
int *closestCenters;
|
||||
float *lowerBound;
|
||||
@@ -201,11 +347,14 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
float *s;
|
||||
float *halfcdist;
|
||||
float *newcdist;
|
||||
MemoryContext kmeansCtx;
|
||||
MemoryContext oldCtx;
|
||||
|
||||
/* Calculate allocation sizes */
|
||||
Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->dim);
|
||||
Size centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->dim);
|
||||
Size newCentersSize = VECTOR_ARRAY_SIZE(numCenters, dimensions);
|
||||
Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->itemsize);
|
||||
Size centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->itemsize);
|
||||
Size newCentersSize = VECTOR_ARRAY_SIZE(numCenters, centers->itemsize);
|
||||
Size aggCentersSize = type == IVFFLAT_TYPE_VECTOR ? 0 : VECTOR_ARRAY_SIZE(numCenters, VECTOR_SIZE(dimensions));
|
||||
Size centerCountsSize = sizeof(int) * numCenters;
|
||||
Size closestCentersSize = sizeof(int) * numSamples;
|
||||
Size lowerBoundSize = sizeof(float) * numSamples * numCenters;
|
||||
@@ -215,7 +364,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
Size newcdistSize = sizeof(float) * numCenters;
|
||||
|
||||
/* Calculate total size */
|
||||
Size totalSize = samplesSize + centersSize + newCentersSize + centerCountsSize + closestCentersSize + lowerBoundSize + upperBoundSize + sSize + halfcdistSize + newcdistSize;
|
||||
Size totalSize = samplesSize + centersSize + newCentersSize + aggCentersSize + centerCountsSize + closestCentersSize + lowerBoundSize + upperBoundSize + sSize + halfcdistSize + newcdistSize;
|
||||
|
||||
/* Check memory requirements */
|
||||
/* Add one to error message to ceil */
|
||||
@@ -234,6 +383,12 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
|
||||
collation = index->rd_indcollation[0];
|
||||
|
||||
/* Use memory context */
|
||||
kmeansCtx = AllocSetContextCreate(CurrentMemoryContext,
|
||||
"Ivfflat kmeans temporary context",
|
||||
ALLOCSET_DEFAULT_SIZES);
|
||||
oldCtx = MemoryContextSwitchTo(kmeansCtx);
|
||||
|
||||
/* Allocate space */
|
||||
/* Use float instead of double to save memory */
|
||||
centerCounts = palloc(centerCountsSize);
|
||||
@@ -244,29 +399,50 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
halfcdist = palloc_extended(halfcdistSize, MCXT_ALLOC_HUGE);
|
||||
newcdist = palloc(newcdistSize);
|
||||
|
||||
newCenters = VectorArrayInit(numCenters, dimensions);
|
||||
for (j = 0; j < numCenters; j++)
|
||||
aggCenters = VectorArrayInit(numCenters, dimensions, VECTOR_SIZE(dimensions));
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
vec = VectorArrayGet(newCenters, j);
|
||||
Vector *vec = (Vector *) VectorArrayGet(aggCenters, j);
|
||||
|
||||
SET_VARSIZE(vec, VECTOR_SIZE(dimensions));
|
||||
vec->dim = dimensions;
|
||||
}
|
||||
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
{
|
||||
/* Use same centers to save memory */
|
||||
newCenters = aggCenters;
|
||||
}
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
{
|
||||
newCenters = VectorArrayInit(numCenters, dimensions, centers->itemsize);
|
||||
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
HalfVector *vec = (HalfVector *) VectorArrayGet(newCenters, j);
|
||||
|
||||
SET_VARSIZE(vec, HALFVEC_SIZE(dimensions));
|
||||
vec->dim = dimensions;
|
||||
}
|
||||
}
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
|
||||
#ifdef IVFFLAT_MEMORY
|
||||
ShowMemoryUsage(totalSize);
|
||||
ShowMemoryUsage(oldCtx, totalSize);
|
||||
#endif
|
||||
|
||||
/* Pick initial centers */
|
||||
InitCenters(index, samples, centers, lowerBound);
|
||||
|
||||
/* Assign each x to its closest initial center c(x) = argmin d(x,c) */
|
||||
for (j = 0; j < numSamples; j++)
|
||||
for (int64 j = 0; j < numSamples; j++)
|
||||
{
|
||||
float minDistance = FLT_MAX;
|
||||
int closestCenter = 0;
|
||||
|
||||
/* Find closest center */
|
||||
for (k = 0; k < numCenters; k++)
|
||||
for (int64 k = 0; k < numCenters; k++)
|
||||
{
|
||||
/* TODO Use Lemma 1 in k-means++ initialization */
|
||||
float distance = lowerBound[j * numCenters + k];
|
||||
@@ -292,13 +468,13 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
CHECK_FOR_INTERRUPTS();
|
||||
|
||||
/* Step 1: For all centers, compute distance */
|
||||
for (j = 0; j < numCenters; j++)
|
||||
for (int64 j = 0; j < numCenters; j++)
|
||||
{
|
||||
vec = VectorArrayGet(centers, j);
|
||||
Datum vec = PointerGetDatum(VectorArrayGet(centers, j));
|
||||
|
||||
for (k = j + 1; k < numCenters; k++)
|
||||
for (int64 k = j + 1; k < numCenters; k++)
|
||||
{
|
||||
float distance = 0.5 * DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, k))));
|
||||
float distance = 0.5 * DatumGetFloat8(FunctionCall2Coll(procinfo, collation, vec, PointerGetDatum(VectorArrayGet(centers, k))));
|
||||
|
||||
halfcdist[j * numCenters + k] = distance;
|
||||
halfcdist[k * numCenters + j] = distance;
|
||||
@@ -306,11 +482,11 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
}
|
||||
|
||||
/* For all centers c, compute s(c) */
|
||||
for (j = 0; j < numCenters; j++)
|
||||
for (int64 j = 0; j < numCenters; j++)
|
||||
{
|
||||
float minDistance = FLT_MAX;
|
||||
|
||||
for (k = 0; k < numCenters; k++)
|
||||
for (int64 k = 0; k < numCenters; k++)
|
||||
{
|
||||
float distance;
|
||||
|
||||
@@ -327,7 +503,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
|
||||
rjreset = iteration != 0;
|
||||
|
||||
for (j = 0; j < numSamples; j++)
|
||||
for (int64 j = 0; j < numSamples; j++)
|
||||
{
|
||||
bool rj;
|
||||
|
||||
@@ -337,8 +513,9 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
|
||||
rj = rjreset;
|
||||
|
||||
for (k = 0; k < numCenters; k++)
|
||||
for (int64 k = 0; k < numCenters; k++)
|
||||
{
|
||||
Datum vec;
|
||||
float dxcx;
|
||||
|
||||
/* Step 3: For all remaining points x and centers c */
|
||||
@@ -351,12 +528,12 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
if (upperBound[j] <= halfcdist[closestCenters[j] * numCenters + k])
|
||||
continue;
|
||||
|
||||
vec = VectorArrayGet(samples, j);
|
||||
vec = PointerGetDatum(VectorArrayGet(samples, j));
|
||||
|
||||
/* Step 3a */
|
||||
if (rj)
|
||||
{
|
||||
dxcx = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, closestCenters[j]))));
|
||||
dxcx = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, vec, PointerGetDatum(VectorArrayGet(centers, closestCenters[j]))));
|
||||
|
||||
/* d(x,c(x)) computed, which is a form of d(x,c) */
|
||||
lowerBound[j * numCenters + closestCenters[j]] = dxcx;
|
||||
@@ -370,7 +547,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
/* Step 3b */
|
||||
if (dxcx > lowerBound[j * numCenters + k] || dxcx > halfcdist[closestCenters[j] * numCenters + k])
|
||||
{
|
||||
float dxc = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, k))));
|
||||
float dxc = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, vec, PointerGetDatum(VectorArrayGet(centers, k))));
|
||||
|
||||
/* d(x,c) calculated */
|
||||
lowerBound[j * numCenters + k] = dxc;
|
||||
@@ -389,66 +566,15 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
}
|
||||
|
||||
/* Step 4: For each center c, let m(c) be mean of all points assigned */
|
||||
for (j = 0; j < numCenters; j++)
|
||||
{
|
||||
vec = VectorArrayGet(newCenters, j);
|
||||
for (k = 0; k < dimensions; k++)
|
||||
vec->x[k] = 0.0;
|
||||
|
||||
centerCounts[j] = 0;
|
||||
}
|
||||
|
||||
for (j = 0; j < numSamples; j++)
|
||||
{
|
||||
int closestCenter;
|
||||
|
||||
vec = VectorArrayGet(samples, j);
|
||||
closestCenter = closestCenters[j];
|
||||
|
||||
/* Increment sum and count of closest center */
|
||||
newCenter = VectorArrayGet(newCenters, closestCenter);
|
||||
for (k = 0; k < dimensions; k++)
|
||||
newCenter->x[k] += vec->x[k];
|
||||
|
||||
centerCounts[closestCenter] += 1;
|
||||
}
|
||||
|
||||
for (j = 0; j < numCenters; j++)
|
||||
{
|
||||
vec = VectorArrayGet(newCenters, j);
|
||||
|
||||
if (centerCounts[j] > 0)
|
||||
{
|
||||
/* Double avoids overflow, but requires more memory */
|
||||
/* TODO Update bounds */
|
||||
for (k = 0; k < dimensions; k++)
|
||||
{
|
||||
if (isinf(vec->x[k]))
|
||||
vec->x[k] = vec->x[k] > 0 ? FLT_MAX : -FLT_MAX;
|
||||
}
|
||||
|
||||
for (k = 0; k < dimensions; k++)
|
||||
vec->x[k] /= centerCounts[j];
|
||||
}
|
||||
else
|
||||
{
|
||||
/* TODO Handle empty centers properly */
|
||||
for (k = 0; k < dimensions; k++)
|
||||
vec->x[k] = RandomDouble();
|
||||
}
|
||||
|
||||
/* Normalize if needed */
|
||||
if (normprocinfo != NULL)
|
||||
ApplyNorm(normprocinfo, collation, vec);
|
||||
}
|
||||
ComputeNewCenters(samples, aggCenters, newCenters, centerCounts, closestCenters, normprocinfo, collation, type);
|
||||
|
||||
/* Step 5 */
|
||||
for (j = 0; j < numCenters; j++)
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
newcdist[j] = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(VectorArrayGet(centers, j)), PointerGetDatum(VectorArrayGet(newCenters, j))));
|
||||
|
||||
for (j = 0; j < numSamples; j++)
|
||||
for (int64 j = 0; j < numSamples; j++)
|
||||
{
|
||||
for (k = 0; k < numCenters; k++)
|
||||
for (int64 k = 0; k < numCenters; k++)
|
||||
{
|
||||
float distance = lowerBound[j * numCenters + k] - newcdist[k];
|
||||
|
||||
@@ -461,32 +587,26 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
|
||||
/* Step 6 */
|
||||
/* We reset r(x) before Step 3 in the next iteration */
|
||||
for (j = 0; j < numSamples; j++)
|
||||
for (int j = 0; j < numSamples; j++)
|
||||
upperBound[j] += newcdist[closestCenters[j]];
|
||||
|
||||
/* Step 7 */
|
||||
for (j = 0; j < numCenters; j++)
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
VectorArraySet(centers, j, VectorArrayGet(newCenters, j));
|
||||
|
||||
if (changes == 0 && iteration != 0)
|
||||
break;
|
||||
}
|
||||
|
||||
VectorArrayFree(newCenters);
|
||||
pfree(centerCounts);
|
||||
pfree(closestCenters);
|
||||
pfree(lowerBound);
|
||||
pfree(upperBound);
|
||||
pfree(s);
|
||||
pfree(halfcdist);
|
||||
pfree(newcdist);
|
||||
MemoryContextSwitchTo(oldCtx);
|
||||
MemoryContextDelete(kmeansCtx);
|
||||
}
|
||||
|
||||
/*
|
||||
* Detect issues with centers
|
||||
*/
|
||||
static void
|
||||
CheckCenters(Relation index, VectorArray centers)
|
||||
CheckCenters(Relation index, VectorArray centers, IvfflatType type)
|
||||
{
|
||||
FmgrInfo *normprocinfo;
|
||||
|
||||
@@ -496,24 +616,48 @@ CheckCenters(Relation index, VectorArray centers)
|
||||
/* Ensure no NaN or infinite values */
|
||||
for (int i = 0; i < centers->length; i++)
|
||||
{
|
||||
Vector *vec = VectorArrayGet(centers, i);
|
||||
|
||||
for (int j = 0; j < vec->dim; j++)
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
{
|
||||
if (isnan(vec->x[j]))
|
||||
elog(ERROR, "NaN detected. Please report a bug.");
|
||||
Vector *vec = (Vector *) VectorArrayGet(centers, i);
|
||||
|
||||
if (isinf(vec->x[j]))
|
||||
elog(ERROR, "Infinite value detected. Please report a bug.");
|
||||
for (int j = 0; j < vec->dim; j++)
|
||||
{
|
||||
if (isnan(vec->x[j]))
|
||||
elog(ERROR, "NaN detected. Please report a bug.");
|
||||
|
||||
if (isinf(vec->x[j]))
|
||||
elog(ERROR, "Infinite value detected. Please report a bug.");
|
||||
}
|
||||
}
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
{
|
||||
HalfVector *vec = (HalfVector *) VectorArrayGet(centers, i);
|
||||
|
||||
for (int j = 0; j < vec->dim; j++)
|
||||
{
|
||||
if (HalfIsNan(vec->x[j]))
|
||||
elog(ERROR, "NaN detected. Please report a bug.");
|
||||
|
||||
if (HalfIsInf(vec->x[j]))
|
||||
elog(ERROR, "Infinite value detected. Please report a bug.");
|
||||
}
|
||||
}
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
}
|
||||
|
||||
/* Ensure no duplicate centers */
|
||||
/* Fine to sort in-place */
|
||||
qsort(centers->items, centers->length, VECTOR_SIZE(centers->dim), CompareVectors);
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
qsort(centers->items, centers->length, centers->itemsize, CompareVectors);
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
qsort(centers->items, centers->length, centers->itemsize, CompareHalfVectors);
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
|
||||
for (int i = 1; i < centers->length; i++)
|
||||
{
|
||||
if (CompareVectors(VectorArrayGet(centers, i), VectorArrayGet(centers, i - 1)) == 0)
|
||||
if (datumIsEqual(PointerGetDatum(VectorArrayGet(centers, i)), PointerGetDatum(VectorArrayGet(centers, i - 1)), false, -1))
|
||||
elog(ERROR, "Duplicate centers detected. Please report a bug.");
|
||||
}
|
||||
|
||||
@@ -539,12 +683,12 @@ CheckCenters(Relation index, VectorArray centers)
|
||||
* We use spherical k-means for inner product and cosine
|
||||
*/
|
||||
void
|
||||
IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers)
|
||||
IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatType type)
|
||||
{
|
||||
if (samples->length <= centers->maxlen)
|
||||
QuickCenters(index, samples, centers);
|
||||
QuickCenters(index, samples, centers, type);
|
||||
else
|
||||
ElkanKmeans(index, samples, centers);
|
||||
ElkanKmeans(index, samples, centers, type);
|
||||
|
||||
CheckCenters(index, centers);
|
||||
CheckCenters(index, centers, type);
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
#include "access/relscan.h"
|
||||
#include "catalog/pg_operator_d.h"
|
||||
#include "catalog/pg_type_d.h"
|
||||
#include "halfvec.h"
|
||||
#include "lib/pairingheap.h"
|
||||
#include "ivfflat.h"
|
||||
#include "miscadmin.h"
|
||||
@@ -177,6 +178,42 @@ GetScanItems(IndexScanDesc scan, Datum value)
|
||||
tuplesort_performsort(so->sortstate);
|
||||
}
|
||||
|
||||
/*
|
||||
* Get scan value
|
||||
*/
|
||||
static Datum
|
||||
GetScanValue(IndexScanDesc scan)
|
||||
{
|
||||
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
|
||||
Datum value;
|
||||
|
||||
if (scan->orderByData->sk_flags & SK_ISNULL)
|
||||
{
|
||||
IvfflatType type = IvfflatGetType(scan->indexRelation);
|
||||
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
value = PointerGetDatum(InitVector(so->dimensions));
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
value = PointerGetDatum(InitHalfVector(so->dimensions));
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
}
|
||||
else
|
||||
{
|
||||
value = scan->orderByData->sk_argument;
|
||||
|
||||
/* Value should not be compressed or toasted */
|
||||
Assert(!VARATT_IS_COMPRESSED(DatumGetPointer(value)));
|
||||
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
|
||||
|
||||
/* Fine if normalization fails */
|
||||
if (so->normprocinfo != NULL)
|
||||
IvfflatNormValue(so->normprocinfo, so->collation, &value, IvfflatGetType(scan->indexRelation));
|
||||
}
|
||||
|
||||
return value;
|
||||
}
|
||||
|
||||
/*
|
||||
* Prepare for an index scan
|
||||
*/
|
||||
@@ -268,7 +305,6 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
if (so->first)
|
||||
{
|
||||
Datum value;
|
||||
IvfflatType type = IvfflatGetType(scan->indexRelation);
|
||||
|
||||
/* Count index scan for stats */
|
||||
pgstat_count_index_scan(scan->indexRelation);
|
||||
@@ -282,26 +318,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
if (!IsMVCCSnapshot(scan->xs_snapshot))
|
||||
elog(ERROR, "non-MVCC snapshots are not supported with ivfflat");
|
||||
|
||||
if (scan->orderByData->sk_flags & SK_ISNULL)
|
||||
{
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
value = PointerGetDatum(InitVector(so->dimensions));
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
}
|
||||
else
|
||||
{
|
||||
value = scan->orderByData->sk_argument;
|
||||
|
||||
/* Value should not be compressed or toasted */
|
||||
Assert(!VARATT_IS_COMPRESSED(DatumGetPointer(value)));
|
||||
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
|
||||
|
||||
/* Fine if normalization fails */
|
||||
if (so->normprocinfo != NULL)
|
||||
IvfflatNormValue(so->normprocinfo, so->collation, &value, type);
|
||||
}
|
||||
|
||||
value = GetScanValue(scan);
|
||||
IvfflatBench("GetScanLists", GetScanLists(scan, value));
|
||||
IvfflatBench("GetScanItems", GetScanItems(scan, value));
|
||||
so->first = false;
|
||||
|
||||
@@ -1,22 +1,27 @@
|
||||
#include "postgres.h"
|
||||
|
||||
#include "access/generic_xlog.h"
|
||||
#include "catalog/pg_type.h"
|
||||
#include "halfutils.h"
|
||||
#include "halfvec.h"
|
||||
#include "ivfflat.h"
|
||||
#include "storage/bufmgr.h"
|
||||
#include "vector.h"
|
||||
#include "utils/syscache.h"
|
||||
|
||||
/*
|
||||
* Allocate a vector array
|
||||
*/
|
||||
VectorArray
|
||||
VectorArrayInit(int maxlen, int dimensions)
|
||||
VectorArrayInit(int maxlen, int dimensions, Size itemsize)
|
||||
{
|
||||
VectorArray res = palloc(sizeof(VectorArrayData));
|
||||
|
||||
res->length = 0;
|
||||
res->maxlen = maxlen;
|
||||
res->dim = dimensions;
|
||||
res->items = palloc_extended(maxlen * VECTOR_SIZE(dimensions), MCXT_ALLOC_ZERO | MCXT_ALLOC_HUGE);
|
||||
res->itemsize = itemsize;
|
||||
res->items = palloc_extended(maxlen * itemsize, MCXT_ALLOC_ZERO | MCXT_ALLOC_HUGE);
|
||||
return res;
|
||||
}
|
||||
|
||||
@@ -30,16 +35,6 @@ VectorArrayFree(VectorArray arr)
|
||||
pfree(arr);
|
||||
}
|
||||
|
||||
/*
|
||||
* Print vector array - useful for debugging
|
||||
*/
|
||||
void
|
||||
PrintVectorArray(char *msg, VectorArray arr)
|
||||
{
|
||||
for (int i = 0; i < arr->length; i++)
|
||||
PrintVector(msg, VectorArrayGet(arr, i));
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the number of lists in the index
|
||||
*/
|
||||
@@ -72,7 +67,29 @@ IvfflatOptionalProcInfo(Relation index, uint16 procnum)
|
||||
IvfflatType
|
||||
IvfflatGetType(Relation index)
|
||||
{
|
||||
return IVFFLAT_TYPE_VECTOR;
|
||||
Oid typid = TupleDescAttr(index->rd_att, 0)->atttypid;
|
||||
HeapTuple tuple;
|
||||
Form_pg_type type;
|
||||
IvfflatType result;
|
||||
|
||||
tuple = SearchSysCache1(TYPEOID, ObjectIdGetDatum(typid));
|
||||
if (!HeapTupleIsValid(tuple))
|
||||
elog(ERROR, "cache lookup failed for type %u", typid);
|
||||
|
||||
type = (Form_pg_type) GETSTRUCT(tuple);
|
||||
if (strcmp(NameStr(type->typname), "vector") == 0)
|
||||
result = IVFFLAT_TYPE_VECTOR;
|
||||
else if (strcmp(NameStr(type->typname), "halfvec") == 0)
|
||||
result = IVFFLAT_TYPE_HALFVEC;
|
||||
else
|
||||
{
|
||||
ReleaseSysCache(tuple);
|
||||
elog(ERROR, "type not supported for ivfflat index");
|
||||
}
|
||||
|
||||
ReleaseSysCache(tuple);
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -100,6 +117,16 @@ IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, IvfflatType ty
|
||||
|
||||
*value = PointerGetDatum(result);
|
||||
}
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
{
|
||||
HalfVector *v = DatumGetHalfVector(*value);
|
||||
HalfVector *result = InitHalfVector(v->dim);
|
||||
|
||||
for (int i = 0; i < v->dim; i++)
|
||||
result->x[i] = Float4ToHalfUnchecked(HalfToFloat4(v->x[i]) / norm);
|
||||
|
||||
*value = PointerGetDatum(result);
|
||||
}
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
|
||||
|
||||
218
src/sparsevec.c
218
src/sparsevec.c
@@ -174,19 +174,16 @@ sparsevec_in(PG_FUNCTION_ARGS)
|
||||
char *lit = PG_GETARG_CSTRING(0);
|
||||
int32 typmod = PG_GETARG_INT32(2);
|
||||
int dim;
|
||||
char *pt;
|
||||
char *pt = lit;
|
||||
char *stringEnd;
|
||||
SparseVector *result;
|
||||
float *rvalues;
|
||||
char *litcopy = pstrdup(lit);
|
||||
char *str = litcopy;
|
||||
int32 *indices;
|
||||
float *values;
|
||||
int maxNnz;
|
||||
int nnz = 0;
|
||||
|
||||
maxNnz = 1;
|
||||
pt = str;
|
||||
while (*pt != '\0')
|
||||
{
|
||||
if (*pt == ',')
|
||||
@@ -203,121 +200,137 @@ sparsevec_in(PG_FUNCTION_ARGS)
|
||||
indices = palloc(maxNnz * sizeof(int32));
|
||||
values = palloc(maxNnz * sizeof(float));
|
||||
|
||||
while (sparsevec_isspace(*str))
|
||||
str++;
|
||||
pt = lit;
|
||||
|
||||
if (*str != '{')
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*pt != '{')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed sparsevec literal: \"%s\"", lit),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit),
|
||||
errdetail("Vector contents must start with \"{\".")));
|
||||
|
||||
str++;
|
||||
pt = strtok(str, ",");
|
||||
stringEnd = pt;
|
||||
pt++;
|
||||
|
||||
while (pt != NULL && *stringEnd != '}')
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*pt == '}')
|
||||
pt++;
|
||||
else
|
||||
{
|
||||
long index;
|
||||
float value;
|
||||
for (;;)
|
||||
{
|
||||
long index;
|
||||
float value;
|
||||
|
||||
/* TODO Better error */
|
||||
if (nnz == maxNnz)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("ran out of buffer: \"%s\"", lit)));
|
||||
/* TODO Better error */
|
||||
if (nnz == maxNnz)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("ran out of buffer: \"%s\"", lit)));
|
||||
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
/* Check for empty string like float4in */
|
||||
if (*pt == '\0')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
|
||||
/* Use similar logic as int2vectorin */
|
||||
errno = 0;
|
||||
index = strtol(pt, &stringEnd, 10);
|
||||
|
||||
if (stringEnd == pt)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
|
||||
if (errno == ERANGE || index < 1 || index > INT_MAX)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("index \"%ld\" is out of range for type sparsevec", index)));
|
||||
|
||||
pt = stringEnd;
|
||||
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*pt != ':')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
/* Check for empty string like float4in */
|
||||
if (*pt == '\0')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
/* Use similar logic as int2vectorin */
|
||||
errno = 0;
|
||||
index = strtol(pt, &stringEnd, 10);
|
||||
errno = 0;
|
||||
|
||||
if (stringEnd == pt)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
/* Use strtof like float4in to avoid a double-rounding problem */
|
||||
/* Postgres sets LC_NUMERIC to C on startup */
|
||||
value = strtof(pt, &stringEnd);
|
||||
|
||||
if (errno == ERANGE || index < 1 || index > INT_MAX)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("index \"%ld\" is out of range for type sparsevec", index)));
|
||||
if (stringEnd == pt)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
|
||||
if (stringEnd == pt)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
/* Check for range error like float4in */
|
||||
if (errno == ERANGE && (value == 0 || isinf(value)))
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("\"%s\" is out of range for type sparsevec", pnstrdup(pt, stringEnd - pt))));
|
||||
|
||||
while (sparsevec_isspace(*stringEnd))
|
||||
stringEnd++;
|
||||
CheckElement(value);
|
||||
|
||||
if (*stringEnd != ':')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
/* Do not store zero values */
|
||||
if (value != 0)
|
||||
{
|
||||
indices[nnz] = index;
|
||||
values[nnz] = value;
|
||||
nnz++;
|
||||
}
|
||||
|
||||
stringEnd++;
|
||||
pt = stringEnd;
|
||||
|
||||
while (sparsevec_isspace(*stringEnd))
|
||||
stringEnd++;
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
errno = 0;
|
||||
pt = stringEnd;
|
||||
value = strtof(pt, &stringEnd);
|
||||
|
||||
if (stringEnd == pt)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
|
||||
/* Check for range error like float4in */
|
||||
if (errno == ERANGE && (value == 0 || isinf(value)))
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("\"%s\" is out of range for type sparsevec", pt)));
|
||||
|
||||
/* TODO Decide whether to store zero values */
|
||||
if (value != 0)
|
||||
{
|
||||
indices[nnz] = index;
|
||||
values[nnz] = value;
|
||||
nnz++;
|
||||
if (*pt == ',')
|
||||
pt++;
|
||||
else if (*pt == '}')
|
||||
{
|
||||
pt++;
|
||||
break;
|
||||
}
|
||||
else
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
}
|
||||
|
||||
if (*stringEnd != '\0' && *stringEnd != '}')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
|
||||
pt = strtok(NULL, ",");
|
||||
}
|
||||
|
||||
if (stringEnd == NULL || *stringEnd != '}')
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*pt != '/')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed sparsevec literal: \"%s\"", lit),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit),
|
||||
errdetail("Unexpected end of input.")));
|
||||
|
||||
stringEnd++;
|
||||
pt++;
|
||||
|
||||
if (*stringEnd != '/')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed sparsevec literal: \"%s\"", lit),
|
||||
errdetail("Unexpected end of input.")));
|
||||
|
||||
stringEnd++;
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
/* Use similar logic as int2vectorin */
|
||||
errno = 0;
|
||||
pt = stringEnd;
|
||||
dim = strtol(pt, &stringEnd, 10);
|
||||
|
||||
if (stringEnd == pt)
|
||||
@@ -325,18 +338,18 @@ sparsevec_in(PG_FUNCTION_ARGS)
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit)));
|
||||
|
||||
/* Only whitespace is allowed after the closing brace */
|
||||
while (sparsevec_isspace(*stringEnd))
|
||||
stringEnd++;
|
||||
pt = stringEnd;
|
||||
|
||||
if (*stringEnd != '\0')
|
||||
/* Only whitespace is allowed after the closing brace */
|
||||
while (sparsevec_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*pt != '\0')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed sparsevec literal: \"%s\"", lit),
|
||||
errmsg("invalid input syntax for type sparsevec: \"%s\"", lit),
|
||||
errdetail("Junk after closing.")));
|
||||
|
||||
pfree(litcopy);
|
||||
|
||||
CheckDim(dim);
|
||||
CheckExpectedDim(typmod, dim);
|
||||
|
||||
@@ -348,7 +361,6 @@ sparsevec_in(PG_FUNCTION_ARGS)
|
||||
rvalues[i] = values[i];
|
||||
|
||||
CheckIndex(result->indices, i, dim);
|
||||
CheckElement(rvalues[i]);
|
||||
}
|
||||
|
||||
PG_RETURN_POINTER(result);
|
||||
@@ -583,7 +595,7 @@ vector_to_sparsevec(PG_FUNCTION_ARGS)
|
||||
* Get the L2 squared distance between sparse vectors
|
||||
*/
|
||||
static double
|
||||
l2_distance_squared_internal(SparseVector * a, SparseVector * b)
|
||||
SparsevecL2SquaredDistance(SparseVector * a, SparseVector * b)
|
||||
{
|
||||
float *ax = SPARSEVEC_VALUES(a);
|
||||
float *bx = SPARSEVEC_VALUES(b);
|
||||
@@ -639,7 +651,7 @@ sparsevec_l2_distance(PG_FUNCTION_ARGS)
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(sqrt(l2_distance_squared_internal(a, b)));
|
||||
PG_RETURN_FLOAT8(sqrt(SparsevecL2SquaredDistance(a, b)));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -655,14 +667,14 @@ sparsevec_l2_squared_distance(PG_FUNCTION_ARGS)
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(l2_distance_squared_internal(a, b));
|
||||
PG_RETURN_FLOAT8(SparsevecL2SquaredDistance(a, b));
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the inner product of two sparse vectors
|
||||
*/
|
||||
static double
|
||||
inner_product_internal(SparseVector * a, SparseVector * b)
|
||||
SparsevecInnerProduct(SparseVector * a, SparseVector * b)
|
||||
{
|
||||
float *ax = SPARSEVEC_VALUES(a);
|
||||
float *bx = SPARSEVEC_VALUES(b);
|
||||
@@ -706,7 +718,7 @@ sparsevec_inner_product(PG_FUNCTION_ARGS)
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(inner_product_internal(a, b));
|
||||
PG_RETURN_FLOAT8(SparsevecInnerProduct(a, b));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -721,7 +733,7 @@ sparsevec_negative_inner_product(PG_FUNCTION_ARGS)
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
PG_RETURN_FLOAT8(-inner_product_internal(a, b));
|
||||
PG_RETURN_FLOAT8(-SparsevecInnerProduct(a, b));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -741,7 +753,7 @@ sparsevec_cosine_distance(PG_FUNCTION_ARGS)
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
similarity = inner_product_internal(a, b);
|
||||
similarity = SparsevecInnerProduct(a, b);
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < a->nnz; i++)
|
||||
|
||||
172
src/vector.c
172
src/vector.c
@@ -6,6 +6,7 @@
|
||||
#include "catalog/pg_type.h"
|
||||
#include "common/shortest_dec.h"
|
||||
#include "fmgr.h"
|
||||
#include "halfutils.h"
|
||||
#include "halfvec.h"
|
||||
#include "hnsw.h"
|
||||
#include "ivfflat.h"
|
||||
@@ -32,6 +33,13 @@
|
||||
#define STATE_DIMS(x) (ARR_DIMS(x)[0] - 1)
|
||||
#define CreateStateDatums(dim) palloc(sizeof(Datum) * (dim + 1))
|
||||
|
||||
/* target_clones requires glibc */
|
||||
#if defined(__x86_64__) && defined(__gnu_linux__) && defined(__has_attribute) && __has_attribute(target_clones) && !defined(__FMA__)
|
||||
#define VECTOR_DISPATCH __attribute__((target_clones("default", "fma")))
|
||||
#else
|
||||
#define VECTOR_DISPATCH
|
||||
#endif
|
||||
|
||||
PG_MODULE_MAGIC;
|
||||
|
||||
/*
|
||||
@@ -41,6 +49,7 @@ PGDLLEXPORT void _PG_init(void);
|
||||
void
|
||||
_PG_init(void)
|
||||
{
|
||||
HalfvecInit();
|
||||
HnswInit();
|
||||
IvfflatInit();
|
||||
}
|
||||
@@ -179,28 +188,32 @@ vector_in(PG_FUNCTION_ARGS)
|
||||
int32 typmod = PG_GETARG_INT32(2);
|
||||
float x[VECTOR_MAX_DIM];
|
||||
int dim = 0;
|
||||
char *pt;
|
||||
char *stringEnd;
|
||||
char *pt = lit;
|
||||
Vector *result;
|
||||
char *litcopy = pstrdup(lit);
|
||||
char *str = litcopy;
|
||||
|
||||
while (vector_isspace(*str))
|
||||
str++;
|
||||
while (vector_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*str != '[')
|
||||
if (*pt != '[')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed vector literal: \"%s\"", lit),
|
||||
errmsg("invalid input syntax for type vector: \"%s\"", lit),
|
||||
errdetail("Vector contents must start with \"[\".")));
|
||||
|
||||
str++;
|
||||
pt = strtok(str, ",");
|
||||
stringEnd = pt;
|
||||
pt++;
|
||||
|
||||
while (pt != NULL && *stringEnd != ']')
|
||||
while (vector_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*pt == ']')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("vector must have at least 1 dimension")));
|
||||
|
||||
for (;;)
|
||||
{
|
||||
float val;
|
||||
char *stringEnd;
|
||||
|
||||
if (dim == VECTOR_MAX_DIM)
|
||||
ereport(ERROR,
|
||||
@@ -216,8 +229,10 @@ vector_in(PG_FUNCTION_ARGS)
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type vector: \"%s\"", lit)));
|
||||
|
||||
/* Use strtof like float4in to avoid a double-rounding problem */
|
||||
errno = 0;
|
||||
|
||||
/* Use strtof like float4in to avoid a double-rounding problem */
|
||||
/* Postgres sets LC_NUMERIC to C on startup */
|
||||
val = strtof(pt, &stringEnd);
|
||||
|
||||
if (stringEnd == pt)
|
||||
@@ -229,56 +244,40 @@ vector_in(PG_FUNCTION_ARGS)
|
||||
if (errno == ERANGE && isinf(val))
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("\"%s\" is out of range for type vector", pt)));
|
||||
errmsg("\"%s\" is out of range for type vector", pnstrdup(pt, stringEnd - pt))));
|
||||
|
||||
CheckElement(val);
|
||||
x[dim++] = val;
|
||||
|
||||
while (vector_isspace(*stringEnd))
|
||||
stringEnd++;
|
||||
pt = stringEnd;
|
||||
|
||||
if (*stringEnd != '\0' && *stringEnd != ']')
|
||||
while (vector_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*pt == ',')
|
||||
pt++;
|
||||
else if (*pt == ']')
|
||||
{
|
||||
pt++;
|
||||
break;
|
||||
}
|
||||
else
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("invalid input syntax for type vector: \"%s\"", lit)));
|
||||
|
||||
pt = strtok(NULL, ",");
|
||||
}
|
||||
|
||||
if (stringEnd == NULL || *stringEnd != ']')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed vector literal: \"%s\"", lit),
|
||||
errdetail("Unexpected end of input.")));
|
||||
|
||||
stringEnd++;
|
||||
|
||||
/* Only whitespace is allowed after the closing brace */
|
||||
while (vector_isspace(*stringEnd))
|
||||
stringEnd++;
|
||||
while (vector_isspace(*pt))
|
||||
pt++;
|
||||
|
||||
if (*stringEnd != '\0')
|
||||
if (*pt != '\0')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed vector literal: \"%s\"", lit),
|
||||
errmsg("invalid input syntax for type vector: \"%s\"", lit),
|
||||
errdetail("Junk after closing right brace.")));
|
||||
|
||||
/* Ensure no consecutive delimiters since strtok skips */
|
||||
for (pt = lit + 1; *pt != '\0'; pt++)
|
||||
{
|
||||
if (pt[-1] == ',' && *pt == ',')
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
|
||||
errmsg("malformed vector literal: \"%s\"", lit)));
|
||||
}
|
||||
|
||||
if (dim < 1)
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_DATA_EXCEPTION),
|
||||
errmsg("vector must have at least 1 dimension")));
|
||||
|
||||
pfree(litcopy);
|
||||
|
||||
CheckDim(dim);
|
||||
CheckExpectedDim(typmod, dim);
|
||||
|
||||
result = InitVector(dim);
|
||||
@@ -565,6 +564,22 @@ halfvec_to_vector(PG_FUNCTION_ARGS)
|
||||
PG_RETURN_POINTER(result);
|
||||
}
|
||||
|
||||
VECTOR_DISPATCH static float
|
||||
VectorL2SquaredDistance(int dim, float *ax, float *bx)
|
||||
{
|
||||
float distance = 0.0;
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < dim; i++)
|
||||
{
|
||||
float diff = ax[i] - bx[i];
|
||||
|
||||
distance += diff * diff;
|
||||
}
|
||||
|
||||
return distance;
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the L2 distance between vectors
|
||||
*/
|
||||
@@ -574,21 +589,10 @@ l2_distance(PG_FUNCTION_ARGS)
|
||||
{
|
||||
Vector *a = PG_GETARG_VECTOR_P(0);
|
||||
Vector *b = PG_GETARG_VECTOR_P(1);
|
||||
float *ax = a->x;
|
||||
float *bx = b->x;
|
||||
float distance = 0.0;
|
||||
float diff;
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < a->dim; i++)
|
||||
{
|
||||
diff = ax[i] - bx[i];
|
||||
distance += diff * diff;
|
||||
}
|
||||
|
||||
PG_RETURN_FLOAT8(sqrt((double) distance));
|
||||
PG_RETURN_FLOAT8(sqrt((double) VectorL2SquaredDistance(a->dim, a->x, b->x)));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -601,21 +605,22 @@ vector_l2_squared_distance(PG_FUNCTION_ARGS)
|
||||
{
|
||||
Vector *a = PG_GETARG_VECTOR_P(0);
|
||||
Vector *b = PG_GETARG_VECTOR_P(1);
|
||||
float *ax = a->x;
|
||||
float *bx = b->x;
|
||||
float distance = 0.0;
|
||||
float diff;
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < a->dim; i++)
|
||||
{
|
||||
diff = ax[i] - bx[i];
|
||||
distance += diff * diff;
|
||||
}
|
||||
PG_RETURN_FLOAT8((double) VectorL2SquaredDistance(a->dim, a->x, b->x));
|
||||
}
|
||||
|
||||
PG_RETURN_FLOAT8((double) distance);
|
||||
VECTOR_DISPATCH static float
|
||||
VectorInnerProduct(int dim, float *ax, float *bx)
|
||||
{
|
||||
float distance = 0.0;
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < dim; i++)
|
||||
distance += ax[i] * bx[i];
|
||||
|
||||
return distance;
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -627,17 +632,10 @@ inner_product(PG_FUNCTION_ARGS)
|
||||
{
|
||||
Vector *a = PG_GETARG_VECTOR_P(0);
|
||||
Vector *b = PG_GETARG_VECTOR_P(1);
|
||||
float *ax = a->x;
|
||||
float *bx = b->x;
|
||||
float distance = 0.0;
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < a->dim; i++)
|
||||
distance += ax[i] * bx[i];
|
||||
|
||||
PG_RETURN_FLOAT8((double) distance);
|
||||
PG_RETURN_FLOAT8((double) VectorInnerProduct(a->dim, a->x, b->x));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -649,17 +647,10 @@ vector_negative_inner_product(PG_FUNCTION_ARGS)
|
||||
{
|
||||
Vector *a = PG_GETARG_VECTOR_P(0);
|
||||
Vector *b = PG_GETARG_VECTOR_P(1);
|
||||
float *ax = a->x;
|
||||
float *bx = b->x;
|
||||
float distance = 0.0;
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < a->dim; i++)
|
||||
distance += ax[i] * bx[i];
|
||||
|
||||
PG_RETURN_FLOAT8((double) distance * -1);
|
||||
PG_RETURN_FLOAT8((double) -VectorInnerProduct(a->dim, a->x, b->x));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -717,18 +708,11 @@ vector_spherical_distance(PG_FUNCTION_ARGS)
|
||||
{
|
||||
Vector *a = PG_GETARG_VECTOR_P(0);
|
||||
Vector *b = PG_GETARG_VECTOR_P(1);
|
||||
float *ax = a->x;
|
||||
float *bx = b->x;
|
||||
float dp = 0.0;
|
||||
double distance;
|
||||
|
||||
CheckDims(a, b);
|
||||
|
||||
/* Auto-vectorized */
|
||||
for (int i = 0; i < a->dim; i++)
|
||||
dp += ax[i] * bx[i];
|
||||
|
||||
distance = (double) dp;
|
||||
distance = (double) VectorInnerProduct(a->dim, a->x, b->x);
|
||||
|
||||
/* Prevent NaN with acos with loss of precision */
|
||||
if (distance > 1)
|
||||
|
||||
@@ -28,6 +28,26 @@ SELECT ARRAY[1,2,3]::numeric[]::vector;
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2,3]'::vector::real[];
|
||||
float4
|
||||
---------
|
||||
{1,2,3}
|
||||
(1 row)
|
||||
|
||||
SELECT '{1,2,3}'::real[]::vector;
|
||||
vector
|
||||
---------
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '{1,2,3}'::real[]::vector(3);
|
||||
vector
|
||||
---------
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '{1,2,3}'::real[]::vector(2);
|
||||
ERROR: expected 2 dimensions, not 3
|
||||
SELECT '{NULL}'::real[]::vector;
|
||||
ERROR: array must not contain nulls
|
||||
SELECT '{NaN}'::real[]::vector;
|
||||
@@ -40,10 +60,26 @@ SELECT '{}'::real[]::vector;
|
||||
ERROR: vector must have at least 1 dimension
|
||||
SELECT '{{1}}'::real[]::vector;
|
||||
ERROR: array must be 1-D
|
||||
SELECT '[1,2,3]'::vector::real[];
|
||||
float4
|
||||
SELECT '{1,2,3}'::double precision[]::vector;
|
||||
vector
|
||||
---------
|
||||
{1,2,3}
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '{1,2,3}'::double precision[]::vector(3);
|
||||
vector
|
||||
---------
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '{1,2,3}'::double precision[]::vector(2);
|
||||
ERROR: expected 2 dimensions, not 3
|
||||
SELECT '{4e38,-4e38}'::double precision[]::vector;
|
||||
ERROR: infinite value not allowed in vector
|
||||
SELECT '{1e-46,-1e-46}'::double precision[]::vector;
|
||||
vector
|
||||
--------
|
||||
[0,-0]
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2,3]'::vector::halfvec;
|
||||
@@ -52,24 +88,88 @@ SELECT '[1,2,3]'::vector::halfvec;
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2,3]'::halfvec::vector;
|
||||
vector
|
||||
SELECT '[1,2,3]'::vector::halfvec(3);
|
||||
halfvec
|
||||
---------
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2,3]'::vector::halfvec(2);
|
||||
ERROR: expected 2 dimensions, not 3
|
||||
SELECT '[1,2,3]'::halfvec::vector(2);
|
||||
ERROR: expected 2 dimensions, not 3
|
||||
SELECT '[65520]'::vector::halfvec;
|
||||
ERROR: infinite value not allowed in halfvec
|
||||
ERROR: "65520" is out of range for type halfvec
|
||||
SELECT '[1e-8]'::vector::halfvec;
|
||||
halfvec
|
||||
---------
|
||||
[0]
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2,3]'::halfvec::vector;
|
||||
vector
|
||||
---------
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2,3]'::halfvec::vector(3);
|
||||
vector
|
||||
---------
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2,3]'::halfvec::vector(2);
|
||||
ERROR: expected 2 dimensions, not 3
|
||||
SELECT '{1,2,3}'::real[]::halfvec;
|
||||
halfvec
|
||||
---------
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '{1,2,3}'::real[]::halfvec(3);
|
||||
halfvec
|
||||
---------
|
||||
[1,2,3]
|
||||
(1 row)
|
||||
|
||||
SELECT '{1,2,3}'::real[]::halfvec(2);
|
||||
ERROR: expected 2 dimensions, not 3
|
||||
SELECT '{65520,-65520}'::real[]::halfvec;
|
||||
ERROR: "65520" is out of range for type halfvec
|
||||
SELECT '{1e-8,-1e-8}'::real[]::halfvec;
|
||||
halfvec
|
||||
---------
|
||||
[0,-0]
|
||||
(1 row)
|
||||
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec;
|
||||
sparsevec
|
||||
-----------------
|
||||
{2:1.5,4:3.5}/5
|
||||
(1 row)
|
||||
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec(5);
|
||||
sparsevec
|
||||
-----------------
|
||||
{2:1.5,4:3.5}/5
|
||||
(1 row)
|
||||
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec(4);
|
||||
ERROR: expected 4 dimensions, not 5
|
||||
SELECT '{2:1.5,4:3.5}/5'::sparsevec::vector;
|
||||
vector
|
||||
-----------------
|
||||
[0,1.5,0,3.5,0]
|
||||
(1 row)
|
||||
|
||||
SELECT '{2:1.5,4:3.5}/5'::sparsevec::vector(5);
|
||||
vector
|
||||
-----------------
|
||||
[0,1.5,0,3.5,0]
|
||||
(1 row)
|
||||
|
||||
SELECT '{2:1.5,4:3.5}/5'::sparsevec::vector(4);
|
||||
ERROR: expected 4 dimensions, not 5
|
||||
SELECT '{}/16001'::sparsevec::vector;
|
||||
ERROR: vector cannot have more than 16000 dimensions
|
||||
SELECT array_agg(n)::vector FROM generate_series(1, 16001) n;
|
||||
ERROR: vector cannot have more than 16000 dimensions
|
||||
SELECT array_to_vector(array_agg(n), 16001, false) FROM generate_series(1, 16001) n;
|
||||
|
||||
@@ -1,15 +1,15 @@
|
||||
CREATE TABLE t (val vector(3), val2 halfvec(3));
|
||||
INSERT INTO t (val, val2) VALUES ('[0,0,0]', '[0,0,0]'), ('[1,2,3]', '[1,2,3]'), ('[1,1,1]', '[1,1,1]'), (NULL, NULL);
|
||||
CREATE TABLE t2 (val vector(3), val2 halfvec(3));
|
||||
CREATE TABLE t (val vector(3), val2 halfvec(3), val3 sparsevec(3));
|
||||
INSERT INTO t (val, val2, val3) VALUES ('[0,0,0]', '[0,0,0]', '{}/3'), ('[1,2,3]', '[1,2,3]', '{1:1,2:2,3:3}/3'), ('[1,1,1]', '[1,1,1]', '{1:1,2:1,3:1}/3'), (NULL, NULL, NULL);
|
||||
CREATE TABLE t2 (val vector(3), val2 halfvec(3), val3 sparsevec(3));
|
||||
\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 | val2
|
||||
---------+---------
|
||||
[0,0,0] | [0,0,0]
|
||||
[1,1,1] | [1,1,1]
|
||||
[1,2,3] | [1,2,3]
|
||||
|
|
||||
val | val2 | val3
|
||||
---------+---------+-----------------
|
||||
[0,0,0] | [0,0,0] | {}/3
|
||||
[1,1,1] | [1,1,1] | {1:1,2:1,3:1}/3
|
||||
[1,2,3] | [1,2,3] | {1:1,2:2,3:3}/3
|
||||
| |
|
||||
(4 rows)
|
||||
|
||||
DROP TABLE t;
|
||||
|
||||
@@ -1,3 +1,21 @@
|
||||
SELECT round(halfvec_norm('[1,1]')::numeric, 5);
|
||||
round
|
||||
---------
|
||||
1.41421
|
||||
(1 row)
|
||||
|
||||
SELECT halfvec_norm('[3,4]');
|
||||
halfvec_norm
|
||||
--------------
|
||||
5
|
||||
(1 row)
|
||||
|
||||
SELECT halfvec_norm('[0,1]');
|
||||
halfvec_norm
|
||||
--------------
|
||||
1
|
||||
(1 row)
|
||||
|
||||
SELECT l2_distance('[0,0]'::halfvec, '[3,4]');
|
||||
l2_distance
|
||||
-------------
|
||||
@@ -12,6 +30,12 @@ SELECT l2_distance('[0,0]'::halfvec, '[0,1]');
|
||||
|
||||
SELECT l2_distance('[1,2]'::halfvec, '[3]');
|
||||
ERROR: different halfvec dimensions 2 and 1
|
||||
SELECT l2_distance('[1,1,1,1,1,1,1,1,1]'::halfvec, '[1,1,1,1,1,1,1,4,5]');
|
||||
l2_distance
|
||||
-------------
|
||||
5
|
||||
(1 row)
|
||||
|
||||
SELECT '[0,0]'::halfvec <-> '[3,4]';
|
||||
?column?
|
||||
----------
|
||||
@@ -32,6 +56,12 @@ SELECT inner_product('[65504]'::halfvec, '[65504]');
|
||||
4290774016
|
||||
(1 row)
|
||||
|
||||
SELECT inner_product('[1,1,1,1,1,1,1,1,1]'::halfvec, '[1,2,3,4,5,6,7,8,9]');
|
||||
inner_product
|
||||
---------------
|
||||
45
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2]'::halfvec <#> '[3,4]';
|
||||
?column?
|
||||
----------
|
||||
|
||||
@@ -55,51 +55,64 @@ ERROR: "65520" is out of range for type halfvec
|
||||
LINE 1: SELECT '[65520,-65520]'::halfvec;
|
||||
^
|
||||
SELECT '[1e-8,-1e-8]'::halfvec;
|
||||
ERROR: "1e-8" is out of range for type halfvec
|
||||
LINE 1: SELECT '[1e-8,-1e-8]'::halfvec;
|
||||
^
|
||||
halfvec
|
||||
---------
|
||||
[0,-0]
|
||||
(1 row)
|
||||
|
||||
SELECT '[4e38,1]'::halfvec;
|
||||
ERROR: "4e38" is out of range for type halfvec
|
||||
LINE 1: SELECT '[4e38,1]'::halfvec;
|
||||
^
|
||||
SELECT '[1e-46,1]'::halfvec;
|
||||
ERROR: "1e-46" is out of range for type halfvec
|
||||
LINE 1: SELECT '[1e-46,1]'::halfvec;
|
||||
^
|
||||
halfvec
|
||||
---------
|
||||
[0,1]
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2,3'::halfvec;
|
||||
ERROR: malformed halfvec literal: "[1,2,3"
|
||||
ERROR: invalid input syntax for type halfvec: "[1,2,3"
|
||||
LINE 1: SELECT '[1,2,3'::halfvec;
|
||||
^
|
||||
DETAIL: Unexpected end of input.
|
||||
SELECT '[1,2,3]9'::halfvec;
|
||||
ERROR: malformed halfvec literal: "[1,2,3]9"
|
||||
ERROR: invalid input syntax for type halfvec: "[1,2,3]9"
|
||||
LINE 1: SELECT '[1,2,3]9'::halfvec;
|
||||
^
|
||||
DETAIL: Junk after closing right brace.
|
||||
SELECT '1,2,3'::halfvec;
|
||||
ERROR: malformed halfvec literal: "1,2,3"
|
||||
ERROR: invalid input syntax for type halfvec: "1,2,3"
|
||||
LINE 1: SELECT '1,2,3'::halfvec;
|
||||
^
|
||||
DETAIL: Vector contents must start with "[".
|
||||
SELECT ''::halfvec;
|
||||
ERROR: malformed halfvec literal: ""
|
||||
ERROR: invalid input syntax for type halfvec: ""
|
||||
LINE 1: SELECT ''::halfvec;
|
||||
^
|
||||
DETAIL: Vector contents must start with "[".
|
||||
SELECT '['::halfvec;
|
||||
ERROR: malformed halfvec literal: "["
|
||||
ERROR: invalid input syntax for type halfvec: "["
|
||||
LINE 1: SELECT '['::halfvec;
|
||||
^
|
||||
DETAIL: Unexpected end of input.
|
||||
SELECT '[ '::halfvec;
|
||||
ERROR: invalid input syntax for type halfvec: "[ "
|
||||
LINE 1: SELECT '[ '::halfvec;
|
||||
^
|
||||
SELECT '[,'::halfvec;
|
||||
ERROR: malformed halfvec literal: "[,"
|
||||
ERROR: invalid input syntax for type halfvec: "[,"
|
||||
LINE 1: SELECT '[,'::halfvec;
|
||||
^
|
||||
DETAIL: Unexpected end of input.
|
||||
SELECT '[]'::halfvec;
|
||||
ERROR: halfvec must have at least 1 dimension
|
||||
LINE 1: SELECT '[]'::halfvec;
|
||||
^
|
||||
SELECT '[ ]'::halfvec;
|
||||
ERROR: halfvec must have at least 1 dimension
|
||||
LINE 1: SELECT '[ ]'::halfvec;
|
||||
^
|
||||
SELECT '[,]'::halfvec;
|
||||
ERROR: invalid input syntax for type halfvec: "[,]"
|
||||
LINE 1: SELECT '[,]'::halfvec;
|
||||
^
|
||||
SELECT '[1,]'::halfvec;
|
||||
ERROR: invalid input syntax for type halfvec: "[1,]"
|
||||
LINE 1: SELECT '[1,]'::halfvec;
|
||||
@@ -109,7 +122,7 @@ ERROR: invalid input syntax for type halfvec: "[1a]"
|
||||
LINE 1: SELECT '[1a]'::halfvec;
|
||||
^
|
||||
SELECT '[1,,3]'::halfvec;
|
||||
ERROR: malformed halfvec literal: "[1,,3]"
|
||||
ERROR: invalid input syntax for type halfvec: "[1,,3]"
|
||||
LINE 1: SELECT '[1,,3]'::halfvec;
|
||||
^
|
||||
SELECT '[1, ,3]'::halfvec;
|
||||
|
||||
@@ -24,4 +24,6 @@ CREATE TABLE t (val varbit(3));
|
||||
CREATE INDEX ON t USING hnsw (val bit_hamming_ops);
|
||||
ERROR: type not supported for hnsw index
|
||||
CREATE INDEX ON t USING hnsw ((val::bit(3)) bit_hamming_ops);
|
||||
CREATE INDEX ON t USING hnsw ((val::bit(64001)) bit_hamming_ops);
|
||||
ERROR: column cannot have more than 64000 dimensions for hnsw index
|
||||
DROP TABLE t;
|
||||
|
||||
26
test/expected/ivfflat_halfvec_cosine.out
Normal file
26
test/expected/ivfflat_halfvec_cosine.out
Normal file
@@ -0,0 +1,26 @@
|
||||
SET enable_seqscan = off;
|
||||
CREATE TABLE t (val halfvec(3));
|
||||
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val halfvec_cosine_ops) WITH (lists = 1);
|
||||
INSERT INTO t (val) VALUES ('[1,2,4]');
|
||||
SELECT * FROM t ORDER BY val <=> '[3,3,3]';
|
||||
val
|
||||
---------
|
||||
[1,1,1]
|
||||
[1,2,3]
|
||||
[1,2,4]
|
||||
(3 rows)
|
||||
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
|
||||
count
|
||||
-------
|
||||
3
|
||||
(1 row)
|
||||
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::halfvec)) t2;
|
||||
count
|
||||
-------
|
||||
3
|
||||
(1 row)
|
||||
|
||||
DROP TABLE t;
|
||||
21
test/expected/ivfflat_halfvec_ip.out
Normal file
21
test/expected/ivfflat_halfvec_ip.out
Normal file
@@ -0,0 +1,21 @@
|
||||
SET enable_seqscan = off;
|
||||
CREATE TABLE t (val halfvec(3));
|
||||
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val halfvec_ip_ops) WITH (lists = 1);
|
||||
INSERT INTO t (val) VALUES ('[1,2,4]');
|
||||
SELECT * FROM t ORDER BY val <#> '[3,3,3]';
|
||||
val
|
||||
---------
|
||||
[1,2,4]
|
||||
[1,2,3]
|
||||
[1,1,1]
|
||||
[0,0,0]
|
||||
(4 rows)
|
||||
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::halfvec)) t2;
|
||||
count
|
||||
-------
|
||||
4
|
||||
(1 row)
|
||||
|
||||
DROP TABLE t;
|
||||
36
test/expected/ivfflat_halfvec_l2.out
Normal file
36
test/expected/ivfflat_halfvec_l2.out
Normal file
@@ -0,0 +1,36 @@
|
||||
SET enable_seqscan = off;
|
||||
CREATE TABLE t (val halfvec(3));
|
||||
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val halfvec_l2_ops) WITH (lists = 1);
|
||||
INSERT INTO t (val) VALUES ('[1,2,4]');
|
||||
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
|
||||
val
|
||||
---------
|
||||
[1,2,3]
|
||||
[1,2,4]
|
||||
[1,1,1]
|
||||
[0,0,0]
|
||||
(4 rows)
|
||||
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <-> (SELECT NULL::halfvec)) t2;
|
||||
count
|
||||
-------
|
||||
4
|
||||
(1 row)
|
||||
|
||||
SELECT COUNT(*) FROM t;
|
||||
count
|
||||
-------
|
||||
5
|
||||
(1 row)
|
||||
|
||||
TRUNCATE t;
|
||||
NOTICE: ivfflat index created with little data
|
||||
DETAIL: This will cause low recall.
|
||||
HINT: Drop the index until the table has more data.
|
||||
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
|
||||
val
|
||||
-----
|
||||
(0 rows)
|
||||
|
||||
DROP TABLE t;
|
||||
@@ -12,14 +12,11 @@ SELECT * FROM t ORDER BY val <-> '[3,3,3]';
|
||||
[0,0,0]
|
||||
(4 rows)
|
||||
|
||||
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector);
|
||||
val
|
||||
---------
|
||||
[0,0,0]
|
||||
[1,1,1]
|
||||
[1,2,3]
|
||||
[1,2,4]
|
||||
(4 rows)
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector)) t2;
|
||||
count
|
||||
-------
|
||||
4
|
||||
(1 row)
|
||||
|
||||
SELECT COUNT(*) FROM t;
|
||||
count
|
||||
|
||||
@@ -1,3 +1,27 @@
|
||||
SELECT round(sparsevec_norm('{1:1,2:1}/2')::numeric, 5);
|
||||
round
|
||||
---------
|
||||
1.41421
|
||||
(1 row)
|
||||
|
||||
SELECT sparsevec_norm('{1:3,2:4}/2');
|
||||
sparsevec_norm
|
||||
----------------
|
||||
5
|
||||
(1 row)
|
||||
|
||||
SELECT sparsevec_norm('{2:1}/2');
|
||||
sparsevec_norm
|
||||
----------------
|
||||
1
|
||||
(1 row)
|
||||
|
||||
SELECT sparsevec_norm('{1:3e37,2:4e37}/2')::real;
|
||||
sparsevec_norm
|
||||
----------------
|
||||
5e+37
|
||||
(1 row)
|
||||
|
||||
SELECT l2_distance('{}/2'::sparsevec, '{1:3,2:4}/2');
|
||||
l2_distance
|
||||
-------------
|
||||
|
||||
@@ -4,26 +4,159 @@ SELECT '{1:1.5,3:3.5}/5'::sparsevec;
|
||||
{1:1.5,3:3.5}/5
|
||||
(1 row)
|
||||
|
||||
SELECT '{1:1.5,3:3.5}/5'::sparsevec::vector;
|
||||
vector
|
||||
-----------------
|
||||
[1.5,0,3.5,0,0]
|
||||
SELECT '{1:-2,3:-4}/5'::sparsevec;
|
||||
sparsevec
|
||||
---------------
|
||||
{1:-2,3:-4}/5
|
||||
(1 row)
|
||||
|
||||
SELECT '{1:1.5,3:3.5}/5'::sparsevec::vector(5);
|
||||
vector
|
||||
-----------------
|
||||
[1.5,0,3.5,0,0]
|
||||
SELECT '{1:2.,3:4.}/5'::sparsevec;
|
||||
sparsevec
|
||||
-------------
|
||||
{1:2,3:4}/5
|
||||
(1 row)
|
||||
|
||||
SELECT '{1:1.5,3:3.5}/5'::sparsevec::vector(4);
|
||||
ERROR: expected 4 dimensions, not 5
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec;
|
||||
SELECT ' { 1 : 1.5 , 3 : 3.5 } / 5 '::sparsevec;
|
||||
sparsevec
|
||||
-----------------
|
||||
{2:1.5,4:3.5}/5
|
||||
{1:1.5,3:3.5}/5
|
||||
(1 row)
|
||||
|
||||
SELECT '{1:1.23456}/1'::sparsevec;
|
||||
sparsevec
|
||||
---------------
|
||||
{1:1.23456}/1
|
||||
(1 row)
|
||||
|
||||
SELECT '{1:hello,2:1}/2'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{1:hello,2:1}/2"
|
||||
LINE 1: SELECT '{1:hello,2:1}/2'::sparsevec;
|
||||
^
|
||||
SELECT '{1:NaN,2:1}/2'::sparsevec;
|
||||
ERROR: NaN not allowed in sparsevec
|
||||
LINE 1: SELECT '{1:NaN,2:1}/2'::sparsevec;
|
||||
^
|
||||
SELECT '{1:Infinity,2:1}/2'::sparsevec;
|
||||
ERROR: infinite value not allowed in sparsevec
|
||||
LINE 1: SELECT '{1:Infinity,2:1}/2'::sparsevec;
|
||||
^
|
||||
SELECT '{1:-Infinity,2:1}/2'::sparsevec;
|
||||
ERROR: infinite value not allowed in sparsevec
|
||||
LINE 1: SELECT '{1:-Infinity,2:1}/2'::sparsevec;
|
||||
^
|
||||
SELECT '{1:1.5e38,2:-1.5e38}/2'::sparsevec;
|
||||
sparsevec
|
||||
--------------------------
|
||||
{1:1.5e+38,2:-1.5e+38}/2
|
||||
(1 row)
|
||||
|
||||
SELECT '{1:1.5e+38,2:-1.5e+38}/2'::sparsevec;
|
||||
sparsevec
|
||||
--------------------------
|
||||
{1:1.5e+38,2:-1.5e+38}/2
|
||||
(1 row)
|
||||
|
||||
SELECT '{1:1.5e-38,2:-1.5e-38}/2'::sparsevec;
|
||||
sparsevec
|
||||
--------------------------
|
||||
{1:1.5e-38,2:-1.5e-38}/2
|
||||
(1 row)
|
||||
|
||||
SELECT '{1:4e38,2:1}/2'::sparsevec;
|
||||
ERROR: "4e38" is out of range for type sparsevec
|
||||
LINE 1: SELECT '{1:4e38,2:1}/2'::sparsevec;
|
||||
^
|
||||
SELECT '{1:-4e38,2:1}/2'::sparsevec;
|
||||
ERROR: "-4e38" is out of range for type sparsevec
|
||||
LINE 1: SELECT '{1:-4e38,2:1}/2'::sparsevec;
|
||||
^
|
||||
SELECT '{1:1e-46,2:1}/2'::sparsevec;
|
||||
ERROR: "1e-46" is out of range for type sparsevec
|
||||
LINE 1: SELECT '{1:1e-46,2:1}/2'::sparsevec;
|
||||
^
|
||||
SELECT '{1:-1e-46,2:1}/2'::sparsevec;
|
||||
ERROR: "-1e-46" is out of range for type sparsevec
|
||||
LINE 1: SELECT '{1:-1e-46,2:1}/2'::sparsevec;
|
||||
^
|
||||
SELECT ''::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: ""
|
||||
LINE 1: SELECT ''::sparsevec;
|
||||
^
|
||||
DETAIL: Vector contents must start with "{".
|
||||
SELECT '{'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{"
|
||||
LINE 1: SELECT '{'::sparsevec;
|
||||
^
|
||||
SELECT '{ '::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{ "
|
||||
LINE 1: SELECT '{ '::sparsevec;
|
||||
^
|
||||
SELECT '{:'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{:"
|
||||
LINE 1: SELECT '{:'::sparsevec;
|
||||
^
|
||||
SELECT '{,'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{,"
|
||||
LINE 1: SELECT '{,'::sparsevec;
|
||||
^
|
||||
SELECT '{}'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{}"
|
||||
LINE 1: SELECT '{}'::sparsevec;
|
||||
^
|
||||
DETAIL: Unexpected end of input.
|
||||
SELECT '{}/'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{}/"
|
||||
LINE 1: SELECT '{}/'::sparsevec;
|
||||
^
|
||||
SELECT '{}/1'::sparsevec;
|
||||
sparsevec
|
||||
-----------
|
||||
{}/1
|
||||
(1 row)
|
||||
|
||||
SELECT '{}/1a'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{}/1a"
|
||||
LINE 1: SELECT '{}/1a'::sparsevec;
|
||||
^
|
||||
DETAIL: Junk after closing.
|
||||
SELECT '{ }/1'::sparsevec;
|
||||
sparsevec
|
||||
-----------
|
||||
{}/1
|
||||
(1 row)
|
||||
|
||||
SELECT '{:}/1'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{:}/1"
|
||||
LINE 1: SELECT '{:}/1'::sparsevec;
|
||||
^
|
||||
SELECT '{,}/1'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{,}/1"
|
||||
LINE 1: SELECT '{,}/1'::sparsevec;
|
||||
^
|
||||
SELECT '{1,}/1'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{1,}/1"
|
||||
LINE 1: SELECT '{1,}/1'::sparsevec;
|
||||
^
|
||||
SELECT '{:1}/1'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{:1}/1"
|
||||
LINE 1: SELECT '{:1}/1'::sparsevec;
|
||||
^
|
||||
SELECT '{1:}/1'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{1:}/1"
|
||||
LINE 1: SELECT '{1:}/1'::sparsevec;
|
||||
^
|
||||
SELECT '{1a:1}/1'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{1a:1}/1"
|
||||
LINE 1: SELECT '{1a:1}/1'::sparsevec;
|
||||
^
|
||||
SELECT '{1:1a}/1'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{1:1a}/1"
|
||||
LINE 1: SELECT '{1:1a}/1'::sparsevec;
|
||||
^
|
||||
SELECT '{1:1,}/1'::sparsevec;
|
||||
ERROR: invalid input syntax for type sparsevec: "{1:1,}/1"
|
||||
LINE 1: SELECT '{1:1,}/1'::sparsevec;
|
||||
^
|
||||
SELECT '{1:0,2:1,3:0}/3'::sparsevec;
|
||||
sparsevec
|
||||
-----------
|
||||
@@ -48,8 +181,6 @@ SELECT '{}/100001'::sparsevec;
|
||||
ERROR: sparsevec cannot have more than 100000 dimensions
|
||||
LINE 1: SELECT '{}/100001'::sparsevec;
|
||||
^
|
||||
SELECT '{}/16001'::sparsevec::vector;
|
||||
ERROR: vector cannot have more than 16000 dimensions
|
||||
SELECT '{0:1}/1'::sparsevec;
|
||||
ERROR: index "0" is out of range for type sparsevec
|
||||
LINE 1: SELECT '{0:1}/1'::sparsevec;
|
||||
@@ -58,5 +189,27 @@ SELECT '{2:1}/1'::sparsevec;
|
||||
ERROR: index must be less than or equal to dimensions
|
||||
LINE 1: SELECT '{2:1}/1'::sparsevec;
|
||||
^
|
||||
SELECT '{}/1'::sparsevec(2);
|
||||
ERROR: expected 2 dimensions, not 1
|
||||
SELECT '{}/3'::sparsevec(3);
|
||||
sparsevec
|
||||
-----------
|
||||
{}/3
|
||||
(1 row)
|
||||
|
||||
SELECT '{}/3'::sparsevec(2);
|
||||
ERROR: expected 2 dimensions, not 3
|
||||
SELECT '{}/3'::sparsevec(3, 2);
|
||||
ERROR: invalid type modifier
|
||||
LINE 1: SELECT '{}/3'::sparsevec(3, 2);
|
||||
^
|
||||
SELECT '{}/3'::sparsevec('a');
|
||||
ERROR: invalid input syntax for type integer: "a"
|
||||
LINE 1: SELECT '{}/3'::sparsevec('a');
|
||||
^
|
||||
SELECT '{}/3'::sparsevec(0);
|
||||
ERROR: dimensions for type sparsevec must be at least 1
|
||||
LINE 1: SELECT '{}/3'::sparsevec(0);
|
||||
^
|
||||
SELECT '{}/3'::sparsevec(100001);
|
||||
ERROR: dimensions for type sparsevec cannot exceed 100000
|
||||
LINE 1: SELECT '{}/3'::sparsevec(100001);
|
||||
^
|
||||
|
||||
@@ -83,39 +83,48 @@ SELECT '[-1e-46,1]'::vector;
|
||||
(1 row)
|
||||
|
||||
SELECT '[1,2,3'::vector;
|
||||
ERROR: malformed vector literal: "[1,2,3"
|
||||
ERROR: invalid input syntax for type vector: "[1,2,3"
|
||||
LINE 1: SELECT '[1,2,3'::vector;
|
||||
^
|
||||
DETAIL: Unexpected end of input.
|
||||
SELECT '[1,2,3]9'::vector;
|
||||
ERROR: malformed vector literal: "[1,2,3]9"
|
||||
ERROR: invalid input syntax for type vector: "[1,2,3]9"
|
||||
LINE 1: SELECT '[1,2,3]9'::vector;
|
||||
^
|
||||
DETAIL: Junk after closing right brace.
|
||||
SELECT '1,2,3'::vector;
|
||||
ERROR: malformed vector literal: "1,2,3"
|
||||
ERROR: invalid input syntax for type vector: "1,2,3"
|
||||
LINE 1: SELECT '1,2,3'::vector;
|
||||
^
|
||||
DETAIL: Vector contents must start with "[".
|
||||
SELECT ''::vector;
|
||||
ERROR: malformed vector literal: ""
|
||||
ERROR: invalid input syntax for type vector: ""
|
||||
LINE 1: SELECT ''::vector;
|
||||
^
|
||||
DETAIL: Vector contents must start with "[".
|
||||
SELECT '['::vector;
|
||||
ERROR: malformed vector literal: "["
|
||||
ERROR: invalid input syntax for type vector: "["
|
||||
LINE 1: SELECT '['::vector;
|
||||
^
|
||||
DETAIL: Unexpected end of input.
|
||||
SELECT '[ '::vector;
|
||||
ERROR: invalid input syntax for type vector: "[ "
|
||||
LINE 1: SELECT '[ '::vector;
|
||||
^
|
||||
SELECT '[,'::vector;
|
||||
ERROR: malformed vector literal: "[,"
|
||||
ERROR: invalid input syntax for type vector: "[,"
|
||||
LINE 1: SELECT '[,'::vector;
|
||||
^
|
||||
DETAIL: Unexpected end of input.
|
||||
SELECT '[]'::vector;
|
||||
ERROR: vector must have at least 1 dimension
|
||||
LINE 1: SELECT '[]'::vector;
|
||||
^
|
||||
SELECT '[ ]'::vector;
|
||||
ERROR: vector must have at least 1 dimension
|
||||
LINE 1: SELECT '[ ]'::vector;
|
||||
^
|
||||
SELECT '[,]'::vector;
|
||||
ERROR: invalid input syntax for type vector: "[,]"
|
||||
LINE 1: SELECT '[,]'::vector;
|
||||
^
|
||||
SELECT '[1,]'::vector;
|
||||
ERROR: invalid input syntax for type vector: "[1,]"
|
||||
LINE 1: SELECT '[1,]'::vector;
|
||||
@@ -125,7 +134,7 @@ ERROR: invalid input syntax for type vector: "[1a]"
|
||||
LINE 1: SELECT '[1a]'::vector;
|
||||
^
|
||||
SELECT '[1,,3]'::vector;
|
||||
ERROR: malformed vector literal: "[1,,3]"
|
||||
ERROR: invalid input syntax for type vector: "[1,,3]"
|
||||
LINE 1: SELECT '[1,,3]'::vector;
|
||||
^
|
||||
SELECT '[1, ,3]'::vector;
|
||||
|
||||
@@ -3,19 +3,50 @@ SELECT ARRAY[1.0,2.0,3.0]::vector;
|
||||
SELECT ARRAY[1,2,3]::float4[]::vector;
|
||||
SELECT ARRAY[1,2,3]::float8[]::vector;
|
||||
SELECT ARRAY[1,2,3]::numeric[]::vector;
|
||||
|
||||
SELECT '[1,2,3]'::vector::real[];
|
||||
|
||||
SELECT '{1,2,3}'::real[]::vector;
|
||||
SELECT '{1,2,3}'::real[]::vector(3);
|
||||
SELECT '{1,2,3}'::real[]::vector(2);
|
||||
SELECT '{NULL}'::real[]::vector;
|
||||
SELECT '{NaN}'::real[]::vector;
|
||||
SELECT '{Infinity}'::real[]::vector;
|
||||
SELECT '{-Infinity}'::real[]::vector;
|
||||
SELECT '{}'::real[]::vector;
|
||||
SELECT '{{1}}'::real[]::vector;
|
||||
SELECT '[1,2,3]'::vector::real[];
|
||||
|
||||
SELECT '{1,2,3}'::double precision[]::vector;
|
||||
SELECT '{1,2,3}'::double precision[]::vector(3);
|
||||
SELECT '{1,2,3}'::double precision[]::vector(2);
|
||||
SELECT '{4e38,-4e38}'::double precision[]::vector;
|
||||
SELECT '{1e-46,-1e-46}'::double precision[]::vector;
|
||||
|
||||
SELECT '[1,2,3]'::vector::halfvec;
|
||||
SELECT '[1,2,3]'::halfvec::vector;
|
||||
SELECT '[1,2,3]'::vector::halfvec(3);
|
||||
SELECT '[1,2,3]'::vector::halfvec(2);
|
||||
SELECT '[1,2,3]'::halfvec::vector(2);
|
||||
SELECT '[65520]'::vector::halfvec;
|
||||
SELECT '[1e-8]'::vector::halfvec;
|
||||
|
||||
SELECT '[1,2,3]'::halfvec::vector;
|
||||
SELECT '[1,2,3]'::halfvec::vector(3);
|
||||
SELECT '[1,2,3]'::halfvec::vector(2);
|
||||
|
||||
SELECT '{1,2,3}'::real[]::halfvec;
|
||||
SELECT '{1,2,3}'::real[]::halfvec(3);
|
||||
SELECT '{1,2,3}'::real[]::halfvec(2);
|
||||
SELECT '{65520,-65520}'::real[]::halfvec;
|
||||
SELECT '{1e-8,-1e-8}'::real[]::halfvec;
|
||||
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec;
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec(5);
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec(4);
|
||||
|
||||
SELECT '{2:1.5,4:3.5}/5'::sparsevec::vector;
|
||||
SELECT '{2:1.5,4:3.5}/5'::sparsevec::vector(5);
|
||||
SELECT '{2:1.5,4:3.5}/5'::sparsevec::vector(4);
|
||||
SELECT '{}/16001'::sparsevec::vector;
|
||||
|
||||
SELECT array_agg(n)::vector FROM generate_series(1, 16001) n;
|
||||
SELECT array_to_vector(array_agg(n), 16001, false) FROM generate_series(1, 16001) n;
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
CREATE TABLE t (val vector(3), val2 halfvec(3));
|
||||
INSERT INTO t (val, val2) VALUES ('[0,0,0]', '[0,0,0]'), ('[1,2,3]', '[1,2,3]'), ('[1,1,1]', '[1,1,1]'), (NULL, NULL);
|
||||
CREATE TABLE t (val vector(3), val2 halfvec(3), val3 sparsevec(3));
|
||||
INSERT INTO t (val, val2, val3) VALUES ('[0,0,0]', '[0,0,0]', '{}/3'), ('[1,2,3]', '[1,2,3]', '{1:1,2:2,3:3}/3'), ('[1,1,1]', '[1,1,1]', '{1:1,2:1,3:1}/3'), (NULL, NULL, NULL);
|
||||
|
||||
CREATE TABLE t2 (val vector(3), val2 halfvec(3));
|
||||
CREATE TABLE t2 (val vector(3), val2 halfvec(3), val3 sparsevec(3));
|
||||
|
||||
\copy t TO 'results/data.bin' WITH (FORMAT binary)
|
||||
\copy t2 FROM 'results/data.bin' WITH (FORMAT binary)
|
||||
|
||||
@@ -1,11 +1,17 @@
|
||||
SELECT round(halfvec_norm('[1,1]')::numeric, 5);
|
||||
SELECT halfvec_norm('[3,4]');
|
||||
SELECT halfvec_norm('[0,1]');
|
||||
|
||||
SELECT l2_distance('[0,0]'::halfvec, '[3,4]');
|
||||
SELECT l2_distance('[0,0]'::halfvec, '[0,1]');
|
||||
SELECT l2_distance('[1,2]'::halfvec, '[3]');
|
||||
SELECT l2_distance('[1,1,1,1,1,1,1,1,1]'::halfvec, '[1,1,1,1,1,1,1,4,5]');
|
||||
SELECT '[0,0]'::halfvec <-> '[3,4]';
|
||||
|
||||
SELECT inner_product('[1,2]'::halfvec, '[3,4]');
|
||||
SELECT inner_product('[1,2]'::halfvec, '[3]');
|
||||
SELECT inner_product('[65504]'::halfvec, '[65504]');
|
||||
SELECT inner_product('[1,1,1,1,1,1,1,1,1]'::halfvec, '[1,2,3,4,5,6,7,8,9]');
|
||||
SELECT '[1,2]'::halfvec <#> '[3,4]';
|
||||
|
||||
SELECT cosine_distance('[1,2]'::halfvec, '[2,4]');
|
||||
|
||||
@@ -17,8 +17,11 @@ SELECT '[1,2,3]9'::halfvec;
|
||||
SELECT '1,2,3'::halfvec;
|
||||
SELECT ''::halfvec;
|
||||
SELECT '['::halfvec;
|
||||
SELECT '[ '::halfvec;
|
||||
SELECT '[,'::halfvec;
|
||||
SELECT '[]'::halfvec;
|
||||
SELECT '[ ]'::halfvec;
|
||||
SELECT '[,]'::halfvec;
|
||||
SELECT '[1,]'::halfvec;
|
||||
SELECT '[1a]'::halfvec;
|
||||
SELECT '[1,,3]'::halfvec;
|
||||
|
||||
@@ -15,4 +15,5 @@ DROP TABLE t;
|
||||
CREATE TABLE t (val varbit(3));
|
||||
CREATE INDEX ON t USING hnsw (val bit_hamming_ops);
|
||||
CREATE INDEX ON t USING hnsw ((val::bit(3)) bit_hamming_ops);
|
||||
CREATE INDEX ON t USING hnsw ((val::bit(64001)) bit_hamming_ops);
|
||||
DROP TABLE t;
|
||||
|
||||
13
test/sql/ivfflat_halfvec_cosine.sql
Normal file
13
test/sql/ivfflat_halfvec_cosine.sql
Normal file
@@ -0,0 +1,13 @@
|
||||
SET enable_seqscan = off;
|
||||
|
||||
CREATE TABLE t (val halfvec(3));
|
||||
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val halfvec_cosine_ops) WITH (lists = 1);
|
||||
|
||||
INSERT INTO t (val) VALUES ('[1,2,4]');
|
||||
|
||||
SELECT * FROM t ORDER BY val <=> '[3,3,3]';
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::halfvec)) t2;
|
||||
|
||||
DROP TABLE t;
|
||||
12
test/sql/ivfflat_halfvec_ip.sql
Normal file
12
test/sql/ivfflat_halfvec_ip.sql
Normal file
@@ -0,0 +1,12 @@
|
||||
SET enable_seqscan = off;
|
||||
|
||||
CREATE TABLE t (val halfvec(3));
|
||||
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val halfvec_ip_ops) WITH (lists = 1);
|
||||
|
||||
INSERT INTO t (val) VALUES ('[1,2,4]');
|
||||
|
||||
SELECT * FROM t ORDER BY val <#> '[3,3,3]';
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::halfvec)) t2;
|
||||
|
||||
DROP TABLE t;
|
||||
16
test/sql/ivfflat_halfvec_l2.sql
Normal file
16
test/sql/ivfflat_halfvec_l2.sql
Normal file
@@ -0,0 +1,16 @@
|
||||
SET enable_seqscan = off;
|
||||
|
||||
CREATE TABLE t (val halfvec(3));
|
||||
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
|
||||
CREATE INDEX ON t USING ivfflat (val halfvec_l2_ops) WITH (lists = 1);
|
||||
|
||||
INSERT INTO t (val) VALUES ('[1,2,4]');
|
||||
|
||||
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <-> (SELECT NULL::halfvec)) t2;
|
||||
SELECT COUNT(*) FROM t;
|
||||
|
||||
TRUNCATE t;
|
||||
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
|
||||
|
||||
DROP TABLE t;
|
||||
@@ -7,7 +7,7 @@ CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 1);
|
||||
INSERT INTO t (val) VALUES ('[1,2,4]');
|
||||
|
||||
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
|
||||
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector);
|
||||
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector)) t2;
|
||||
SELECT COUNT(*) FROM t;
|
||||
|
||||
TRUNCATE t;
|
||||
|
||||
@@ -1,3 +1,8 @@
|
||||
SELECT round(sparsevec_norm('{1:1,2:1}/2')::numeric, 5);
|
||||
SELECT sparsevec_norm('{1:3,2:4}/2');
|
||||
SELECT sparsevec_norm('{2:1}/2');
|
||||
SELECT sparsevec_norm('{1:3e37,2:4e37}/2')::real;
|
||||
|
||||
SELECT l2_distance('{}/2'::sparsevec, '{1:3,2:4}/2');
|
||||
SELECT l2_distance('{}/2'::sparsevec, '{2:1}/2');
|
||||
SELECT '{}/2'::sparsevec <-> '{1:3,2:4}/2';
|
||||
|
||||
@@ -1,19 +1,48 @@
|
||||
SELECT '{1:1.5,3:3.5}/5'::sparsevec;
|
||||
SELECT '{1:1.5,3:3.5}/5'::sparsevec::vector;
|
||||
SELECT '{1:1.5,3:3.5}/5'::sparsevec::vector(5);
|
||||
SELECT '{1:1.5,3:3.5}/5'::sparsevec::vector(4);
|
||||
SELECT '[0,1.5,0,3.5,0]'::vector::sparsevec;
|
||||
|
||||
SELECT '{1:-2,3:-4}/5'::sparsevec;
|
||||
SELECT '{1:2.,3:4.}/5'::sparsevec;
|
||||
SELECT ' { 1 : 1.5 , 3 : 3.5 } / 5 '::sparsevec;
|
||||
SELECT '{1:1.23456}/1'::sparsevec;
|
||||
SELECT '{1:hello,2:1}/2'::sparsevec;
|
||||
SELECT '{1:NaN,2:1}/2'::sparsevec;
|
||||
SELECT '{1:Infinity,2:1}/2'::sparsevec;
|
||||
SELECT '{1:-Infinity,2:1}/2'::sparsevec;
|
||||
SELECT '{1:1.5e38,2:-1.5e38}/2'::sparsevec;
|
||||
SELECT '{1:1.5e+38,2:-1.5e+38}/2'::sparsevec;
|
||||
SELECT '{1:1.5e-38,2:-1.5e-38}/2'::sparsevec;
|
||||
SELECT '{1:4e38,2:1}/2'::sparsevec;
|
||||
SELECT '{1:-4e38,2:1}/2'::sparsevec;
|
||||
SELECT '{1:1e-46,2:1}/2'::sparsevec;
|
||||
SELECT '{1:-1e-46,2:1}/2'::sparsevec;
|
||||
SELECT ''::sparsevec;
|
||||
SELECT '{'::sparsevec;
|
||||
SELECT '{ '::sparsevec;
|
||||
SELECT '{:'::sparsevec;
|
||||
SELECT '{,'::sparsevec;
|
||||
SELECT '{}'::sparsevec;
|
||||
SELECT '{}/'::sparsevec;
|
||||
SELECT '{}/1'::sparsevec;
|
||||
SELECT '{}/1a'::sparsevec;
|
||||
SELECT '{ }/1'::sparsevec;
|
||||
SELECT '{:}/1'::sparsevec;
|
||||
SELECT '{,}/1'::sparsevec;
|
||||
SELECT '{1,}/1'::sparsevec;
|
||||
SELECT '{:1}/1'::sparsevec;
|
||||
SELECT '{1:}/1'::sparsevec;
|
||||
SELECT '{1a:1}/1'::sparsevec;
|
||||
SELECT '{1:1a}/1'::sparsevec;
|
||||
SELECT '{1:1,}/1'::sparsevec;
|
||||
SELECT '{1:0,2:1,3:0}/3'::sparsevec;
|
||||
|
||||
SELECT '{2:1,1:1}/2'::sparsevec;
|
||||
|
||||
SELECT '{}/5'::sparsevec;
|
||||
SELECT '{}/-1'::sparsevec;
|
||||
SELECT '{}/100001'::sparsevec;
|
||||
SELECT '{}/16001'::sparsevec::vector;
|
||||
|
||||
SELECT '{0:1}/1'::sparsevec;
|
||||
SELECT '{2:1}/1'::sparsevec;
|
||||
|
||||
SELECT '{}/1'::sparsevec(2);
|
||||
SELECT '{}/3'::sparsevec(3);
|
||||
SELECT '{}/3'::sparsevec(2);
|
||||
SELECT '{}/3'::sparsevec(3, 2);
|
||||
SELECT '{}/3'::sparsevec('a');
|
||||
SELECT '{}/3'::sparsevec(0);
|
||||
SELECT '{}/3'::sparsevec(100001);
|
||||
|
||||
@@ -19,8 +19,11 @@ SELECT '[1,2,3]9'::vector;
|
||||
SELECT '1,2,3'::vector;
|
||||
SELECT ''::vector;
|
||||
SELECT '['::vector;
|
||||
SELECT '[ '::vector;
|
||||
SELECT '[,'::vector;
|
||||
SELECT '[]'::vector;
|
||||
SELECT '[ ]'::vector;
|
||||
SELECT '[,]'::vector;
|
||||
SELECT '[1,]'::vector;
|
||||
SELECT '[1a]'::vector;
|
||||
SELECT '[1,,3]'::vector;
|
||||
|
||||
@@ -30,18 +30,19 @@ sub test_recall
|
||||
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]);
|
||||
my %expected_set = map { $_ => 1 } @expected_ids;
|
||||
|
||||
foreach (@expected_ids)
|
||||
foreach (@actual_ids)
|
||||
{
|
||||
if (exists($actual_set{$_}))
|
||||
if (exists($expected_set{$_}))
|
||||
{
|
||||
$correct++;
|
||||
}
|
||||
$total++;
|
||||
}
|
||||
|
||||
$total += $limit;
|
||||
}
|
||||
|
||||
cmp_ok($correct / $total, ">=", $min, $operator);
|
||||
@@ -81,7 +82,12 @@ for my $i (0 .. $#operators)
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
WITH top AS (
|
||||
SELECT v $operator '$_' AS distance FROM tst ORDER BY distance LIMIT $limit
|
||||
)
|
||||
SELECT i FROM tst WHERE (v $operator '$_') <= (SELECT MAX(distance) FROM top)
|
||||
));
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
@@ -98,8 +104,16 @@ for my $i (0 .. $#operators)
|
||||
test_recall(1, 0.71, $operator);
|
||||
test_recall(10, 0.95, $operator);
|
||||
}
|
||||
# Account for equal distances
|
||||
test_recall(100, 0.9925, $operator);
|
||||
|
||||
# Test probes equals lists
|
||||
if ($operator eq "<=>")
|
||||
{
|
||||
test_recall(100, 0.9925, $operator);
|
||||
}
|
||||
else
|
||||
{
|
||||
test_recall(100, 1.00, $operator);
|
||||
}
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
|
||||
@@ -119,8 +133,16 @@ for my $i (0 .. $#operators)
|
||||
test_recall(1, 0.71, $operator);
|
||||
test_recall(10, 0.95, $operator);
|
||||
}
|
||||
# Account for equal distances
|
||||
test_recall(100, 0.9925, $operator);
|
||||
|
||||
# Test probes equals lists
|
||||
if ($operator eq "<=>")
|
||||
{
|
||||
test_recall(100, 0.9925, $operator);
|
||||
}
|
||||
else
|
||||
{
|
||||
test_recall(100, 1.00, $operator);
|
||||
}
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
}
|
||||
|
||||
@@ -98,6 +98,7 @@ for my $i (0 .. $#operators)
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
# Test approximate results
|
||||
my $min = $operator eq "<#>" ? 0.80 : 0.99;
|
||||
test_recall($min, $operator);
|
||||
|
||||
|
||||
@@ -85,7 +85,7 @@ for my $i (0 .. $#operators)
|
||||
# Handle ties
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
WITH top AS (
|
||||
SELECT v $operator $_ AS distance FROM tst ORDER BY v $operator $_ LIMIT $limit
|
||||
SELECT v $operator $_ AS distance FROM tst ORDER BY distance LIMIT $limit
|
||||
)
|
||||
SELECT i FROM tst WHERE (v $operator $_) <= (SELECT MAX(distance) FROM top)
|
||||
));
|
||||
|
||||
@@ -94,7 +94,7 @@ for my $i (0 .. $#operators)
|
||||
));
|
||||
|
||||
# Test approximate results
|
||||
my $min = $operator eq "<#>" ? 0.95 : 0.98;
|
||||
my $min = $operator eq "<#>" ? 0.93 : 0.98;
|
||||
test_recall($min, $operator);
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
|
||||
116
test/t/023_hnsw_bit_insert_recall.pl
Normal file
116
test/t/023_hnsw_bit_insert_recall.pl
Normal file
@@ -0,0 +1,116 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
my $node;
|
||||
my @queries = ();
|
||||
my @expected;
|
||||
my $limit = 20;
|
||||
my $dim = 52;
|
||||
my $max = 2**$dim;
|
||||
|
||||
sub test_recall
|
||||
{
|
||||
my ($min, $operator) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = 100;
|
||||
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;
|
||||
SET hnsw.ef_search = 100;
|
||||
SELECT i FROM tst ORDER BY v $operator $queries[$i] LIMIT $limit;
|
||||
));
|
||||
my @actual_ids = split("\n", $actual);
|
||||
|
||||
my @expected_ids = split("\n", $expected[$i]);
|
||||
my %expected_set = map { $_ => 1 } @expected_ids;
|
||||
|
||||
foreach (@actual_ids)
|
||||
{
|
||||
if (exists($expected_set{$_}))
|
||||
{
|
||||
$correct++;
|
||||
}
|
||||
}
|
||||
|
||||
$total += $limit;
|
||||
}
|
||||
|
||||
cmp_ok($correct / $total, ">=", $min, $operator);
|
||||
}
|
||||
|
||||
# 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 serial, v bit($dim));");
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my $r = int(rand() * $max);
|
||||
push(@queries, "${r}::bigint::bit($dim)");
|
||||
}
|
||||
|
||||
# Check each index type
|
||||
my @operators = ("<~>", "<\%>");
|
||||
my @opclasses = ("bit_hamming_ops", "bit_jaccard_ops");
|
||||
|
||||
for my $i (0 .. $#operators)
|
||||
{
|
||||
my $operator = $operators[$i];
|
||||
my $opclass = $opclasses[$i];
|
||||
|
||||
# Add index
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v $opclass);");
|
||||
|
||||
# Use concurrent inserts
|
||||
$node->pgbench(
|
||||
"--no-vacuum --client=10 --transactions=1000",
|
||||
0,
|
||||
[qr{actually processed}],
|
||||
[qr{^$}],
|
||||
"concurrent INSERTs",
|
||||
{
|
||||
"023_hnsw_bit_insert_recall_$opclass" => "INSERT INTO tst (v) VALUES ((random() * $max)::bigint::bit($dim));"
|
||||
}
|
||||
);
|
||||
|
||||
# Get exact results
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
# Handle ties
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
SET enable_indexscan = off;
|
||||
WITH top AS (
|
||||
SELECT v $operator $_ AS distance FROM tst ORDER BY distance LIMIT $limit
|
||||
)
|
||||
SELECT i FROM tst WHERE (v $operator $_) <= (SELECT MAX(distance) FROM top)
|
||||
));
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
# Test approximate results
|
||||
my $min = $operator eq "<\%>" ? 0.95 : 0.98;
|
||||
test_recall($min, $operator);
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
$node->safe_psql("postgres", "TRUNCATE tst;");
|
||||
}
|
||||
|
||||
done_testing();
|
||||
113
test/t/024_hnsw_halfvec_insert_recall.pl
Normal file
113
test/t/024_hnsw_halfvec_insert_recall.pl
Normal file
@@ -0,0 +1,113 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
my $node;
|
||||
my @queries = ();
|
||||
my @expected;
|
||||
my $limit = 20;
|
||||
my $dim = 10;
|
||||
my $array_sql = join(",", ('random()') x $dim);
|
||||
|
||||
sub test_recall
|
||||
{
|
||||
my ($min, $operator) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v $operator '$queries[0]' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan/);
|
||||
|
||||
for my $i (0 .. $#queries)
|
||||
{
|
||||
my $actual = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
|
||||
));
|
||||
my @actual_ids = split("\n", $actual);
|
||||
my %actual_set = map { $_ => 1 } @actual_ids;
|
||||
|
||||
my @expected_ids = split("\n", $expected[$i]);
|
||||
|
||||
foreach (@expected_ids)
|
||||
{
|
||||
if (exists($actual_set{$_}))
|
||||
{
|
||||
$correct++;
|
||||
}
|
||||
$total++;
|
||||
}
|
||||
}
|
||||
|
||||
cmp_ok($correct / $total, ">=", $min, $operator);
|
||||
}
|
||||
|
||||
# Initialize node
|
||||
$node = 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 serial, v halfvec($dim));");
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my @r = ();
|
||||
for (1 .. $dim)
|
||||
{
|
||||
push(@r, rand());
|
||||
}
|
||||
push(@queries, "[" . join(",", @r) . "]");
|
||||
}
|
||||
|
||||
# Check each index type
|
||||
my @operators = ("<->", "<#>", "<=>");
|
||||
my @opclasses = ("halfvec_l2_ops", "halfvec_ip_ops", "halfvec_cosine_ops");
|
||||
|
||||
for my $i (0 .. $#operators)
|
||||
{
|
||||
my $operator = $operators[$i];
|
||||
my $opclass = $opclasses[$i];
|
||||
|
||||
# Add index
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v $opclass);");
|
||||
|
||||
# Use concurrent inserts
|
||||
$node->pgbench(
|
||||
"--no-vacuum --client=10 --transactions=1000",
|
||||
0,
|
||||
[qr{actually processed}],
|
||||
[qr{^$}],
|
||||
"concurrent INSERTs",
|
||||
{
|
||||
"024_hnsw_halfvec_insert_recall_$opclass" => "INSERT INTO tst (v) VALUES (ARRAY[$array_sql]);"
|
||||
}
|
||||
);
|
||||
|
||||
# Get exact results
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
SET enable_indexscan = off;
|
||||
SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;
|
||||
));
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
# Test approximate results
|
||||
my $min = $operator eq "<#>" ? 0.94 : 0.98;
|
||||
test_recall($min, $operator);
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
$node->safe_psql("postgres", "TRUNCATE tst;");
|
||||
}
|
||||
|
||||
done_testing();
|
||||
109
test/t/025_hnsw_sparsevec_insert_recall.pl
Normal file
109
test/t/025_hnsw_sparsevec_insert_recall.pl
Normal file
@@ -0,0 +1,109 @@
|
||||
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 serial, v sparsevec(3));");
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my $r1 = rand();
|
||||
my $r2 = rand();
|
||||
my $r3 = rand();
|
||||
push(@queries, "{1:$r1,2:$r2,3:$r3}/3");
|
||||
}
|
||||
|
||||
# Check each index type
|
||||
my @operators = ("<->", "<#>", "<=>");
|
||||
my @opclasses = ("sparsevec_l2_ops", "sparsevec_ip_ops", "sparsevec_cosine_ops");
|
||||
|
||||
for my $i (0 .. $#operators)
|
||||
{
|
||||
my $operator = $operators[$i];
|
||||
my $opclass = $opclasses[$i];
|
||||
|
||||
# Add index
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v $opclass);");
|
||||
|
||||
# Use concurrent inserts
|
||||
$node->pgbench(
|
||||
"--no-vacuum --client=10 --transactions=1000",
|
||||
0,
|
||||
[qr{actually processed}],
|
||||
[qr{^$}],
|
||||
"concurrent INSERTs",
|
||||
{
|
||||
"025_hnsw_sparsevec_insert_recall_$opclass" => "INSERT INTO tst (v) VALUES (ARRAY[random(), random(), random()]::vector::sparsevec);"
|
||||
}
|
||||
);
|
||||
|
||||
# Get exact results
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
SET enable_indexscan = off;
|
||||
SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;
|
||||
));
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
# Test approximate results
|
||||
my $min = $operator eq "<#>" ? 0.80 : 0.99;
|
||||
test_recall($min, $operator);
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
$node->safe_psql("postgres", "TRUNCATE tst;");
|
||||
}
|
||||
|
||||
done_testing();
|
||||
58
test/t/026_hnsw_bit_duplicates.pl
Normal file
58
test/t/026_hnsw_bit_duplicates.pl
Normal file
@@ -0,0 +1,58 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
# Initialize node
|
||||
my $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 (v bit(3));");
|
||||
|
||||
sub insert_vectors
|
||||
{
|
||||
for my $i (1 .. 20)
|
||||
{
|
||||
$node->safe_psql("postgres", "INSERT INTO tst VALUES ('111');");
|
||||
}
|
||||
}
|
||||
|
||||
sub test_duplicates
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = 1;
|
||||
SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <~> '111') t;
|
||||
));
|
||||
is($res, 10);
|
||||
}
|
||||
|
||||
# Test duplicates with build
|
||||
insert_vectors();
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v bit_hamming_ops);");
|
||||
test_duplicates();
|
||||
|
||||
# Reset
|
||||
$node->safe_psql("postgres", "TRUNCATE tst;");
|
||||
|
||||
# Test duplicates with inserts
|
||||
insert_vectors();
|
||||
test_duplicates();
|
||||
|
||||
# Test fallback path for inserts
|
||||
$node->pgbench(
|
||||
"--no-vacuum --client=5 --transactions=100",
|
||||
0,
|
||||
[qr{actually processed}],
|
||||
[qr{^$}],
|
||||
"concurrent INSERTs",
|
||||
{
|
||||
"026_hnsw_bit_duplicates" => "INSERT INTO tst VALUES ('111');"
|
||||
}
|
||||
);
|
||||
|
||||
done_testing();
|
||||
58
test/t/027_hnsw_halfvec_duplicates.pl
Normal file
58
test/t/027_hnsw_halfvec_duplicates.pl
Normal file
@@ -0,0 +1,58 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
# Initialize node
|
||||
my $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 (v halfvec(3));");
|
||||
|
||||
sub insert_vectors
|
||||
{
|
||||
for my $i (1 .. 20)
|
||||
{
|
||||
$node->safe_psql("postgres", "INSERT INTO tst VALUES ('[1,1,1]');");
|
||||
}
|
||||
}
|
||||
|
||||
sub test_duplicates
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = 1;
|
||||
SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '[1,1,1]') t;
|
||||
));
|
||||
is($res, 10);
|
||||
}
|
||||
|
||||
# Test duplicates with build
|
||||
insert_vectors();
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v halfvec_l2_ops);");
|
||||
test_duplicates();
|
||||
|
||||
# Reset
|
||||
$node->safe_psql("postgres", "TRUNCATE tst;");
|
||||
|
||||
# Test duplicates with inserts
|
||||
insert_vectors();
|
||||
test_duplicates();
|
||||
|
||||
# Test fallback path for inserts
|
||||
$node->pgbench(
|
||||
"--no-vacuum --client=5 --transactions=100",
|
||||
0,
|
||||
[qr{actually processed}],
|
||||
[qr{^$}],
|
||||
"concurrent INSERTs",
|
||||
{
|
||||
"027_hnsw_halfvec_duplicates" => "INSERT INTO tst VALUES ('[1,1,1]');"
|
||||
}
|
||||
);
|
||||
|
||||
done_testing();
|
||||
58
test/t/028_hnsw_sparsevec_duplicates.pl
Normal file
58
test/t/028_hnsw_sparsevec_duplicates.pl
Normal file
@@ -0,0 +1,58 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
# Initialize node
|
||||
my $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 (v sparsevec(3));");
|
||||
|
||||
sub insert_vectors
|
||||
{
|
||||
for my $i (1 .. 20)
|
||||
{
|
||||
$node->safe_psql("postgres", "INSERT INTO tst VALUES ('{1:1,2:1,3:1}/3');");
|
||||
}
|
||||
}
|
||||
|
||||
sub test_duplicates
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = 1;
|
||||
SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '{1:1,2:1,3:1}/3') t;
|
||||
));
|
||||
is($res, 10);
|
||||
}
|
||||
|
||||
# Test duplicates with build
|
||||
insert_vectors();
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v sparsevec_l2_ops);");
|
||||
test_duplicates();
|
||||
|
||||
# Reset
|
||||
$node->safe_psql("postgres", "TRUNCATE tst;");
|
||||
|
||||
# Test duplicates with inserts
|
||||
insert_vectors();
|
||||
test_duplicates();
|
||||
|
||||
# Test fallback path for inserts
|
||||
$node->pgbench(
|
||||
"--no-vacuum --client=5 --transactions=100",
|
||||
0,
|
||||
[qr{actually processed}],
|
||||
[qr{^$}],
|
||||
"concurrent INSERTs",
|
||||
{
|
||||
"028_hnsw_sparsevec_duplicates" => "INSERT INTO tst VALUES ('{1:1,2:1,3:1}/3');"
|
||||
}
|
||||
);
|
||||
|
||||
done_testing();
|
||||
100
test/t/029_hnsw_bit_vacuum_recall.pl
Normal file
100
test/t/029_hnsw_bit_vacuum_recall.pl
Normal file
@@ -0,0 +1,100 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
my $node;
|
||||
my @queries = ();
|
||||
my @expected;
|
||||
my $limit = 20;
|
||||
my $dim = 52;
|
||||
my $max = 2**$dim;
|
||||
|
||||
sub test_recall
|
||||
{
|
||||
my ($min, $ef_search, $test_name) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = $ef_search;
|
||||
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <~> $queries[0] LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan/);
|
||||
|
||||
for my $i (0 .. $#queries)
|
||||
{
|
||||
my $actual = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = $ef_search;
|
||||
SELECT i FROM tst ORDER BY v <~> $queries[$i] LIMIT $limit;
|
||||
));
|
||||
my @actual_ids = split("\n", $actual);
|
||||
|
||||
my @expected_ids = split("\n", $expected[$i]);
|
||||
my %expected_set = map { $_ => 1 } @expected_ids;
|
||||
|
||||
foreach (@actual_ids)
|
||||
{
|
||||
if (exists($expected_set{$_}))
|
||||
{
|
||||
$correct++;
|
||||
}
|
||||
}
|
||||
|
||||
$total += $limit;
|
||||
}
|
||||
|
||||
cmp_ok($correct / $total, ">=", $min, $test_name);
|
||||
}
|
||||
|
||||
# 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 bit($dim));");
|
||||
$node->safe_psql("postgres", "ALTER TABLE tst SET (autovacuum_enabled = false);");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, (random() * $max)::bigint::bit($dim) FROM generate_series(1, 10000) i;"
|
||||
);
|
||||
|
||||
# Add index
|
||||
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v bit_hamming_ops) WITH (m = 4, ef_construction = 8);");
|
||||
|
||||
# Delete data
|
||||
$node->safe_psql("postgres", "DELETE FROM tst WHERE i > 2500;");
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my $r = int(rand() * $max);
|
||||
push(@queries, "${r}::bigint::bit($dim)");
|
||||
}
|
||||
|
||||
# Get exact results
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
SET enable_indexscan = off;
|
||||
WITH top AS (
|
||||
SELECT v <~> $_ AS distance FROM tst ORDER BY distance LIMIT $limit
|
||||
)
|
||||
SELECT i FROM tst WHERE (v <~> $_) <= (SELECT MAX(distance) FROM top)
|
||||
));
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
test_recall(0.35, 100, "before vacuum");
|
||||
|
||||
# TODO Test concurrent inserts with vacuum
|
||||
$node->safe_psql("postgres", "VACUUM tst;");
|
||||
|
||||
test_recall(0.80, 100, "after vacuum");
|
||||
|
||||
done_testing();
|
||||
97
test/t/030_hnsw_halfvec_vacuum_recall.pl
Normal file
97
test/t/030_hnsw_halfvec_vacuum_recall.pl
Normal file
@@ -0,0 +1,97 @@
|
||||
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, $ef_search, $test_name) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = $ef_search;
|
||||
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$queries[0]' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan/);
|
||||
|
||||
for my $i (0 .. $#queries)
|
||||
{
|
||||
my $actual = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = $ef_search;
|
||||
SELECT i FROM tst ORDER BY v <-> '$queries[$i]' LIMIT $limit;
|
||||
));
|
||||
my @actual_ids = split("\n", $actual);
|
||||
my %actual_set = map { $_ => 1 } @actual_ids;
|
||||
|
||||
my @expected_ids = split("\n", $expected[$i]);
|
||||
|
||||
foreach (@expected_ids)
|
||||
{
|
||||
if (exists($actual_set{$_}))
|
||||
{
|
||||
$correct++;
|
||||
}
|
||||
$total++;
|
||||
}
|
||||
}
|
||||
|
||||
cmp_ok($correct / $total, ">=", $min, $test_name);
|
||||
}
|
||||
|
||||
# Initialize node
|
||||
$node = 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 halfvec(3));");
|
||||
$node->safe_psql("postgres", "ALTER TABLE tst SET (autovacuum_enabled = false);");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 10000) i;"
|
||||
);
|
||||
|
||||
# Add index
|
||||
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v halfvec_l2_ops) WITH (m = 4, ef_construction = 8);");
|
||||
|
||||
# Delete data
|
||||
$node->safe_psql("postgres", "DELETE FROM tst WHERE i > 2500;");
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my $r1 = rand();
|
||||
my $r2 = rand();
|
||||
my $r3 = rand();
|
||||
push(@queries, "[$r1,$r2,$r3]");
|
||||
}
|
||||
|
||||
# Get exact results
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
SET enable_indexscan = off;
|
||||
SELECT i FROM tst ORDER BY v <-> '$_' LIMIT $limit;
|
||||
));
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
test_recall(0.18, $limit, "before vacuum");
|
||||
test_recall(0.95, 100, "before vacuum");
|
||||
|
||||
# TODO Test concurrent inserts with vacuum
|
||||
$node->safe_psql("postgres", "VACUUM tst;");
|
||||
|
||||
test_recall(0.95, $limit, "after vacuum");
|
||||
|
||||
done_testing();
|
||||
97
test/t/031_hnsw_sparsevec_vacuum_recall.pl
Normal file
97
test/t/031_hnsw_sparsevec_vacuum_recall.pl
Normal file
@@ -0,0 +1,97 @@
|
||||
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, $ef_search, $test_name) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = $ef_search;
|
||||
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$queries[0]' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan/);
|
||||
|
||||
for my $i (0 .. $#queries)
|
||||
{
|
||||
my $actual = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = $ef_search;
|
||||
SELECT i FROM tst ORDER BY v <-> '$queries[$i]' LIMIT $limit;
|
||||
));
|
||||
my @actual_ids = split("\n", $actual);
|
||||
my %actual_set = map { $_ => 1 } @actual_ids;
|
||||
|
||||
my @expected_ids = split("\n", $expected[$i]);
|
||||
|
||||
foreach (@expected_ids)
|
||||
{
|
||||
if (exists($actual_set{$_}))
|
||||
{
|
||||
$correct++;
|
||||
}
|
||||
$total++;
|
||||
}
|
||||
}
|
||||
|
||||
cmp_ok($correct / $total, ">=", $min, $test_name);
|
||||
}
|
||||
|
||||
# Initialize node
|
||||
$node = 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 sparsevec(3));");
|
||||
$node->safe_psql("postgres", "ALTER TABLE tst SET (autovacuum_enabled = false);");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()]::vector::sparsevec(3) FROM generate_series(1, 10000) i;"
|
||||
);
|
||||
|
||||
# Add index
|
||||
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v sparsevec_l2_ops) WITH (m = 4, ef_construction = 8);");
|
||||
|
||||
# Delete data
|
||||
$node->safe_psql("postgres", "DELETE FROM tst WHERE i > 2500;");
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my $r1 = rand();
|
||||
my $r2 = rand();
|
||||
my $r3 = rand();
|
||||
push(@queries, "{1:$r1,2:$r2,3:$r3}/3");
|
||||
}
|
||||
|
||||
# Get exact results
|
||||
@expected = ();
|
||||
foreach (@queries)
|
||||
{
|
||||
my $res = $node->safe_psql("postgres", qq(
|
||||
SET enable_indexscan = off;
|
||||
SELECT i FROM tst ORDER BY v <-> '$_' LIMIT $limit;
|
||||
));
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
test_recall(0.18, $limit, "before vacuum");
|
||||
test_recall(0.95, 100, "before vacuum");
|
||||
|
||||
# TODO Test concurrent inserts with vacuum
|
||||
$node->safe_psql("postgres", "VACUUM tst;");
|
||||
|
||||
test_recall(0.95, $limit, "after vacuum");
|
||||
|
||||
done_testing();
|
||||
154
test/t/032_ivfflat_halfvec_build_recall.pl
Normal file
154
test/t/032_ivfflat_halfvec_build_recall.pl
Normal file
@@ -0,0 +1,154 @@
|
||||
use strict;
|
||||
use warnings;
|
||||
use PostgresNode;
|
||||
use TestLib;
|
||||
use Test::More;
|
||||
|
||||
my $node;
|
||||
my @queries = ();
|
||||
my @expected;
|
||||
my $limit = 20;
|
||||
my $dim = 10;
|
||||
my $array_sql = join(",", ('random()') x $dim);
|
||||
|
||||
sub test_recall
|
||||
{
|
||||
my ($probes, $min, $operator) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET ivfflat.probes = $probes;
|
||||
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v $operator '$queries[0]' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan using idx on tst/);
|
||||
|
||||
for my $i (0 .. $#queries)
|
||||
{
|
||||
my $actual = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET ivfflat.probes = $probes;
|
||||
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
|
||||
));
|
||||
my @actual_ids = split("\n", $actual);
|
||||
|
||||
my @expected_ids = split("\n", $expected[$i]);
|
||||
my %expected_set = map { $_ => 1 } @expected_ids;
|
||||
|
||||
foreach (@actual_ids)
|
||||
{
|
||||
if (exists($expected_set{$_}))
|
||||
{
|
||||
$correct++;
|
||||
}
|
||||
}
|
||||
|
||||
$total += $limit;
|
||||
}
|
||||
|
||||
cmp_ok($correct / $total, ">=", $min, $operator);
|
||||
}
|
||||
|
||||
# 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 halfvec($dim));");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
|
||||
);
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my @r = ();
|
||||
for (1 .. $dim)
|
||||
{
|
||||
push(@r, rand());
|
||||
}
|
||||
push(@queries, "[" . join(",", @r) . "]");
|
||||
}
|
||||
|
||||
# Check each index type
|
||||
my @operators = ("<->", "<#>", "<=>");
|
||||
my @opclasses = ("halfvec_l2_ops", "halfvec_ip_ops", "halfvec_cosine_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", qq(
|
||||
WITH top AS (
|
||||
SELECT v $operator '$_' AS distance FROM tst ORDER BY distance LIMIT $limit
|
||||
)
|
||||
SELECT i FROM tst WHERE (v $operator '$_') <= (SELECT MAX(distance) FROM top)
|
||||
));
|
||||
push(@expected, $res);
|
||||
}
|
||||
|
||||
# Build index serially
|
||||
$node->safe_psql("postgres", qq(
|
||||
SET max_parallel_maintenance_workers = 0;
|
||||
CREATE INDEX idx ON tst USING ivfflat (v $opclass);
|
||||
));
|
||||
|
||||
# Test approximate results
|
||||
if ($operator ne "<#>")
|
||||
{
|
||||
# TODO Fix test (uniform random vectors all have similar inner product)
|
||||
test_recall(1, 0.35, $operator);
|
||||
test_recall(10, 0.95, $operator);
|
||||
}
|
||||
|
||||
# Test probes equals lists
|
||||
if ($operator eq "<=>")
|
||||
{
|
||||
test_recall(100, 0.98, $operator);
|
||||
}
|
||||
else
|
||||
{
|
||||
test_recall(100, 1.00, $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;
|
||||
CREATE INDEX idx ON tst USING ivfflat (v $opclass);
|
||||
));
|
||||
is($ret, 0, $stderr);
|
||||
like($stderr, qr/using \d+ parallel workers/);
|
||||
|
||||
# Test approximate results
|
||||
if ($operator ne "<#>")
|
||||
{
|
||||
# TODO Fix test (uniform random vectors all have similar inner product)
|
||||
test_recall(1, 0.35, $operator);
|
||||
test_recall(10, 0.95, $operator);
|
||||
}
|
||||
|
||||
# Test probes equals lists
|
||||
if ($operator eq "<=>")
|
||||
{
|
||||
test_recall(100, 0.98, $operator);
|
||||
}
|
||||
else
|
||||
{
|
||||
test_recall(100, 1.00, $operator);
|
||||
}
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
}
|
||||
|
||||
done_testing();
|
||||
Reference in New Issue
Block a user