liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-06 04:30:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Removed type-specific code from IVFFlat - #527

Browse Source
This commit is contained in:
Andrew Kane
2024-04-24 22:19:53 -07:00
parent 52bfedddc2
commit 8f00d679d6
9 changed files with 217 additions and 314 deletions
Download Patch File Download Diff File Expand all files Collapse all files

98
src/ivfutils.c
View File

@@ -1,8 +1,11 @@
#include "postgres.h"
#include "access/generic_xlog.h"
#include "bitvec.h"
#include "catalog/pg_type.h"
#include "fmgr.h"
#include "halfutils.h"
#include "halfvec.h"
#include "ivfflat.h"
#include "storage/bufmgr.h"
@@ -239,21 +242,96 @@ ivfflat_bit_max_dims(PG_FUNCTION_ARGS)
PG_RETURN_INT32(IVFFLAT_MAX_DIM * 32);
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_support);
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_vector_update_center);
Datum
ivfflat_halfvec_support(PG_FUNCTION_ARGS)
ivfflat_vector_update_center(PG_FUNCTION_ARGS)
{
PG_RETURN_INT32(IVFFLAT_TYPE_HALFVEC);
Vector *vec = PG_GETARG_VECTOR_P(0);
int dimensions = PG_GETARG_INT32(1);
float *x = (float *) PG_GETARG_POINTER(2);
SET_VARSIZE(vec, VECTOR_SIZE(dimensions));
vec->dim = dimensions;
for (int k = 0; k < dimensions; k++)
vec->x[k] = x[k];
PG_RETURN_VOID();
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_bit_support);
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_update_center);
Datum
ivfflat_bit_support(PG_FUNCTION_ARGS)
ivfflat_halfvec_update_center(PG_FUNCTION_ARGS)
{
Oid typid = PG_GETARG_OID(0);
HalfVector *vec = PG_GETARG_HALFVEC_P(0);
int dimensions = PG_GETARG_INT32(1);
float *x = (float *) PG_GETARG_POINTER(2);
if (typid == BITOID)
PG_RETURN_INT32(IVFFLAT_TYPE_BIT);
else
PG_RETURN_INT32(IVFFLAT_TYPE_UNSUPPORTED);
SET_VARSIZE(vec, HALFVEC_SIZE(dimensions));
vec->dim = dimensions;
for (int k = 0; k < dimensions; k++)
vec->x[k] = Float4ToHalfUnchecked(x[k]);
PG_RETURN_VOID();
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_bit_update_center);
Datum
ivfflat_bit_update_center(PG_FUNCTION_ARGS)
{
VarBit *vec = PG_GETARG_VARBIT_P(0);
int dimensions = PG_GETARG_INT32(1);
float *x = (float *) PG_GETARG_POINTER(2);
unsigned char *nx = VARBITS(vec);
SET_VARSIZE(vec, VARBITTOTALLEN(dimensions));
VARBITLEN(vec) = dimensions;
for (uint32 k = 0; k < VARBITBYTES(vec); k++)
nx[k] = 0;
for (int k = 0; k < dimensions; k++)
nx[k / 8] |= (x[k] > 0.5 ? 1 : 0) << (7 - (k % 8));
PG_RETURN_VOID();
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_vector_sum_center);
Datum
ivfflat_vector_sum_center(PG_FUNCTION_ARGS)
{
Vector *vec = PG_GETARG_VECTOR_P(0);
float *x = (float *) PG_GETARG_POINTER(1);
for (int k = 0; k < vec->dim; k++)
x[k] += vec->x[k];
PG_RETURN_VOID();
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_sum_center);
Datum
ivfflat_halfvec_sum_center(PG_FUNCTION_ARGS)
{
HalfVector *vec = PG_GETARG_HALFVEC_P(0);
float *x = (float *) PG_GETARG_POINTER(1);
for (int k = 0; k < vec->dim; k++)
x[k] += HalfToFloat4(vec->x[k]);
PG_RETURN_VOID();
}
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_bit_sum_center);
Datum
ivfflat_bit_sum_center(PG_FUNCTION_ARGS)
{
VarBit *vec = PG_GETARG_VARBIT_P(0);
float *x = (float *) PG_GETARG_POINTER(1);
for (int k = 0; k < VARBITLEN(vec); k++)
x[k] += (float) (((VARBITS(vec)[k / 8]) >> (7 - (k % 8))) & 0x01);
PG_RETURN_VOID();
}