MB-69655: Fix vector normalization to handle multi-vectors correctly (#2260)

- When indexing multi-vector fields (e.g., `[[3,0,0], [0,4,0]]`) with
`cosine` similarity, normalization was incorrectly applied to the entire
flattened array instead of each sub-vector independently, resulting in
degraded similarity scores.
- Added `NormalizeMultiVector(vec, dims)` that normalizes each
sub-vector separately, fixing scores for multi-vector documents (e.g.,
score now correctly returns 1.0 instead of 0.6 for exact matches).

Author: CascadingRadium

mapping/mapping_vectors.go

@@ -20,6 +20,7 @@ package mapping
 import (
 	"fmt"
 	"reflect"
+	"slices"
 
 	"github.com/blevesearch/bleve/v2/document"
 	"github.com/blevesearch/bleve/v2/util"
@@ -151,8 +152,10 @@ func (fm *FieldMapping) processVector(propertyMightBeVector interface{},
 		vectorIndexOptimizedFor = index.DefaultIndexOptimization
 	}
 	// normalize raw vector if similarity is cosine
+	// Since the vector can be multi-vector (flattened array of multiple vectors),
+	// we use NormalizeMultiVector to normalize each sub-vector independently.
 	if similarity == index.CosineSimilarity {
-		vector = NormalizeVector(vector)
+		vector = NormalizeMultiVector(vector, fm.Dims)
 	}
 
 	fieldName := getFieldName(pathString, path, fm)
@@ -186,7 +189,8 @@ func (fm *FieldMapping) processVectorBase64(propertyMightBeVectorBase64 interfac
 	if err != nil || len(decodedVector) != fm.Dims {
 		return
 	}
-	// normalize raw vector if similarity is cosine
+	// normalize raw vector if similarity is cosine, multi-vector is not supported
+	// for base64 encoded vectors, so we use NormalizeVector directly.
 	if similarity == index.CosineSimilarity {
 		decodedVector = NormalizeVector(decodedVector)
 	}
@@ -292,11 +296,33 @@ func validateVectorFieldAlias(field *FieldMapping, path []string,
 	return nil
 }
 
+// NormalizeVector normalizes a single vector to unit length.
+// It makes a copy of the input vector to avoid modifying it in-place.
 func NormalizeVector(vec []float32) []float32 {
 	// make a copy of the vector to avoid modifying the original
 	// vector in-place
-	vecCopy := make([]float32, len(vec))
-	copy(vecCopy, vec)
+	vecCopy := slices.Clone(vec)
 	// normalize the vector copy using in-place normalization provided by faiss
 	return faiss.NormalizeVector(vecCopy)
 }
+
+// NormalizeMultiVector normalizes each sub-vector of size `dims` independently.
+// For a flattened array containing multiple vectors, each sub-vector is
+// normalized separately to unit length.
+// It makes a copy of the input vector to avoid modifying it in-place.
+func NormalizeMultiVector(vec []float32, dims int) []float32 {
+	if len(vec) == 0 || dims <= 0 || len(vec)%dims != 0 {
+		return vec
+	}
+	// Single vector - delegate to NormalizeVector
+	if len(vec) == dims {
+		return NormalizeVector(vec)
+	}
+	// Multi-vector - make a copy to avoid modifying the original
+	result := slices.Clone(vec)
+	// Normalize each sub-vector in-place
+	for i := 0; i < len(result); i += dims {
+		faiss.NormalizeVector(result[i : i+dims])
+	}
+	return result
+}

mapping/mapping_vectors_test.go

@@ -18,6 +18,7 @@
 package mapping
 
 import (
+	"math"
 	"reflect"
 	"strings"
 	"testing"
@@ -1069,3 +1070,120 @@ func TestNormalizeVector(t *testing.T) {
 		}
 	}
 }
+
+func TestNormalizeMultiVectors(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    []float32
+		dims     int
+		expected []float32
+	}{
+		{
+			name:     "single vector - already normalized",
+			input:    []float32{1, 0, 0},
+			dims:     3,
+			expected: []float32{1, 0, 0},
+		},
+		{
+			name:     "single vector - needs normalization",
+			input:    []float32{3, 0, 0},
+			dims:     3,
+			expected: []float32{1, 0, 0},
+		},
+		{
+			name:     "two vectors - X and Y directions",
+			input:    []float32{3, 0, 0, 0, 4, 0},
+			dims:     3,
+			expected: []float32{1, 0, 0, 0, 1, 0},
+		},
+		{
+			name:     "three vectors",
+			input:    []float32{3, 0, 0, 0, 4, 0, 0, 0, 5},
+			dims:     3,
+			expected: []float32{1, 0, 0, 0, 1, 0, 0, 0, 1},
+		},
+		{
+			name:     "two 2D vectors",
+			input:    []float32{3, 4, 5, 12},
+			dims:     2,
+			expected: []float32{0.6, 0.8, 0.38461538, 0.92307693},
+		},
+		{
+			name:     "empty vector",
+			input:    []float32{},
+			dims:     3,
+			expected: []float32{},
+		},
+		{
+			name:     "zero dims",
+			input:    []float32{1, 2, 3},
+			dims:     0,
+			expected: []float32{1, 2, 3},
+		},
+		{
+			name:     "negative dims",
+			input:    []float32{1, 2, 3},
+			dims:     -1,
+			expected: []float32{1, 2, 3},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Make a copy of input to verify original is not modified
+			inputCopy := make([]float32, len(tt.input))
+			copy(inputCopy, tt.input)
+
+			result := NormalizeMultiVector(tt.input, tt.dims)
+
+			// Check result matches expected
+			if len(result) != len(tt.expected) {
+				t.Errorf("length mismatch: expected %d, got %d", len(tt.expected), len(result))
+				return
+			}
+
+			for i := range result {
+				if !floatApproxEqual(result[i], tt.expected[i], 1e-5) {
+					t.Errorf("value mismatch at index %d: expected %v, got %v",
+						i, tt.expected[i], result[i])
+				}
+			}
+
+			// Verify original input was not modified
+			if !reflect.DeepEqual(tt.input, inputCopy) {
+				t.Errorf("original input was modified: was %v, now %v", inputCopy, tt.input)
+			}
+
+			// For valid multi-vectors, verify each sub-vector has unit magnitude
+			if tt.dims > 0 && len(tt.input) > 0 && len(tt.input)%tt.dims == 0 {
+				numVecs := len(result) / tt.dims
+				for i := 0; i < numVecs; i++ {
+					subVec := result[i*tt.dims : (i+1)*tt.dims]
+					mag := magnitude(subVec)
+					// Allow for zero vectors (magnitude 0) or unit vectors (magnitude 1)
+					if mag > 1e-6 && !floatApproxEqual(mag, 1.0, 1e-5) {
+						t.Errorf("sub-vector %d has magnitude %v, expected 1.0", i, mag)
+					}
+				}
+			}
+		})
+	}
+}
+
+// Helper to compute magnitude of a vector
+func magnitude(v []float32) float32 {
+	var sum float32
+	for _, x := range v {
+		sum += x * x
+	}
+	return float32(math.Sqrt(float64(sum)))
+}
+
+// Helper for approximate float comparison
+func floatApproxEqual(a, b, epsilon float32) bool {
+	diff := a - b
+	if diff < 0 {
+		diff = -diff
+	}
+	return diff < epsilon
+}

