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

DRY halfvec distance functions

Browse Source
This commit is contained in:
Andrew Kane
2024-04-07 19:03:20 -07:00
parent d861a0304e
commit 3bd67fef54
1 changed files with 35 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
src/halfvec.c
View File

@@ -801,20 +801,15 @@ vector_to_halfvec(PG_FUNCTION_ARGS)
} }
/* /*
* Get the L2 distance between half vectors * Get the L2 squared distance between half vectors
*/ */
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_l2_distance); static double
Datum l2_distance_squared_internal(HalfVector * a, HalfVector * b)
halfvec_l2_distance(PG_FUNCTION_ARGS)
{ {
HalfVector *a = PG_GETARG_HALFVEC_P(0);
HalfVector *b = PG_GETARG_HALFVEC_P(1);
half *ax = a->x; half *ax = a->x;
half *bx = b->x; half *bx = b->x;
float distance = 0.0; float distance = 0.0;
CheckDims(a, b);
/* Auto-vectorized */ /* Auto-vectorized */
for (int i = 0; i < a->dim; i++) for (int i = 0; i < a->dim; i++)
{ {
@@ -823,7 +818,22 @@ halfvec_l2_distance(PG_FUNCTION_ARGS)
distance += diff * diff; distance += diff * diff;
} }
PG_RETURN_FLOAT8(sqrt((double) distance)); return (double) distance;
}
/*
* Get the L2 distance between half vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(halfvec_l2_distance);
Datum
halfvec_l2_distance(PG_FUNCTION_ARGS)
{
HalfVector *a = PG_GETARG_HALFVEC_P(0);
HalfVector *b = PG_GETARG_HALFVEC_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8(sqrt(l2_distance_squared_internal(a, b)));
} }
/* /*
@@ -835,21 +845,27 @@ halfvec_l2_squared_distance(PG_FUNCTION_ARGS)
{ {
HalfVector *a = PG_GETARG_HALFVEC_P(0); HalfVector *a = PG_GETARG_HALFVEC_P(0);
HalfVector *b = PG_GETARG_HALFVEC_P(1); HalfVector *b = PG_GETARG_HALFVEC_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8(l2_distance_squared_internal(a, b));
}
/*
* Get the inner product of two half vectors
*/
static double
inner_product_internal(HalfVector * a, HalfVector * b)
{
half *ax = a->x; half *ax = a->x;
half *bx = b->x; half *bx = b->x;
float distance = 0.0; float distance = 0.0;
CheckDims(a, b);
/* Auto-vectorized */ /* Auto-vectorized */
for (int i = 0; i < a->dim; i++) for (int i = 0; i < a->dim; i++)
{ distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);
float diff = HalfToFloat4(ax[i]) - HalfToFloat4(bx[i]);
distance += diff * diff; return (double) distance;
}
PG_RETURN_FLOAT8((double) distance);
} }
/* /*
@@ -861,17 +877,10 @@ halfvec_inner_product(PG_FUNCTION_ARGS)
{ {
HalfVector *a = PG_GETARG_HALFVEC_P(0); HalfVector *a = PG_GETARG_HALFVEC_P(0);
HalfVector *b = PG_GETARG_HALFVEC_P(1); HalfVector *b = PG_GETARG_HALFVEC_P(1);
half *ax = a->x;
half *bx = b->x;
float distance = 0.0;
CheckDims(a, b); CheckDims(a, b);
/* Auto-vectorized */ PG_RETURN_FLOAT8(inner_product_internal(a, b));
for (int i = 0; i < a->dim; i++)
distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);
PG_RETURN_FLOAT8((double) distance);
} }
/* /*
@@ -883,17 +892,10 @@ halfvec_negative_inner_product(PG_FUNCTION_ARGS)
{ {
HalfVector *a = PG_GETARG_HALFVEC_P(0); HalfVector *a = PG_GETARG_HALFVEC_P(0);
HalfVector *b = PG_GETARG_HALFVEC_P(1); HalfVector *b = PG_GETARG_HALFVEC_P(1);
half *ax = a->x;
half *bx = b->x;
float distance = 0.0;
CheckDims(a, b); CheckDims(a, b);
/* Auto-vectorized */ PG_RETURN_FLOAT8(-inner_product_internal(a, b));
for (int i = 0; i < a->dim; i++)
distance += HalfToFloat4(ax[i]) * HalfToFloat4(bx[i]);
PG_RETURN_FLOAT8((double) distance * -1);
} }
/* /*