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.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

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",
     ],

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,34 +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: 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.
   // NOTE: This is always nil when using the L2Squared distance metric.
-  repeated float centroid_dot_products = 6;
+  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 = 7;
+  float centroid_norm = 8;
 }
 
 // UnQuantizedVectorSet trivially implements the QuantizedVectorSet interface,

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

@@ -126,37 +126,41 @@ func (q *RaBitQuantizer) QuantizeInSet(
 
 // NewSet implements the Quantizer interface
 func (q *RaBitQuantizer) NewSet(capacity int, centroid vector.T) QuantizedVectorSet {
-	codeWidth := RaBitQCodeSetWidth(q.GetDims())
-	dataBuffer := make([]uint64, 0, capacity*codeWidth)
+	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.quantizedDotProductStorage[:0]
 
-		// L2Squared doesn't use these, so don't make extra calculations.
+		// L2Squared doesn't use this.
 		if q.distanceMetric != vecpb.L2SquaredDistance {
 			quantized.CentroidDotProducts = quantized.centroidDotProductStorage[:0]
-			quantized.CentroidNorm = num32.Norm(centroid)
 		}
-		return &quantized.RaBitQuantizedVectorSet
+		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)
+		}
 	}
 
-	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 these, so don't make extra allocation or calculation.
+	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
 }
 
