mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-22 12:07:34 +08:00
Compare commits
3 Commits
valgrind-s
...
index-type
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f26bd7054d | ||
|
|
4305e723cc | ||
|
|
571a277af9 |
118
.github/workflows/build.yml
vendored
118
.github/workflows/build.yml
vendored
@@ -1,6 +1,122 @@
|
||||
name: build
|
||||
on: [push, pull_request]
|
||||
jobs:
|
||||
ubuntu:
|
||||
runs-on: ${{ matrix.os }}
|
||||
if: ${{ !startsWith(github.ref_name, 'mac') && !startsWith(github.ref_name, 'windows') }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
- postgres: 17
|
||||
os: ubuntu-22.04
|
||||
- postgres: 16
|
||||
os: ubuntu-22.04
|
||||
- postgres: 15
|
||||
os: ubuntu-22.04
|
||||
- postgres: 14
|
||||
os: ubuntu-22.04
|
||||
- postgres: 13
|
||||
os: ubuntu-20.04
|
||||
- postgres: 12
|
||||
os: ubuntu-20.04
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: ankane/setup-postgres@v1
|
||||
with:
|
||||
postgres-version: ${{ matrix.postgres }}
|
||||
dev-files: true
|
||||
- run: make
|
||||
env:
|
||||
PG_CFLAGS: -DUSE_ASSERT_CHECKING -Wall -Wextra -Werror -Wno-unused-parameter -Wno-sign-compare
|
||||
- run: |
|
||||
export PG_CONFIG=`which pg_config`
|
||||
sudo --preserve-env=PG_CONFIG make install
|
||||
- run: make installcheck
|
||||
- if: ${{ failure() }}
|
||||
run: cat regression.diffs
|
||||
- run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install libipc-run-perl
|
||||
- run: make prove_installcheck
|
||||
mac:
|
||||
runs-on: ${{ matrix.os }}
|
||||
if: ${{ !startsWith(github.ref_name, 'windows') }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
- postgres: 16
|
||||
os: macos-14
|
||||
- postgres: 14
|
||||
os: macos-12
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: ankane/setup-postgres@v1
|
||||
with:
|
||||
postgres-version: ${{ matrix.postgres }}
|
||||
- run: make
|
||||
env:
|
||||
PG_CFLAGS: -DUSE_ASSERT_CHECKING -Wall -Wextra -Werror -Wno-unused-parameter
|
||||
- run: make install
|
||||
- run: make installcheck
|
||||
- if: ${{ failure() }}
|
||||
run: cat regression.diffs
|
||||
# Homebrew Postgres does not enable TAP tests, so need to download
|
||||
- run: |
|
||||
brew install cpanm
|
||||
cpanm --notest IPC::Run
|
||||
wget -q https://github.com/postgres/postgres/archive/refs/tags/$TAG.tar.gz
|
||||
tar xf $TAG.tar.gz
|
||||
mv postgres-$TAG postgres
|
||||
env:
|
||||
TAG: ${{ matrix.postgres == 16 && 'REL_16_2' || 'REL_14_11' }}
|
||||
- run: make prove_installcheck PROVE_FLAGS="-I ./postgres/src/test/perl -I ./test/perl"
|
||||
env:
|
||||
PERL5LIB: /Users/runner/perl5/lib/perl5
|
||||
- run: make clean && $(brew --prefix llvm@15)/bin/scan-build --status-bugs make
|
||||
env:
|
||||
PG_CFLAGS: -DUSE_ASSERT_CHECKING
|
||||
windows:
|
||||
runs-on: windows-latest
|
||||
if: ${{ !startsWith(github.ref_name, 'mac') }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: ankane/setup-postgres@v1
|
||||
with:
|
||||
postgres-version: 14
|
||||
- run: |
|
||||
call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvars64.bat" && ^
|
||||
cd %TEMP% && ^
|
||||
nmake /NOLOGO /F Makefile.win && ^
|
||||
nmake /NOLOGO /F Makefile.win install && ^
|
||||
nmake /NOLOGO /F Makefile.win installcheck && ^
|
||||
nmake /NOLOGO /F Makefile.win clean && ^
|
||||
nmake /NOLOGO /F Makefile.win uninstall
|
||||
shell: cmd
|
||||
i386:
|
||||
if: ${{ !startsWith(github.ref_name, 'mac') && !startsWith(github.ref_name, 'windows') }}
|
||||
runs-on: ubuntu-latest
|
||||
container:
|
||||
image: debian:12
|
||||
options: --platform linux/386
|
||||
steps:
|
||||
- run: apt-get update && apt-get install -y build-essential git libipc-run-perl postgresql-15 postgresql-server-dev-15 sudo
|
||||
- run: service postgresql start
|
||||
- run: |
|
||||
git clone https://github.com/${{ github.repository }}.git pgvector
|
||||
cd pgvector
|
||||
git fetch origin ${{ github.ref }}
|
||||
git reset --hard FETCH_HEAD
|
||||
make
|
||||
make install
|
||||
chown -R postgres .
|
||||
sudo -u postgres make installcheck
|
||||
sudo -u postgres make prove_installcheck
|
||||
env:
|
||||
PG_CFLAGS: -DUSE_ASSERT_CHECKING -Wall -Wextra -Werror -Wno-unused-parameter -Wno-sign-compare
|
||||
- if: ${{ failure() }}
|
||||
run: cat pgvector/regression.diffs
|
||||
valgrind:
|
||||
if: ${{ !startsWith(github.ref_name, 'mac') && !startsWith(github.ref_name, 'windows') }}
|
||||
runs-on: ubuntu-latest
|
||||
@@ -12,4 +128,4 @@ jobs:
|
||||
check-ub: yes
|
||||
- run: make OPTFLAGS=""
|
||||
- run: sudo --preserve-env=PG_CONFIG make install
|
||||
- run: make prove_installcheck PROVE_TESTS=test/t/031_hnsw_sparsevec_vacuum_recall.pl
|
||||
- run: make installcheck
|
||||
|
||||
@@ -327,7 +327,7 @@ CREATE FUNCTION jaccard_distance(bit, bit) RETURNS float8
|
||||
CREATE FUNCTION bit_ivfflat_support(internal) RETURNS internal
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION bit_hnsw_support(internal) RETURNS internal
|
||||
CREATE FUNCTION bit_hnsw_max_dims(internal) RETURNS internal
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
-- bit operators
|
||||
@@ -355,13 +355,13 @@ CREATE OPERATOR CLASS bit_hamming_ops
|
||||
FOR TYPE bit USING hnsw AS
|
||||
OPERATOR 1 <~> (bit, bit) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 hamming_distance(bit, bit),
|
||||
FUNCTION 4 bit_hnsw_support(internal);
|
||||
FUNCTION 4 bit_hnsw_max_dims(internal);
|
||||
|
||||
CREATE OPERATOR CLASS bit_jaccard_ops
|
||||
FOR TYPE bit USING hnsw AS
|
||||
OPERATOR 1 <%> (bit, bit) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 jaccard_distance(bit, bit),
|
||||
FUNCTION 4 bit_hnsw_support(internal);
|
||||
FUNCTION 4 bit_hnsw_max_dims(internal);
|
||||
|
||||
-- halfvec type
|
||||
|
||||
@@ -473,7 +473,7 @@ CREATE FUNCTION halfvec_avg(double precision[]) RETURNS halfvec
|
||||
CREATE FUNCTION halfvec_ivfflat_support(internal) RETURNS internal
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION halfvec_hnsw_support(internal) RETURNS internal
|
||||
CREATE FUNCTION halfvec_hnsw_max_dims(internal) RETURNS internal
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
-- halfvec aggregates
|
||||
@@ -663,13 +663,13 @@ CREATE OPERATOR CLASS halfvec_l2_ops
|
||||
FOR TYPE halfvec USING hnsw AS
|
||||
OPERATOR 1 <-> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 halfvec_l2_squared_distance(halfvec, halfvec),
|
||||
FUNCTION 4 halfvec_hnsw_support(internal);
|
||||
FUNCTION 4 halfvec_hnsw_max_dims(internal);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_ip_ops
|
||||
FOR TYPE halfvec USING hnsw AS
|
||||
OPERATOR 1 <#> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec),
|
||||
FUNCTION 4 halfvec_hnsw_support(internal);
|
||||
FUNCTION 4 halfvec_hnsw_max_dims(internal);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_cosine_ops
|
||||
FOR TYPE halfvec USING hnsw AS
|
||||
@@ -677,13 +677,13 @@ CREATE OPERATOR CLASS halfvec_cosine_ops
|
||||
FUNCTION 1 halfvec_negative_inner_product(halfvec, halfvec),
|
||||
FUNCTION 2 l2_norm(halfvec),
|
||||
FUNCTION 3 l2_normalize(halfvec),
|
||||
FUNCTION 4 halfvec_hnsw_support(internal);
|
||||
FUNCTION 4 halfvec_hnsw_max_dims(internal);
|
||||
|
||||
CREATE OPERATOR CLASS halfvec_l1_ops
|
||||
FOR TYPE halfvec USING hnsw AS
|
||||
OPERATOR 1 <+> (halfvec, halfvec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 l1_distance(halfvec, halfvec),
|
||||
FUNCTION 4 halfvec_hnsw_support(internal);
|
||||
FUNCTION 4 halfvec_hnsw_max_dims(internal);
|
||||
|
||||
--- sparsevec type
|
||||
|
||||
@@ -779,7 +779,10 @@ CREATE FUNCTION halfvec_to_sparsevec(halfvec, integer, boolean) RETURNS sparseve
|
||||
CREATE FUNCTION sparsevec_to_halfvec(sparsevec, integer, boolean) RETURNS halfvec
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION sparsevec_hnsw_support(internal) RETURNS internal
|
||||
CREATE FUNCTION sparsevec_hnsw_max_dims(internal) RETURNS internal
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
CREATE FUNCTION sparsevec_hnsw_check_value(internal) RETURNS internal
|
||||
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
|
||||
|
||||
-- sparsevec casts
|
||||
@@ -872,13 +875,15 @@ CREATE OPERATOR CLASS sparsevec_l2_ops
|
||||
FOR TYPE sparsevec USING hnsw AS
|
||||
OPERATOR 1 <-> (sparsevec, sparsevec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 sparsevec_l2_squared_distance(sparsevec, sparsevec),
|
||||
FUNCTION 4 sparsevec_hnsw_support(internal);
|
||||
FUNCTION 4 sparsevec_hnsw_max_dims(internal),
|
||||
FUNCTION 5 sparsevec_hnsw_check_value(internal);
|
||||
|
||||
CREATE OPERATOR CLASS sparsevec_ip_ops
|
||||
FOR TYPE sparsevec USING hnsw AS
|
||||
OPERATOR 1 <#> (sparsevec, sparsevec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 sparsevec_negative_inner_product(sparsevec, sparsevec),
|
||||
FUNCTION 4 sparsevec_hnsw_support(internal);
|
||||
FUNCTION 4 sparsevec_hnsw_max_dims(internal),
|
||||
FUNCTION 5 sparsevec_hnsw_check_value(internal);
|
||||
|
||||
CREATE OPERATOR CLASS sparsevec_cosine_ops
|
||||
FOR TYPE sparsevec USING hnsw AS
|
||||
@@ -886,10 +891,12 @@ CREATE OPERATOR CLASS sparsevec_cosine_ops
|
||||
FUNCTION 1 sparsevec_negative_inner_product(sparsevec, sparsevec),
|
||||
FUNCTION 2 l2_norm(sparsevec),
|
||||
FUNCTION 3 l2_normalize(sparsevec),
|
||||
FUNCTION 4 sparsevec_hnsw_support(internal);
|
||||
FUNCTION 4 sparsevec_hnsw_max_dims(internal),
|
||||
FUNCTION 5 sparsevec_hnsw_check_value(internal);
|
||||
|
||||
CREATE OPERATOR CLASS sparsevec_l1_ops
|
||||
FOR TYPE sparsevec USING hnsw AS
|
||||
OPERATOR 1 <+> (sparsevec, sparsevec) FOR ORDER BY float_ops,
|
||||
FUNCTION 1 l1_distance(sparsevec, sparsevec),
|
||||
FUNCTION 4 sparsevec_hnsw_support(internal);
|
||||
FUNCTION 4 sparsevec_hnsw_max_dims(internal),
|
||||
FUNCTION 5 sparsevec_hnsw_check_value(internal);
|
||||
|
||||
@@ -194,7 +194,7 @@ hnswhandler(PG_FUNCTION_ARGS)
|
||||
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
|
||||
|
||||
amroutine->amstrategies = 0;
|
||||
amroutine->amsupport = 4;
|
||||
amroutine->amsupport = 5;
|
||||
#if PG_VERSION_NUM >= 130000
|
||||
amroutine->amoptsprocnum = 0;
|
||||
#endif
|
||||
|
||||
17
src/hnsw.h
17
src/hnsw.h
@@ -23,7 +23,8 @@
|
||||
#define HNSW_DISTANCE_PROC 1
|
||||
#define HNSW_NORM_PROC 2
|
||||
#define HNSW_NORMALIZE_PROC 3
|
||||
#define HNSW_TYPE_SUPPORT_PROC 4
|
||||
#define HNSW_MAX_DIMS_PROC 4
|
||||
#define HNSW_CHECK_VALUE_PROC 5
|
||||
|
||||
#define HNSW_VERSION 1
|
||||
#define HNSW_MAGIC_NUMBER 0xA953A953
|
||||
@@ -58,15 +59,6 @@
|
||||
#define HNSW_UPDATE_ENTRY_GREATER 1
|
||||
#define HNSW_UPDATE_ENTRY_ALWAYS 2
|
||||
|
||||
typedef enum HnswType
|
||||
{
|
||||
HNSW_TYPE_VECTOR,
|
||||
HNSW_TYPE_HALFVEC,
|
||||
HNSW_TYPE_BIT,
|
||||
HNSW_TYPE_SPARSEVEC,
|
||||
HNSW_TYPE_UNSUPPORTED
|
||||
} HnswType;
|
||||
|
||||
/* Build phases */
|
||||
/* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
|
||||
#define PROGRESS_HNSW_PHASE_LOAD 2
|
||||
@@ -254,7 +246,6 @@ typedef struct HnswBuildState
|
||||
Relation index;
|
||||
IndexInfo *indexInfo;
|
||||
ForkNumber forkNum;
|
||||
HnswType type;
|
||||
|
||||
/* Settings */
|
||||
int dimensions;
|
||||
@@ -269,6 +260,7 @@ typedef struct HnswBuildState
|
||||
FmgrInfo *procinfo;
|
||||
FmgrInfo *normprocinfo;
|
||||
FmgrInfo *normalizeprocinfo;
|
||||
FmgrInfo *checkvalueprocinfo;
|
||||
Oid collation;
|
||||
|
||||
/* Variables */
|
||||
@@ -381,10 +373,9 @@ typedef struct HnswVacuumState
|
||||
int HnswGetM(Relation index);
|
||||
int HnswGetEfConstruction(Relation index);
|
||||
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum);
|
||||
HnswType HnswGetType(Relation index);
|
||||
Datum HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum value);
|
||||
bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
|
||||
void HnswCheckValue(Datum value, HnswType type);
|
||||
void HnswCheckValue(FmgrInfo *procinfo, Oid collation, Datum value);
|
||||
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
|
||||
void HnswInitPage(Buffer buf, Page page);
|
||||
void HnswInit(void);
|
||||
|
||||
@@ -488,7 +488,8 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
|
||||
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
|
||||
|
||||
/* Check value */
|
||||
HnswCheckValue(value, buildstate->type);
|
||||
if (buildstate->checkvalueprocinfo != NULL)
|
||||
HnswCheckValue(buildstate->checkvalueprocinfo, buildstate->collation, value);
|
||||
|
||||
/* Normalize if needed */
|
||||
if (buildstate->normprocinfo != NULL)
|
||||
@@ -675,18 +676,14 @@ HnswSharedMemoryAlloc(Size size, void *state)
|
||||
* Get max dimensions
|
||||
*/
|
||||
static int
|
||||
GetMaxDimensions(HnswType type)
|
||||
GetMaxDimensions(Relation index)
|
||||
{
|
||||
int maxDimensions = HNSW_MAX_DIM;
|
||||
FmgrInfo *procinfo = HnswOptionalProcInfo(index, HNSW_MAX_DIMS_PROC);
|
||||
|
||||
if (type == HNSW_TYPE_HALFVEC)
|
||||
maxDimensions *= 2;
|
||||
else if (type == HNSW_TYPE_BIT)
|
||||
maxDimensions *= 32;
|
||||
else if (type == HNSW_TYPE_SPARSEVEC)
|
||||
maxDimensions = INT_MAX;
|
||||
if (procinfo == NULL)
|
||||
return HNSW_MAX_DIM;
|
||||
|
||||
return maxDimensions;
|
||||
return DatumGetInt32(FunctionCall1(procinfo, PointerGetDatum(NULL)));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -701,13 +698,15 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
|
||||
buildstate->index = index;
|
||||
buildstate->indexInfo = indexInfo;
|
||||
buildstate->forkNum = forkNum;
|
||||
buildstate->type = HnswGetType(index);
|
||||
|
||||
buildstate->m = HnswGetM(index);
|
||||
buildstate->efConstruction = HnswGetEfConstruction(index);
|
||||
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
|
||||
|
||||
maxDimensions = GetMaxDimensions(buildstate->type);
|
||||
if (TupleDescAttr(index->rd_att, 0)->atttypid == VARBITOID)
|
||||
elog(ERROR, "type not supported for hnsw index");
|
||||
|
||||
maxDimensions = GetMaxDimensions(index);
|
||||
|
||||
/* Require column to have dimensions to be indexed */
|
||||
if (buildstate->dimensions < 0)
|
||||
@@ -726,6 +725,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
|
||||
buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
||||
buildstate->normalizeprocinfo = HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC);
|
||||
buildstate->checkvalueprocinfo = HnswOptionalProcInfo(index, HNSW_CHECK_VALUE_PROC);
|
||||
buildstate->collation = index->rd_indcollation[0];
|
||||
|
||||
InitGraph(&buildstate->graphData, NULL, maintenance_work_mem * 1024L);
|
||||
|
||||
@@ -612,6 +612,7 @@ static void
|
||||
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid)
|
||||
{
|
||||
Datum value;
|
||||
FmgrInfo *checkvalueprocinfo = HnswOptionalProcInfo(index, HNSW_CHECK_VALUE_PROC);
|
||||
FmgrInfo *normprocinfo;
|
||||
Oid collation = index->rd_indcollation[0];
|
||||
|
||||
@@ -619,7 +620,8 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
|
||||
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
|
||||
|
||||
/* Check value */
|
||||
HnswCheckValue(value, HnswGetType(index));
|
||||
if (checkvalueprocinfo != NULL)
|
||||
HnswCheckValue(checkvalueprocinfo, collation, value);
|
||||
|
||||
/* Normalize if needed */
|
||||
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
||||
|
||||
@@ -152,27 +152,6 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
|
||||
return index_getprocinfo(index, 1, procnum);
|
||||
}
|
||||
|
||||
/*
|
||||
* Get type
|
||||
*/
|
||||
HnswType
|
||||
HnswGetType(Relation index)
|
||||
{
|
||||
FmgrInfo *procinfo = HnswOptionalProcInfo(index, HNSW_TYPE_SUPPORT_PROC);
|
||||
Oid typid = TupleDescAttr(index->rd_att, 0)->atttypid;
|
||||
HnswType result;
|
||||
|
||||
if (procinfo == NULL)
|
||||
return HNSW_TYPE_VECTOR;
|
||||
|
||||
result = (HnswType) DatumGetInt32(FunctionCall1(procinfo, ObjectIdGetDatum(typid)));
|
||||
|
||||
if (result == HNSW_TYPE_UNSUPPORTED)
|
||||
elog(ERROR, "type not supported for hnsw index");
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/*
|
||||
* Normalize value
|
||||
*/
|
||||
@@ -198,15 +177,9 @@ HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value)
|
||||
* Check if a value can be indexed
|
||||
*/
|
||||
void
|
||||
HnswCheckValue(Datum value, HnswType type)
|
||||
HnswCheckValue(FmgrInfo *procinfo, Oid collation, Datum value)
|
||||
{
|
||||
if (type == HNSW_TYPE_SPARSEVEC)
|
||||
{
|
||||
SparseVector *vec = DatumGetSparseVector(value);
|
||||
|
||||
if (vec->nnz > HNSW_MAX_NNZ)
|
||||
elog(ERROR, "sparsevec cannot have more than %d non-zero elements for hnsw index", HNSW_MAX_NNZ);
|
||||
}
|
||||
FunctionCall1Coll(procinfo, collation, value);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -1303,28 +1276,35 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
||||
}
|
||||
}
|
||||
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_hnsw_support);
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_hnsw_max_dims);
|
||||
Datum
|
||||
halfvec_hnsw_support(PG_FUNCTION_ARGS)
|
||||
halfvec_hnsw_max_dims(PG_FUNCTION_ARGS)
|
||||
{
|
||||
PG_RETURN_INT32(HNSW_TYPE_HALFVEC);
|
||||
PG_RETURN_INT32(HNSW_MAX_DIM * 2);
|
||||
};
|
||||
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(bit_hnsw_support);
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(bit_hnsw_max_dims);
|
||||
Datum
|
||||
bit_hnsw_support(PG_FUNCTION_ARGS)
|
||||
bit_hnsw_max_dims(PG_FUNCTION_ARGS)
|
||||
{
|
||||
Oid typid = PG_GETARG_OID(0);
|
||||
|
||||
if (typid == BITOID)
|
||||
PG_RETURN_INT32(HNSW_TYPE_BIT);
|
||||
else
|
||||
PG_RETURN_INT32(HNSW_TYPE_UNSUPPORTED);
|
||||
PG_RETURN_INT32(HNSW_MAX_DIM * 32);
|
||||
};
|
||||
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(sparsevec_hnsw_support);
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(sparsevec_hnsw_max_dims);
|
||||
Datum
|
||||
sparsevec_hnsw_support(PG_FUNCTION_ARGS)
|
||||
sparsevec_hnsw_max_dims(PG_FUNCTION_ARGS)
|
||||
{
|
||||
PG_RETURN_INT32(HNSW_TYPE_SPARSEVEC);
|
||||
PG_RETURN_INT32(SPARSEVEC_MAX_DIM);
|
||||
};
|
||||
|
||||
PGDLLEXPORT PG_FUNCTION_INFO_V1(sparsevec_hnsw_check_value);
|
||||
Datum
|
||||
sparsevec_hnsw_check_value(PG_FUNCTION_ARGS)
|
||||
{
|
||||
SparseVector *vec = PG_GETARG_SPARSEVEC_P(0);
|
||||
|
||||
if (vec->nnz > HNSW_MAX_NNZ)
|
||||
elog(ERROR, "sparsevec cannot have more than %d non-zero elements for hnsw index", HNSW_MAX_NNZ);
|
||||
|
||||
PG_RETURN_VOID();
|
||||
}
|
||||
|
||||
203
src/ivfkmeans.c
203
src/ivfkmeans.c
@@ -89,31 +89,33 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
|
||||
}
|
||||
|
||||
/*
|
||||
* Norm centers
|
||||
* Apply norm to vector
|
||||
*/
|
||||
static void
|
||||
NormCenters(FmgrInfo *normalizeprocinfo, Oid collation, VectorArray centers)
|
||||
static inline void
|
||||
ApplyNorm(FmgrInfo *normprocinfo, Oid collation, Datum value, IvfflatType type)
|
||||
{
|
||||
MemoryContext normCtx = AllocSetContextCreate(CurrentMemoryContext,
|
||||
"Ivfflat norm temporary context",
|
||||
ALLOCSET_DEFAULT_SIZES);
|
||||
MemoryContext oldCtx = MemoryContextSwitchTo(normCtx);
|
||||
double norm = DatumGetFloat8(FunctionCall1Coll(normprocinfo, collation, value));
|
||||
|
||||
for (int j = 0; j < centers->length; j++)
|
||||
/* TODO Handle zero norm */
|
||||
if (norm > 0)
|
||||
{
|
||||
Datum center = PointerGetDatum(VectorArrayGet(centers, j));
|
||||
Datum newCenter = IvfflatNormValue(normalizeprocinfo, collation, center);
|
||||
Size size = VARSIZE_ANY(DatumGetPointer(newCenter));
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
{
|
||||
Vector *vec = DatumGetVector(value);
|
||||
|
||||
if (size > centers->itemsize)
|
||||
elog(ERROR, "safety check failed");
|
||||
for (int i = 0; i < vec->dim; i++)
|
||||
vec->x[i] /= norm;
|
||||
}
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
{
|
||||
HalfVector *vec = DatumGetHalfVector(value);
|
||||
|
||||
memcpy(DatumGetPointer(center), DatumGetPointer(newCenter), size);
|
||||
MemoryContextReset(normCtx);
|
||||
for (int i = 0; i < vec->dim; i++)
|
||||
vec->x[i] = Float4ToHalfUnchecked(HalfToFloat4(vec->x[i]) / norm);
|
||||
}
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
}
|
||||
|
||||
MemoryContextSwitchTo(oldCtx);
|
||||
MemoryContextDelete(normCtx);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -143,22 +145,6 @@ CompareBitVectors(const void *a, const void *b)
|
||||
return DirectFunctionCall2(bitcmp, VarBitPGetDatum((VarBit *) a), VarBitPGetDatum((VarBit *) b));
|
||||
}
|
||||
|
||||
/*
|
||||
* Sort vector array
|
||||
*/
|
||||
static void
|
||||
SortVectorArray(VectorArray arr, IvfflatType type)
|
||||
{
|
||||
if (type == IVFFLAT_TYPE_VECTOR)
|
||||
qsort(arr->items, arr->length, arr->itemsize, CompareVectors);
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
qsort(arr->items, arr->length, arr->itemsize, CompareHalfVectors);
|
||||
else if (type == IVFFLAT_TYPE_BIT)
|
||||
qsort(arr->items, arr->length, arr->itemsize, CompareBitVectors);
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
}
|
||||
|
||||
/*
|
||||
* Quick approach if we have little data
|
||||
*/
|
||||
@@ -168,12 +154,18 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers, IvfflatTy
|
||||
int dimensions = centers->dim;
|
||||
Oid collation = index->rd_indcollation[0];
|
||||
FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
|
||||
FmgrInfo *normalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORMALIZE_PROC);
|
||||
|
||||
/* Copy existing vectors while avoiding duplicates */
|
||||
if (samples->length > 0)
|
||||
{
|
||||
SortVectorArray(samples, type);
|
||||
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 if (type == IVFFLAT_TYPE_BIT)
|
||||
qsort(samples->items, samples->length, samples->itemsize, CompareBitVectors);
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
|
||||
for (int i = 0; i < samples->length; i++)
|
||||
{
|
||||
@@ -225,12 +217,12 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers, IvfflatTy
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
|
||||
/* Normalize if needed (only needed for random centers) */
|
||||
if (normprocinfo != NULL)
|
||||
ApplyNorm(normprocinfo, collation, center, type);
|
||||
|
||||
centers->length++;
|
||||
}
|
||||
|
||||
/* Fine if existing vectors are normalized twice */
|
||||
if (normprocinfo != NULL)
|
||||
NormCenters(normalizeprocinfo, collation, centers);
|
||||
}
|
||||
|
||||
#ifdef IVFFLAT_MEMORY
|
||||
@@ -251,14 +243,27 @@ ShowMemoryUsage(MemoryContext context, Size estimatedSize)
|
||||
#endif
|
||||
|
||||
/*
|
||||
* Sum centers
|
||||
* Compute new centers
|
||||
*/
|
||||
static void
|
||||
SumCenters(VectorArray samples, VectorArray aggCenters, int *closestCenters, IvfflatType type)
|
||||
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++)
|
||||
@@ -294,68 +299,6 @@ SumCenters(VectorArray samples, VectorArray aggCenters, int *closestCenters, Ivf
|
||||
}
|
||||
else
|
||||
elog(ERROR, "Unsupported type");
|
||||
}
|
||||
|
||||
/*
|
||||
* Set new centers
|
||||
*/
|
||||
static void
|
||||
SetNewCenters(VectorArray aggCenters, VectorArray newCenters, IvfflatType type)
|
||||
{
|
||||
int dimensions = aggCenters->dim;
|
||||
int numCenters = aggCenters->maxlen;
|
||||
|
||||
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]);
|
||||
}
|
||||
}
|
||||
else if (type == IVFFLAT_TYPE_BIT)
|
||||
{
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, j);
|
||||
VarBit *newCenter = (VarBit *) VectorArrayGet(newCenters, j);
|
||||
unsigned char *nx = VARBITS(newCenter);
|
||||
|
||||
for (uint32 k = 0; k < VARBITBYTES(newCenter); k++)
|
||||
nx[k] = 0;
|
||||
|
||||
for (int k = 0; k < dimensions; k++)
|
||||
nx[k / 8] |= (aggCenter->x[k] > 0.5) << (7 - (k % 8));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
* Compute new centers
|
||||
*/
|
||||
static void
|
||||
ComputeNewCenters(VectorArray samples, VectorArray aggCenters, VectorArray newCenters, int *centerCounts, int *closestCenters, FmgrInfo *normprocinfo, FmgrInfo *normalizeprocinfo, 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 */
|
||||
SumCenters(samples, aggCenters, closestCenters, type);
|
||||
|
||||
/* Increment count of closest center */
|
||||
for (int j = 0; j < numSamples; j++)
|
||||
@@ -388,11 +331,43 @@ ComputeNewCenters(VectorArray samples, VectorArray aggCenters, VectorArray newCe
|
||||
}
|
||||
|
||||
/* Set new centers if different from agg centers */
|
||||
SetNewCenters(aggCenters, newCenters, type);
|
||||
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]);
|
||||
}
|
||||
}
|
||||
else if (type == IVFFLAT_TYPE_BIT)
|
||||
{
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, j);
|
||||
VarBit *newCenter = (VarBit *) VectorArrayGet(newCenters, j);
|
||||
unsigned char *nx = VARBITS(newCenter);
|
||||
|
||||
for (uint32 k = 0; k < VARBITBYTES(newCenter); k++)
|
||||
nx[k] = 0;
|
||||
|
||||
for (int k = 0; k < dimensions; k++)
|
||||
nx[k / 8] |= (aggCenter->x[k] > 0.5) << (7 - (k % 8));
|
||||
}
|
||||
}
|
||||
|
||||
/* Normalize if needed */
|
||||
if (normprocinfo != NULL)
|
||||
NormCenters(normalizeprocinfo, collation, newCenters);
|
||||
{
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
Datum newCenter = PointerGetDatum(VectorArrayGet(newCenters, j));
|
||||
|
||||
ApplyNorm(normprocinfo, collation, newCenter, type);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -408,7 +383,6 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
|
||||
{
|
||||
FmgrInfo *procinfo;
|
||||
FmgrInfo *normprocinfo;
|
||||
FmgrInfo *normalizeprocinfo;
|
||||
Oid collation;
|
||||
int dimensions = centers->dim;
|
||||
int numCenters = centers->maxlen;
|
||||
@@ -456,7 +430,6 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
|
||||
/* Set support functions */
|
||||
procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
|
||||
normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
|
||||
normalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORMALIZE_PROC);
|
||||
collation = index->rd_indcollation[0];
|
||||
|
||||
/* Use memory context */
|
||||
@@ -476,8 +449,6 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
|
||||
newcdist = palloc(newcdistSize);
|
||||
|
||||
aggCenters = VectorArrayInit(numCenters, dimensions, VECTOR_SIZE(dimensions));
|
||||
aggCenters->length = numCenters;
|
||||
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
Vector *vec = (Vector *) VectorArrayGet(aggCenters, j);
|
||||
@@ -494,7 +465,6 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
|
||||
else if (type == IVFFLAT_TYPE_HALFVEC)
|
||||
{
|
||||
newCenters = VectorArrayInit(numCenters, dimensions, centers->itemsize);
|
||||
newCenters->length = numCenters;
|
||||
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
@@ -507,7 +477,6 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
|
||||
else if (type == IVFFLAT_TYPE_BIT)
|
||||
{
|
||||
newCenters = VectorArrayInit(numCenters, dimensions, centers->itemsize);
|
||||
newCenters->length = numCenters;
|
||||
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
{
|
||||
@@ -658,7 +627,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
|
||||
}
|
||||
|
||||
/* Step 4: For each center c, let m(c) be mean of all points assigned */
|
||||
ComputeNewCenters(samples, aggCenters, newCenters, centerCounts, closestCenters, normprocinfo, normalizeprocinfo, collation, type);
|
||||
ComputeNewCenters(samples, aggCenters, newCenters, centerCounts, closestCenters, normprocinfo, collation, type);
|
||||
|
||||
/* Step 5 */
|
||||
for (int j = 0; j < numCenters; j++)
|
||||
@@ -741,7 +710,13 @@ CheckCenters(Relation index, VectorArray centers, IvfflatType type)
|
||||
if (type != IVFFLAT_TYPE_BIT)
|
||||
{
|
||||
/* Ensure no duplicate centers */
|
||||
SortVectorArray(centers, type);
|
||||
/* Fine to sort in-place */
|
||||
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++)
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user