MB-71383: Expose vector field stats in scorch

go.mod

@@ -5,14 +5,14 @@ go 1.25.0
 require (
 	github.com/RoaringBitmap/roaring/v2 v2.14.5
 	github.com/bits-and-blooms/bitset v1.24.2
-	github.com/blevesearch/bleve_index_api v1.3.10
+	github.com/blevesearch/bleve_index_api v1.3.11
 	github.com/blevesearch/geo v0.2.5
 	github.com/blevesearch/go-faiss v1.0.34
 	github.com/blevesearch/go-metrics v0.0.0-20201227073835-cf1acfcdf475
 	github.com/blevesearch/go-porterstemmer v1.0.3
 	github.com/blevesearch/goleveldb v1.0.1
 	github.com/blevesearch/gtreap v0.1.1
-	github.com/blevesearch/scorch_segment_api/v2 v2.4.6
+	github.com/blevesearch/scorch_segment_api/v2 v2.4.7
 	github.com/blevesearch/segment v0.9.1
 	github.com/blevesearch/snowball v0.6.1
 	github.com/blevesearch/snowballstem v0.9.0
@@ -25,7 +25,7 @@ require (
 	github.com/blevesearch/zapx/v14 v14.4.3
 	github.com/blevesearch/zapx/v15 v15.4.3
 	github.com/blevesearch/zapx/v16 v16.3.4
-	github.com/blevesearch/zapx/v17 v17.0.12-0.20260421144655-b9ffe761cd37
+	github.com/blevesearch/zapx/v17 v17.0.12-0.20260421145725-c120519962c1
 	github.com/couchbase/moss v0.2.0
 	github.com/spf13/cobra v1.10.2
 	go.etcd.io/bbolt v1.4.0

go.sum

@@ -2,8 +2,8 @@ github.com/RoaringBitmap/roaring/v2 v2.14.5 h1:ckd0o545JqDPeVJDgeFoaM21eBixUnlWf
 github.com/RoaringBitmap/roaring/v2 v2.14.5/go.mod h1:eq4wdNXxtJIS/oikeCzdX1rBzek7ANzbth041hrU8Q4=
 github.com/bits-and-blooms/bitset v1.24.2 h1:M7/NzVbsytmtfHbumG+K2bremQPMJuqv1JD3vOaFxp0=
 github.com/bits-and-blooms/bitset v1.24.2/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8=
-github.com/blevesearch/bleve_index_api v1.3.10 h1:a7G+IOMa2xuO6f8vtutbTsqjVLpLuCuH3uoTZHkGiYg=
-github.com/blevesearch/bleve_index_api v1.3.10/go.mod h1:xvd48t5XMeeioWQ5/jZvgLrV98flT2rdvEJ3l/ki4Ko=
+github.com/blevesearch/bleve_index_api v1.3.11 h1:x29vbV8OjWfLcrDVd7Lr1q+BkLNS0JWNEig0MCVnKH4=
+github.com/blevesearch/bleve_index_api v1.3.11/go.mod h1:xvd48t5XMeeioWQ5/jZvgLrV98flT2rdvEJ3l/ki4Ko=
 github.com/blevesearch/geo v0.2.5 h1:yJg9FX1oRwLnjXSXF+ECHfXFTF4diF02Ca/qUGVjJhE=
 github.com/blevesearch/geo v0.2.5/go.mod h1:Jhq7WE2K6mJTx1xS44M2pUO6Io+wjCSHh1+co3YOgH4=
 github.com/blevesearch/go-faiss v1.0.34 h1:cFE1jRkjJfk7qMMsqXBqGEivbYQz/tjSf5yyoH50xbY=
@@ -19,8 +19,8 @@ github.com/blevesearch/gtreap v0.1.1/go.mod h1:QaQyDRAT51sotthUWAH4Sj08awFSSWzgY
 github.com/blevesearch/mmap-go v1.0.2/go.mod h1:ol2qBqYaOUsGdm7aRMRrYGgPvnwLe6Y+7LMvAB5IbSA=
 github.com/blevesearch/mmap-go v1.2.0 h1:l33nNKPFcBjJUMwem6sAYJPUzhUCABoK9FxZDGiFNBI=
 github.com/blevesearch/mmap-go v1.2.0/go.mod h1:Vd6+20GBhEdwJnU1Xohgt88XCD/CTWcqbCNxkZpyBo0=
-github.com/blevesearch/scorch_segment_api/v2 v2.4.6 h1:6D0ZarXRbBIIartND0QMhpzH6YR0eDWbRs7k+nS+zd8=
-github.com/blevesearch/scorch_segment_api/v2 v2.4.6/go.mod h1:Ry0cjO/wbmjBU0Vxf/+TW6IGcXYCdWMwAicvxUWNNeU=
+github.com/blevesearch/scorch_segment_api/v2 v2.4.7 h1:GlMzW08hcsM3DnLUxhyF/1PcDal1qtvvIuytuph5djw=
+github.com/blevesearch/scorch_segment_api/v2 v2.4.7/go.mod h1://IJ7tG3QCf0cWW/aVSXqy77tc1AvLu3fcJLYEvOAFs=
 github.com/blevesearch/segment v0.9.1 h1:+dThDy+Lvgj5JMxhmOVlgFfkUtZV2kw49xax4+jTfSU=
 github.com/blevesearch/segment v0.9.1/go.mod h1:zN21iLm7+GnBHWTao9I+Au/7MBiL8pPFtJBJTsk6kQw=
 github.com/blevesearch/snowball v0.6.1 h1:cDYjn/NCH+wwt2UdehaLpr2e4BwLIjN4V/TdLsL+B5A=
@@ -45,8 +45,8 @@ github.com/blevesearch/zapx/v15 v15.4.3 h1:iJiMJOHrz216jyO6lS0m9RTCEkprUnzvqAI2l
 github.com/blevesearch/zapx/v15 v15.4.3/go.mod h1:1pssev/59FsuWcgSnTa0OeEpOzmhtmr/0/11H0Z8+Nw=
 github.com/blevesearch/zapx/v16 v16.3.4 h1:hDAqA8qusZTNbPEL7//w5P65UZ2de6yhSeUaTbp0Po0=
 github.com/blevesearch/zapx/v16 v16.3.4/go.mod h1:zqkPPqs9GS9FzVWzCO3Wf1X044yWAV17+4zb+FTiEHg=
-github.com/blevesearch/zapx/v17 v17.0.12-0.20260421144655-b9ffe761cd37 h1:LaROcKdceOI37F6OMc9pySbdXCAQVqOgcMQk47yosqA=
-github.com/blevesearch/zapx/v17 v17.0.12-0.20260421144655-b9ffe761cd37/go.mod h1:62wlIX0vYZoLoLLKmix4zQvyCevvUt7RLuvcV5D3/N0=
+github.com/blevesearch/zapx/v17 v17.0.12-0.20260421145725-c120519962c1 h1:1qM+d5vKedxmdL7rIldvQfgh68NevZXMNE7aQwkj5cU=
+github.com/blevesearch/zapx/v17 v17.0.12-0.20260421145725-c120519962c1/go.mod h1:be77zp3wB5sTGTWo/6KwCEHnPRyOZYkIeQEr3YIO55E=
 github.com/couchbase/ghistogram v0.1.0 h1:b95QcQTCzjTUocDXp/uMgSNQi8oj1tGwnJ4bODWZnps=
 github.com/couchbase/ghistogram v0.1.0/go.mod h1:s1Jhy76zqfEecpNWJfWUiKZookAFaiGOEoyzgHt9i7k=
 github.com/couchbase/moss v0.2.0 h1:VCYrMzFwEryyhRSeI+/b3tRBSeTpi/8gn5Kf6dxqn+o=

index/scorch/scorch.go

@@ -840,6 +840,20 @@ func (s *Scorch) StatsMap() map[string]interface{} {
 			m["field:"+fieldName+":"+statName] = val
 		}
 	}
+
+	aggVectorStats := newFieldStats()
+	for _, segmentSnapshot := range indexSnapshot.Segments() {
+		if vsr, ok := segmentSnapshot.Segment().(segment.VectorFieldStatsReporter); ok {
+			segStats := newFieldStats()
+			vsr.UpdateVectorFieldStats(segStats)
+			aggVectorStats.Aggregate(segStats)
+		}
+	}
+	for statName, stats := range aggVectorStats.Fetch() {
+		for fieldName, val := range stats {
+			m["field:"+fieldName+":"+statName] = val
+		}
+	}
 	return m
 }
 

