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

Added support for IVFFlat [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-09-23 18:40:49 -07:00
parent 6105a01881
commit 0a2803130f
3 changed files with 228 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
src/ivfutils.c
View File

@@ -7,6 +7,7 @@
#include "halfutils.h"
#include "halfvec.h"
#include "ivfflat.h"
#include "minivec.h"
#include "storage/bufmgr.h"
/*
@@ -231,6 +232,7 @@ IvfflatUpdateList(Relation index, ListInfo listInfo,
PGDLLEXPORT Datum l2_normalize(PG_FUNCTION_ARGS);
PGDLLEXPORT Datum halfvec_l2_normalize(PG_FUNCTION_ARGS);
PGDLLEXPORT Datum minivec_l2_normalize(PG_FUNCTION_ARGS);
PGDLLEXPORT Datum sparsevec_l2_normalize(PG_FUNCTION_ARGS);
static Size
@@ -245,6 +247,12 @@ HalfvecItemSize(int dimensions)
return HALFVEC_SIZE(dimensions);
}
static Size
MinivecItemSize(int dimensions)
{
return MINIVEC_SIZE(dimensions);
}
static Size
BitItemSize(int dimensions)
{
@@ -275,6 +283,18 @@ HalfvecUpdateCenter(Pointer v, int dimensions, float *x)
vec->x[k] = Float4ToHalfUnchecked(x[k]);
}
static void
MinivecUpdateCenter(Pointer v, int dimensions, float *x)
{
MiniVector *vec = (MiniVector *) v;
SET_VARSIZE(vec, MINIVEC_SIZE(dimensions));
vec->dim = dimensions;
for (int k = 0; k < dimensions; k++)
vec->x[k] = Float4ToFp8Unchecked(x[k]);
}
static void
BitUpdateCenter(Pointer v, int dimensions, float *x)
{
@@ -309,6 +329,15 @@ HalfvecSumCenter(Pointer v, float *x)
x[k] += HalfToFloat4(vec->x[k]);
}
static void
MinivecSumCenter(Pointer v, float *x)
{
MiniVector *vec = (MiniVector *) v;
for (int k = 0; k < vec->dim; k++)
x[k] += Fp8ToFloat4(vec->x[k]);
}
static void
BitSumCenter(Pointer v, float *x)
{
@@ -357,6 +386,21 @@ ivfflat_halfvec_support(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(&typeInfo);
};
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(ivfflat_minivec_support);
Datum
ivfflat_minivec_support(PG_FUNCTION_ARGS)
{
static const IvfflatTypeInfo typeInfo = {
.maxDimensions = IVFFLAT_MAX_DIM * 4,
.normalize = minivec_l2_normalize,
.itemSize = MinivecItemSize,
.updateCenter = MinivecUpdateCenter,
.sumCenter = MinivecSumCenter
};
PG_RETURN_POINTER(&typeInfo);
};
FUNCTION_PREFIX PG_FUNCTION_INFO_V1(ivfflat_bit_support);
Datum
ivfflat_bit_support(PG_FUNCTION_ARGS)