Added support functions for max dimensions for ivfflat

This commit is contained in:
Andrew Kane
2024-04-24 15:27:10 -07:00
parent e81547847f
commit 69c3e719f7
6 changed files with 55 additions and 18 deletions

View File

@@ -323,16 +323,14 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
* Get max dimensions
*/
static int
GetMaxDimensions(IvfflatType type)
GetMaxDimensions(Relation index)
{
int maxDimensions = IVFFLAT_MAX_DIM;
FmgrInfo *procinfo = IvfflatOptionalProcInfo(index, IVFFLAT_MAX_DIMS_PROC);
if (type == IVFFLAT_TYPE_HALFVEC)
maxDimensions *= 2;
else if (type == IVFFLAT_TYPE_BIT)
maxDimensions *= 32;
if (procinfo == NULL)
return IVFFLAT_MAX_DIM;
return maxDimensions;
return DatumGetInt32(FunctionCall1(procinfo, PointerGetDatum(NULL)));
}
/*
@@ -367,7 +365,11 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->lists = IvfflatGetLists(index);
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
maxDimensions = GetMaxDimensions(buildstate->type);
/* Disallow varbit since require fixed dimensions */
if (TupleDescAttr(index->rd_att, 0)->atttypid == VARBITOID)
elog(ERROR, "type not supported for ivfflat index");
maxDimensions = GetMaxDimensions(index);
/* Require column to have dimensions to be indexed */
if (buildstate->dimensions < 0)

View File

@@ -188,7 +188,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0;
amroutine->amsupport = 6;
amroutine->amsupport = 7;
#if PG_VERSION_NUM >= 130000
amroutine->amoptsprocnum = 0;
#endif

View File

@@ -29,7 +29,8 @@
#define IVFFLAT_KMEANS_DISTANCE_PROC 3
#define IVFFLAT_KMEANS_NORM_PROC 4
#define IVFFLAT_NORMALIZE_PROC 5
#define IVFFLAT_TYPE_SUPPORT_PROC 6
#define IVFFLAT_MAX_DIMS_PROC 6
#define IVFFLAT_TYPE_SUPPORT_PROC 7
#define IVFFLAT_VERSION 1
#define IVFFLAT_MAGIC_NUMBER 0x14FF1A7

View File

@@ -246,6 +246,20 @@ IvfflatUpdateList(Relation index, ListInfo listInfo,
}
}
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_max_dims);
Datum
ivfflat_halfvec_max_dims(PG_FUNCTION_ARGS)
{
PG_RETURN_INT32(IVFFLAT_MAX_DIM * 2);
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_bit_max_dims);
Datum
ivfflat_bit_max_dims(PG_FUNCTION_ARGS)
{
PG_RETURN_INT32(IVFFLAT_MAX_DIM * 32);
};
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_support);
Datum
ivfflat_halfvec_support(PG_FUNCTION_ARGS)