quantize: add GetCentroidDistances method to Quantizer interface

GetCentroidDistances returns the exact distance of each vector in a quantized
set to the centroid of that set, according to the quantizer's distance metric
(e.g. L2Squared or Cosine). It also supports distance to the spherical centroid
for Cosine and InnerProduct metrics. This method will be used in a later PR by
the partition split code.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

pkg/sql/vecindex/cspann/quantize/BUILD.bazel

@@ -71,5 +71,6 @@ go_test(
         "//pkg/util/vector",
         "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_stretchr_testify//require",
+        "@org_gonum_v1_gonum//floats/scalar",
     ],
 )

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

@@ -35,9 +35,9 @@ message RaBitQuantizedVectorSet {
   // 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: This is always the mean centroid, regardless of the distance metric.
+  // The caller is responsible for converting this to a spherical centroid when
+  // that's needed.
   // NOTE: The centroid should be treated as immutable.
   repeated float centroid = 1;
   // Codes is a set of RaBitQ quantization codes, with one code per quantized
@@ -55,7 +55,11 @@ message RaBitQuantizedVectorSet {
   repeated float quantized_dot_products = 5;
   // CentroidDotProducts is a slice of the exact inner products between the
   // original full-size vectors and the centroid.
+  // NOTE: This is always nil when using the L2Squared distance metric.
   repeated float centroid_dot_products = 6;
+  // CentroidNorm is the L2 norm of the mean centroid.
+  // NOTE: This is always nil when using the L2Squared distance metric.
+  float centroid_norm = 7;
 }
 
 // UnQuantizedVectorSet trivially implements the QuantizedVectorSet interface,

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

@@ -61,6 +61,18 @@ type Quantizer interface {
 		distances []float32,
 		errorBounds []float32,
 	)
+
+	// GetCentroidDistances returns the exact distance of each vector in
+	// "quantizedSet" from that set's centroid, according to the quantizer's
+	// distance metric (e.g. L2Squared or Cosine). By default, it returns
+	// distances to the mean centroid. However, if "spherical" is true and the
+	// distance metric is Cosine or InnerProduct, then it returns distances to
+	// the spherical centroid instead.
+	//
+	// The caller is responsible for allocating the "distances" slice with length
+	// equal to the number of quantized vectors in "quantizedSet". The centroid
+	// distances will be copied into that slice.
+	GetCentroidDistances(quantizedSet QuantizedVectorSet, distances []float32, spherical bool)
 }
 
 // QuantizedVectorSet is the compressed form of an original set of full-size

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

@@ -41,6 +41,9 @@ func TestQuantizers(t *testing.T) {
 			case "estimate-distances":
 				return state.estimateDistances(t, d)
 
+			case "get-centroid-distances":
+				return state.getCentroidDistances(t, d)
+
 			case "calculate-recall":
 				return state.calculateRecall(t, d)
 
@@ -172,6 +175,73 @@ func (s *testState) estimateDistances(t *testing.T, d *datadriven.TestData) stri
 	return buf.String()
 }
 
