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

Reduced memory usage for few samples for IVFFlat index builds [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2026-07-10 17:33:05 -07:00
parent 73356ecfa7
commit f341cea329
1 changed files with 33 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
src/ivfkmeans.c
View File

@@ -8,6 +8,7 @@
#include "fmgr.h" #include "fmgr.h"
#include "ivfflat.h" #include "ivfflat.h"
#include "miscadmin.h" #include "miscadmin.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#include "utils/relcache.h" #include "utils/relcache.h"
@@ -105,16 +106,44 @@ NormCenters(const IvfflatTypeInfo * typeInfo, Oid collation, VectorArray centers
} }
/* /*
* Quick approach if we have no data * Check if value exists
*/
static bool
VectorArrayExists(VectorArray arr, Pointer val)
{
Datum d = PointerGetDatum(val);
for (int i = 0; i < arr->length; i++)
{
if (datumIsEqual(d, PointerGetDatum(VectorArrayGet(arr, i)), false, -1))
return true;
}
return false;
}
/*
* Quick approach if we have little data
*/ */
static void static void
RandomCenters(Relation index, VectorArray centers, const IvfflatTypeInfo * typeInfo) QuickCenters(Relation index, VectorArray samples, VectorArray centers, const IvfflatTypeInfo * typeInfo)
{ {
int dimensions = centers->dim; int dimensions = centers->dim;
FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC); FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
float *x = (float *) palloc(sizeof(float) * dimensions); float *x = (float *) palloc(sizeof(float) * dimensions);
/* Fill with unique samples (already normalized) */
for (int i = 0; i < samples->length; i++)
{
Pointer sample = VectorArrayGet(samples, i);
if (!VectorArrayExists(centers, sample))
{
VectorArraySet(centers, centers->length, sample);
centers->length++;
}
}
/* Fill with random data */ /* Fill with random data */
while (centers->length < centers->maxlen) while (centers->length < centers->maxlen)
{ {
@@ -548,8 +577,8 @@ IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, const Iv
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
MemoryContext oldCtx = MemoryContextSwitchTo(kmeansCtx); MemoryContext oldCtx = MemoryContextSwitchTo(kmeansCtx);
if (samples->length == 0) if (samples->length <= centers->maxlen)
RandomCenters(index, centers, typeInfo); QuickCenters(index, samples, centers, typeInfo);
else else
ElkanKmeans(index, samples, centers, typeInfo, memoryUsed); ElkanKmeans(index, samples, centers, typeInfo, memoryUsed);