diff --git a/go.mod b/go.mod index c54f38e4a..a81abb787 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 799131ab5..32e149993 100644 --- a/go.sum +++ b/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= diff --git a/index/scorch/scorch.go b/index/scorch/scorch.go index a18cafc3f..e6b960206 100644 --- a/index/scorch/scorch.go +++ b/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 } diff --git a/index/scorch/scorch_test.go b/index/scorch/scorch_test.go index cb7a877cb..6ad8e2513 100644 --- a/index/scorch/scorch_test.go +++ b/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) + } +}