+func (s *testState) getCentroidDistances(t *testing.T, d *datadriven.TestData) string {
+	var dims int
+	var err error
+	var vectors vector.Set
+	for _, arg := range d.CmdArgs {
+		switch arg.Key {
+		case "dims":
+			require.Len(t, arg.Vals, 1)
+			dims, err = strconv.Atoi(arg.Vals[0])
+			require.NoError(t, err)
+
+			// Parse the input vectors.
+			vectors = vector.MakeSet(dims)
+			for _, line := range strings.Split(d.Input, "\n") {
+				line = strings.TrimSpace(line)
+				if len(line) == 0 {
+					continue
+				}
+
+				vec, err := vector.ParseVector(line)
+				require.NoError(t, err)
+				vectors.Add(vec)
+			}
+
+		default:
+			t.Fatalf("unknown arg: %s", arg.Key)
+		}
+	}
+
+	var buf bytes.Buffer
+	doTest := func(metric vecpb.DistanceMetric) {
+		rabitQ := quantize.NewRaBitQuantizer(dims, 42, metric)
+		quantizedSet := rabitQ.Quantize(&s.Workspace, vectors).(*quantize.RaBitQuantizedVectorSet)
+
+		buf.WriteString("  Centroid = ")
+		utils.WriteVector(&buf, quantizedSet.Centroid, 4)
+		buf.WriteByte('\n')
+
+		buf.WriteString("  Mean Centroid Distances = ")
+		distances := make([]float32, quantizedSet.GetCount())
+		rabitQ.GetCentroidDistances(quantizedSet, distances, false /* spherical */)
+		utils.WriteVector(&buf, distances, 4)
+		buf.WriteByte('\n')
+
+		buf.WriteString("  Spherical Centroid Distances = ")
+		rabitQ.GetCentroidDistances(quantizedSet, distances, true /* spherical */)
+		utils.WriteVector(&buf, distances, 4)
+		buf.WriteByte('\n')
+	}
+
+	buf.WriteString("L2Squared\n")
+	doTest(vecpb.L2SquaredDistance)
+
+	buf.WriteString("InnerProduct\n")
+	doTest(vecpb.InnerProductDistance)
+
+	// For cosine distance, normalize the input vectors.
+	for i := range vectors.Count {
+		num32.Normalize(vectors.At(i))
+	}
+
+	buf.WriteString("Cosine\n")
+	doTest(vecpb.CosineDistance)
+
+	return buf.String()
+}
+
 func (s *testState) calculateRecall(t *testing.T, d *datadriven.TestData) string {
 	var datasetName string
 	randomize := false

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

@@ -54,9 +54,10 @@ type RaBitQuantizer struct {
 // to the statically-allocated arrays in this struct.
 type raBitQuantizedVector struct {
 	RaBitQuantizedVectorSet
-	codeCountStorage        [1]uint32
-	centroidDistanceStorage [1]float32
-	dotProductStorage       [1]float32
+	codeCountStorage           [1]uint32
+	centroidDistanceStorage    [1]float32
+	quantizedDotProductStorage [1]float32
+	centroidDotProductStorage  [1]float32
 }
 
 var _ Quantizer = (*RaBitQuantizer)(nil)
@@ -134,7 +135,13 @@ func (q *RaBitQuantizer) NewQuantizedVectorSet(capacity int, centroid vector.T)
 		quantized.Codes = MakeRaBitQCodeSetFromRawData(dataBuffer, codeWidth)
 		quantized.CodeCounts = quantized.codeCountStorage[:0]
 		quantized.CentroidDistances = quantized.centroidDistanceStorage[:0]
-		quantized.QuantizedDotProducts = quantized.dotProductStorage[:0]
+		quantized.QuantizedDotProducts = quantized.quantizedDotProductStorage[:0]
+
+		// L2Squared doesn't use these, so don't make extra calculations.
+		if q.distanceMetric != vecpb.L2SquaredDistance {
+			quantized.CentroidDotProducts = quantized.centroidDotProductStorage[:0]
+			quantized.CentroidNorm = num32.Norm(centroid)
+		}
 		return &quantized.RaBitQuantizedVectorSet
 	}
 
@@ -145,9 +152,10 @@ func (q *RaBitQuantizer) NewQuantizedVectorSet(capacity int, centroid vector.T)
 		CentroidDistances:    make([]float32, 0, capacity),
 		QuantizedDotProducts: make([]float32, 0, capacity),
 	}
-	// L2Squared doesn't use this, so don't make extra allocation.
+	// L2Squared doesn't use these, so don't make extra allocation or calculation.
 	if q.distanceMetric != vecpb.L2SquaredDistance {
 		vs.CentroidDotProducts = make([]float32, 0, capacity)
+		vs.CentroidNorm = num32.Norm(centroid)
 	}
 	return vs
 }
