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

Improved k-means code [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-04-11 16:54:43 -07:00
parent 66a29dbdf3
commit 626bc053e5
1 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/ivfkmeans.c
View File

@@ -148,23 +148,25 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers, IvfflatTy
/* Fill remaining with random data */ /* Fill remaining with random data */
while (centers->length < centers->maxlen) while (centers->length < centers->maxlen)
{ {
Datum center = PointerGetDatum(VectorArrayGet(centers, centers->length));
if (type == IVFFLAT_TYPE_VECTOR) if (type == IVFFLAT_TYPE_VECTOR)
{ {
Vector *vec = (Vector *) VectorArrayGet(centers, centers->length); Vector *vec = DatumGetVector(center);
SET_VARSIZE(vec, VECTOR_SIZE(dimensions)); SET_VARSIZE(vec, VECTOR_SIZE(dimensions));
vec->dim = dimensions; vec->dim = dimensions;
for (int j = 0; j < dimensions; j++) for (int j = 0; j < dimensions; j++)
vec->x[j] = RandomDouble(); vec->x[j] = RandomDouble();
/* Normalize if needed (only needed for random centers) */
if (normprocinfo != NULL)
ApplyNorm(normprocinfo, collation, PointerGetDatum(vec), type);
} }
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++;
} }
} }
@@ -434,7 +436,8 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
for (int64 j = 0; j < numCenters; j++) for (int64 j = 0; j < numCenters; j++)
{ {
Vector *vec = (Vector *) VectorArrayGet(newCenters, j); Datum center = PointerGetDatum(VectorArrayGet(newCenters, j));
Vector *vec = DatumGetVector(center);
if (centerCounts[j] > 0) if (centerCounts[j] > 0)
{ {
@@ -458,7 +461,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
/* Normalize if needed */ /* Normalize if needed */
if (normprocinfo != NULL) if (normprocinfo != NULL)
ApplyNorm(normprocinfo, collation, PointerGetDatum(vec), type); ApplyNorm(normprocinfo, collation, center, type);
} }
/* Step 5 */ /* Step 5 */