index/scorch/scorch_test.go

@@ -31,6 +31,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/RoaringBitmap/roaring/v2"
 	"github.com/blevesearch/bleve/v2/analysis"
 	"github.com/blevesearch/bleve/v2/analysis/analyzer/keyword"
 	"github.com/blevesearch/bleve/v2/analysis/analyzer/standard"
@@ -39,6 +40,7 @@ import (
 	"github.com/blevesearch/bleve/v2/index/scorch/mergeplan"
 	"github.com/blevesearch/bleve/v2/mapping"
 	index "github.com/blevesearch/bleve_index_api"
+	segment "github.com/blevesearch/scorch_segment_api/v2"
 )
 
 func init() {
@@ -3294,3 +3296,253 @@ func TestPersistenceWithoutExclude(t *testing.T) {
 		t.Fatalf("expected internal value to be nil, got %s", val)
 	}
 }
+
+// mockSegmentBase satisfies segment.Segment but does NOT implement
+// VectorFieldStatsReporter. Both mock types embed this so the stubs are
+// not duplicated, while keeping the interface sets distinct.
+type mockSegmentBase struct {
+	fields []string
+}
+
+func (m *mockSegmentBase) Dictionary(_ string) (segment.TermDictionary, error) { return nil, nil }
+func (m *mockSegmentBase) VisitStoredFields(_ uint64, _ segment.StoredFieldValueVisitor) error {
+	return nil
+}
+func (m *mockSegmentBase) DocID(_ uint64) ([]byte, error) { return nil, nil }
+func (m *mockSegmentBase) Count() uint64                  { return 0 }
+func (m *mockSegmentBase) DocNumbers(_ []string) (*roaring.Bitmap, error) {
+	return roaring.New(), nil
+}
+func (m *mockSegmentBase) Fields() []string        { return m.fields }
+func (m *mockSegmentBase) Close() error            { return nil }
+func (m *mockSegmentBase) Size() int               { return 0 }
+func (m *mockSegmentBase) AddRef()                 {}
+func (m *mockSegmentBase) DecRef() error           { return nil }
+func (m *mockSegmentBase) BytesRead() uint64       { return 0 }
+func (m *mockSegmentBase) BytesWritten() uint64    { return 0 }
+func (m *mockSegmentBase) ResetBytesRead(_ uint64) {}
+func (m *mockSegmentBase) Ancestors(_ uint64, prealloc []index.AncestorID) []index.AncestorID {
+	return prealloc
+}
+
+// mockVectorSegment adds VectorFieldStatsReporter on top of the base.
+// inGPU controls whether the index is reported as residing in GPU or CPU memory.
+type mockVectorSegment struct {
+	mockSegmentBase
+	inGPU bool
+}
+
+func (m *mockVectorSegment) UpdateVectorFieldStats(stats segment.FieldStats) {
+	for _, f := range m.fields {
+		if m.inGPU {
+			stats.Store("num_vector_indexes_in_gpu", f, 1)
+		} else {
+			stats.Store("num_vector_indexes_in_cpu", f, 1)
+		}
+	}
+}
+
+// mockPlainSegment is a segment that does NOT implement VectorFieldStatsReporter.
+// It is used to verify that non-vector segments are silently skipped.
+type mockPlainSegment struct {
+	mockSegmentBase
+}
+
+// makeSegmentSnapshot wraps a segment in a SegmentSnapshot without any static
+// field stats (stats == nil), matching the state of a live in-memory segment
+// before it has been persisted.
+func makeSegmentSnapshot(id uint64, seg segment.Segment) *SegmentSnapshot {
+	return &SegmentSnapshot{
+		id:         id,
+		segment:    seg,
+		cachedDocs: &cachedDocs{cache: nil},
+		cachedMeta: &cachedMeta{meta: nil},
+	}
+}
+
+// TestVectorFieldStatsAggregation verifies that StatsMap correctly aggregates
+// num_vector_indexes_in_gpu and num_vector_indexes_in_cpu across multiple segments.
+//
+// Setup:
+//   - seg1: field "vec" -> index in GPU memory
+//   - seg2: field "vec" -> index in GPU memory
+//   - seg3: field "vec" -> index in CPU memory
+//   - seg4: plain segment (no VectorFieldStatsReporter) -> must be ignored
+//
+// Expected:
+//
+//	field:vec:num_vector_indexes_in_gpu = 2
+//	field:vec:num_vector_indexes_in_cpu = 1
+func TestVectorFieldStatsAggregation(t *testing.T) {
+	cfg := CreateConfig("TestVectorFieldStatsAggregation")
+	if err := InitTest(cfg); err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		if err := DestroyTest(cfg); err != nil {
+			t.Log(err)
+		}
+	}()
+
+	analysisQueue := index.NewAnalysisQueue(1)
+	idx, err := NewScorch(Name, cfg, analysisQueue)
+	if err != nil {
+		t.Fatal(err)
+	}
+	s := idx.(*Scorch)
+	if err = s.Open(); err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		if err := s.Close(); err != nil {
+			t.Log(err)
+		}
+	}()
+
+	seg1 := &mockVectorSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}, inGPU: true}
+	seg2 := &mockVectorSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}, inGPU: true}
+	seg3 := &mockVectorSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}, inGPU: false}
+	seg4 := &mockPlainSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}}
+
+	s.rootLock.Lock()
+	s.root.segment = append(s.root.segment,
+		makeSegmentSnapshot(100, seg1),
+		makeSegmentSnapshot(101, seg2),
+		makeSegmentSnapshot(102, seg3),
+		makeSegmentSnapshot(103, seg4),
+	)
+	s.rootLock.Unlock()
+
+	m := s.StatsMap()
+	if m == nil {
+		t.Fatal("StatsMap returned nil")
+	}
+
+	checkUint64Stat(t, m, "field:vec:num_vector_indexes_in_gpu", 2)
+	checkUint64Stat(t, m, "field:vec:num_vector_indexes_in_cpu", 1)
+}
+
+// TestVectorFieldStatsMultipleFields verifies that stats are tracked independently
+// per field when a segment exposes more than one vector field.
+func TestVectorFieldStatsMultipleFields(t *testing.T) {
+	cfg := CreateConfig("TestVectorFieldStatsMultipleFields")
+	if err := InitTest(cfg); err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		if err := DestroyTest(cfg); err != nil {
+			t.Log(err)
+		}
+	}()
+
+	analysisQueue := index.NewAnalysisQueue(1)
+	idx, err := NewScorch(Name, cfg, analysisQueue)
+	if err != nil {
+		t.Fatal(err)
+	}
+	s := idx.(*Scorch)
+	if err = s.Open(); err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		if err := s.Close(); err != nil {
+			t.Log(err)
+		}
+	}()
+
+	// seg1: fieldA in GPU, fieldB in GPU
+	// seg2: fieldA in CPU
+	// seg3: fieldB in GPU
+	seg1 := &mockVectorSegment{mockSegmentBase: mockSegmentBase{fields: []string{"fieldA", "fieldB"}}, inGPU: true}
+	seg2 := &mockVectorSegment{mockSegmentBase: mockSegmentBase{fields: []string{"fieldA"}}, inGPU: false}
+	seg3 := &mockVectorSegment{mockSegmentBase: mockSegmentBase{fields: []string{"fieldB"}}, inGPU: true}
+
+	s.rootLock.Lock()
+	s.root.segment = append(s.root.segment,
+		makeSegmentSnapshot(200, seg1),
+		makeSegmentSnapshot(201, seg2),
+		makeSegmentSnapshot(202, seg3),
+	)
+	s.rootLock.Unlock()
+
+	m := s.StatsMap()
+	if m == nil {
+		t.Fatal("StatsMap returned nil")
+	}
+
+	// fieldA: 1 in GPU (seg1), 1 in CPU (seg2)
+	checkUint64Stat(t, m, "field:fieldA:num_vector_indexes_in_gpu", 1)
+	checkUint64Stat(t, m, "field:fieldA:num_vector_indexes_in_cpu", 1)
+
+	// fieldB: 2 in GPU (seg1 + seg3), 0 in CPU
+	checkUint64Stat(t, m, "field:fieldB:num_vector_indexes_in_gpu", 2)
+	if _, ok := m["field:fieldB:num_vector_indexes_in_cpu"]; ok {
+		t.Errorf("expected no num_vector_indexes_in_cpu stat for fieldB, but got one")
+	}
+}
+
+// TestVectorFieldStatsNoVectorSegments verifies that when no segment implements
+// VectorFieldStatsReporter, the vector stat keys are absent from StatsMap.
+func TestVectorFieldStatsNoVectorSegments(t *testing.T) {
+	cfg := CreateConfig("TestVectorFieldStatsNoVectorSegments")
+	if err := InitTest(cfg); err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		if err := DestroyTest(cfg); err != nil {
+			t.Log(err)
+		}
+	}()
+
+	analysisQueue := index.NewAnalysisQueue(1)
+	idx, err := NewScorch(Name, cfg, analysisQueue)
+	if err != nil {
+		t.Fatal(err)
+	}
+	s := idx.(*Scorch)
+	if err = s.Open(); err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		if err := s.Close(); err != nil {
+			t.Log(err)
+		}
+	}()
+
+	s.rootLock.Lock()
+	s.root.segment = append(s.root.segment,
+		makeSegmentSnapshot(300, &mockPlainSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}}),
+	)
+	s.rootLock.Unlock()
+
+	m := s.StatsMap()
+	if m == nil {
+		t.Fatal("StatsMap returned nil")
+	}
+
+	for _, key := range []string{
+		"field:vec:num_vector_indexes_in_gpu",
+		"field:vec:num_vector_indexes_in_cpu",
+	} {
+		if _, ok := m[key]; ok {
+			t.Errorf("expected key %q to be absent for non-vector segments, but it was present", key)
+		}
+	}
+}
+
+func checkUint64Stat(t *testing.T, m map[string]interface{}, key string, want uint64) {
+	t.Helper()
+	v, ok := m[key]
+	if !ok {
+		t.Errorf("expected stat %q to be present in StatsMap, but it was missing", key)
+		return
+	}
+	got, ok := v.(uint64)
+	if !ok {
+		t.Errorf("stat %q: expected uint64, got %T (%v)", key, v, v)
+		return
+	}
+	if got != want {
+		t.Errorf("stat %q: got %d, want %d", key, got, want)
+	}
+}