Compare commits

..

2 Commits

Author SHA1 Message Date
Andrew Kane
4c675198b3 Updated readme [skip ci] 2024-03-28 00:06:54 -07:00
Andrew Kane
d49d053e84 Added subvector function 2024-03-25 21:51:03 -07:00
18 changed files with 138 additions and 53 deletions

View File

@@ -1,3 +1,7 @@
## 0.7.0 (unreleased)
- Added `subvector` function
## 0.6.2 (2024-03-18)
- Reduced lock contention with parallel HNSW index builds

View File

@@ -10,15 +10,21 @@ TESTS = $(wildcard test/sql/*.sql)
REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS))
REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION)
OPTFLAGS =
# Since runtime dispatch not supported
ifeq ($(shell uname -s), Darwin)
ifeq ($(shell uname -m), x86_64)
OPTFLAGS = -march=native
# Mac ARM doesn't support -march=native
ifeq ($(shell uname -s), Darwin)
ifeq ($(shell uname -p), arm)
# no difference with -march=armv8.5-a
OPTFLAGS =
endif
endif
# PowerPC doesn't support -march=native
ifneq ($(filter ppc64%, $(shell uname -m)), )
OPTFLAGS =
endif
# For auto-vectorization:
# - GCC (needs -ftree-vectorize OR -O3) - https://gcc.gnu.org/projects/tree-ssa/vectorization.html
# - Clang (could use pragma instead) - https://llvm.org/docs/Vectorizers.html

View File

@@ -671,8 +671,6 @@ ALTER TABLE items ALTER COLUMN embedding SET STORAGE PLAIN;
Results are limited by the size of the dynamic candidate list (`hnsw.ef_search`). There may be even less results due to dead tuples or filtering conditions in the query. We recommend setting `hnsw.ef_search` to at least twice the `LIMIT` of the query. If you need more than 500 results, use an IVFFlat index instead.
Also, note that `NULL` vectors are not indexed (as well as zero vectors for cosine distance).
#### Why are there less results for a query after adding an IVFFlat index?
The index was likely created with too little data for the number of lists. Drop the index until the table has more data.
@@ -683,13 +681,11 @@ DROP INDEX index_name;
Results can also be limited by the number of probes (`ivfflat.probes`).
Also, note that `NULL` vectors are not indexed (as well as zero vectors for cosine distance).
## Reference
### Vector Type
Each vector takes `4 * dimensions + 8` bytes of storage. Each element is a single-precision floating-point number (like the `real` type in Postgres), and all elements must be finite (no `NaN`, `Infinity` or `-Infinity`). Vectors can have up to 16,000 dimensions.
Each vector takes `4 * dimensions + 8` bytes of storage. Each element is a single precision floating-point number (like the `real` type in Postgres), and all elements must be finite (no `NaN`, `Infinity` or `-Infinity`). Vectors can have up to 16,000 dimensions.
### Vector Operators
@@ -710,10 +706,11 @@ cosine_distance(vector, vector) → double precision | cosine distance |
inner_product(vector, vector) → double precision | inner product |
l2_distance(vector, vector) → double precision | Euclidean distance |
l1_distance(vector, vector) → double precision | taxicab distance | 0.5.0
subvector(vector, integer, integer) → vector | subvector | 0.7.0 [unreleased]
vector_dims(vector) → integer | number of dimensions |
vector_norm(vector) → double precision | Euclidean norm |
### Vector Aggregate Functions
### Aggregate Functions
Function | Description | Added
--- | --- | ---

View File

@@ -0,0 +1,5 @@
-- complain if script is sourced in psql, rather than via CREATE EXTENSION
\echo Use "ALTER EXTENSION vector UPDATE TO '0.7.0'" to load this file. \quit
CREATE FUNCTION subvector(vector, int, int) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

View File

@@ -58,6 +58,9 @@ CREATE FUNCTION vector_sub(vector, vector) RETURNS vector
CREATE FUNCTION vector_mul(vector, vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION subvector(vector, int, int) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- private functions
CREATE FUNCTION vector_lt(vector, vector) RETURNS bool

View File

@@ -262,6 +262,7 @@ typedef struct HnswBuildState
HnswGraph *graph;
double ml;
int maxLevel;
Vector *normvec;
/* Memory */
MemoryContext graphCtx;
@@ -366,7 +367,7 @@ typedef struct HnswVacuumState
int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index);
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum);
bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value);
bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page);
void HnswInit(void);

View File

