pgvector.git / summary / log / commit / refs

commit    fb6683ade7cf07e6be59bfdcf5a427cf7e20f6ef
Author:   Andrew Kane <andrew@ankane.org>
Date:     Wed Apr 24 22:56:50 2024 +0000

    Improved SetNewCenters [skip ci]


 src/ivfkmeans.c | 54 +++++++++++++++++++++++++++++++-----------------------
 1 file changed, 31 insertions(+), 23 deletions(-)

diff --git a/src/ivfkmeans.c b/src/ivfkmeans.c
index 548f7ffc..4f2865df 100644
--- a/src/ivfkmeans.c
+++ b/src/ivfkmeans.c
@@ -308,40 +308,48 @@ SumCenters(VectorArray samples, VectorArray aggCenters, int *closestCenters, Ivf
 	}
 }
 
+static void
+HalfvecSetNewCenter(const void *v, Vector * aggCenter)
+{
+	HalfVector *newCenter = (HalfVector *) v;
+
+	for (int k = 0; k < aggCenter->dim; k++)
+		newCenter->x[k] = Float4ToHalfUnchecked(aggCenter->x[k]);
+}
+
+static void
+BitSetNewCenter(const void *v, Vector * aggCenter)
+{
+	VarBit	   *newCenter = (VarBit *) v;
+	unsigned char *nx = VARBITS(newCenter);
+
+	for (uint32 k = 0; k < VARBITBYTES(newCenter); k++)
+		nx[k] = 0;
+
+	for (int k = 0; k < aggCenter->dim; k++)
+		nx[k / 8] |= (aggCenter->x[k] > 0.5) << (7 - (k % 8));
+}
+
 /*
  * Set new centers
  */
 static void
 SetNewCenters(VectorArray aggCenters, VectorArray newCenters, IvfflatType type)
 {
-	int			dimensions = aggCenters->dim;
-	int			numCenters = aggCenters->maxlen;
+	void		(*setNewCenter) (const void *v, Vector * aggCenter);
 
 	if (type == IVFFLAT_TYPE_HALFVEC)
-	{
-		for (int j = 0; j < numCenters; j++)
-		{
-			Vector	   *aggCenter = (Vector *) VectorArrayGet(aggCenters, j);
-			HalfVector *newCenter = (HalfVector *) VectorArrayGet(newCenters, j);
-
-			for (int k = 0; k < dimensions; k++)
-				newCenter->x[k] = Float4ToHalfUnchecked(aggCenter->x[k]);
-		}
-	}
+		setNewCenter = HalfvecSetNewCenter;
 	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);
+		setNewCenter = BitSetNewCenter;
+	else
+		return;
 
-			for (uint32 k = 0; k < VARBITBYTES(newCenter); k++)
-				nx[k] = 0;
+	for (int j = 0; j < aggCenters->length; j++)
+	{
+		Vector	   *aggCenter = (Vector *) VectorArrayGet(aggCenters, j);
 
-			for (int k = 0; k < dimensions; k++)
-				nx[k / 8] |= (aggCenter->x[k] > 0.5) << (7 - (k % 8));
-		}
+		setNewCenter(VectorArrayGet(newCenters, j), aggCenter);
 	}
 }
 


[parent: 558953ca6b2a]