admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-28 22:45:26 +08:00
Code Issues Packages Projects Releases Wiki Activity
milvus/internal/querynodev2/segments/mock_loader.go
congqixia 21ed1fabfd
feat: support reopen segment for data/schema changes (#46359) 
issue: #46358

This PR implements segment reopening functionality on query nodes,
enabling the application of data or schema changes to already-loaded
segments without requiring a full reload.

### Core (C++)

**New SegmentLoadInfo class**
(`internal/core/src/segcore/SegmentLoadInfo.h/cpp`):
- Encapsulates segment load configuration with structured access
- Implements `ComputeDiff()` to calculate differences between old and
new load states
- Tracks indexes, binlogs, and column groups that need to be loaded or
dropped
- Provides `ConvertFieldIndexInfoToLoadIndexInfo()` for index loading

**ChunkedSegmentSealedImpl modifications**:
- Added `Reopen(const SegmentLoadInfo&)` method to apply incremental
changes based on computed diff
- Refactored `LoadColumnGroups()` and `LoadColumnGroup()` to support
selective loading via field ID map
- Extracted `LoadBatchIndexes()` and `LoadBatchFieldData()` for reusable
batch loading logic
- Added `LoadManifest()` for manifest-based loading path
- Updated all methods to use `SegmentLoadInfo` wrapper instead of direct
proto access

**SegmentGrowingImpl modifications**:
- Added `Reopen()` stub method for interface compliance

**C API additions** (`segment_c.h/cpp`):
- Added `ReopenSegment()` function exposing reopen to Go layer

### Go Side

**QueryNode handlers** (`internal/querynodev2/`):
- Added `HandleReopen()` in handlers.go
- Added `ReopenSegments()` RPC in services.go

**Segment interface** (`internal/querynodev2/segments/`):
- Extended `Segment` interface with `Reopen()` method
- Implemented `Reopen()` in LocalSegment
- Added `Reopen()` to segment loader

**Segcore wrapper** (`internal/util/segcore/`):
- Added `Reopen()` method in segment.go
- Added `ReopenSegmentRequest` in requests.go

### Proto

- Added new fields to support reopen in `query_coord.proto`

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2025-12-17 15:49:16 +08:00

513 lines
18 KiB
Go
Raw Blame History

// Code generated by mockery v2.53.3. DO NOT EDIT.
package segments
import (
context "context"
commonpb "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
datapb "github.com/milvus-io/milvus/pkg/v2/proto/datapb"
mock "github.com/stretchr/testify/mock"
pkoracle "github.com/milvus-io/milvus/internal/querynodev2/pkoracle"
querypb "github.com/milvus-io/milvus/pkg/v2/proto/querypb"
storage "github.com/milvus-io/milvus/internal/storage"
typeutil "github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
// MockLoader is an autogenerated mock type for the Loader type
type MockLoader struct {
mock.Mock
}
type MockLoader_Expecter struct {
mock *mock.Mock
}
func (_m *MockLoader) EXPECT() *MockLoader_Expecter {
return &MockLoader_Expecter{mock: &_m.Mock}
}
// Load provides a mock function with given fields: ctx, collectionID, segmentType, version, _a4
func (_m *MockLoader) Load(ctx context.Context, collectionID int64, segmentType commonpb.SegmentState, version int64, _a4 ...*querypb.SegmentLoadInfo) ([]Segment, error) {
_va := make([]interface{}, len(_a4))
for _i := range _a4 {
_va[_i] = _a4[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx, collectionID, segmentType, version)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
if len(ret) == 0 {
panic("no return value specified for Load")
}
var r0 []Segment
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, int64, commonpb.SegmentState, int64, ...*querypb.SegmentLoadInfo) ([]Segment, error)); ok {
return rf(ctx, collectionID, segmentType, version, _a4...)
}
if rf, ok := ret.Get(0).(func(context.Context, int64, commonpb.SegmentState, int64, ...*querypb.SegmentLoadInfo) []Segment); ok {
r0 = rf(ctx, collectionID, segmentType, version, _a4...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]Segment)
}
}
if rf, ok := ret.Get(1).(func(context.Context, int64, commonpb.SegmentState, int64, ...*querypb.SegmentLoadInfo) error); ok {
r1 = rf(ctx, collectionID, segmentType, version, _a4...)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockLoader_Load_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Load'
type MockLoader_Load_Call struct {
*mock.Call
}
// Load is a helper method to define mock.On call
// - ctx context.Context
// - collectionID int64
// - segmentType commonpb.SegmentState
// - version int64
// - _a4 ...*querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) Load(ctx interface{}, collectionID interface{}, segmentType interface{}, version interface{}, _a4 ...interface{}) *MockLoader_Load_Call {
return &MockLoader_Load_Call{Call: _e.mock.On("Load",
append([]interface{}{ctx, collectionID, segmentType, version}, _a4...)...)}
}
func (_c *MockLoader_Load_Call) Run(run func(ctx context.Context, collectionID int64, segmentType commonpb.SegmentState, version int64, _a4 ...*querypb.SegmentLoadInfo)) *MockLoader_Load_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]*querypb.SegmentLoadInfo, len(args)-4)
for i, a := range args[4:] {
if a != nil {
variadicArgs[i] = a.(*querypb.SegmentLoadInfo)
}
}
run(args[0].(context.Context), args[1].(int64), args[2].(commonpb.SegmentState), args[3].(int64), variadicArgs...)
})
return _c
}
func (_c *MockLoader_Load_Call) Return(_a0 []Segment, _a1 error) *MockLoader_Load_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockLoader_Load_Call) RunAndReturn(run func(context.Context, int64, commonpb.SegmentState, int64, ...*querypb.SegmentLoadInfo) ([]Segment, error)) *MockLoader_Load_Call {
_c.Call.Return(run)
return _c
}
// LoadBM25Stats provides a mock function with given fields: ctx, collectionID, infos
func (_m *MockLoader) LoadBM25Stats(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) (*typeutil.ConcurrentMap[int64, map[int64]*storage.BM25Stats], error) {
_va := make([]interface{}, len(infos))
for _i := range infos {
_va[_i] = infos[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx, collectionID)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
if len(ret) == 0 {
panic("no return value specified for LoadBM25Stats")
}
var r0 *typeutil.ConcurrentMap[int64, map[int64]*storage.BM25Stats]
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, int64, ...*querypb.SegmentLoadInfo) (*typeutil.ConcurrentMap[int64, map[int64]*storage.BM25Stats], error)); ok {
return rf(ctx, collectionID, infos...)
}
if rf, ok := ret.Get(0).(func(context.Context, int64, ...*querypb.SegmentLoadInfo) *typeutil.ConcurrentMap[int64, map[int64]*storage.BM25Stats]); ok {
r0 = rf(ctx, collectionID, infos...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*typeutil.ConcurrentMap[int64, map[int64]*storage.BM25Stats])
}
}
if rf, ok := ret.Get(1).(func(context.Context, int64, ...*querypb.SegmentLoadInfo) error); ok {
r1 = rf(ctx, collectionID, infos...)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockLoader_LoadBM25Stats_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadBM25Stats'
type MockLoader_LoadBM25Stats_Call struct {
*mock.Call
}
// LoadBM25Stats is a helper method to define mock.On call
// - ctx context.Context
// - collectionID int64
// - infos ...*querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) LoadBM25Stats(ctx interface{}, collectionID interface{}, infos ...interface{}) *MockLoader_LoadBM25Stats_Call {
return &MockLoader_LoadBM25Stats_Call{Call: _e.mock.On("LoadBM25Stats",
append([]interface{}{ctx, collectionID}, infos...)...)}
}
func (_c *MockLoader_LoadBM25Stats_Call) Run(run func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo)) *MockLoader_LoadBM25Stats_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]*querypb.SegmentLoadInfo, len(args)-2)
for i, a := range args[2:] {
if a != nil {
variadicArgs[i] = a.(*querypb.SegmentLoadInfo)
}
}
run(args[0].(context.Context), args[1].(int64), variadicArgs...)
})
return _c
}
func (_c *MockLoader_LoadBM25Stats_Call) Return(_a0 *typeutil.ConcurrentMap[int64, map[int64]*storage.BM25Stats], _a1 error) *MockLoader_LoadBM25Stats_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockLoader_LoadBM25Stats_Call) RunAndReturn(run func(context.Context, int64, ...*querypb.SegmentLoadInfo) (*typeutil.ConcurrentMap[int64, map[int64]*storage.BM25Stats], error)) *MockLoader_LoadBM25Stats_Call {
_c.Call.Return(run)
return _c
}
// LoadBloomFilterSet provides a mock function with given fields: ctx, collectionID, infos
func (_m *MockLoader) LoadBloomFilterSet(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) ([]*pkoracle.BloomFilterSet, error) {
_va := make([]interface{}, len(infos))
for _i := range infos {
_va[_i] = infos[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx, collectionID)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
if len(ret) == 0 {
panic("no return value specified for LoadBloomFilterSet")
}
var r0 []*pkoracle.BloomFilterSet
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, int64, ...*querypb.SegmentLoadInfo) ([]*pkoracle.BloomFilterSet, error)); ok {
return rf(ctx, collectionID, infos...)
}
if rf, ok := ret.Get(0).(func(context.Context, int64, ...*querypb.SegmentLoadInfo) []*pkoracle.BloomFilterSet); ok {
r0 = rf(ctx, collectionID, infos...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*pkoracle.BloomFilterSet)
}
}
if rf, ok := ret.Get(1).(func(context.Context, int64, ...*querypb.SegmentLoadInfo) error); ok {
r1 = rf(ctx, collectionID, infos...)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockLoader_LoadBloomFilterSet_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadBloomFilterSet'
type MockLoader_LoadBloomFilterSet_Call struct {
*mock.Call
}
// LoadBloomFilterSet is a helper method to define mock.On call
// - ctx context.Context
// - collectionID int64
// - infos ...*querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) LoadBloomFilterSet(ctx interface{}, collectionID interface{}, infos ...interface{}) *MockLoader_LoadBloomFilterSet_Call {
return &MockLoader_LoadBloomFilterSet_Call{Call: _e.mock.On("LoadBloomFilterSet",
append([]interface{}{ctx, collectionID}, infos...)...)}
}
func (_c *MockLoader_LoadBloomFilterSet_Call) Run(run func(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo)) *MockLoader_LoadBloomFilterSet_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]*querypb.SegmentLoadInfo, len(args)-2)
for i, a := range args[2:] {
if a != nil {
variadicArgs[i] = a.(*querypb.SegmentLoadInfo)
}
}
run(args[0].(context.Context), args[1].(int64), variadicArgs...)
})
return _c
}
func (_c *MockLoader_LoadBloomFilterSet_Call) Return(_a0 []*pkoracle.BloomFilterSet, _a1 error) *MockLoader_LoadBloomFilterSet_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockLoader_LoadBloomFilterSet_Call) RunAndReturn(run func(context.Context, int64, ...*querypb.SegmentLoadInfo) ([]*pkoracle.BloomFilterSet, error)) *MockLoader_LoadBloomFilterSet_Call {
_c.Call.Return(run)
return _c
}
// LoadDeltaLogs provides a mock function with given fields: ctx, segment, deltaLogs
func (_m *MockLoader) LoadDeltaLogs(ctx context.Context, segment Segment, deltaLogs []*datapb.FieldBinlog) error {
ret := _m.Called(ctx, segment, deltaLogs)
if len(ret) == 0 {
panic("no return value specified for LoadDeltaLogs")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, Segment, []*datapb.FieldBinlog) error); ok {
r0 = rf(ctx, segment, deltaLogs)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockLoader_LoadDeltaLogs_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadDeltaLogs'
type MockLoader_LoadDeltaLogs_Call struct {
*mock.Call
}
// LoadDeltaLogs is a helper method to define mock.On call
// - ctx context.Context
// - segment Segment
// - deltaLogs []*datapb.FieldBinlog
func (_e *MockLoader_Expecter) LoadDeltaLogs(ctx interface{}, segment interface{}, deltaLogs interface{}) *MockLoader_LoadDeltaLogs_Call {
return &MockLoader_LoadDeltaLogs_Call{Call: _e.mock.On("LoadDeltaLogs", ctx, segment, deltaLogs)}
}
func (_c *MockLoader_LoadDeltaLogs_Call) Run(run func(ctx context.Context, segment Segment, deltaLogs []*datapb.FieldBinlog)) *MockLoader_LoadDeltaLogs_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(Segment), args[2].([]*datapb.FieldBinlog))
})
return _c
}
func (_c *MockLoader_LoadDeltaLogs_Call) Return(_a0 error) *MockLoader_LoadDeltaLogs_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockLoader_LoadDeltaLogs_Call) RunAndReturn(run func(context.Context, Segment, []*datapb.FieldBinlog) error) *MockLoader_LoadDeltaLogs_Call {
_c.Call.Return(run)
return _c
}
// LoadIndex provides a mock function with given fields: ctx, segment, info, version
func (_m *MockLoader) LoadIndex(ctx context.Context, segment Segment, info *querypb.SegmentLoadInfo, version int64) error {
ret := _m.Called(ctx, segment, info, version)
if len(ret) == 0 {
panic("no return value specified for LoadIndex")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, Segment, *querypb.SegmentLoadInfo, int64) error); ok {
r0 = rf(ctx, segment, info, version)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockLoader_LoadIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadIndex'
type MockLoader_LoadIndex_Call struct {
*mock.Call
}
// LoadIndex is a helper method to define mock.On call
// - ctx context.Context
// - segment Segment
// - info *querypb.SegmentLoadInfo
// - version int64
func (_e *MockLoader_Expecter) LoadIndex(ctx interface{}, segment interface{}, info interface{}, version interface{}) *MockLoader_LoadIndex_Call {
return &MockLoader_LoadIndex_Call{Call: _e.mock.On("LoadIndex", ctx, segment, info, version)}
}
func (_c *MockLoader_LoadIndex_Call) Run(run func(ctx context.Context, segment Segment, info *querypb.SegmentLoadInfo, version int64)) *MockLoader_LoadIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(Segment), args[2].(*querypb.SegmentLoadInfo), args[3].(int64))
})
return _c
}
func (_c *MockLoader_LoadIndex_Call) Return(_a0 error) *MockLoader_LoadIndex_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockLoader_LoadIndex_Call) RunAndReturn(run func(context.Context, Segment, *querypb.SegmentLoadInfo, int64) error) *MockLoader_LoadIndex_Call {
_c.Call.Return(run)
return _c
}
// LoadJSONIndex provides a mock function with given fields: ctx, segment, info
func (_m *MockLoader) LoadJSONIndex(ctx context.Context, segment Segment, info *querypb.SegmentLoadInfo) error {
ret := _m.Called(ctx, segment, info)
if len(ret) == 0 {
panic("no return value specified for LoadJSONIndex")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, Segment, *querypb.SegmentLoadInfo) error); ok {
r0 = rf(ctx, segment, info)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockLoader_LoadJSONIndex_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadJSONIndex'
type MockLoader_LoadJSONIndex_Call struct {
*mock.Call
}
// LoadJSONIndex is a helper method to define mock.On call
// - ctx context.Context
// - segment Segment
// - info *querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) LoadJSONIndex(ctx interface{}, segment interface{}, info interface{}) *MockLoader_LoadJSONIndex_Call {
return &MockLoader_LoadJSONIndex_Call{Call: _e.mock.On("LoadJSONIndex", ctx, segment, info)}
}
func (_c *MockLoader_LoadJSONIndex_Call) Run(run func(ctx context.Context, segment Segment, info *querypb.SegmentLoadInfo)) *MockLoader_LoadJSONIndex_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(Segment), args[2].(*querypb.SegmentLoadInfo))
})
return _c
}
func (_c *MockLoader_LoadJSONIndex_Call) Return(_a0 error) *MockLoader_LoadJSONIndex_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockLoader_LoadJSONIndex_Call) RunAndReturn(run func(context.Context, Segment, *querypb.SegmentLoadInfo) error) *MockLoader_LoadJSONIndex_Call {
_c.Call.Return(run)
return _c
}
// LoadLazySegment provides a mock function with given fields: ctx, segment, loadInfo
func (_m *MockLoader) LoadLazySegment(ctx context.Context, segment Segment, loadInfo *querypb.SegmentLoadInfo) error {
ret := _m.Called(ctx, segment, loadInfo)
if len(ret) == 0 {
panic("no return value specified for LoadLazySegment")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, Segment, *querypb.SegmentLoadInfo) error); ok {
r0 = rf(ctx, segment, loadInfo)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockLoader_LoadLazySegment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadLazySegment'
type MockLoader_LoadLazySegment_Call struct {
*mock.Call
}
// LoadLazySegment is a helper method to define mock.On call
// - ctx context.Context
// - segment Segment
// - loadInfo *querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) LoadLazySegment(ctx interface{}, segment interface{}, loadInfo interface{}) *MockLoader_LoadLazySegment_Call {
return &MockLoader_LoadLazySegment_Call{Call: _e.mock.On("LoadLazySegment", ctx, segment, loadInfo)}
}
func (_c *MockLoader_LoadLazySegment_Call) Run(run func(ctx context.Context, segment Segment, loadInfo *querypb.SegmentLoadInfo)) *MockLoader_LoadLazySegment_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(Segment), args[2].(*querypb.SegmentLoadInfo))
})
return _c
}
func (_c *MockLoader_LoadLazySegment_Call) Return(_a0 error) *MockLoader_LoadLazySegment_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockLoader_LoadLazySegment_Call) RunAndReturn(run func(context.Context, Segment, *querypb.SegmentLoadInfo) error) *MockLoader_LoadLazySegment_Call {
_c.Call.Return(run)
return _c
}
// ReopenSegments provides a mock function with given fields: ctx, loadInfos
func (_m *MockLoader) ReopenSegments(ctx context.Context, loadInfos []*querypb.SegmentLoadInfo) error {
ret := _m.Called(ctx, loadInfos)
if len(ret) == 0 {
panic("no return value specified for ReopenSegments")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, []*querypb.SegmentLoadInfo) error); ok {
r0 = rf(ctx, loadInfos)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockLoader_ReopenSegments_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ReopenSegments'
type MockLoader_ReopenSegments_Call struct {
*mock.Call
}
// ReopenSegments is a helper method to define mock.On call
// - ctx context.Context
// - loadInfos []*querypb.SegmentLoadInfo
func (_e *MockLoader_Expecter) ReopenSegments(ctx interface{}, loadInfos interface{}) *MockLoader_ReopenSegments_Call {
return &MockLoader_ReopenSegments_Call{Call: _e.mock.On("ReopenSegments", ctx, loadInfos)}
}
func (_c *MockLoader_ReopenSegments_Call) Run(run func(ctx context.Context, loadInfos []*querypb.SegmentLoadInfo)) *MockLoader_ReopenSegments_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].([]*querypb.SegmentLoadInfo))
})
return _c
}
func (_c *MockLoader_ReopenSegments_Call) Return(_a0 error) *MockLoader_ReopenSegments_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockLoader_ReopenSegments_Call) RunAndReturn(run func(context.Context, []*querypb.SegmentLoadInfo) error) *MockLoader_ReopenSegments_Call {
_c.Call.Return(run)
return _c
}
// NewMockLoader creates a new instance of MockLoader. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewMockLoader(t interface {
mock.TestingT
Cleanup(func())
}) *MockLoader {
mock := &MockLoader{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}
Reference in New Issue View Git Blame Copy Permalink