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

Switched to support function for normalizing centers for k-means

Browse Source
This commit is contained in:
Andrew Kane
2024-04-23 15:39:58 -07:00
parent 0da6213a60
commit 9cd789fe06
1 changed files with 25 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
src/ivfkmeans.c
View File

@@ -89,33 +89,27 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
} }
/* /*
* Apply norm to vector * Norm centers
*/ */
static inline void static void
ApplyNorm(FmgrInfo *normprocinfo, Oid collation, Datum value, IvfflatType type) NormCenters(FmgrInfo *normalizeprocinfo, Oid collation, VectorArray centers)
{ {
double norm = DatumGetFloat8(FunctionCall1Coll(normprocinfo, collation, value)); MemoryContext normCtx = AllocSetContextCreate(CurrentMemoryContext,
"Ivfflat norm temporary context",
ALLOCSET_DEFAULT_SIZES);
MemoryContext oldCtx = MemoryContextSwitchTo(normCtx);
/* TODO Handle zero norm */ for (int j = 0; j < centers->maxlen; j++)
if (norm > 0)
{ {
if (type == IVFFLAT_TYPE_VECTOR) Datum center = PointerGetDatum(VectorArrayGet(centers, j));
{ Datum newCenter = IvfflatNormValue(normalizeprocinfo, collation, center);
Vector *vec = DatumGetVector(value);
for (int i = 0; i < vec->dim; i++) memcpy(DatumGetPointer(center), DatumGetPointer(newCenter), VARSIZE_ANY(DatumGetPointer(newCenter)));
vec->x[i] /= norm; MemoryContextReset(normCtx);
}
else if (type == IVFFLAT_TYPE_HALFVEC)
{
HalfVector *vec = DatumGetHalfVector(value);
for (int i = 0; i < vec->dim; i++)
vec->x[i] = Float4ToHalfUnchecked(HalfToFloat4(vec->x[i]) / norm);
}
else
elog(ERROR, "Unsupported type");
} }
MemoryContextSwitchTo(oldCtx);
MemoryContextDelete(normCtx);
} }
/* /*
@@ -154,6 +148,7 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers, IvfflatTy
int dimensions = centers->dim; int dimensions = centers->dim;
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC); FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
FmgrInfo *normalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORMALIZE_PROC);
/* Copy existing vectors while avoiding duplicates */ /* Copy existing vectors while avoiding duplicates */
if (samples->length > 0) if (samples->length > 0)
@@ -217,12 +212,12 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers, IvfflatTy
else else
elog(ERROR, "Unsupported type"); elog(ERROR, "Unsupported type");
/* Normalize if needed (only needed for random centers) */
if (normprocinfo != NULL)
ApplyNorm(normprocinfo, collation, center, type);
centers->length++; centers->length++;
} }
/* Fine if existing vectors are normalized twice */
if (normprocinfo != NULL)
NormCenters(normalizeprocinfo, collation, centers);
} }
#ifdef IVFFLAT_MEMORY #ifdef IVFFLAT_MEMORY
@@ -246,7 +241,7 @@ ShowMemoryUsage(MemoryContext context, Size estimatedSize)
* Compute new centers * Compute new centers
*/ */
static void static void
ComputeNewCenters(VectorArray samples, VectorArray aggCenters, VectorArray newCenters, int *centerCounts, int *closestCenters, FmgrInfo *normprocinfo, Oid collation, IvfflatType type) ComputeNewCenters(VectorArray samples, VectorArray aggCenters, VectorArray newCenters, int *centerCounts, int *closestCenters, FmgrInfo *normprocinfo, FmgrInfo *normalizeprocinfo, Oid collation, IvfflatType type)
{ {
int dimensions = aggCenters->dim; int dimensions = aggCenters->dim;
int numCenters = aggCenters->maxlen; int numCenters = aggCenters->maxlen;
@@ -360,14 +355,7 @@ ComputeNewCenters(VectorArray samples, VectorArray aggCenters, VectorArray newCe
/* Normalize if needed */ /* Normalize if needed */
if (normprocinfo != NULL) if (normprocinfo != NULL)
{ NormCenters(normalizeprocinfo, collation, newCenters);
for (int j = 0; j < numCenters; j++)
{
Datum newCenter = PointerGetDatum(VectorArrayGet(newCenters, j));
ApplyNorm(normprocinfo, collation, newCenter, type);
}
}
} }
/* /*
@@ -383,6 +371,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
FmgrInfo *normalizeprocinfo;
Oid collation; Oid collation;
int dimensions = centers->dim; int dimensions = centers->dim;
int numCenters = centers->maxlen; int numCenters = centers->maxlen;
@@ -430,6 +419,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
/* Set support functions */ /* Set support functions */
procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC); procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC); normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
normalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORMALIZE_PROC);
collation = index->rd_indcollation[0]; collation = index->rd_indcollation[0];
/* Use memory context */ /* Use memory context */
@@ -627,7 +617,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
} }
/* Step 4: For each center c, let m(c) be mean of all points assigned */ /* Step 4: For each center c, let m(c) be mean of all points assigned */
ComputeNewCenters(samples, aggCenters, newCenters, centerCounts, closestCenters, normprocinfo, collation, type); ComputeNewCenters(samples, aggCenters, newCenters, centerCounts, closestCenters, normprocinfo, normalizeprocinfo, collation, type);
/* Step 5 */ /* Step 5 */
for (int j = 0; j < numCenters; j++) for (int j = 0; j < numCenters; j++)