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

Added support for bit to IVFFlat

Browse Source
This commit is contained in:
Andrew Kane
2024-04-16 17:12:27 -07:00
parent 819b6cf312
commit 04af15c9d6
14 changed files with 315 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
src/ivfkmeans.c
View File

@@ -3,10 +3,12 @@
#include <float.h>
#include <math.h>
#include "bitvector.h"
#include "halfutils.h"
#include "halfvec.h"
#include "ivfflat.h"
#include "miscadmin.h"
#include "utils/builtins.h"
#include "utils/datum.h"
#include "utils/memutils.h"
#include "vector.h"
@@ -134,6 +136,15 @@ CompareHalfVectors(const void *a, const void *b)
return halfvec_cmp_internal((HalfVector *) a, (HalfVector *) b);
}
/*
* Compare bit vectors
*/
static int
CompareBitVectors(const void *a, const void *b)
{
return DirectFunctionCall2(bitcmp, VarBitPGetDatum((VarBit *) a), VarBitPGetDatum((VarBit *) b));
}
/*
* Quick approach if we have little data
*/
@@ -151,6 +162,8 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers, IvfflatTy
qsort(samples->items, samples->length, samples->itemsize, CompareVectors);
else if (type == IVFFLAT_TYPE_HALFVEC)
qsort(samples->items, samples->length, samples->itemsize, CompareHalfVectors);
else if (type == IVFFLAT_TYPE_BIT)
qsort(samples->items, samples->length, samples->itemsize, CompareBitVectors);
else
elog(ERROR, "Unsupported type");
@@ -191,6 +204,16 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers, IvfflatTy
for (int j = 0; j < dimensions; j++)
vec->x[j] = Float4ToHalfUnchecked((float) RandomDouble());
}
else if (type == IVFFLAT_TYPE_BIT)
{
VarBit *vec = DatumGetVarBitP(center);
SET_VARSIZE(vec, VARBITTOTALLEN(dimensions));
VARBITLEN(vec) = dimensions;
for (int j = 0; j < dimensions; j++)
VARBITS(vec)[j / dimensions] |= (RandomDouble() > 0.5 ? 1 : 0) << (7 - (j % 8));
}
else
elog(ERROR, "Unsupported type");
@@ -263,6 +286,17 @@ ComputeNewCenters(VectorArray samples, VectorArray aggCenters, VectorArray newCe
aggCenter->x[k] += HalfToFloat4(vec->x[k]);
}
}
else if (type == IVFFLAT_TYPE_BIT)
{
for (int j = 0; j < numSamples; j++)
{
Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, closestCenters[j]);
VarBit *vec = (VarBit *) VectorArrayGet(samples, j);
for (int k = 0; k < dimensions; k++)
aggCenter->x[k] += (float) (((VARBITS(vec)[k / 8]) >> (7 - (k % 8))) & 0x01);
}
}
else
elog(ERROR, "Unsupported type");
@@ -308,6 +342,21 @@ ComputeNewCenters(VectorArray samples, VectorArray aggCenters, VectorArray newCe
newCenter->x[k] = Float4ToHalfUnchecked(aggCenter->x[k]);
}
}
else if (type == IVFFLAT_TYPE_BIT)
{
for (int j = 0; j < numCenters; j++)
{
Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, j);
VarBit *newCenter = (VarBit *) VectorArrayGet(newCenters, j);
unsigned char *nx = VARBITS(newCenter);
for (uint32 k = 0; k < VARBITBYTES(newCenter); k++)
nx[k] = 0;
for (int k = 0; k < dimensions; k++)
nx[k / 8] |= (aggCenter->x[k] > 0.5) << (7 - (k % 8));
}
}
/* Normalize if needed */
if (normprocinfo != NULL)
@@ -425,6 +474,18 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers, IvfflatTyp
vec->dim = dimensions;
}
}
else if (type == IVFFLAT_TYPE_BIT)
{
newCenters = VectorArrayInit(numCenters, dimensions, centers->itemsize);
for (int j = 0; j < numCenters; j++)
{
VarBit *vec = (VarBit *) VectorArrayGet(newCenters, j);
SET_VARSIZE(vec, VARBITTOTALLEN(dimensions));
VARBITLEN(vec) = dimensions;
}
}
else
elog(ERROR, "Unsupported type");
@@ -642,7 +703,7 @@ CheckCenters(Relation index, VectorArray centers, IvfflatType type)
elog(ERROR, "Infinite value detected. Please report a bug.");
}
}
else
else if (type != IVFFLAT_TYPE_BIT)
elog(ERROR, "Unsupported type");
}
@@ -652,6 +713,8 @@ CheckCenters(Relation index, VectorArray centers, IvfflatType type)
qsort(centers->items, centers->length, centers->itemsize, CompareVectors);
else if (type == IVFFLAT_TYPE_HALFVEC)
qsort(centers->items, centers->length, centers->itemsize, CompareHalfVectors);
else if (type == IVFFLAT_TYPE_BIT)
qsort(centers->items, centers->length, centers->itemsize, CompareBitVectors);
else
elog(ERROR, "Unsupported type");