milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Assign metric type in load segment request (#24917) 

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
pull/24921/head
yihao.dai 2023-06-16 16:50:41 +08:00 committed by GitHub
parent c5b1533fdc
commit c7dc1c067a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 150 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
internal/querycoordv2/task/executor.go
View File

@ -254,9 +254,21 @@ func (ex *Executor) loadSegment(task *SegmentTask, step int) error {
return err
}
// TODO: improve this, queryCoord should keep index and schema in memory to save RPC.
indexInfo, err := ex.broker.DescribeIndex(ctx, task.CollectionID())
if err != nil {
log.Warn("fail to get index meta of collection")
return err
}
metricType, err := getMetricType(indexInfo, schema)
if err != nil {
log.Warn("failed to get metric type", zap.Error(err))
return err
}
loadMeta := packLoadMeta(
ex.meta.GetLoadType(task.CollectionID()),
"",
metricType,
task.CollectionID(),
partitions...,
)

79
internal/querycoordv2/task/task_test.go
View File

@ -358,10 +358,6 @@ func (suite *TaskSuite) TestUnsubscribeChannelTask() {
}
}
func (suite *TaskSuite) expectationsForLoadSegments() {
}
func (suite *TaskSuite) TestLoadSegmentTask() {
ctx := context.Background()
timeout := 10 * time.Second
@ -375,6 +371,9 @@ func (suite *TaskSuite) TestLoadSegmentTask() {
// Expect
suite.broker.EXPECT().GetCollectionSchema(mock.Anything, suite.collection).Return(&schemapb.CollectionSchema{
Name: "TestLoadSegmentTask",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
},
}, nil)
for _, segment := range suite.loadSegments {
suite.broker.EXPECT().GetSegmentInfo(mock.Anything, segment).Return(&datapb.GetSegmentInfoResponse{Infos: []*datapb.SegmentInfo{
@ -387,6 +386,18 @@ func (suite *TaskSuite) TestLoadSegmentTask() {
}, nil)
suite.broker.EXPECT().GetIndexInfo(mock.Anything, suite.collection, segment).Return(nil, nil)
}
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{
CollectionID: suite.collection,
FieldID: 100,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.MetricTypeKey,
Value: "L2",
},
},
},
}, nil)
suite.cluster.EXPECT().LoadSegments(mock.Anything, targetNode, mock.Anything).Return(utils.WrapStatus(commonpb.ErrorCode_Success, ""), nil)
// Test load segment task
@ -457,6 +468,9 @@ func (suite *TaskSuite) TestLoadSegmentTaskFailed() {
// Expect
suite.broker.EXPECT().GetCollectionSchema(mock.Anything, suite.collection).Return(&schemapb.CollectionSchema{
Name: "TestLoadSegmentTask",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
},
}, nil)
for _, segment := range suite.loadSegments {
suite.broker.EXPECT().GetSegmentInfo(mock.Anything, segment).Return(&datapb.GetSegmentInfoResponse{Infos: []*datapb.SegmentInfo{
@ -469,6 +483,18 @@ func (suite *TaskSuite) TestLoadSegmentTaskFailed() {
}, nil)
suite.broker.EXPECT().GetIndexInfo(mock.Anything, suite.collection, segment).Return(nil, errors.New("index not ready"))
}
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{
CollectionID: suite.collection,
FieldID: 100,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.MetricTypeKey,
Value: "L2",
},
},
},
}, nil)
// Test load segment task
suite.dist.ChannelDistManager.Update(targetNode, meta.DmChannelFromVChannel(&datapb.VchannelInfo{
@ -650,6 +676,9 @@ func (suite *TaskSuite) TestMoveSegmentTask() {
// Expect
suite.broker.EXPECT().GetCollectionSchema(mock.Anything, suite.collection).Return(&schemapb.CollectionSchema{
Name: "TestMoveSegmentTask",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
},
}, nil)
for _, segment := range suite.moveSegments {
suite.broker.EXPECT().GetSegmentInfo(mock.Anything, segment).Return(&datapb.GetSegmentInfoResponse{Infos: []*datapb.SegmentInfo{
@ -662,6 +691,18 @@ func (suite *TaskSuite) TestMoveSegmentTask() {
}, nil)
suite.broker.EXPECT().GetIndexInfo(mock.Anything, suite.collection, segment).Return(nil, nil)
}
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{
CollectionID: suite.collection,
FieldID: 100,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.MetricTypeKey,
Value: "L2",
},
},
},
}, nil)
suite.cluster.EXPECT().LoadSegments(mock.Anything, leader, mock.Anything).Return(utils.WrapStatus(commonpb.ErrorCode_Success, ""), nil)
suite.cluster.EXPECT().ReleaseSegments(mock.Anything, leader, mock.Anything).Return(utils.WrapStatus(commonpb.ErrorCode_Success, ""), nil)
@ -747,6 +788,9 @@ func (suite *TaskSuite) TestTaskCanceled() {
// Expect
suite.broker.EXPECT().GetCollectionSchema(mock.Anything, suite.collection).Return(&schemapb.CollectionSchema{
Name: "TestSubscribeChannelTask",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
},
}, nil)
for _, segment := range suite.loadSegments {
suite.broker.EXPECT().GetSegmentInfo(mock.Anything, segment).Return(&datapb.GetSegmentInfoResponse{Infos: []*datapb.SegmentInfo{
@ -759,6 +803,18 @@ func (suite *TaskSuite) TestTaskCanceled() {
}, nil)
suite.broker.EXPECT().GetIndexInfo(mock.Anything, suite.collection, segment).Return(nil, nil)
}
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{
CollectionID: suite.collection,
FieldID: 100,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.MetricTypeKey,
Value: "L2",
},
},
},
}, nil)
suite.cluster.EXPECT().LoadSegments(mock.Anything, targetNode, mock.Anything).Return(utils.WrapStatus(commonpb.ErrorCode_Success, ""), nil)
// Test load segment task
@ -823,6 +879,9 @@ func (suite *TaskSuite) TestSegmentTaskStale() {
// Expect
suite.broker.EXPECT().GetCollectionSchema(mock.Anything, suite.collection).Return(&schemapb.CollectionSchema{
Name: "TestSegmentTaskStale",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
},
}, nil)
for _, segment := range suite.loadSegments {
suite.broker.EXPECT().GetSegmentInfo(mock.Anything, segment).Return(&datapb.GetSegmentInfoResponse{Infos: []*datapb.SegmentInfo{
@ -835,6 +894,18 @@ func (suite *TaskSuite) TestSegmentTaskStale() {
}, nil)
suite.broker.EXPECT().GetIndexInfo(mock.Anything, suite.collection, segment).Return(nil, nil)
}
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{
CollectionID: suite.collection,
FieldID: 100,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.MetricTypeKey,
Value: "L2",
},
},
},
}, nil)
suite.cluster.EXPECT().LoadSegments(mock.Anything, targetNode, mock.Anything).Return(utils.WrapStatus(commonpb.ErrorCode_Success, ""), nil)
// Test load segment task