search_knn_test.go

@@ -1856,6 +1856,143 @@ func TestKNNMerger(t *testing.T) {
 	})
 }
 
+// TestMultiVectorCosineNormalization verifies that multi-vector fields are
+// normalized correctly with cosine similarity. Each sub-vector in a multi-vector
+// should be independently normalized, producing correct similarity scores.
+func TestMultiVectorCosineNormalization(t *testing.T) {
+	tmpIndexPath := createTmpIndexPath(t)
+	defer cleanupTmpIndexPath(t, tmpIndexPath)
+
+	const dims = 3
+
+	// Create index with cosine similarity
+	indexMapping := NewIndexMapping()
+	vecFieldMapping := mapping.NewVectorFieldMapping()
+	vecFieldMapping.Dims = dims
+	vecFieldMapping.Similarity = index.CosineSimilarity
+	indexMapping.DefaultMapping.AddFieldMappingsAt("vec", vecFieldMapping)
+
+	// Multi-vector field
+	vecFieldMappingNested := mapping.NewVectorFieldMapping()
+	vecFieldMappingNested.Dims = dims
+	vecFieldMappingNested.Similarity = index.CosineSimilarity
+	indexMapping.DefaultMapping.AddFieldMappingsAt("vec_nested", vecFieldMappingNested)
+
+	idx, err := New(tmpIndexPath, indexMapping)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := idx.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	docsString := []string{
+		`{"vec": [3, 0, 0]}`,
+		`{"vec": [0, 4, 0]}`,
+		`{"vec_nested": [[3, 0, 0], [0, 4, 0]]}`,
+	}
+
+	for i, docStr := range docsString {
+		var doc map[string]interface{}
+		err = json.Unmarshal([]byte(docStr), &doc)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = idx.Index(fmt.Sprintf("doc%d", i+1), doc)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	// Query for X direction [1,0,0]
+	searchReq := NewSearchRequest(query.NewMatchNoneQuery())
+	searchReq.AddKNN("vec", []float32{1, 0, 0}, 3, 1.0)
+	res, err := idx.Search(searchReq)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res.Hits) != 2 {
+		t.Fatalf("expected 2 hits, got %d", len(res.Hits))
+	}
+	// Hit 1 should be doc1 with score 1.0 (perfect match)
+	if res.Hits[0].ID != "doc1" {
+		t.Fatalf("expected doc1 as first hit, got %s", res.Hits[0].ID)
+	}
+	if math.Abs(float64(res.Hits[0].Score-1.0)) > 1e-6 {
+		t.Fatalf("expected score 1.0, got %f", res.Hits[0].Score)
+	}
+	// Hit 2 should be doc2 with a score of 0.0 (orthogonal)
+	if res.Hits[1].ID != "doc2" {
+		t.Fatalf("expected doc2 as second hit, got %s", res.Hits[1].ID)
+	}
+	if math.Abs(float64(res.Hits[1].Score-0.0)) > 1e-6 {
+		t.Fatalf("expected score 0.0, got %f", res.Hits[1].Score)
+	}
+
+	// Query for Y direction [0,1,0]
+	searchReq = NewSearchRequest(query.NewMatchNoneQuery())
+	searchReq.AddKNN("vec", []float32{0, 1, 0}, 3, 1.0)
+	res, err = idx.Search(searchReq)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res.Hits) != 2 {
+		t.Fatalf("expected 2 hits, got %d", len(res.Hits))
+	}
+	// Hit 1 should be doc2 with score 1.0 (perfect match)
+	if res.Hits[0].ID != "doc2" {
+		t.Fatalf("expected doc2 as first hit, got %s", res.Hits[0].ID)
+	}
+	if math.Abs(float64(res.Hits[0].Score-1.0)) > 1e-6 {
+		t.Fatalf("expected score 1.0, got %f", res.Hits[0].Score)
+	}
+	// Hit 2 should be doc1 with a score of 0.0 (orthogonal)
+	if res.Hits[1].ID != "doc1" {
+		t.Fatalf("expected doc1 as second hit, got %s", res.Hits[1].ID)
+	}
+	if math.Abs(float64(res.Hits[1].Score-0.0)) > 1e-6 {
+		t.Fatalf("expected score 0.0, got %f", res.Hits[1].Score)
+	}
+
+	// Now test querying the nested multi-vector field
+	searchReq = NewSearchRequest(query.NewMatchNoneQuery())
+	searchReq.AddKNN("vec_nested", []float32{1, 0, 0}, 3, 1.0)
+	res, err = idx.Search(searchReq)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res.Hits) != 1 {
+		t.Fatalf("expected 1 hit, got %d", len(res.Hits))
+	}
+	// Hit should be doc3 with score 1.0 (perfect match on first sub-vector)
+	if res.Hits[0].ID != "doc3" {
+		t.Fatalf("expected doc3 as first hit, got %s", res.Hits[0].ID)
+	}
+	if math.Abs(float64(res.Hits[0].Score-1.0)) > 1e-6 {
+		t.Fatalf("expected score 1.0, got %f", res.Hits[0].Score)
+	}
+	// Query for Y direction [0,1,0] on nested field
+	searchReq = NewSearchRequest(query.NewMatchNoneQuery())
+	searchReq.AddKNN("vec_nested", []float32{0, 1, 0}, 3, 1.0)
+	res, err = idx.Search(searchReq)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res.Hits) != 1 {
+		t.Fatalf("expected 1 hit, got %d", len(res.Hits))
+	}
+	// Hit should be doc3 with score 1.0 (perfect match on second sub-vector)
+	if res.Hits[0].ID != "doc3" {
+		t.Fatalf("expected doc3 as first hit, got %s", res.Hits[0].ID)
+	}
+	if math.Abs(float64(res.Hits[0].Score-1.0)) > 1e-6 {
+		t.Fatalf("expected score 1.0, got %f", res.Hits[0].Score)
+	}
+}
+
 func TestNumVecsStat(t *testing.T) {
 
 	dataset, _, err := readDatasetAndQueries(testInputCompressedFile)