quantize: set zero dot product for centroid data vector

Previously, the dot product between a data unit vector and its quantized
form was not being set in the case where the data vector is equal to the
centroid. This could cause an issue when a RaBitQuantizedSet is reused
and the dot product memory is not zero. Fix this buglet and update the
code to scribble undefined memory.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

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

@@ -557,9 +557,11 @@ func (q *RaBitQuantizer) quantizeHelper(
 		// Store the inverted dot product, which will be used to make distance
 		// estimates. The dot product is only zero in the case where the data vector
 		// is equal to the centroid vector. That case is handled separately in
-		// EstimatedDistances.
+		// EstimateDistances.
 		if dotProduct != 0 {
 			dotProducts[i] = 1.0 / dotProduct
+		} else {
+			dotProducts[i] = 0
 		}
 	}
 }

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

@@ -150,19 +150,29 @@ func TestRaBitQuantizerEdge(t *testing.T) {
 
 	t.Run("add centroid to set", func(t *testing.T) {
 		quantizer := NewRaBitQuantizer(2, 42, vecpb.L2SquaredDistance)
+		quantizedSet := quantizer.NewSet(4, []float32{3, 9}).(*RaBitQuantizedVectorSet)
 		vectors := vector.MakeSetFromRawData([]float32{1, 5, 5, 13}, 2)
-		quantizedSet := quantizer.Quantize(&workspace, vectors).(*RaBitQuantizedVectorSet)
-		require.Equal(t, []float32{3, 9}, quantizedSet.Centroid)
+		quantizer.QuantizeInSet(&workspace, quantizedSet, vectors)
 
 		// Add centroid to the set along with another vector.
 		vectors = vector.MakeSetFromRawData([]float32{1, 5, 3, 9}, 2)
 		quantizer.QuantizeInSet(&workspace, quantizedSet, vectors)
+		require.Equal(t, float32(0), quantizedSet.QuantizedDotProducts[3],
+			"dot product for centroid should be zero")
+
+		// Estimate distances from a query vector not in the set.
 		distances := make([]float32, 4)
 		errorBounds := make([]float32, 4)
 		quantizer.EstimateDistances(
 			&workspace, quantizedSet, vector.T{3, 2}, distances, errorBounds)
 		require.Equal(t, []float32{22.33, 115.67, 22.33, 49}, testutils.RoundFloats(distances, 2))
 		require.Equal(t, []float32{44.27, 44.27, 44.27, 0}, testutils.RoundFloats(errorBounds, 2))
+
+		// Estimate distances when the query vector is the centroid.
+		quantizer.EstimateDistances(
+			&workspace, quantizedSet, vector.T{3, 9}, distances, errorBounds)
+		require.Equal(t, []float32{20, 20, 20, 0}, testutils.RoundFloats(distances, 2))
+		require.Equal(t, []float32{0, 0, 0, 0}, testutils.RoundFloats(errorBounds, 2))
 	})
 
 	t.Run("query vector is centroid", func(t *testing.T) {

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

@@ -98,6 +98,11 @@ func (cs *RaBitQCodeSet) AddUndefined(count int) {
 	cs.Data = slices.Grow(cs.Data, count*cs.Width)
 	cs.Count += count
 	cs.Data = cs.Data[:cs.Count*cs.Width]
+	if buildutil.CrdbTestBuild {
+		for i := len(cs.Data) - count*cs.Width; i < len(cs.Data); i++ {
+			cs.Data[i] = 0xBADF00D
+		}
+	}
 }
 
 // ReplaceWithLast removes the code at the given offset from the set, replacing
@@ -148,16 +153,7 @@ func (vs *RaBitQuantizedVectorSet) Clear(centroid vector.T) {
 		if vs.Centroid == nil {
 			panic(errors.New("Clear cannot be called on an uninitialized vector set"))
 		}
-		for i := range len(vs.CodeCounts) {
-			vs.CodeCounts[i] = 0xBADF00D
-		}
-		for i := range len(vs.CentroidDistances) {
-			vs.CentroidDistances[i] = math.Pi
-		}
-		for i := range len(vs.QuantizedDotProducts) {
-			vs.QuantizedDotProducts[i] = math.Pi
-		}
-		// RaBitQCodeSet.Clear takes care of scribbling memory for vs.Codes.
+		vs.scribble(0, len(vs.CodeCounts))
 	}
 
 	// Recompute the centroid norm for Cosine and InnerProduct metrics, but only
@@ -193,4 +189,28 @@ func (vs *RaBitQuantizedVectorSet) AddUndefined(count int) {
 		vs.CentroidDotProducts = slices.Grow(vs.CentroidDotProducts, count)
 		vs.CentroidDotProducts = vs.CentroidDotProducts[:newCount]
 	}
+	if buildutil.CrdbTestBuild {
+		vs.scribble(newCount-count, newCount)
+	}
+}
+
+// scribble writes garbage values to undefined vector set values. This is only
+// called in test builds to make detecting bugs easier.
+func (vs *RaBitQuantizedVectorSet) scribble(start, end int) {
+	for i := start; i < end; i++ {
+		vs.CodeCounts[i] = 0xBADF00D
+	}
+	for i := start; i < end; i++ {
+		vs.CentroidDistances[i] = math.Pi
+	}
+	for i := start; i < end; i++ {
+		vs.QuantizedDotProducts[i] = math.Pi
+	}
+	if vs.Metric != vecpb.L2SquaredDistance {
+		for i := start; i < end; i++ {
+			vs.CentroidDotProducts[i] = math.Pi
+		}
+	}
+	// RaBitQCodeSet Clear and AddUndefined methods take care of scribbling
+	// memory for vs.Codes.
 }