@@ -412,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 {

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

@@ -135,6 +135,7 @@ func (vs *RaBitQuantizedVectorSet) ReplaceWithLast(offset int) {
 // Clone implements the QuantizedVectorSet interface.
 func (vs *RaBitQuantizedVectorSet) Clone() QuantizedVectorSet {
 	return &RaBitQuantizedVectorSet{
+		Metric:               vs.Metric,
 		Centroid:             vs.Centroid, // Centroid is immutable
 		Codes:                vs.Codes.Clone(),
 		CodeCounts:           slices.Clone(vs.CodeCounts),
@@ -167,14 +168,16 @@ 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)
+	if vs.Metric != vecpb.L2SquaredDistance {
+		if &vs.Centroid[0] != &centroid[0] {
+			vs.CentroidNorm = num32.Norm(centroid)
+		}
 	}
 }
 
 // AddUndefined adds the given number of quantized vectors to this set. The new
 // quantized vector information should be set to defined values before use.
-func (vs *RaBitQuantizedVectorSet) AddUndefined(count int, distanceMetric vecpb.DistanceMetric) {
+func (vs *RaBitQuantizedVectorSet) AddUndefined(count int) {
 	newCount := len(vs.CodeCounts) + count
 	vs.Codes.AddUndefined(count)
 	vs.CodeCounts = slices.Grow(vs.CodeCounts, count)
@@ -183,7 +186,7 @@ func (vs *RaBitQuantizedVectorSet) AddUndefined(count int, distanceMetric vecpb.
 	vs.CentroidDistances = vs.CentroidDistances[:newCount]
 	vs.QuantizedDotProducts = slices.Grow(vs.QuantizedDotProducts, count)
 	vs.QuantizedDotProducts = vs.QuantizedDotProducts[:newCount]
-	if distanceMetric != vecpb.L2SquaredDistance {
+	if vs.Metric != vecpb.L2SquaredDistance {
 		// L2Squared doesn't need this.
 		vs.CentroidDotProducts = slices.Grow(vs.CentroidDotProducts, count)
 		vs.CentroidDotProducts = vs.CentroidDotProducts[:newCount]

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

@@ -62,10 +62,11 @@ func TestRaBitCodeSet(t *testing.T) {
 
 func TestRaBitQuantizedVectorSet(t *testing.T) {
 	var quantizedSet RaBitQuantizedVectorSet
+	quantizedSet.Metric = vecpb.L2SquaredDistance
 	quantizedSet.Centroid = []float32{1, 2, 3}
 	quantizedSet.Codes.Width = 3
 
-	quantizedSet.AddUndefined(5, vecpb.L2SquaredDistance)
+	quantizedSet.AddUndefined(5)
 	copy(quantizedSet.Codes.At(4), []uint64{1, 2, 3})
 	quantizedSet.CodeCounts[4] = 15
 	quantizedSet.CentroidDistances[4] = 1.23
@@ -107,7 +108,8 @@ func TestRaBitQuantizedVectorSet(t *testing.T) {
 	// Test InnerProduct distance metric, which uses the CentroidDotProducts
 	// field (L2Squared does not use it).
 	quantizedSet.Clear(quantizedSet.Centroid)
-	quantizedSet.AddUndefined(2, vecpb.InnerProductDistance)
+	quantizedSet.Metric = vecpb.InnerProductDistance
+	quantizedSet.AddUndefined(2)
 	copy(quantizedSet.Codes.At(1), []uint64{1, 2, 3})
 	quantizedSet.CodeCounts[1] = 15
 	quantizedSet.CentroidDistances[1] = 1.23

pkg/sql/vecindex/vecencoding/encoding.go

@@ -202,22 +202,16 @@ func EncodeMetadataValue(metadata cspann.PartitionMetadata) []byte {
 }
 
 // EncodeRaBitQVector encodes a RaBitQ vector into the given byte slice.
-func EncodeRaBitQVector(
-	appendTo []byte,
-	codeCount uint32,
-	centroidDistance float32,
-	quantizedDotProduct float32,
-	centroidDotProduct float32,
-	code quantize.RaBitQCode,
-	metric vecpb.DistanceMetric,
+func EncodeRaBitQVectorFromSet(
+	appendTo []byte, vectorSet *quantize.RaBitQuantizedVectorSet, offset int,
 ) []byte {
-	appendTo = encoding.EncodeUint32Ascending(appendTo, codeCount)
-	appendTo = encoding.EncodeUntaggedFloat32Value(appendTo, centroidDistance)
-	appendTo = encoding.EncodeUntaggedFloat32Value(appendTo, quantizedDotProduct)
-	if metric != vecpb.L2SquaredDistance {
-		appendTo = encoding.EncodeUntaggedFloat32Value(appendTo, centroidDotProduct)
+	appendTo = encoding.EncodeUint32Ascending(appendTo, vectorSet.CodeCounts[offset])
+	appendTo = encoding.EncodeUntaggedFloat32Value(appendTo, vectorSet.CentroidDistances[offset])
+	appendTo = encoding.EncodeUntaggedFloat32Value(appendTo, vectorSet.QuantizedDotProducts[offset])
+	if vectorSet.Metric != vecpb.L2SquaredDistance {
+		appendTo = encoding.EncodeUntaggedFloat32Value(appendTo, vectorSet.CentroidDotProducts[offset])
 	}
-	for _, c := range code {
+	for _, c := range vectorSet.Codes.At(offset) {
 		appendTo = encoding.EncodeUint64Ascending(appendTo, c)
 	}
 	return appendTo
@@ -318,7 +312,7 @@ func DecodeMetadataValue(encMetadata []byte) (metadata cspann.PartitionMetadata,
 // RaBitQuantizedVectorSet. The vector set must have been initialized with the
 // correct number of dimensions. It returns the remainder of the input buffer.
 func DecodeRaBitQVectorToSet(
-	encVector []byte, vectorSet *quantize.RaBitQuantizedVectorSet, metric vecpb.DistanceMetric,
+	encVector []byte, vectorSet *quantize.RaBitQuantizedVectorSet,
 ) ([]byte, error) {
 	encVector, codeCount, err := encoding.DecodeUint32Ascending(encVector)
 	if err != nil {
@@ -332,7 +326,7 @@ func DecodeRaBitQVectorToSet(
 	if err != nil {
 		return nil, err
 	}
-	if metric != vecpb.L2SquaredDistance {
+	if vectorSet.Metric != vecpb.L2SquaredDistance {
 		var centroidDotProduct float32
 		encVector, centroidDotProduct, err = encoding.DecodeUntaggedFloat32Value(encVector)
 		if err != nil {

pkg/sql/vecindex/vecencoding/encoding_test.go

@@ -94,15 +94,7 @@ func testEncodeDecodeRoundTripImpl(t *testing.T, rnd *rand.Rand, set vector.Set)
 				buf, err = vecencoding.EncodeUnquantizerVector(buf, set.At(i))
 				require.NoError(t, err)
 			case *quantize.RaBitQuantizedVectorSet:
-				var centroidDotProduct float32
-				if distMetric != vecpb.L2SquaredDistance {
-					centroidDotProduct = quantizedSet.CentroidDotProducts[i]
-				}
-				buf = vecencoding.EncodeRaBitQVector(buf,
-					quantizedSet.CodeCounts[i], quantizedSet.CentroidDistances[i],
-					quantizedSet.QuantizedDotProducts[i], centroidDotProduct,
-					quantizedSet.Codes.At(i), distMetric,
-				)
+				buf = vecencoding.EncodeRaBitQVectorFromSet(buf, quantizedSet, i)
 			}
 		}
 
@@ -131,9 +123,7 @@ func testEncodeDecodeRoundTripImpl(t *testing.T, rnd *rand.Rand, set vector.Set)
 			decodedSet = quantizer.NewSet(set.Count, decodedMetadata.Centroid)
 			for range set.Count {
 				remainder, err = vecencoding.DecodeRaBitQVectorToSet(
-					remainder,
-					decodedSet.(*quantize.RaBitQuantizedVectorSet),
-					distMetric,
+					remainder, decodedSet.(*quantize.RaBitQuantizedVectorSet),
 				)
 				require.NoError(t, err)
 			}
@@ -188,6 +178,7 @@ func testingAssertPartitionsEqual(t *testing.T, l, r *cspann.Partition) {
 	case *quantize.RaBitQuantizedVectorSet:
 		rightSet, ok := q2.(*quantize.RaBitQuantizedVectorSet)
 		require.True(t, ok, "quantized set types do not match")
+		require.Equal(t, leftSet.Metric, rightSet.Metric)
 		require.Equal(t, leftSet.CodeCounts, rightSet.CodeCounts, "code counts do not match")
 		require.Equal(t, leftSet.Codes, rightSet.Codes, "codes do not match")
 		require.Equal(t, leftSet.QuantizedDotProducts, rightSet.QuantizedDotProducts,

pkg/sql/vecindex/vecstore/BUILD.bazel

@@ -28,7 +28,6 @@ go_library(
         "//pkg/sql/vecindex/cspann/quantize",
         "//pkg/sql/vecindex/cspann/workspace",
         "//pkg/sql/vecindex/vecencoding",
-        "//pkg/sql/vecindex/vecpb",
         "//pkg/util/log",
         "//pkg/util/unique",
         "//pkg/util/vector",

pkg/sql/vecindex/vecstore/codec.go

@@ -10,7 +10,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann/quantize"
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann/workspace"
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecencoding"
-	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecpb"
 	"github.com/cockroachdb/cockroach/pkg/util/vector"
 	"github.com/cockroachdb/errors"
 )
@@ -56,9 +55,7 @@ func (sc *storeCodec) DecodeVector(encodedVector []byte) ([]byte, error) {
 			encodedVector, sc.tmpVectorSet.(*quantize.UnQuantizedVectorSet))
 	case *quantize.RaBitQuantizer:
 		return vecencoding.DecodeRaBitQVectorToSet(
-			encodedVector,
-			sc.tmpVectorSet.(*quantize.RaBitQuantizedVectorSet),
-			sc.quantizer.GetDistanceMetric(),
+			encodedVector, sc.tmpVectorSet.(*quantize.RaBitQuantizedVectorSet),
 		)
 	}
 	return nil, errors.Errorf("unknown quantizer type %T", sc.quantizer)
@@ -75,21 +72,7 @@ func (sc *storeCodec) EncodeVector(w *workspace.T, v vector.T, centroid vector.T
 	case *quantize.UnQuantizedVectorSet:
 		return vecencoding.EncodeUnquantizerVector([]byte{}, t.Vectors.At(0))
 	case *quantize.RaBitQuantizedVectorSet:
-		metric := sc.quantizer.GetDistanceMetric()
-		var centroidDotProduct float32
-		if metric != vecpb.L2SquaredDistance {
-			// CentroidDotProducts is only defined for non-L2 distance metrics.
-			centroidDotProduct = t.CentroidDotProducts[0]
-		}
-		return vecencoding.EncodeRaBitQVector(
-			[]byte{},
-			t.CodeCounts[0],
-			t.CentroidDistances[0],
-			t.QuantizedDotProducts[0],
-			centroidDotProduct,
-			t.Codes.At(0),
-			metric,
-		), nil
+		return vecencoding.EncodeRaBitQVectorFromSet([]byte{}, t, 0), nil
 	default:
 		return nil, errors.Errorf("unknown quantizer type %T", t)
 	}