Merge #148427

148427: cspann: add GetCentroidDistances method to Quantizer interface r=drewkimball a=andy-kimball

#### quantize: track distance metric in RaBitQuantizedSet

Include the distance metric in RaBitQuantizedSet, since its needed for
methods like Clear and AddUndefined. Also, clean up the encoding methods
for RaBitQVector a bit.

#### quantize: rename NewQuantizedVectorSet to NewSet

Rename the NewQuantizedVectorSet method on the Quantizer interface to
NewSet. "QuantizedVector" is already implied by context and consistent
with other methdods like "QuantizeInSet".

#### 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.

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

Author: craig[bot]

pkg/sql/vecindex/cspann/partition.go

@@ -282,6 +282,6 @@ func (p *Partition) Clear() int {
 // CreateEmptyPartition returns an empty partition for the given quantizer and
 // level.
 func CreateEmptyPartition(quantizer quantize.Quantizer, metadata PartitionMetadata) *Partition {
-	quantizedSet := quantizer.NewQuantizedVectorSet(0, metadata.Centroid)
+	quantizedSet := quantizer.NewSet(0, metadata.Centroid)
 	return NewPartition(metadata, quantizer, quantizedSet, []ChildKey(nil), []ValueBytes(nil))
 }

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

@@ -8,6 +8,7 @@ proto_library(
     strip_import_prefix = "/pkg",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/sql/vecindex/vecpb:vecpb_proto",
         "//pkg/util/vector:vector_proto",
         "@com_github_gogo_protobuf//gogoproto:gogo_proto",
     ],
@@ -20,6 +21,7 @@ go_proto_library(
     proto = ":quantize_proto",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/sql/vecindex/vecpb",
         "//pkg/util/vector",
         "@com_github_gogo_protobuf//gogoproto",
     ],
@@ -71,5 +73,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

@@ -7,6 +7,7 @@ syntax = "proto3";
 package cockroach.sql.vecindex.quantize;
 option go_package = "github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann/quantize";
 
+import "sql/vecindex/vecpb/vec.proto";
 import "util/vector/vector.proto";
 import "gogoproto/gogo.proto";
 
@@ -32,30 +33,37 @@ message RaBitQCodeSet {
 message RaBitQuantizedVectorSet {
   option (gogoproto.equal) = true;
 
+  // Metric specifies the metric used to compute similarity between vectors in
+  // the set.
+  cockroach.sql.vecindex.vecpb.DistanceMetric metric = 1;
   // 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;
+  repeated float centroid = 2;
   // Codes is a set of RaBitQ quantization codes, with one code per quantized
   // vector in the set.
-  RaBitQCodeSet codes = 2 [(gogoproto.nullable) = false];
+  RaBitQCodeSet codes = 3 [(gogoproto.nullable) = false];
   // CodeCounts records the count of "1" bits in each of the quantization codes.
-  repeated uint32 code_counts = 3;
+  repeated uint32 code_counts = 4;
   // CentroidDistances is a slice of the exact Euclidean distances (non-squared)
   // of the original full-size vectors from the centroid.
-  repeated float centroid_distances = 4;
+  repeated float centroid_distances = 5;
   // QuantizedDotProducts is a slice of the exact inner products between the
   // original full-size vectors and their corresponding quantized vectors.
   // NOTE: These values have been inverted (1/inner_product) to avoid expensive
   // division during distance estimation.
-  repeated float quantized_dot_products = 5;
+  repeated float quantized_dot_products = 6;
   // CentroidDotProducts is a slice of the exact inner products between the
   // original full-size vectors and the centroid.
-  repeated float centroid_dot_products = 6;
+  // NOTE: This is always nil when using the L2Squared distance metric.
+  repeated float centroid_dot_products = 7;
+  // CentroidNorm is the L2 norm of the mean centroid.
+  // NOTE: This is always nil when using the L2Squared distance metric.
+  float centroid_norm = 8;
 }
 
 // UnQuantizedVectorSet trivially implements the QuantizedVectorSet interface,

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

@@ -41,9 +41,9 @@ type Quantizer interface {
 	//       vectors.
 	QuantizeInSet(w *workspace.T, quantizedSet QuantizedVectorSet, vectors vector.Set)
 
-	// NewQuantizedVectorSet returns a new empty vector set preallocated to the
-	// number of vectors specified.
-	NewQuantizedVectorSet(capacity int, centroid vector.T) QuantizedVectorSet
+	// NewSet returns a new empty vector set preallocated to the number of vectors
+	// specified.
+	NewSet(capacity int, centroid vector.T) QuantizedVectorSet
 
 	// EstimateDistances returns the estimated distances of the query vector from
 	// each data vector represented in the given quantized vector set, as well as
@@ -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)
 
@@ -100,7 +103,7 @@ func (s *testState) estimateDistances(t *testing.T, d *datadriven.TestData) stri
 
 		// Now add the vectors one-by-one and ensure that distances and error
 		// bounds stay the same.
-		quantizedSet = quantizer.NewQuantizedVectorSet(vectors.Count, centroid)
+		quantizedSet = quantizer.NewSet(vectors.Count, centroid)
 		for i := range vectors.Count {
 			quantizer.QuantizeInSet(&s.Workspace, quantizedSet, vectors.Slice(i, 1))
 		}
@@ -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)
@@ -111,7 +112,7 @@ func (q *RaBitQuantizer) Quantize(w *workspace.T, vectors vector.Set) QuantizedV
 		centroid = vectors.Centroid(make(vector.T, vectors.Dims))
 	}
 
