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

Switched to static const for IVFFlat type info

Browse Source
This commit is contained in:
Andrew Kane
2024-04-25 12:30:49 -07:00
parent 91cf4d223e
commit ec640f3b57
4 changed files with 38 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/ivfbuild.c
View File

@@ -325,18 +325,14 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
static void static void
InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, IndexInfo *indexInfo) InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, IndexInfo *indexInfo)
{ {
IvfflatTypeInfo *typeInfo = &buildstate->typeInfo;
buildstate->heap = heap; buildstate->heap = heap;
buildstate->index = index; buildstate->index = index;
buildstate->indexInfo = indexInfo; buildstate->indexInfo = indexInfo;
buildstate->typeInfo = IvfflatGetTypeInfo(index);
buildstate->lists = IvfflatGetLists(index); buildstate->lists = IvfflatGetLists(index);
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod; buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
typeInfo->dimensions = buildstate->dimensions;
IvfflatGetTypeInfo(typeInfo, index);
/* Disallow varbit since require fixed dimensions */ /* Disallow varbit since require fixed dimensions */
if (TupleDescAttr(index->rd_att, 0)->atttypid == VARBITOID) if (TupleDescAttr(index->rd_att, 0)->atttypid == VARBITOID)
elog(ERROR, "type not supported for ivfflat index"); elog(ERROR, "type not supported for ivfflat index");
@@ -345,8 +341,8 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
if (buildstate->dimensions < 0) if (buildstate->dimensions < 0)
elog(ERROR, "column does not have dimensions"); elog(ERROR, "column does not have dimensions");
if (buildstate->dimensions > typeInfo->maxDimensions) if (buildstate->dimensions > buildstate->typeInfo->maxDimensions)
elog(ERROR, "column cannot have more than %d dimensions for ivfflat index", typeInfo->maxDimensions); elog(ERROR, "column cannot have more than %d dimensions for ivfflat index", buildstate->typeInfo->maxDimensions);
buildstate->reltuples = 0; buildstate->reltuples = 0;
buildstate->indtuples = 0; buildstate->indtuples = 0;
@@ -370,7 +366,8 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual); buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, typeInfo->itemsize); /* TODO Fix item size */
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, VECTOR_SIZE(buildstate->dimensions));
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists); buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
@@ -440,7 +437,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
} }
/* Calculate centers */ /* Calculate centers */
IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers, &buildstate->typeInfo)); IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers, buildstate->typeInfo));
/* Free samples before we allocate more memory */ /* Free samples before we allocate more memory */
VectorArrayFree(buildstate->samples); VectorArrayFree(buildstate->samples);

6
src/ivfflat.h
View File

@@ -165,7 +165,7 @@ typedef struct IvfflatBuildState
Relation heap; Relation heap;
Relation index; Relation index;
IndexInfo *indexInfo; IndexInfo *indexInfo;
IvfflatTypeInfo typeInfo; const IvfflatTypeInfo *typeInfo;
/* Settings */ /* Settings */
int dimensions; int dimensions;
@@ -279,7 +279,7 @@ typedef IvfflatScanOpaqueData * IvfflatScanOpaque;
/* Methods */ /* Methods */
VectorArray VectorArrayInit(int maxlen, int dimensions, Size itemsize); VectorArray VectorArrayInit(int maxlen, int dimensions, Size itemsize);
void VectorArrayFree(VectorArray arr); void VectorArrayFree(VectorArray arr);
void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTypeInfo * typeInfo); void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, const IvfflatTypeInfo * typeInfo);
FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum); FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum);
Datum IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum value); Datum IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum value);
bool IvfflatCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value); bool IvfflatCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
@@ -292,7 +292,7 @@ Buffer IvfflatNewBuffer(Relation index, ForkNumber forkNum);
void IvfflatInitPage(Buffer buf, Page page); void IvfflatInitPage(Buffer buf, Page page);
void IvfflatInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state); void IvfflatInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
void IvfflatInit(void); void IvfflatInit(void);
void IvfflatGetTypeInfo(IvfflatTypeInfo * typeInfo, Relation index); const IvfflatTypeInfo *IvfflatGetTypeInfo(Relation index);
PGDLLEXPORT void IvfflatParallelBuildMain(dsm_segment *seg, shm_toc *toc); PGDLLEXPORT void IvfflatParallelBuildMain(dsm_segment *seg, shm_toc *toc);
/* Index access methods */ /* Index access methods */

14
src/ivfkmeans.c
View File

