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

Added casts between halfvec and sparsevec

Browse Source
This commit is contained in:
Andrew Kane
2024-04-19 18:03:07 -07:00
parent fd4fbd238c
commit fb77671d05
6 changed files with 132 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
src/halfvec.c
View File

@@ -11,6 +11,7 @@
#include "lib/stringinfo.h"
#include "libpq/pqformat.h"
#include "port.h" /* for strtof() */
#include "sparsevec.h"
#include "utils/array.h"
#include "utils/builtins.h"
#include "utils/float.h"
@@ -1174,3 +1175,26 @@ halfvec_avg(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(result);
}
/*
* Convert sparse vector to half vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(sparsevec_to_halfvec);
Datum
sparsevec_to_halfvec(PG_FUNCTION_ARGS)
{
SparseVector *svec = PG_GETARG_SPARSEVEC_P(0);
int32 typmod = PG_GETARG_INT32(1);
HalfVector *result;
int dim = svec->dim;
float *values = SPARSEVEC_VALUES(svec);
CheckDim(dim);
CheckExpectedDim(typmod, dim);
result = InitHalfVector(dim);
for (int i = 0; i < svec->nnz; i++)
result->x[svec->indices[i] - 1] = Float4ToHalf(values[i]);
PG_RETURN_POINTER(result);
}

45
src/sparsevec.c
View File

@@ -4,6 +4,8 @@
#include <math.h>
#include "fmgr.h"
#include "halfutils.h"
#include "halfvec.h"
#include "libpq/pqformat.h"
#include "sparsevec.h"
#include "utils/array.h"
@@ -612,6 +614,49 @@ vector_to_sparsevec(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(result);
}
/*
* Convert half vector to sparse vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_to_sparsevec);
Datum
halfvec_to_sparsevec(PG_FUNCTION_ARGS)
{
HalfVector *vec = PG_GETARG_HALFVEC_P(0);
int32 typmod = PG_GETARG_INT32(1);
SparseVector *result;
int dim = vec->dim;
int nnz = 0;
float *values;
int j = 0;
CheckDim(dim);
CheckExpectedDim(typmod, dim);
for (int i = 0; i < dim; i++)
{
if (!HalfIsZero(vec->x[i]))
nnz++;
}
result = InitSparseVector(dim, nnz);
values = SPARSEVEC_VALUES(result);
for (int i = 0; i < dim; i++)
{
if (!HalfIsZero(vec->x[i]))
{
/* Safety check */
if (j >= result->nnz)
elog(ERROR, "safety check failed");
result->indices[j] = i + 1;
values[j] = HalfToFloat4(vec->x[i]);
j++;
}
}
PG_RETURN_POINTER(result);
}
/*
* Get the L2 squared distance between sparse vectors
*/