-	quantizedSet := q.NewQuantizedVectorSet(vectors.Count, centroid)
+	quantizedSet := q.NewSet(vectors.Count, centroid)
 	q.quantizeHelper(w, quantizedSet.(*RaBitQuantizedVectorSet), vectors)
 	return quantizedSet
 }
@@ -123,32 +124,43 @@ func (q *RaBitQuantizer) QuantizeInSet(
 	q.quantizeHelper(w, quantizedSet.(*RaBitQuantizedVectorSet), vectors)
 }
 
-// NewQuantizedVectorSet implements the Quantizer interface
-func (q *RaBitQuantizer) NewQuantizedVectorSet(capacity int, centroid vector.T) QuantizedVectorSet {
-	codeWidth := RaBitQCodeSetWidth(q.GetDims())
-	dataBuffer := make([]uint64, 0, capacity*codeWidth)
+// NewSet implements the Quantizer interface
+func (q *RaBitQuantizer) NewSet(capacity int, centroid vector.T) QuantizedVectorSet {
+	var vs *RaBitQuantizedVectorSet
+
 	if capacity <= 1 {
 		// Special case capacity of zero or one by using in-line storage.
-		var quantized raBitQuantizedVector
-		quantized.Centroid = centroid
-		quantized.Codes = MakeRaBitQCodeSetFromRawData(dataBuffer, codeWidth)
+		quantized := &raBitQuantizedVector{}
 		quantized.CodeCounts = quantized.codeCountStorage[:0]
 		quantized.CentroidDistances = quantized.centroidDistanceStorage[:0]
-		quantized.QuantizedDotProducts = quantized.dotProductStorage[:0]
-		return &quantized.RaBitQuantizedVectorSet
-	}
+		quantized.QuantizedDotProducts = quantized.quantizedDotProductStorage[:0]
 
-	vs := &RaBitQuantizedVectorSet{
-		Centroid:             centroid,
-		Codes:                MakeRaBitQCodeSetFromRawData(dataBuffer, codeWidth),
-		CodeCounts:           make([]uint32, 0, capacity),
-		CentroidDistances:    make([]float32, 0, capacity),
-		QuantizedDotProducts: make([]float32, 0, capacity),
+		// L2Squared doesn't use this.
+		if q.distanceMetric != vecpb.L2SquaredDistance {
+			quantized.CentroidDotProducts = quantized.centroidDotProductStorage[:0]
+		}
+		vs = &quantized.RaBitQuantizedVectorSet
+	} else {
+		vs = &RaBitQuantizedVectorSet{
+			CodeCounts:           make([]uint32, 0, capacity),
+			CentroidDistances:    make([]float32, 0, capacity),
+			QuantizedDotProducts: make([]float32, 0, capacity),
+		}
+		// L2Squared doesn't use these, so don't make extra allocation or calculation.
+		if q.distanceMetric != vecpb.L2SquaredDistance {
+			vs.CentroidDotProducts = make([]float32, 0, capacity)
+		}
 	}
-	// L2Squared doesn't use this, so don't make extra allocation.
+
+	vs.Metric = q.distanceMetric
+	vs.Centroid = centroid
+	codeWidth := RaBitQCodeSetWidth(q.GetDims())
+	dataBuffer := make([]uint64, 0, capacity*codeWidth)
+	vs.Codes = MakeRaBitQCodeSetFromRawData(dataBuffer, codeWidth)
 	if q.distanceMetric != vecpb.L2SquaredDistance {
-		vs.CentroidDotProducts = make([]float32, 0, capacity)
+		vs.CentroidNorm = num32.Norm(centroid)
 	}
+
 	return vs
 }
 
@@ -180,28 +192,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 +201,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 +362,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(
@@ -383,7 +416,7 @@ func (q *RaBitQuantizer) quantizeHelper(
 	// Extend any existing slices in the vector set.
 	count := vectors.Count
 	oldCount := qs.GetCount()
-	qs.AddUndefined(count, q.distanceMetric)
+	qs.AddUndefined(count)
 
 	// L2Squared doesn't use this, so don't store it.
 	if q.distanceMetric != vecpb.L2SquaredDistance {