@@ -489,7 +489,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!HnswNormValue(buildstate->normprocinfo, buildstate->collation, &value))
if (!HnswNormValue(buildstate->normprocinfo, buildstate->collation, &value, buildstate->normvec))
return false;
}
@@ -703,6 +703,9 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
/* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions);
buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw build graph context",
#if PG_VERSION_NUM >= 150000
@@ -726,6 +729,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
static void
FreeBuildState(HnswBuildState * buildstate)
{
pfree(buildstate->normvec);
MemoryContextDelete(buildstate->graphCtx);
MemoryContextDelete(buildstate->tmpCtx);
}

View File

@@ -622,7 +622,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
if (normprocinfo != NULL)
{
if (!HnswNormValue(normprocinfo, collation, &value))
if (!HnswNormValue(normprocinfo, collation, &value, NULL))
return;
}

View File

@@ -84,7 +84,7 @@ GetScanValue(IndexScanDesc scan)
/* Fine if normalization fails */
if (so->normprocinfo != NULL)
HnswNormValue(so->normprocinfo, so->collation, &value);
HnswNormValue(so->normprocinfo, so->collation, &value, NULL);
}
return value;

View File

@@ -158,14 +158,16 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
* if it's different than the original value
*/
bool
HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value)
HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
{
double norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
if (norm > 0)
{
Vector *v = DatumGetVector(*value);
Vector *result = InitVector(v->dim);
if (result == NULL)
result = InitVector(v->dim);
for (int i = 0; i < v->dim; i++)
result->x[i] = v->x[i] / norm;

View File

@@ -57,7 +57,7 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
*/
if (buildstate->kmeansnormprocinfo != NULL)
{
if (!IvfflatNormValue(buildstate->kmeansnormprocinfo, buildstate->collation, &value))
if (!IvfflatNormValue(buildstate->kmeansnormprocinfo, buildstate->collation, &value, buildstate->normvec))
return;
}
@@ -153,7 +153,7 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!IvfflatNormValue(buildstate->normprocinfo, buildstate->collation, &value))
if (!IvfflatNormValue(buildstate->normprocinfo, buildstate->collation, &value, buildstate->normvec))
return;
}
@@ -356,6 +356,9 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions);
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
/* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions);
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Ivfflat build temporary context",
ALLOCSET_DEFAULT_SIZES);
@@ -377,6 +380,7 @@ FreeBuildState(IvfflatBuildState * buildstate)
{
VectorArrayFree(buildstate->centers);
pfree(buildstate->listInfo);
pfree(buildstate->normvec);
#ifdef IVFFLAT_KMEANS_DEBUG
pfree(buildstate->listSums);

View File

@@ -172,6 +172,7 @@ typedef struct IvfflatBuildState
VectorArray samples;
VectorArray centers;
ListInfo *listInfo;
Vector *normvec;
#ifdef IVFFLAT_KMEANS_DEBUG
double inertia;
@@ -266,7 +267,7 @@ void VectorArrayFree(VectorArray arr);
void PrintVectorArray(char *msg, VectorArray arr);
void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);
FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum);
bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value);
bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
int IvfflatGetLists(Relation index);
void IvfflatGetMetaPageInfo(Relation index, int *lists, int *dimensions);
void IvfflatUpdateList(Relation index, ListInfo listInfo, BlockNumber insertPage, BlockNumber originalInsertPage, BlockNumber startPage, ForkNumber forkNum);

View File

@@ -85,7 +85,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
if (normprocinfo != NULL)
{
if (!IvfflatNormValue(normprocinfo, index->rd_indcollation[0], &value))
if (!IvfflatNormValue(normprocinfo, index->rd_indcollation[0], &value, NULL))
return;
}

View File

@@ -293,7 +293,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
/* Fine if normalization fails */
if (so->normprocinfo != NULL)
IvfflatNormValue(so->normprocinfo, so->collation, &value);
IvfflatNormValue(so->normprocinfo, so->collation, &value, NULL);
}
IvfflatBench("GetScanLists", GetScanLists(scan, value));

View File

