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

Improved memory context for k-means [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-04-25 14:22:24 -07:00
parent ebc76114ad
commit 8daa581f42
1 changed files with 9 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
src/ivfkmeans.c
View File

@@ -142,8 +142,6 @@ RandomCenters(Relation index, VectorArray centers, const IvfflatTypeInfo * typeI
if (normprocinfo != NULL) if (normprocinfo != NULL)
NormCenters(typeInfo, collation, centers); NormCenters(typeInfo, collation, centers);
pfree(x);
} }
#ifdef IVFFLAT_MEMORY #ifdef IVFFLAT_MEMORY
@@ -279,8 +277,6 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, const Ivff
float *s; float *s;
float *halfcdist; float *halfcdist;
float *newcdist; float *newcdist;
MemoryContext kmeansCtx;
MemoryContext oldCtx;
/* Calculate allocation sizes */ /* Calculate allocation sizes */
Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->itemsize); Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->itemsize);
@@ -315,12 +311,6 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, const Ivff
normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC); normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
collation = index->rd_indcollation[0]; collation = index->rd_indcollation[0];
/* Use memory context */
kmeansCtx = AllocSetContextCreate(CurrentMemoryContext,
"Ivfflat kmeans temporary context",
ALLOCSET_DEFAULT_SIZES);
oldCtx = MemoryContextSwitchTo(kmeansCtx);
/* Allocate space */ /* Allocate space */
/* Use float instead of double to save memory */ /* Use float instead of double to save memory */
agg = palloc(aggSize); agg = palloc(aggSize);
@@ -337,7 +327,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, const Ivff
newCenters->length = numCenters; newCenters->length = numCenters;
#ifdef IVFFLAT_MEMORY #ifdef IVFFLAT_MEMORY
ShowMemoryUsage(oldCtx, totalSize); ShowMemoryUsage(MemoryContextGetParent(CurrentMemoryContext));
#endif #endif
/* Pick initial centers */ /* Pick initial centers */
@@ -505,9 +495,6 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, const Ivff
if (changes == 0 && iteration != 0) if (changes == 0 && iteration != 0)
break; break;
} }
MemoryContextSwitchTo(oldCtx);
MemoryContextDelete(kmeansCtx);
} }
/* /*
@@ -556,8 +543,6 @@ CheckCenters(Relation index, VectorArray centers, const IvfflatTypeInfo * typeIn
elog(ERROR, "Zero norm detected. Please report a bug."); elog(ERROR, "Zero norm detected. Please report a bug.");
} }
} }
pfree(scratch);
} }
/* /*
@@ -567,10 +552,18 @@ CheckCenters(Relation index, VectorArray centers, const IvfflatTypeInfo * typeIn
void void
IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, const IvfflatTypeInfo * typeInfo) IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, const IvfflatTypeInfo * typeInfo)
{ {
MemoryContext kmeansCtx = AllocSetContextCreate(CurrentMemoryContext,
"Ivfflat kmeans temporary context",
ALLOCSET_DEFAULT_SIZES);
MemoryContext oldCtx = MemoryContextSwitchTo(kmeansCtx);
if (samples->length == 0) if (samples->length == 0)
RandomCenters(index, centers, typeInfo); RandomCenters(index, centers, typeInfo);
else else
ElkanKmeans(index, samples, centers, typeInfo); ElkanKmeans(index, samples, centers, typeInfo);
CheckCenters(index, centers, typeInfo); CheckCenters(index, centers, typeInfo);
MemoryContextSwitchTo(oldCtx);
MemoryContextDelete(kmeansCtx);
} }