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

Use memory context for k-means [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-04-11 09:34:05 -07:00
parent e1647473c9
commit d45f561d75
1 changed files with 10 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
src/ivfkmeans.c
View File

@@ -5,10 +5,7 @@
#include "ivfflat.h" #include "ivfflat.h"
#include "miscadmin.h" #include "miscadmin.h"
#ifdef IVFFLAT_MEMORY
#include "utils/memutils.h" #include "utils/memutils.h"
#endif
/* /*
* Initialize with kmeans++ * Initialize with kmeans++
@@ -201,6 +198,8 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
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->dim); Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->dim);
@@ -234,6 +233,12 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
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 */
centerCounts = palloc(centerCountsSize); centerCounts = palloc(centerCountsSize);
@@ -472,14 +477,8 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
break; break;
} }
VectorArrayFree(newCenters); MemoryContextSwitchTo(oldCtx);
pfree(centerCounts); MemoryContextDelete(kmeansCtx);
pfree(closestCenters);
pfree(lowerBound);
pfree(upperBound);
pfree(s);
pfree(halfcdist);
pfree(newcdist);
} }
/* /*