@@ -180,28 +188,7 @@ func (q *RaBitQuantizer) EstimateDistances(
 
 	if queryCentroidDistance == 0 {
 		// The query vector is the centroid.
-		switch q.distanceMetric {
-		case vecpb.L2SquaredDistance:
-			// The distance from the query to the data vectors are just the centroid
-			// distances that have already been calculated, but just need to be
-			// squared.
-			num32.MulTo(distances, raBitSet.CentroidDistances, raBitSet.CentroidDistances)
-
-		case vecpb.InnerProductDistance:
-			// The dot products between the centroid and the data vectors have
-			// already been computed, just need to negate them.
-			num32.ScaleTo(distances, -1, raBitSet.CentroidDotProducts)
-
-		case vecpb.CosineDistance:
-			// All vectors have been normalized, so cosine distance = 1 - dot product.
-			num32.ScaleTo(distances, -1, raBitSet.CentroidDotProducts)
-			num32.AddConst(1, distances)
-
-		default:
-			panic(errors.AssertionFailedf(
-				"RaBitQuantizer does not support distance metric %s", q.distanceMetric))
-		}
-
+		q.GetCentroidDistances(quantizedSet, distances, false /* spherical */)
 		num32.Zero(errorBounds)
 		return
 	}
@@ -210,7 +197,7 @@ func (q *RaBitQuantizer) EstimateDistances(
 	var squaredCentroidNorm, queryCentroidDotProduct float32
 	if q.distanceMetric != vecpb.L2SquaredDistance {
 		queryCentroidDotProduct = num32.Dot(queryVector, raBitSet.Centroid)
-		squaredCentroidNorm = num32.SquaredNorm(raBitSet.Centroid)
+		squaredCentroidNorm = raBitSet.CentroidNorm * raBitSet.CentroidNorm
 	}
 
 	tempQueryUnitVector := tempQueryDiff
@@ -371,6 +358,48 @@ func (q *RaBitQuantizer) EstimateDistances(
 	}
 }
 
+// GetCentroidDistances implements the Quantizer interface.
+func (q *RaBitQuantizer) GetCentroidDistances(
+	quantizedSet QuantizedVectorSet, distances []float32, spherical bool,
+) {
+	raBitSet := quantizedSet.(*RaBitQuantizedVectorSet)
+
+	switch q.distanceMetric {
+	case vecpb.L2SquaredDistance:
+		// The distance from the query to the data vectors are just the centroid
+		// distances that have already been calculated, but just need to be
+		// squared.
+		num32.MulTo(distances, raBitSet.CentroidDistances, raBitSet.CentroidDistances)
+
+	case vecpb.InnerProductDistance:
+		// Need to negate precomputed centroid dot products to compute inner
+		// product distance.
+		multiplier := float32(-1)
+		if spherical && raBitSet.CentroidNorm != 0 {
+			// Convert the mean centroid dot product into a spherical centroid
+			// dot product by dividing by the centroid's norm.
+			multiplier /= raBitSet.CentroidNorm
+		}
+		num32.ScaleTo(distances, multiplier, raBitSet.CentroidDotProducts)
+
+	case vecpb.CosineDistance:
+		// Cosine distance = 1 - dot product when vectors are normalized. The
+		// precomputed centroid dot products were computed with normalized data
+		// vectors, but the centroid was not normalized. Do that now by dividing
+		// the dot products by the centroid's norm. Also negate the result.
+		multiplier := float32(-1)
+		if raBitSet.CentroidNorm != 0 {
+			multiplier /= raBitSet.CentroidNorm
+		}
+		num32.ScaleTo(distances, multiplier, raBitSet.CentroidDotProducts)
+		num32.AddConst(1, distances)
+
+	default:
+		panic(errors.AssertionFailedf(
+			"RaBitQuantizer does not support distance metric %s", q.distanceMetric))
+	}
+}
+
 // quantizeHelper quantizes the given set of vectors and adds the quantization
 // information to the provided quantized vector set.
 func (q *RaBitQuantizer) quantizeHelper(

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

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/num32"
 	"github.com/cockroachdb/cockroach/pkg/util/vector"
 	"github.com/stretchr/testify/require"
+	"gonum.org/v1/gonum/floats/scalar"
 )
 
 // Basic tests.
@@ -77,20 +78,21 @@ func TestRaBitQuantizerSimple(t *testing.T) {
 	})
 
 	t.Run("empty quantized set with capacity", func(t *testing.T) {
-		quantizer := NewRaBitQuantizer(65, 42, vecpb.L2SquaredDistance)
+		quantizer := NewRaBitQuantizer(65, 42, vecpb.InnerProductDistance)
 		centroid := make([]float32, 65)
 		for i := range centroid {
 			centroid[i] = float32(i)
 		}
-		quantizedSet := quantizer.NewQuantizedVectorSet(
-			5, centroid).(*RaBitQuantizedVectorSet)
+		quantizedSet := quantizer.NewQuantizedVectorSet(5, centroid).(*RaBitQuantizedVectorSet)
 		require.Equal(t, centroid, quantizedSet.Centroid)
 		require.Equal(t, 0, quantizedSet.Codes.Count)
 		require.Equal(t, 2, quantizedSet.Codes.Width)
 		require.Equal(t, 10, cap(quantizedSet.Codes.Data))
 		require.Equal(t, 5, cap(quantizedSet.CodeCounts))
 		require.Equal(t, 5, cap(quantizedSet.CentroidDistances))
 		require.Equal(t, 5, cap(quantizedSet.QuantizedDotProducts))
+		require.Equal(t, 5, cap(quantizedSet.CentroidDotProducts))
+		require.Equal(t, float64(299.07), scalar.Round(float64(quantizedSet.CentroidNorm), 2))
 	})
 }
 

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

