quantize: recompute norm when centroid is updated

When RaBitQuantizedVectorSet.Clear was called with a new centroid, the norm
was not being recomputed. This commit fixes that bug.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

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

@@ -145,6 +145,9 @@ func (vs *RaBitQuantizedVectorSet) Clone() QuantizedVectorSet {
 // Clear implements the QuantizedVectorSet interface
 func (vs *RaBitQuantizedVectorSet) Clear(centroid vector.T) {
 	if buildutil.CrdbTestBuild {
+		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
 		}
@@ -157,18 +160,21 @@ func (vs *RaBitQuantizedVectorSet) Clear(centroid vector.T) {
 		// RaBitQCodeSet.Clear takes care of scribbling memory for vs.Codes.
 	}
 
+	// Recompute the centroid norm for Cosine and InnerProduct metrics, but only
+	// if a new centroid is provided.
+	if vs.Metric != vecpb.L2SquaredDistance {
+		if &vs.Centroid[0] != &centroid[0] {
+			vs.CentroidNorm = num32.Norm(centroid)
+		}
+	}
+
 	// vs.Centroid is immutable, so do not try to reuse its memory.
 	vs.Centroid = centroid
 	vs.Codes.Clear()
 	vs.CodeCounts = vs.CodeCounts[:0]
 	vs.CentroidDistances = vs.CentroidDistances[:0]
 	vs.QuantizedDotProducts = vs.QuantizedDotProducts[:0]
 	vs.CentroidDotProducts = vs.CentroidDotProducts[:0]
-	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

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

@@ -105,10 +105,15 @@ func TestRaBitQuantizedVectorSet(t *testing.T) {
 	require.Equal(t, []float32{10}, cloned.CentroidDistances)
 	require.Equal(t, []float32{10}, cloned.QuantizedDotProducts)
 
+	// Clear the set and ensure that norm is not updated.
+	quantizedSet.Clear(quantizedSet.Centroid)
+	require.Equal(t, float32(0), quantizedSet.CentroidNorm)
+
 	// Test InnerProduct distance metric, which uses the CentroidDotProducts
 	// field (L2Squared does not use it).
-	quantizedSet.Clear(quantizedSet.Centroid)
 	quantizedSet.Metric = vecpb.InnerProductDistance
+	quantizedSet.Clear(quantizedSet.Centroid)
+	require.Equal(t, float32(0), quantizedSet.CentroidNorm)
 	quantizedSet.AddUndefined(2)
 	copy(quantizedSet.Codes.At(1), []uint64{1, 2, 3})
 	quantizedSet.CodeCounts[1] = 15
@@ -124,4 +129,8 @@ func TestRaBitQuantizedVectorSet(t *testing.T) {
 	require.Len(t, cloned.CentroidDotProducts, 2)
 	cloned.Clear(quantizedSet.Centroid)
 	require.Len(t, cloned.CentroidDotProducts, 0)
+
+	// Update the centroid and ensure that norm is updated.
+	quantizedSet.Clear([]float32{2, 3, 6})
+	require.Equal(t, float32(7), quantizedSet.CentroidNorm)
 }