11
internal/querynodev2/services.go
View File

@ -433,6 +433,11 @@ func (node *QueryNode) LoadPartitions(ctx context.Context, req *querypb.LoadPart
if err != nil {
return merr.Status(err), nil
}
// check metric type
if metricType == "" {
err := fmt.Errorf("empty metric type, collection = %d", req.GetCollectionID())
return merr.Status(err), nil
}
node.manager.Collection.Put(req.GetCollectionID(), req.GetSchema(), &segcorepb.CollectionIndexMeta{
IndexMetas: fieldIndexMetas,
MaxIndexRowCount: maxIndexRecordPerSegment,
@ -480,6 +485,12 @@ func (node *QueryNode) LoadSegments(ctx context.Context, req *querypb.LoadSegmen
return node.loadDeltaLogs(ctx, req), nil
}
// check metric type
if req.GetLoadMeta().GetMetricType() == "" {
err := fmt.Errorf("empty metric type, collection = %d", req.GetCollectionID())
return merr.Status(err), nil
}
node.manager.Collection.Put(req.GetCollectionID(), req.GetSchema(), nil, req.GetLoadMeta())
// Delegates request to workers

53
internal/querynodev2/services_test.go
View File

@ -494,6 +494,9 @@ func (suite *ServiceSuite) TestLoadSegments_Int64() {
Schema: schema,
DeltaPositions: []*msgpb.MsgPosition{{Timestamp: 20000}},
NeedTransfer: true,
LoadMeta: &querypb.LoadMetaInfo{
MetricType: defaultMetricType,
},
}
// LoadSegment
@ -507,13 +510,14 @@ func (suite *ServiceSuite) TestLoadSegments_VarChar() {
suite.TestWatchDmChannelsVarchar()
// data
schema := segments.GenTestCollectionSchema(suite.collectionName, schemapb.DataType_VarChar)
LoadMeta := &querypb.LoadMetaInfo{
loadMeta := &querypb.LoadMetaInfo{
LoadType: querypb.LoadType_LoadCollection,
CollectionID: suite.collectionID,
PartitionIDs: suite.partitionIDs,
MetricType: defaultMetricType,
}
suite.node.manager.Collection = segments.NewCollectionManager()
suite.node.manager.Collection.Put(suite.collectionID, schema, nil, LoadMeta)
suite.node.manager.Collection.Put(suite.collectionID, schema, nil, loadMeta)
req := &querypb.LoadSegmentsRequest{
Base: &commonpb.MsgBase{
MsgID: rand.Int63(),
@ -525,6 +529,7 @@ func (suite *ServiceSuite) TestLoadSegments_VarChar() {
Schema: schema,
DeltaPositions: []*msgpb.MsgPosition{{Timestamp: 20000}},
NeedTransfer: true,
LoadMeta: loadMeta,
}
// LoadSegment
@ -549,6 +554,9 @@ func (suite *ServiceSuite) TestLoadDeltaInt64() {
Schema: schema,
NeedTransfer: true,
LoadScope: querypb.LoadScope_Delta,
LoadMeta: &querypb.LoadMetaInfo{
MetricType: defaultMetricType,
},
}
// LoadSegment
@ -573,6 +581,9 @@ func (suite *ServiceSuite) TestLoadDeltaVarchar() {
Schema: schema,
NeedTransfer: true,
LoadScope: querypb.LoadScope_Delta,
LoadMeta: &querypb.LoadMetaInfo{
MetricType: defaultMetricType,
},
}
// LoadSegment
@ -595,6 +606,9 @@ func (suite *ServiceSuite) TestLoadSegments_Failed() {
Infos: suite.genSegmentLoadInfos(schema),
Schema: schema,
NeedTransfer: true,
LoadMeta: &querypb.LoadMetaInfo{
MetricType: defaultMetricType,
},
}
// Delegator not found
@ -614,6 +628,14 @@ func (suite *ServiceSuite) TestLoadSegments_Failed() {
status, err = suite.node.LoadSegments(ctx, req)
suite.NoError(err)
suite.Equal(commonpb.ErrorCode_NotReadyServe, status.GetErrorCode())
// no metric type
req.LoadMeta = nil
req.Base.TargetID = paramtable.GetNodeID()
suite.node.UpdateStateCode(commonpb.StateCode_Healthy)
status, err = suite.node.LoadSegments(ctx, req)
suite.NoError(err)
suite.Equal(commonpb.ErrorCode_UnexpectedError, status.GetErrorCode())
}
func (suite *ServiceSuite) TestLoadSegments_Transfer() {
@ -637,6 +659,9 @@ func (suite *ServiceSuite) TestLoadSegments_Transfer() {
Infos: suite.genSegmentLoadInfos(schema),
Schema: schema,
NeedTransfer: true,
LoadMeta: &querypb.LoadMetaInfo{
MetricType: defaultMetricType,
},
}
// LoadSegment
@ -658,6 +683,9 @@ func (suite *ServiceSuite) TestLoadSegments_Transfer() {
Infos: suite.genSegmentLoadInfos(schema),
Schema: schema,
NeedTransfer: true,
LoadMeta: &querypb.LoadMetaInfo{
MetricType: defaultMetricType,
},
}
// LoadSegment
@ -684,8 +712,12 @@ func (suite *ServiceSuite) TestLoadSegments_Transfer() {
Infos: suite.genSegmentLoadInfos(schema),
Schema: schema,
NeedTransfer: true,
LoadMeta: &querypb.LoadMetaInfo{
MetricType: defaultMetricType,
},
}
// LoadSegment
// LoadSegment
status, err := suite.node.LoadSegments(ctx, req)
suite.NoError(err)
@ -1540,6 +1572,23 @@ func (suite *ServiceSuite) TestLoadPartition() {
suite.NoError(err)
suite.Equal(commonpb.ErrorCode_UnexpectedError, status.GetErrorCode())
// empty metric type
req.IndexInfoList = []*indexpb.IndexInfo{
{
CollectionID: suite.collectionID,
FieldID: vecField.GetFieldID(),
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.MetricTypeKey,
Value: "",
},
},
},
}
status, err = suite.node.LoadPartitions(ctx, req)
suite.NoError(err)
suite.Equal(commonpb.ErrorCode_UnexpectedError, status.ErrorCode)
// collection not exist and schema is not nil
req.IndexInfoList = []*indexpb.IndexInfo{
{