@@ -11,6 +11,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecpb"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
+	"github.com/cockroachdb/cockroach/pkg/util/num32"
 	"github.com/cockroachdb/cockroach/pkg/util/vector"
 	"github.com/cockroachdb/errors"
 )
@@ -140,6 +141,7 @@ func (vs *RaBitQuantizedVectorSet) Clone() QuantizedVectorSet {
 		CentroidDistances:    slices.Clone(vs.CentroidDistances),
 		QuantizedDotProducts: slices.Clone(vs.QuantizedDotProducts),
 		CentroidDotProducts:  slices.Clone(vs.CentroidDotProducts),
+		CentroidNorm:         vs.CentroidNorm,
 	}
 }
 
@@ -165,6 +167,9 @@ func (vs *RaBitQuantizedVectorSet) Clear(centroid vector.T) {
 	vs.CentroidDistances = vs.CentroidDistances[:0]
 	vs.QuantizedDotProducts = vs.QuantizedDotProducts[:0]
 	vs.CentroidDotProducts = vs.CentroidDotProducts[:0]
+	if &vs.Centroid[0] != &centroid[0] {
+		vs.CentroidNorm = num32.Norm(centroid)
+	}
 }
 
 // AddUndefined adds the given number of quantized vectors to this set. The new

pkg/sql/vecindex/cspann/quantize/testdata/estimate-distances.ddt

@@ -103,7 +103,28 @@ Cosine
   (-1, 0): exact is 2, estimate is 2 ± 0.7071
   (1, 0): exact is 0, estimate is 0 ± 0.7071
 
-# Query vector is equal to the centroid.
+# Query vector is equal to the centroid at the origin.
+estimate-distances query=[0,0]
+[-2,0]
+[2,0]
+----
+L2Squared
+  Query = (0, 0)
+  Centroid = (0, 0)
+  (-2, 0): exact is 4, estimate is 4
+  (2, 0): exact is 4, estimate is 4
+InnerProduct
+  Query = (0, 0)
+  Centroid = (0, 0)
+  (-2, 0): exact is 0, estimate is 0
+  (2, 0): exact is 0, estimate is 0
+Cosine
+  Query = (0, 0)
+  Centroid = (0, 0)
+  (-1, 0): exact is 1, estimate is 1
+  (1, 0): exact is 1, estimate is 1
+
+# Query vector is equal to the centroid at a non-origin point.
 estimate-distances query=[2,2]
 [0,2]
 [4,2]