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

Mark type-specific code [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-04-11 16:44:10 -07:00
parent 000cc13c29
commit 1c26da6ef5
3 changed files with 32 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/ivfbuild.c
View File

@@ -437,7 +437,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
} }
/* Calculate centers */ /* Calculate centers */
IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers)); IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers, buildstate->type));
/* Free samples before we allocate more memory */ /* Free samples before we allocate more memory */
VectorArrayFree(buildstate->samples); VectorArrayFree(buildstate->samples);

2
src/ivfflat.h
View File

@@ -270,7 +270,7 @@ typedef IvfflatScanOpaqueData * IvfflatScanOpaque;
/* Methods */ /* Methods */
VectorArray VectorArrayInit(int maxlen, int dimensions, Size itemsize); VectorArray VectorArrayInit(int maxlen, int dimensions, Size itemsize);
void VectorArrayFree(VectorArray arr); void VectorArrayFree(VectorArray arr);
void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers); void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatType type);
FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum); FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum);
IvfflatType IvfflatGetType(Relation index); IvfflatType IvfflatGetType(Relation index);
bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, IvfflatType type); bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, IvfflatType type);

47
src/ivfkmeans.c
View File

@@ -112,7 +112,7 @@ CompareVectors(const void *a, const void *b)
* Quick approach if we have little data * Quick approach if we have little data
*/ */
static void static void
QuickCenters(Relation index, VectorArray samples, VectorArray centers) QuickCenters(Relation index, VectorArray samples, VectorArray centers, IvfflatType type)
{ {
int dimensions = centers->dim; int dimensions = centers->dim;
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
@@ -121,7 +121,11 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
/* Copy existing vectors while avoiding duplicates */ /* Copy existing vectors while avoiding duplicates */
if (samples->length > 0) if (samples->length > 0)
{ {
qsort(samples->items, samples->length, samples->itemsize, CompareVectors); if (type == IVFFLAT_TYPE_VECTOR)
qsort(samples->items, samples->length, samples->itemsize, CompareVectors);
else
elog(ERROR, "Unsupported type");
for (int i = 0; i < samples->length; i++) for (int i = 0; i < samples->length; i++)
{ {
Datum vec = PointerGetDatum(VectorArrayGet(samples, i)); Datum vec = PointerGetDatum(VectorArrayGet(samples, i));
@@ -137,17 +141,22 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
/* Fill remaining with random data */ /* Fill remaining with random data */
while (centers->length < centers->maxlen) while (centers->length < centers->maxlen)
{ {
Vector *vec = (Vector *) VectorArrayGet(centers, centers->length); if (type == IVFFLAT_TYPE_VECTOR)
{
Vector *vec = (Vector *) VectorArrayGet(centers, centers->length);
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) */ /* Normalize if needed (only needed for random centers) */
if (normprocinfo != NULL) if (normprocinfo != NULL)
ApplyNorm(normprocinfo, collation, vec); ApplyNorm(normprocinfo, collation, vec);
}
else
elog(ERROR, "Unsupported type");
centers->length++; centers->length++;
} }
@@ -179,7 +188,7 @@ ShowMemoryUsage(Size estimatedSize)
* https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf * https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf
*/ */
static void static void
ElkanKmeans(Relation index, VectorArray samples, VectorArray centers) ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatType type)
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
@@ -483,7 +492,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
* Detect issues with centers * Detect issues with centers
*/ */
static void static void
CheckCenters(Relation index, VectorArray centers) CheckCenters(Relation index, VectorArray centers, IvfflatType type)
{ {
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
@@ -507,7 +516,11 @@ CheckCenters(Relation index, VectorArray centers)
/* Ensure no duplicate centers */ /* Ensure no duplicate centers */
/* Fine to sort in-place */ /* Fine to sort in-place */
qsort(centers->items, centers->length, centers->itemsize, CompareVectors); if (type == IVFFLAT_TYPE_VECTOR)
qsort(centers->items, centers->length, centers->itemsize, CompareVectors);
else
elog(ERROR, "Unsupported type");
for (int i = 1; i < centers->length; i++) for (int i = 1; i < centers->length; i++)
{ {
if (datumIsEqual(PointerGetDatum(VectorArrayGet(centers, i)), PointerGetDatum(VectorArrayGet(centers, i - 1)), false, -1)) if (datumIsEqual(PointerGetDatum(VectorArrayGet(centers, i)), PointerGetDatum(VectorArrayGet(centers, i - 1)), false, -1))
@@ -536,12 +549,12 @@ CheckCenters(Relation index, VectorArray centers)
* We use spherical k-means for inner product and cosine * We use spherical k-means for inner product and cosine
*/ */
void void
IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers) IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatType type)
{ {
if (samples->length <= centers->maxlen) if (samples->length <= centers->maxlen)
QuickCenters(index, samples, centers); QuickCenters(index, samples, centers, type);
else else
ElkanKmeans(index, samples, centers); ElkanKmeans(index, samples, centers, type);
CheckCenters(index, centers); CheckCenters(index, centers, type);
} }