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

Added IVFFLAT_MEMORY flag to show memory usage [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2023-12-24 09:27:09 -05:00
parent 9a782d29f8
commit a508b120c1
1 changed files with 25 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
src/ivfkmeans.c
View File

@@ -6,6 +6,10 @@
#include "ivfflat.h" #include "ivfflat.h"
#include "miscadmin.h" #include "miscadmin.h"
#ifdef IVFFLAT_MEMORY
#include "utils/memutils.h"
#endif
/* /*
* Initialize with kmeans++ * Initialize with kmeans++
* *
@@ -151,6 +155,23 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
} }
} }
#ifdef IVFFLAT_MEMORY
/*
* Show memory usage
*/
static void
ShowMemoryUsage(Size estimatedSize)
{
#if PG_VERSION_NUM >= 130000
elog(INFO, "total memory: %zu MB",
MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#else
MemoryContextStats(CurrentMemoryContext);
#endif
elog(INFO, "estimated memory: %zu MB", estimatedSize / (1024 * 1024));
}
#endif
/* /*
* Use Elkan for performance. This requires distance function to satisfy triangle inequality. * Use Elkan for performance. This requires distance function to satisfy triangle inequality.
* *
@@ -231,6 +252,10 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
vec->dim = dimensions; vec->dim = dimensions;
} }
#ifdef IVFFLAT_MEMORY
ShowMemoryUsage(totalSize);
#endif
/* Pick initial centers */ /* Pick initial centers */
InitCenters(index, samples, centers, lowerBound); InitCenters(index, samples, centers, lowerBound);