Added support for halfvec to IVFFlat

This commit is contained in:
Andrew Kane
2024-04-11 19:56:39 -07:00
parent a4531ca51f
commit 8d9400bae3
17 changed files with 439 additions and 39 deletions

View File

@@ -2,6 +2,7 @@
#include "access/generic_xlog.h"
#include "catalog/pg_type.h"
#include "halfvec.h"
#include "ivfflat.h"
#include "storage/bufmgr.h"
#include "vector.h"
@@ -77,6 +78,8 @@ IvfflatGetType(Relation index)
type = (Form_pg_type) GETSTRUCT(tuple);
if (strcmp(NameStr(type->typname), "vector") == 0)
result = IVFFLAT_TYPE_VECTOR;
else if (strcmp(NameStr(type->typname), "halfvec") == 0)
result = IVFFLAT_TYPE_HALFVEC;
else
{
ReleaseSysCache(tuple);
@@ -113,6 +116,16 @@ IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, IvfflatType ty
*value = PointerGetDatum(result);
}
else if (type == IVFFLAT_TYPE_HALFVEC)
{
HalfVector *v = DatumGetHalfVector(*value);
HalfVector *result = InitHalfVector(v->dim);
for (int i = 0; i < v->dim; i++)
result->x[i] = Float4ToHalfUnchecked(HalfToFloat4(v->x[i]) / norm);
*value = PointerGetDatum(result);
}
else
elog(ERROR, "Unsupported type");