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

Improved initialization of new centers [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-04-24 16:45:16 -07:00
parent 25b98540c9
commit 15ee38456f
1 changed files with 23 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
src/ivfkmeans.c
View File

@@ -474,6 +474,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
float *newcdist; float *newcdist;
MemoryContext kmeansCtx; MemoryContext kmeansCtx;
MemoryContext oldCtx; MemoryContext oldCtx;
void (*initNewCenter) (Pointer v, int dimensions);
/* Calculate allocation sizes */ /* Calculate allocation sizes */
Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->itemsize); Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->itemsize);
@@ -525,50 +526,36 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
halfcdist = palloc_extended(halfcdistSize, MCXT_ALLOC_HUGE); halfcdist = palloc_extended(halfcdistSize, MCXT_ALLOC_HUGE);
newcdist = palloc(newcdistSize); newcdist = palloc(newcdistSize);
aggCenters = VectorArrayInit(numCenters, dimensions, VECTOR_SIZE(dimensions)); /* Initialize new centers */
aggCenters->length = numCenters; if (type == IVFFLAT_TYPE_VECTOR)
initNewCenter = VectorInitNewCenter;
else if (type == IVFFLAT_TYPE_HALFVEC)
initNewCenter = HalfvecInitNewCenter;
else if (type == IVFFLAT_TYPE_BIT)
initNewCenter = BitInitNewCenter;
else
elog(ERROR, "Unsupported type");
newCenters = VectorArrayInit(numCenters, dimensions, centers->itemsize);
newCenters->length = numCenters;
for (int j = 0; j < numCenters; j++) for (int j = 0; j < numCenters; j++)
{ initNewCenter(VectorArrayGet(newCenters, j), dimensions);
Vector *vec = (Vector *) VectorArrayGet(aggCenters, j);
SET_VARSIZE(vec, VECTOR_SIZE(dimensions));
vec->dim = dimensions;
}
/* Initialize agg centers */
if (type == IVFFLAT_TYPE_VECTOR) if (type == IVFFLAT_TYPE_VECTOR)
{ {
/* Use same centers to save memory */ /* Use same centers to save memory */
newCenters = aggCenters; aggCenters = newCenters;
}
else if (type == IVFFLAT_TYPE_HALFVEC)
{
newCenters = VectorArrayInit(numCenters, dimensions, centers->itemsize);
newCenters->length = numCenters;
for (int j = 0; j < numCenters; j++)
{
HalfVector *vec = (HalfVector *) VectorArrayGet(newCenters, j);
SET_VARSIZE(vec, HALFVEC_SIZE(dimensions));
vec->dim = dimensions;
}
}
else if (type == IVFFLAT_TYPE_BIT)
{
newCenters = VectorArrayInit(numCenters, dimensions, centers->itemsize);
newCenters->length = numCenters;
for (int j = 0; j < numCenters; j++)
{
VarBit *vec = (VarBit *) VectorArrayGet(newCenters, j);
SET_VARSIZE(vec, VARBITTOTALLEN(dimensions));
VARBITLEN(vec) = dimensions;
}
} }
else else
elog(ERROR, "Unsupported type"); {
aggCenters = VectorArrayInit(numCenters, dimensions, VECTOR_SIZE(dimensions));
aggCenters->length = numCenters;
for (int j = 0; j < numCenters; j++)
VectorInitNewCenter(VectorArrayGet(aggCenters, j), dimensions);
}
#ifdef IVFFLAT_MEMORY #ifdef IVFFLAT_MEMORY
ShowMemoryUsage(oldCtx, totalSize); ShowMemoryUsage(oldCtx, totalSize);