Merge #147027

147027: vecindex: fix bugs and clean up r=drewkimball a=andy-kimball

#### quantize: fix calculation when vectors are added incrementally

The quantization code is mis-calculating the centroid dot product when
vectors are added incrementally to a quantized set that uses a non-Euclidean
distance metric. Rather than appending the dot products to the end of the
CentroidDotProducts slice, it was overwriting the initial elements of the
slice. Fix the bug and add more tests that ensure it won't happen again.

#### kmeans: fix small issue with convergence calculation

Update the convergence calculation to match the logic in scikit-learn, which
sums the shifts across all centroids before comparing to the tolerance value.
Also improve comments and add an interesting test case.

#### quantize: remove the QuantizedVectorSet.GetCentroid method

Remove GetCentroid method from the QuantizedVectorSet interface. Having
it present on the interface forces every quantizer to store it, even
those (like UnQuantizer) that don't need it. It's also redundant in most
contexts with Partition.Centroid(), so it's cleaner to have code use
that consistently.

Co-authored-by: Andrew Kimball <[email protected]>

Author: craig[bot]

pkg/sql/vecindex/cspann/commontest/utils.go

@@ -72,7 +72,7 @@ func ValidatePartitionsEqual(t *testing.T, l, r *cspann.Partition) {
 	require.Equal(t, l.Level(), r.Level(), "levels do not match")
 	require.Equal(t, l.ChildKeys(), r.ChildKeys(), "childKeys do not match")
 	require.Equal(t, l.ValueBytes(), r.ValueBytes(), "valueBytes do not match")
-	require.Equal(t, q1.GetCentroid(), q2.GetCentroid(), "centroids do not match")
+	require.Equal(t, l.Centroid, r.Centroid, "centroids do not match")
 	require.Equal(t, q1.GetCount(), q2.GetCount(), "counts do not match")
 	if eq, ok := q1.(equaler); ok {
 		require.True(t, eq.Equal(q2))

pkg/sql/vecindex/cspann/kmeans.go

@@ -78,6 +78,10 @@ func (km *BalancedKmeans) ComputeCentroids(
 	tolerance := km.calculateTolerance(vectors)
 
 	// Pick 2 centroids to start, using the K-means++ algorithm.
+	// TOOD(andyk): We should consider adding an outer loop here to generate new
+	// random centroids in the case where more than 2/3 the vectors are assigned
+	// to one of the partitions. If we're that unbalanced, it might just be
+	// because we picked bad starting centroids and retrying could correct that.
 	tempLeftCentroid := km.Workspace.AllocVector(vectors.Dims)
 	defer km.Workspace.FreeVector(tempLeftCentroid)
 	newLeftCentroid := tempLeftCentroid
@@ -109,12 +113,15 @@ func (km *BalancedKmeans) ComputeCentroids(
 		}
 		calcPartitionCentroid(km.DistanceMetric, vectors, rightOffsets, newRightCentroid)
 
-		// Check if algorithm has converged.
-		// TODO(andyk): Is there anything better than using Euclidean distances to
-		// check for convergence?
+		// Check for convergence using the scikit-learn algorithm.
+		// NOTE: This uses Euclidean distance, even when using spherical centroids
+		// with Cosine or InnerProduct distances. This approach mirrors the
+		// spherecluster library. Since spherical centroids are always normalized
+		// (unit vectors), the squared Euclidean distance is 2x the Cosine or
+		// InnerProduct distance, so it's a reasonable convergence check.
 		leftCentroidShift := num32.L2SquaredDistance(leftCentroid, newLeftCentroid)
 		rightCentroidShift := num32.L2SquaredDistance(rightCentroid, newRightCentroid)
-		if leftCentroidShift <= tolerance && rightCentroidShift <= tolerance {
+		if leftCentroidShift+rightCentroidShift <= tolerance {
 			break
 		}
 

pkg/sql/vecindex/cspann/kmeans_test.go

@@ -34,8 +34,8 @@ func TestBalancedKMeans(t *testing.T) {
 	}
 
 	workspace := &workspace.T{}
-	rng := rand.New(rand.NewSource(42))
-	dataset := testutils.LoadDataset(t, testutils.ImagesDataset)
+	images := testutils.LoadDataset(t, testutils.ImagesDataset)
+	laion := testutils.LoadDataset(t, testutils.LaionDataset)
 
 	testCases := []struct {
 		desc           string
@@ -46,6 +46,7 @@ func TestBalancedKMeans(t *testing.T) {
 		leftCentroid   vector.T
 		rightCentroid  vector.T
 		skipPinTest    bool
+		testUnbalanced bool
 	}{
 		{
 			desc:           "partition vector set with only 2 elements",
@@ -128,8 +129,8 @@ func TestBalancedKMeans(t *testing.T) {
 				9, -14, 20,
 				5, 9, 4,
 			}, 3),
-			leftOffsets:  []uint64{0, 3},
-			rightOffsets: []uint64{1, 2, 4},
+			leftOffsets:  []uint64{1, 2},
+			rightOffsets: []uint64{0, 3, 4},
 		},
 		{
 			desc:           "cosine distance",
@@ -147,7 +148,7 @@ func TestBalancedKMeans(t *testing.T) {
 		{
 			desc:           "high-dimensional unit vectors, Euclidean distance",
 			distanceMetric: vecdist.L2Squared,
-			vectors:        dataset.Slice(0, 100),
+			vectors:        images.Slice(0, 100),
 			// It's challenging to test pinLeftCentroid for this case, due to the
 			// inherent randomness of the K-means++ algorithm. The other test cases
 			// should be sufficient to test that, however.
@@ -156,22 +157,39 @@ func TestBalancedKMeans(t *testing.T) {
 		{
 			desc:           "high-dimensional unit vectors, InnerProduct distance",
 			distanceMetric: vecdist.InnerProduct,
-			vectors:        dataset.Slice(0, 100),
+			vectors:        images.Slice(0, 100),
 			skipPinTest:    true,
 		},
 		{
 			desc:           "high-dimensional unit vectors, Cosine distance",
 			distanceMetric: vecdist.Cosine,
-			vectors:        dataset.Slice(0, 100),
+			vectors:        images.Slice(0, 100),
 			skipPinTest:    true,
 		},
+		{
+			// Note that laion.Slice(0, 100) actually fails the check that vectors
+			// in the left partition are closer to the left centroid than vectors
+			// in the right partition. This is because K-means++ happens to pick
+			// bad centroids that result in > 2/3rd the vectors being closer to
+			// the right centroid. In that case, the BalancedKMeans class will
+			// deliberately move vectors to the left partition, even though they
+			// are closer to the right partition, in order to obey the balancing
+			// constraint.
+			desc:           "different dataset, InnerProduct distance",
+			distanceMetric: vecdist.InnerProduct,
+			vectors:        laion.Slice(0, 100),
+			skipPinTest:    true,
+			testUnbalanced: true,
+		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
+			// Re-initialize rng for each iteration so that order of test cases
+			// doesn't matter.
 			kmeans := BalancedKmeans{
 				Workspace:      workspace,
-				Rand:           rng,
+				Rand:           rand.New(rand.NewSource(42)),
 				DistanceMetric: tc.distanceMetric,
 			}
 
@@ -217,7 +235,11 @@ func TestBalancedKMeans(t *testing.T) {
 			// partition than those in the right partition.
 			leftMean := calcMeanDistance(tc.distanceMetric, tc.vectors, leftCentroid, leftOffsets)
 			rightMean := calcMeanDistance(tc.distanceMetric, tc.vectors, leftCentroid, rightOffsets)
-			require.LessOrEqual(t, leftMean, rightMean)
+			if tc.testUnbalanced {
+				require.Greater(t, leftMean, rightMean)
+			} else {
+				require.LessOrEqual(t, leftMean, rightMean)
+			}
 
 			if !tc.skipPinTest {
 				// Check that pinning the left centroid returns the same right centroid.

pkg/sql/vecindex/cspann/memstore/memstore.go

@@ -695,17 +695,19 @@ func Load(data []byte) (*Store, error) {
 
 		var quantizer quantize.Quantizer
 		var quantizedSet quantize.QuantizedVectorSet
+		var centroid vector.T
 		if partitionProto.RaBitQ != nil {
 			quantizer = raBitQuantizer
 			quantizedSet = partitionProto.RaBitQ
+			centroid = partitionProto.RaBitQ.Centroid
 		} else {
 			quantizer = unquantizer
 			quantizedSet = partitionProto.UnQuantized
 		}
 
 		metadata := cspann.PartitionMetadata{
 			Level:    partitionProto.Metadata.Level,
-			Centroid: quantizedSet.GetCentroid(),
+			Centroid: centroid,
 			StateDetails: cspann.PartitionStateDetails{
 				State:   partitionProto.Metadata.State,
 				Target1: partitionProto.Metadata.Target1,

pkg/sql/vecindex/cspann/memstore/memstore_test.go

@@ -253,7 +253,6 @@ func TestInMemoryStoreMarshalling(t *testing.T) {
 		cspann.MakeReadyPartitionMetadata(1, centroid),
 		unquantizer,
 		&quantize.UnQuantizedVectorSet{
-			Centroid: centroid,
 			Vectors: vector.Set{
 				Dims:  2,
 				Count: 3,
@@ -270,7 +269,6 @@ func TestInMemoryStoreMarshalling(t *testing.T) {
 		cspann.MakeReadyPartitionMetadata(2, centroid),
 		raBitQuantizer,
 		&quantize.UnQuantizedVectorSet{
-			Centroid: centroid,
 			Vectors: vector.Set{
 				Dims:  2,
 				Count: 3,

pkg/sql/vecindex/cspann/partition.go

@@ -127,7 +127,7 @@ func (p *Partition) QuantizedSet() quantize.QuantizedVectorSet {
 // Centroid is the full-sized centroid vector for this partition.
 // NOTE: The centroid is immutable and therefore this method is thread-safe.
 func (p *Partition) Centroid() vector.T {
-	return p.quantizedSet.GetCentroid()
+	return p.metadata.Centroid
 }
 
 // ChildKeys point to the location of the full-size vectors that are quantized
@@ -271,7 +271,7 @@ func (p *Partition) Find(childKey ChildKey) int {
 // vectors that were cleared. The centroid stays the same.
 func (p *Partition) Clear() int {
 	count := len(p.childKeys)
-	p.quantizedSet.Clear(p.quantizedSet.GetCentroid())
+	p.quantizedSet.Clear(p.metadata.Centroid)
 	clear(p.childKeys)
 	p.childKeys = p.childKeys[:0]
 	clear(p.valueBytes)

pkg/sql/vecindex/cspann/partition_test.go

@@ -55,7 +55,7 @@ func TestPartition(t *testing.T) {
 		quantizedSet := quantizer.Quantize(&workspace, vectors)
 		childKeys := []ChildKey{childKey10, childKey20, childKey30}
 		valueBytes := []ValueBytes{valueBytes10, valueBytes20, valueBytes30}
-		metadata := MakeReadyPartitionMetadata(1, quantizedSet.GetCentroid())
+		metadata := MakeReadyPartitionMetadata(1, vectors.Centroid(make(vector.T, vectors.Dims)))
 		return NewPartition(metadata, quantizer, quantizedSet, childKeys, valueBytes)
 	}
 

pkg/sql/vecindex/cspann/quantize/quantize.proto

@@ -64,15 +64,6 @@ message RaBitQuantizedVectorSet {
 message UnQuantizedVectorSet {
   option (gogoproto.equal) = true;
 
-  // Centroid is the average of vectors in the set, representing its "center of
-  // mass". Note that the centroid is computed when a vector set is created and
-  // is not updated when vectors are added or removed.
-  // NOTE: By default, this is the mean centroid for the L2Squared distance
-  // metric, but is the spherical centroid for InnerProduct and Cosine metrics.
-  // The spherical centroid is simply the mean centroid, but normalized.
-  // NOTE: The centroid should be treated as immutable.
-  repeated float centroid = 1;
-  reserved 2; // CentroidDistances
   // Vectors is the set of original full-size vectors.
-  cockroach.util.vector.Set vectors = 3 [(gogoproto.nullable) = false];
+  cockroach.util.vector.Set vectors = 1 [(gogoproto.nullable) = false];
 }

pkg/sql/vecindex/cspann/quantize/quantizer.go

@@ -70,16 +70,6 @@ type QuantizedVectorSet interface {
 	// GetCount returns the number of quantized vectors in the set.
 	GetCount() int
 
-	// GetCentroid returns the full-size centroid vector for the set. The
-	// centroid is the average of the vectors across all dimensions.
-	// NOTE: By default, this is the mean centroid for the L2Squared distance
-	// metric, but is the spherical centroid for InnerProduct and Cosine metrics.
-	// The spherical centroid is simply the mean centroid, but normalized.
-	// NOTE: This centroid is calculated once, when the set is first created. It
-	// is not updated when quantized vectors are added to or removed from the set.
-	// Since it is immutable, this method is thread-safe.
-	GetCentroid() vector.T
-
 	// ReplaceWithLast removes the quantized vector at the given offset from the
 	// set, replacing it with the last quantized vector in the set. The modified
 	// set has one less element and the last quantized vector's position changes.

pkg/sql/vecindex/cspann/quantize/quantizer_test.go

@@ -88,33 +88,63 @@ func (s *testState) estimateDistances(t *testing.T, d *datadriven.TestData) stri
 		}
 	}
 
+	addVectors := func(
+		quantizer quantize.Quantizer, centroid vector.T, vectors vector.Set,
+	) (quantizedSet quantize.QuantizedVectorSet, distances, errorBounds []float32) {
+		// Quantize all vectors at once.
+		quantizedSet = quantizer.Quantize(&s.Workspace, vectors)
+		distances = make([]float32, quantizedSet.GetCount())
+		errorBounds = make([]float32, quantizedSet.GetCount())
+		quantizer.EstimateDistances(
+			&s.Workspace, quantizedSet, queryVector, distances, errorBounds)
+
+		// Now add the vectors one-by-one and ensure that distances and error
+		// bounds stay the same.
+		quantizedSet = quantizer.NewQuantizedVectorSet(vectors.Count, centroid)
+		for i := range vectors.Count {
+			quantizer.QuantizeInSet(&s.Workspace, quantizedSet, vectors.Slice(i, 1))
+		}
+		distances2 := make([]float32, quantizedSet.GetCount())
+		errorBounds2 := make([]float32, quantizedSet.GetCount())
+		quantizer.EstimateDistances(
+			&s.Workspace, quantizedSet, queryVector, distances2, errorBounds2)
+		require.Equal(t, distances2, distances)
+		require.Equal(t, errorBounds2, errorBounds)
+
+		return quantizedSet, distances, errorBounds
+	}
+
 	var buf bytes.Buffer
 	doTest := func(metric vecdist.Metric, prec int) {
-		unquantizer := quantize.NewUnQuantizer(len(queryVector), metric)
-		unQuantizedSet := unquantizer.Quantize(&s.Workspace, vectors)
-		exact := make([]float32, unQuantizedSet.GetCount())
-		errorBounds := make([]float32, unQuantizedSet.GetCount())
-		unquantizer.EstimateDistances(
-			&s.Workspace, unQuantizedSet, queryVector, exact, errorBounds)
+		centroid := vectors.Centroid(make(vector.T, vectors.Dims))
+		if metric == vecdist.InnerProduct || metric == vecdist.Cosine {
+			// Use spherical centroid for inner product and cosine distances,
+			// which is the mean centroid, but normalized.
+			num32.Normalize(centroid)
+		}
+
+		// Test UnQuantizer.
+		unQuantizer := quantize.NewUnQuantizer(len(queryVector), metric)
+		quantizedSet, exact, errorBounds := addVectors(unQuantizer, centroid, vectors)
+		unQuantizedSet := quantizedSet.(*quantize.UnQuantizedVectorSet)
 		for _, error := range errorBounds {
+			// ErrorBounds should always be zero for UnQuantizer.
 			require.Zero(t, error)
 		}
 
+		// Test RaBitQuantizer.
 		rabitQ := quantize.NewRaBitQuantizer(len(queryVector), 42, metric)
-		rabitQSet := rabitQ.Quantize(&s.Workspace, vectors)
-		estimated := make([]float32, rabitQSet.GetCount())
-		rabitQ.EstimateDistances(
-			&s.Workspace, rabitQSet, queryVector, estimated, errorBounds)
+		quantizedSet, estimated, errorBounds := addVectors(rabitQ, centroid, vectors)
+		rabitQSet := quantizedSet.(*quantize.RaBitQuantizedVectorSet)
 
-		// UnQuantizer and RaBitQuantizer should have calculated same centroid.
-		require.Equal(t, unQuantizedSet.GetCentroid(), rabitQSet.GetCentroid())
+		require.Equal(t, unQuantizedSet.GetCount(), rabitQSet.GetCount())
 
 		buf.WriteString("  Query = ")
 		utils.WriteVector(&buf, queryVector, 4)
 		buf.WriteByte('\n')
 
 		buf.WriteString("  Centroid = ")
-		utils.WriteVector(&buf, rabitQSet.GetCentroid(), 4)
+		utils.WriteVector(&buf, rabitQSet.Centroid, 4)
 		buf.WriteByte('\n')
 
 		for i := range vectors.Count {