@@ -120,7 +120,7 @@ NormCenters(FmgrInfo *normalizeprocinfo, Oid collation, VectorArray centers)
* Quick approach if we have no data * Quick approach if we have no data
*/ */
static void static void
RandomCenters(Relation index, VectorArray centers, IvfflatTypeInfo * typeInfo) RandomCenters(Relation index, VectorArray centers, const IvfflatTypeInfo * typeInfo)
{ {
int dimensions = centers->dim; int dimensions = centers->dim;
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
@@ -168,7 +168,7 @@ ShowMemoryUsage(MemoryContext context, Size estimatedSize)
* Sum centers * Sum centers
*/ */
static void static void
SumCenters(VectorArray samples, float *agg, int *closestCenters, IvfflatTypeInfo * typeInfo) SumCenters(VectorArray samples, float *agg, int *closestCenters, const IvfflatTypeInfo * typeInfo)
{ {
for (int j = 0; j < samples->length; j++) for (int j = 0; j < samples->length; j++)
{ {
@@ -182,7 +182,7 @@ SumCenters(VectorArray samples, float *agg, int *closestCenters, IvfflatTypeInfo
* Update centers * Update centers
*/ */
static void static void
UpdateCenters(float *agg, VectorArray centers, IvfflatTypeInfo * typeInfo) UpdateCenters(float *agg, VectorArray centers, const IvfflatTypeInfo * typeInfo)
{ {
for (int j = 0; j < centers->length; j++) for (int j = 0; j < centers->length; j++)
{ {
@@ -196,7 +196,7 @@ UpdateCenters(float *agg, VectorArray centers, IvfflatTypeInfo * typeInfo)
* Compute new centers * Compute new centers
*/ */
static void static void
ComputeNewCenters(VectorArray samples, float *agg, VectorArray newCenters, int *centerCounts, int *closestCenters, FmgrInfo *normprocinfo, FmgrInfo *normalizeprocinfo, Oid collation, IvfflatTypeInfo * typeInfo) ComputeNewCenters(VectorArray samples, float *agg, VectorArray newCenters, int *centerCounts, int *closestCenters, FmgrInfo *normprocinfo, FmgrInfo *normalizeprocinfo, Oid collation, const IvfflatTypeInfo * typeInfo)
{ {
int dimensions = newCenters->dim; int dimensions = newCenters->dim;
int numCenters = newCenters->length; int numCenters = newCenters->length;
@@ -263,7 +263,7 @@ ComputeNewCenters(VectorArray samples, float *agg, VectorArray newCenters, int *
* https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf * https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf
*/ */
static void static void
ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTypeInfo * typeInfo) ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, const IvfflatTypeInfo * typeInfo)
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
@@ -517,7 +517,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
* Detect issues with centers * Detect issues with centers
*/ */
static void static void
CheckCenters(Relation index, VectorArray centers, IvfflatTypeInfo * typeInfo) CheckCenters(Relation index, VectorArray centers, const IvfflatTypeInfo * typeInfo)
{ {
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
float *scratch = palloc(sizeof(float) * centers->dim); float *scratch = palloc(sizeof(float) * centers->dim);
@@ -568,7 +568,7 @@ CheckCenters(Relation index, VectorArray centers, IvfflatTypeInfo * typeInfo)
* We use spherical k-means for inner product and cosine * We use spherical k-means for inner product and cosine
*/ */
void void
IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTypeInfo * typeInfo) IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, const IvfflatTypeInfo * typeInfo)
{ {
if (samples->length == 0) if (samples->length == 0)
RandomCenters(index, centers, typeInfo); RandomCenters(index, centers, typeInfo);

43
src/ivfutils.c
View File

@@ -298,46 +298,47 @@ BitSumCenter(Pointer v, float *x)
/* /*
* Get type info * Get type info
*/ */
void const IvfflatTypeInfo *
IvfflatGetTypeInfo(IvfflatTypeInfo * typeInfo, Relation index) IvfflatGetTypeInfo(Relation index)
{ {
FmgrInfo *procinfo = IvfflatOptionalProcInfo(index, IVFFLAT_TYPE_INFO_PROC); FmgrInfo *procinfo = IvfflatOptionalProcInfo(index, IVFFLAT_TYPE_INFO_PROC);
if (procinfo == NULL) if (procinfo == NULL)
{ {
typeInfo->maxDimensions = IVFFLAT_MAX_DIM; static const IvfflatTypeInfo typeInfo = {
typeInfo->itemsize = VECTOR_SIZE(typeInfo->dimensions); .maxDimensions = IVFFLAT_MAX_DIM,
typeInfo->updateCenter = VectorUpdateCenter; .updateCenter = VectorUpdateCenter,
typeInfo->sumCenter = VectorSumCenter; .sumCenter = VectorSumCenter
};
return (&typeInfo);
} }
else else
FunctionCall1(procinfo, PointerGetDatum(typeInfo)); return (const IvfflatTypeInfo *) DatumGetPointer(FunctionCall0Coll(procinfo, InvalidOid));
} }
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_support); PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_halfvec_support);
Datum Datum
ivfflat_halfvec_support(PG_FUNCTION_ARGS) ivfflat_halfvec_support(PG_FUNCTION_ARGS)
{ {
IvfflatTypeInfo *typeInfo = (IvfflatTypeInfo *) PG_GETARG_POINTER(0); static const IvfflatTypeInfo typeInfo = {
.maxDimensions = IVFFLAT_MAX_DIM * 2,
.updateCenter = HalfvecUpdateCenter,
.sumCenter = HalfvecSumCenter
};
typeInfo->maxDimensions = IVFFLAT_MAX_DIM * 2; PG_RETURN_POINTER(&typeInfo);
typeInfo->itemsize = HALFVEC_SIZE(typeInfo->dimensions);
typeInfo->updateCenter = HalfvecUpdateCenter;
typeInfo->sumCenter = HalfvecSumCenter;
PG_RETURN_VOID();
}; };
PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_bit_support); PGDLLEXPORT PG_FUNCTION_INFO_V1(ivfflat_bit_support);
Datum Datum
ivfflat_bit_support(PG_FUNCTION_ARGS) ivfflat_bit_support(PG_FUNCTION_ARGS)
{ {
IvfflatTypeInfo *typeInfo = (IvfflatTypeInfo *) PG_GETARG_POINTER(0); static const IvfflatTypeInfo typeInfo = {
.maxDimensions = IVFFLAT_MAX_DIM * 32,
.updateCenter = BitUpdateCenter,
.sumCenter = BitSumCenter
};
typeInfo->maxDimensions = IVFFLAT_MAX_DIM * 32; PG_RETURN_POINTER(&typeInfo);
typeInfo->itemsize = VARBITTOTALLEN(typeInfo->dimensions);
typeInfo->updateCenter = BitUpdateCenter;
typeInfo->sumCenter = BitSumCenter;
PG_RETURN_VOID();
}; };