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

Added quantize_binary and subvector functions for halfvec

Browse Source
This commit is contained in:
Andrew Kane
2024-04-03 14:53:03 -07:00
parent 253acbccf4
commit aaa2d644ce
8 changed files with 136 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
src/halfvec.c
View File

@@ -2,6 +2,7 @@
#include <math.h>
#include "bitvector.h"
#include "catalog/pg_type.h"
#include "common/shortest_dec.h"
#include "fmgr.h"
@@ -967,3 +968,53 @@ halfvec_norm(PG_FUNCTION_ARGS)
PG_RETURN_FLOAT8(sqrt(norm));
}
/*
* Quantize a half vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_quantize_binary);
Datum
halfvec_quantize_binary(PG_FUNCTION_ARGS)
{
HalfVector *a = PG_GETARG_HALFVEC_P(0);
half *ax = a->x;
VarBit *result = InitBitVector(a->dim);
unsigned char *rx = VARBITS(result);
for (int i = 0; i < a->dim; i++)
rx[i / 8] |= (HalfToFloat4(ax[i]) > 0) << (7 - (i % 8));
PG_RETURN_VARBIT_P(result);
}
/*
* Get a subvector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_subvector);
Datum
halfvec_subvector(PG_FUNCTION_ARGS)
{
HalfVector *a = PG_GETARG_HALFVEC_P(0);
int32 start = PG_GETARG_INT32(1);
int32 count = PG_GETARG_INT32(2);
int32 end = start + count;
half *ax = a->x;
HalfVector *result;
int dim;
/* Indexing starts at 1, like substring */
if (start < 1)
start = 1;
if (end > a->dim)
end = a->dim + 1;
dim = end - start;
CheckDim(dim);
result = InitHalfVector(dim);
for (int i = 0; i < dim; i++)
result->x[i] = ax[start - 1 + i];
PG_RETURN_POINTER(result);
}