From a8d6a5260c079ba3dedfa91fc2f5dde153fafc72 Mon Sep 17 00:00:00 2001 From: Gautham Krithiwas Date: Mon, 20 Apr 2026 21:37:32 +0530 Subject: [PATCH 1/4] MB-71383: Expose GPU field stats in scorch --- go.mod | 4 ++-- go.sum | 8 ++++---- index/scorch/scorch.go | 14 ++++++++++++++ 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/go.mod b/go.mod index 204679f44..71bbd5a21 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,7 @@ require ( 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-0.20260420151438-c4d745b7f7a1 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.11 + github.com/blevesearch/zapx/v17 v17.0.12-0.20260420152705-bc365ef0a9be 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 00c6134ea..450a29ffe 100644 --- a/go.sum +++ b/go.sum @@ -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-0.20260420151438-c4d745b7f7a1 h1:iTv1okrVaN5HZbZER7p6ZAlg00tHfYZKEtf8WiS948g= +github.com/blevesearch/scorch_segment_api/v2 v2.4.7-0.20260420151438-c4d745b7f7a1/go.mod h1:Ry0cjO/wbmjBU0Vxf/+TW6IGcXYCdWMwAicvxUWNNeU= 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.11 h1:3OKEe8NpD4n14GW+GY/op5nh8/x8dmYNwgLdRgvv5go= -github.com/blevesearch/zapx/v17 v17.0.11/go.mod h1:62wlIX0vYZoLoLLKmix4zQvyCevvUt7RLuvcV5D3/N0= +github.com/blevesearch/zapx/v17 v17.0.12-0.20260420152705-bc365ef0a9be h1:w1wrZMbCxe/VB7FwbCphRASQNtdEBed1I/M1DCCUXzY= +github.com/blevesearch/zapx/v17 v17.0.12-0.20260420152705-bc365ef0a9be/go.mod h1:eBLrg5egrPa84AstavdGkx83LLX5yoFTRZbVa3qSNFU= 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..7aacf2c93 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 } } + + aggGPUStats := newFieldStats() + for _, segmentSnapshot := range indexSnapshot.Segments() { + if gsr, ok := segmentSnapshot.Segment().(segment.GPUFieldStatsReporter); ok { + segStats := newFieldStats() + gsr.UpdateGPUFieldStats(segStats) + aggGPUStats.Aggregate(segStats) + } + } + for statName, stats := range aggGPUStats.Fetch() { + for fieldName, val := range stats { + m["field:"+fieldName+":"+statName] = val + } + } return m } From ec549685aedcb1e838bba8f1aa269556a4f66ecf Mon Sep 17 00:00:00 2001 From: Gautham Krithiwas Date: Mon, 20 Apr 2026 22:43:42 +0530 Subject: [PATCH 2/4] add unit tests --- index/scorch/scorch_test.go | 255 ++++++++++++++++++++++++++++++++++++ 1 file changed, 255 insertions(+) diff --git a/index/scorch/scorch_test.go b/index/scorch/scorch_test.go index cb7a877cb..1c01484e2 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,256 @@ func TestPersistenceWithoutExclude(t *testing.T) { t.Fatalf("expected internal value to be nil, got %s", val) } } + +// mockSegmentBase satisfies segment.Segment but does NOT implement +// GPUFieldStatsReporter. 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 +} + +// mockGPUSegment adds GPUFieldStatsReporter on top of the base. +// inGPU controls whether the segment reports GPU or CPU memory. +type mockGPUSegment struct { + mockSegmentBase + inGPU bool +} + +func (m *mockGPUSegment) UpdateGPUFieldStats(stats segment.FieldStats) { + for _, f := range m.fields { + if m.inGPU { + stats.Store("num_gpu_segments_in_gpu_memory", f, 1) + } else { + stats.Store("num_gpu_segments_in_cpu_memory", f, 1) + } + } +} + +// mockPlainSegment is a segment that does NOT implement GPUFieldStatsReporter. +// It is used to verify that non-GPU 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}, + } +} + +// TestGPUFieldStatsAggregation verifies that StatsMap correctly aggregates +// num_gpu_segments_in_gpu_memory and num_gpu_segments_in_cpu_memory across +// multiple segments. +// +// Setup: +// - seg1: field "vec" -> in GPU +// - seg2: field "vec" -> in GPU +// - seg3: field "vec" -> fell back to CPU +// - seg4: plain segment (no GPUFieldStatsReporter) -> must be ignored +// +// Expected: +// +// field:vec:num_gpu_segments_in_gpu_memory = 2 +// field:vec:num_gpu_segments_in_cpu_memory = 1 +func TestGPUFieldStatsAggregation(t *testing.T) { + cfg := CreateConfig("TestGPUFieldStatsAggregation") + 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 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}, inGPU: true} + seg2 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}, inGPU: true} + seg3 := &mockGPUSegment{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_gpu_segments_in_gpu_memory", 2) + checkUint64Stat(t, m, "field:vec:num_gpu_segments_in_cpu_memory", 1) + // plain segment must not contribute — value must still be 2, not 3 + checkUint64Stat(t, m, "field:vec:num_gpu_segments_in_gpu_memory", 2) +} + +// TestGPUFieldStatsMultipleFields verifies that stats are tracked independently +// per field when a segment exposes more than one GPU-backed vector field. +func TestGPUFieldStatsMultipleFields(t *testing.T) { + cfg := CreateConfig("TestGPUFieldStatsMultipleFields") + 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 fell back to CPU + // seg3: fieldB in GPU + seg1 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"fieldA", "fieldB"}}, inGPU: true} + seg2 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"fieldA"}}, inGPU: false} + seg3 := &mockGPUSegment{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 GPU (seg1), 1 CPU (seg2) + checkUint64Stat(t, m, "field:fieldA:num_gpu_segments_in_gpu_memory", 1) + checkUint64Stat(t, m, "field:fieldA:num_gpu_segments_in_cpu_memory", 1) + + // fieldB: 2 GPU (seg1 + seg3), 0 CPU + checkUint64Stat(t, m, "field:fieldB:num_gpu_segments_in_gpu_memory", 2) + if _, ok := m["field:fieldB:num_gpu_segments_in_cpu_memory"]; ok { + t.Errorf("expected no cpu_memory stat for fieldB, but got one") + } +} + +// TestGPUFieldStatsNoGPUSegments verifies that when no segment implements +// GPUFieldStatsReporter, the GPU stat keys are absent from StatsMap. +func TestGPUFieldStatsNoGPUSegments(t *testing.T) { + cfg := CreateConfig("TestGPUFieldStatsNoGPUSegments") + 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_gpu_segments_in_gpu_memory", + "field:vec:num_gpu_segments_in_cpu_memory", + } { + if _, ok := m[key]; ok { + t.Errorf("expected key %q to be absent for non-GPU 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) + } +} From dc2faefd39c488bafd98593990fbd6307c186ed8 Mon Sep 17 00:00:00 2001 From: Gautham Krithiwas Date: Tue, 21 Apr 2026 18:00:37 +0530 Subject: [PATCH 3/4] updated zapx and segment api --- go.mod | 4 +- go.sum | 8 ++-- index/scorch/scorch.go | 10 ++-- index/scorch/scorch_test.go | 95 ++++++++++++++++++------------------- 4 files changed, 57 insertions(+), 60 deletions(-) diff --git a/go.mod b/go.mod index 71bbd5a21..8a09e2a24 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,7 @@ require ( 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.7-0.20260420151438-c4d745b7f7a1 + github.com/blevesearch/scorch_segment_api/v2 v2.4.7-0.20260421120500-9a856e8cc071 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.20260420152705-bc365ef0a9be + github.com/blevesearch/zapx/v17 v17.0.12-0.20260421122854-c2c9e44a8b13 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 450a29ffe..adb566c66 100644 --- a/go.sum +++ b/go.sum @@ -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.7-0.20260420151438-c4d745b7f7a1 h1:iTv1okrVaN5HZbZER7p6ZAlg00tHfYZKEtf8WiS948g= -github.com/blevesearch/scorch_segment_api/v2 v2.4.7-0.20260420151438-c4d745b7f7a1/go.mod h1:Ry0cjO/wbmjBU0Vxf/+TW6IGcXYCdWMwAicvxUWNNeU= +github.com/blevesearch/scorch_segment_api/v2 v2.4.7-0.20260421120500-9a856e8cc071 h1:HzdMdWfzVkh4IAP3L8XgqGlJBWeghrzuvDfc8ZjrVd8= +github.com/blevesearch/scorch_segment_api/v2 v2.4.7-0.20260421120500-9a856e8cc071/go.mod h1:Ry0cjO/wbmjBU0Vxf/+TW6IGcXYCdWMwAicvxUWNNeU= 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.20260420152705-bc365ef0a9be h1:w1wrZMbCxe/VB7FwbCphRASQNtdEBed1I/M1DCCUXzY= -github.com/blevesearch/zapx/v17 v17.0.12-0.20260420152705-bc365ef0a9be/go.mod h1:eBLrg5egrPa84AstavdGkx83LLX5yoFTRZbVa3qSNFU= +github.com/blevesearch/zapx/v17 v17.0.12-0.20260421122854-c2c9e44a8b13 h1:CXyAplQG4HloJO8Ld9cPJ7IPHMRYDO149xOLBk7tzIw= +github.com/blevesearch/zapx/v17 v17.0.12-0.20260421122854-c2c9e44a8b13/go.mod h1:449u8H4UH3VjFm8DJh+QMnYGDB7FiR/IpGCzzKCgQ7U= 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 7aacf2c93..e6b960206 100644 --- a/index/scorch/scorch.go +++ b/index/scorch/scorch.go @@ -841,15 +841,15 @@ func (s *Scorch) StatsMap() map[string]interface{} { } } - aggGPUStats := newFieldStats() + aggVectorStats := newFieldStats() for _, segmentSnapshot := range indexSnapshot.Segments() { - if gsr, ok := segmentSnapshot.Segment().(segment.GPUFieldStatsReporter); ok { + if vsr, ok := segmentSnapshot.Segment().(segment.VectorFieldStatsReporter); ok { segStats := newFieldStats() - gsr.UpdateGPUFieldStats(segStats) - aggGPUStats.Aggregate(segStats) + vsr.UpdateVectorFieldStats(segStats) + aggVectorStats.Aggregate(segStats) } } - for statName, stats := range aggGPUStats.Fetch() { + for statName, stats := range aggVectorStats.Fetch() { for fieldName, val := range stats { m["field:"+fieldName+":"+statName] = val } diff --git a/index/scorch/scorch_test.go b/index/scorch/scorch_test.go index 1c01484e2..6ad8e2513 100644 --- a/index/scorch/scorch_test.go +++ b/index/scorch/scorch_test.go @@ -3298,7 +3298,7 @@ func TestPersistenceWithoutExclude(t *testing.T) { } // mockSegmentBase satisfies segment.Segment but does NOT implement -// GPUFieldStatsReporter. Both mock types embed this so the stubs are +// VectorFieldStatsReporter. Both mock types embed this so the stubs are // not duplicated, while keeping the interface sets distinct. type mockSegmentBase struct { fields []string @@ -3325,25 +3325,25 @@ func (m *mockSegmentBase) Ancestors(_ uint64, prealloc []index.AncestorID) []ind return prealloc } -// mockGPUSegment adds GPUFieldStatsReporter on top of the base. -// inGPU controls whether the segment reports GPU or CPU memory. -type mockGPUSegment struct { +// 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 *mockGPUSegment) UpdateGPUFieldStats(stats segment.FieldStats) { +func (m *mockVectorSegment) UpdateVectorFieldStats(stats segment.FieldStats) { for _, f := range m.fields { if m.inGPU { - stats.Store("num_gpu_segments_in_gpu_memory", f, 1) + stats.Store("num_vector_indexes_in_gpu", f, 1) } else { - stats.Store("num_gpu_segments_in_cpu_memory", f, 1) + stats.Store("num_vector_indexes_in_cpu", f, 1) } } } -// mockPlainSegment is a segment that does NOT implement GPUFieldStatsReporter. -// It is used to verify that non-GPU segments are silently skipped. +// 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 } @@ -3360,22 +3360,21 @@ func makeSegmentSnapshot(id uint64, seg segment.Segment) *SegmentSnapshot { } } -// TestGPUFieldStatsAggregation verifies that StatsMap correctly aggregates -// num_gpu_segments_in_gpu_memory and num_gpu_segments_in_cpu_memory across -// multiple segments. +// TestVectorFieldStatsAggregation verifies that StatsMap correctly aggregates +// num_vector_indexes_in_gpu and num_vector_indexes_in_cpu across multiple segments. // // Setup: -// - seg1: field "vec" -> in GPU -// - seg2: field "vec" -> in GPU -// - seg3: field "vec" -> fell back to CPU -// - seg4: plain segment (no GPUFieldStatsReporter) -> must be ignored +// - 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_gpu_segments_in_gpu_memory = 2 -// field:vec:num_gpu_segments_in_cpu_memory = 1 -func TestGPUFieldStatsAggregation(t *testing.T) { - cfg := CreateConfig("TestGPUFieldStatsAggregation") +// 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) } @@ -3400,9 +3399,9 @@ func TestGPUFieldStatsAggregation(t *testing.T) { } }() - seg1 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}, inGPU: true} - seg2 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}, inGPU: true} - seg3 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"vec"}}, inGPU: false} + 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() @@ -3419,16 +3418,14 @@ func TestGPUFieldStatsAggregation(t *testing.T) { t.Fatal("StatsMap returned nil") } - checkUint64Stat(t, m, "field:vec:num_gpu_segments_in_gpu_memory", 2) - checkUint64Stat(t, m, "field:vec:num_gpu_segments_in_cpu_memory", 1) - // plain segment must not contribute — value must still be 2, not 3 - checkUint64Stat(t, m, "field:vec:num_gpu_segments_in_gpu_memory", 2) + checkUint64Stat(t, m, "field:vec:num_vector_indexes_in_gpu", 2) + checkUint64Stat(t, m, "field:vec:num_vector_indexes_in_cpu", 1) } -// TestGPUFieldStatsMultipleFields verifies that stats are tracked independently -// per field when a segment exposes more than one GPU-backed vector field. -func TestGPUFieldStatsMultipleFields(t *testing.T) { - cfg := CreateConfig("TestGPUFieldStatsMultipleFields") +// 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) } @@ -3454,11 +3451,11 @@ func TestGPUFieldStatsMultipleFields(t *testing.T) { }() // seg1: fieldA in GPU, fieldB in GPU - // seg2: fieldA fell back to CPU + // seg2: fieldA in CPU // seg3: fieldB in GPU - seg1 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"fieldA", "fieldB"}}, inGPU: true} - seg2 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"fieldA"}}, inGPU: false} - seg3 := &mockGPUSegment{mockSegmentBase: mockSegmentBase{fields: []string{"fieldB"}}, inGPU: true} + 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, @@ -3473,21 +3470,21 @@ func TestGPUFieldStatsMultipleFields(t *testing.T) { t.Fatal("StatsMap returned nil") } - // fieldA: 1 GPU (seg1), 1 CPU (seg2) - checkUint64Stat(t, m, "field:fieldA:num_gpu_segments_in_gpu_memory", 1) - checkUint64Stat(t, m, "field:fieldA:num_gpu_segments_in_cpu_memory", 1) + // 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 GPU (seg1 + seg3), 0 CPU - checkUint64Stat(t, m, "field:fieldB:num_gpu_segments_in_gpu_memory", 2) - if _, ok := m["field:fieldB:num_gpu_segments_in_cpu_memory"]; ok { - t.Errorf("expected no cpu_memory stat for fieldB, but got one") + // 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") } } -// TestGPUFieldStatsNoGPUSegments verifies that when no segment implements -// GPUFieldStatsReporter, the GPU stat keys are absent from StatsMap. -func TestGPUFieldStatsNoGPUSegments(t *testing.T) { - cfg := CreateConfig("TestGPUFieldStatsNoGPUSegments") +// 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) } @@ -3524,11 +3521,11 @@ func TestGPUFieldStatsNoGPUSegments(t *testing.T) { } for _, key := range []string{ - "field:vec:num_gpu_segments_in_gpu_memory", - "field:vec:num_gpu_segments_in_cpu_memory", + "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-GPU segments, but it was present", key) + t.Errorf("expected key %q to be absent for non-vector segments, but it was present", key) } } } From 7574f04c29d2ab195f7006d47d979d37a96153d0 Mon Sep 17 00:00:00 2001 From: Abhinav Dangeti Date: Tue, 21 Apr 2026 09:13:48 -0600 Subject: [PATCH 4/4] Use latest of zapx@v17 --- go.mod | 6 +++--- go.sum | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/go.mod b/go.mod index 8a09e2a24..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.7-0.20260421120500-9a856e8cc071 + 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.20260421122854-c2c9e44a8b13 + 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 adb566c66..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.7-0.20260421120500-9a856e8cc071 h1:HzdMdWfzVkh4IAP3L8XgqGlJBWeghrzuvDfc8ZjrVd8= -github.com/blevesearch/scorch_segment_api/v2 v2.4.7-0.20260421120500-9a856e8cc071/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.20260421122854-c2c9e44a8b13 h1:CXyAplQG4HloJO8Ld9cPJ7IPHMRYDO149xOLBk7tzIw= -github.com/blevesearch/zapx/v17 v17.0.12-0.20260421122854-c2c9e44a8b13/go.mod h1:449u8H4UH3VjFm8DJh+QMnYGDB7FiR/IpGCzzKCgQ7U= +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=