@@ -75,14 +75,16 @@ IvfflatOptionalProcInfo(Relation index, uint16 procnum)
* if it's different than the original value
*/
bool
IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value)
IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
{
double norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
if (norm > 0)
{
Vector *v = DatumGetVector(*value);
Vector *result = InitVector(v->dim);
if (result == NULL)
result = InitVector(v->dim);
for (int i = 0; i < v->dim; i++)
result->x[i] = v->x[i] / norm;

View File

@@ -29,15 +29,6 @@
#define STATE_DIMS(x) (ARR_DIMS(x)[0] - 1)
#define CreateStateDatums(dim) palloc(sizeof(Datum) * (dim + 1))
#if defined(__x86_64__) && defined(__gnu_linux__) && defined(__has_attribute) && __has_attribute(target_clones)
#define RUNTIME_DISPATCH __attribute__((target_clones("default", "avx", "fma", "avx512f")))
#elif defined(__aarch64__) && defined(__gnu_linux__) && defined(__has_attribute) && __has_attribute(target_clones)
/* TODO Fix error: target does not support function version dispatcher */
#define RUNTIME_DISPATCH __attribute__((target_clones("default", "arch=armv8.5-a")))
#else
#define RUNTIME_DISPATCH
#endif
PG_MODULE_MAGIC;
/*
@@ -541,23 +532,6 @@ vector_to_float4(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(result);
}
RUNTIME_DISPATCH
static float
l2_squared_distance_impl(int16 dim, float *ax, float *bx)
{
float distance = 0.0;
/* Auto-vectorized */
for (int16 i = 0; i < dim; i++)
{
float diff = ax[i] - bx[i];
distance += diff * diff;
}
return distance;
}
/*
* Get the L2 distance between vectors
*/
@@ -567,11 +541,19 @@ l2_distance(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
float distance;
float *ax = a->x;
float *bx = b->x;
float distance = 0.0;
float diff;
CheckDims(a, b);
distance = l2_squared_distance_impl(a->dim, a->x, b->x);
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++)
{
diff = ax[i] - bx[i];
distance += diff * diff;
}
PG_RETURN_FLOAT8(sqrt((double) distance));
}
@@ -586,11 +568,19 @@ vector_l2_squared_distance(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
float distance;
float *ax = a->x;
float *bx = b->x;
float distance = 0.0;
float diff;
CheckDims(a, b);
distance = l2_squared_distance_impl(a->dim, a->x, b->x);
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++)
{
diff = ax[i] - bx[i];
distance += diff * diff;
}
PG_RETURN_FLOAT8((double) distance);
}
@@ -870,6 +860,37 @@ vector_mul(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(result);
}
/*
* Get a subset of a vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(subvector);
Datum
subvector(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
int32 start = PG_GETARG_INT32(1);
int32 count = PG_GETARG_INT32(2);
int32 end = start + count;
float *ax = a->x;
Vector *result;
int dim;
if (start < 1)
start = 1;
if (end > a->dim)
end = a->dim + 1;
dim = end - start;
CheckDim(dim);
result = InitVector(dim);
for (int i = 0; i < dim; i++)
result->x[i] = ax[start - 1 + i];
PG_RETURN_POINTER(result);
}
/*
* Internal helper to compare vectors
*/

View File

@@ -208,6 +208,34 @@ SELECT l1_distance('[3e38]', '[-3e38]');
Infinity
(1 row)
SELECT subvector('[1,2,3,4,5]', 1, 3);
subvector
-----------
[1,2,3]
(1 row)
SELECT subvector('[1,2,3,4,5]', 3, 2);
subvector
-----------
[3,4]
(1 row)
SELECT subvector('[1,2,3,4,5]', -1, 3);
subvector
-----------
[1]
(1 row)
SELECT subvector('[1,2,3,4,5]', 3, 9);
subvector
-----------
[3,4,5]
(1 row)
SELECT subvector('[1,2,3,4,5]', 1, 0);
ERROR: vector must have at least 1 dimension
SELECT subvector('[1,2,3,4,5]', -1, 2);
ERROR: vector must have at least 1 dimension
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
avg
-----------

View File

@@ -48,6 +48,13 @@ SELECT l1_distance('[0,0]', '[0,1]');
SELECT l1_distance('[1,2]', '[3]');
SELECT l1_distance('[3e38]', '[-3e38]');
SELECT subvector('[1,2,3,4,5]', 1, 3);
SELECT subvector('[1,2,3,4,5]', 3, 2);
SELECT subvector('[1,2,3,4,5]', -1, 3);
SELECT subvector('[1,2,3,4,5]', 3, 9);
SELECT subvector('[1,2,3,4,5]', 1, 0);
SELECT subvector('[1,2,3,4,5]', -1, 2);
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v;
SELECT avg(v) FROM unnest(ARRAY[]::vector[]) v;