2024-01-31 20:45:04 +08:00
|
|
|
// Licensed to the LF AI & Data foundation under one
|
|
|
|
// or more contributor license agreements. See the NOTICE file
|
|
|
|
// distributed with this work for additional information
|
|
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
|
|
// to you under the Apache License, Version 2.0 (the
|
|
|
|
// "License"); you may not use this file except in compliance
|
|
|
|
// with the License. You may obtain a copy of the License at
|
|
|
|
//
|
|
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
//
|
|
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
// See the License for the specific language governing permissions and
|
|
|
|
// limitations under the License.
|
|
|
|
|
|
|
|
package importv2
|
|
|
|
|
|
|
|
import (
|
|
|
|
"github.com/samber/lo"
|
2024-05-23 11:29:41 +08:00
|
|
|
"go.uber.org/zap"
|
2024-01-31 20:45:04 +08:00
|
|
|
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
2024-05-23 11:29:41 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util/conc"
|
2024-01-31 20:45:04 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
type TaskType int
|
|
|
|
|
|
|
|
const (
|
2024-06-07 14:17:20 +08:00
|
|
|
PreImportTaskType TaskType = 0
|
|
|
|
ImportTaskType TaskType = 1
|
|
|
|
L0PreImportTaskType TaskType = 2
|
|
|
|
L0ImportTaskType TaskType = 3
|
2024-01-31 20:45:04 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
var ImportTaskTypeName = map[TaskType]string{
|
|
|
|
0: "PreImportTask",
|
|
|
|
1: "ImportTask",
|
2024-06-07 14:17:20 +08:00
|
|
|
2: "L0PreImportTaskType",
|
|
|
|
3: "L0ImportTaskType",
|
2024-01-31 20:45:04 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (t TaskType) String() string {
|
|
|
|
return ImportTaskTypeName[t]
|
|
|
|
}
|
|
|
|
|
|
|
|
type TaskFilter func(task Task) bool
|
|
|
|
|
2024-03-01 18:31:02 +08:00
|
|
|
func WithStates(states ...datapb.ImportTaskStateV2) TaskFilter {
|
2024-01-31 20:45:04 +08:00
|
|
|
return func(task Task) bool {
|
|
|
|
for _, state := range states {
|
|
|
|
if task.GetState() == state {
|
|
|
|
return true
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return false
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func WithType(taskType TaskType) TaskFilter {
|
|
|
|
return func(task Task) bool {
|
|
|
|
return task.GetType() == taskType
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
type UpdateAction func(task Task)
|
|
|
|
|
2024-03-01 18:31:02 +08:00
|
|
|
func UpdateState(state datapb.ImportTaskStateV2) UpdateAction {
|
2024-01-31 20:45:04 +08:00
|
|
|
return func(t Task) {
|
|
|
|
switch t.GetType() {
|
|
|
|
case PreImportTaskType:
|
|
|
|
t.(*PreImportTask).PreImportTask.State = state
|
|
|
|
case ImportTaskType:
|
|
|
|
t.(*ImportTask).ImportTaskV2.State = state
|
2024-06-07 14:17:20 +08:00
|
|
|
case L0PreImportTaskType:
|
|
|
|
t.(*L0PreImportTask).PreImportTask.State = state
|
|
|
|
case L0ImportTaskType:
|
|
|
|
t.(*L0ImportTask).ImportTaskV2.State = state
|
2024-01-31 20:45:04 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func UpdateReason(reason string) UpdateAction {
|
|
|
|
return func(t Task) {
|
|
|
|
switch t.GetType() {
|
|
|
|
case PreImportTaskType:
|
|
|
|
t.(*PreImportTask).PreImportTask.Reason = reason
|
|
|
|
case ImportTaskType:
|
|
|
|
t.(*ImportTask).ImportTaskV2.Reason = reason
|
2024-06-07 14:17:20 +08:00
|
|
|
case L0PreImportTaskType:
|
|
|
|
t.(*L0PreImportTask).PreImportTask.Reason = reason
|
|
|
|
case L0ImportTaskType:
|
|
|
|
t.(*L0ImportTask).ImportTaskV2.Reason = reason
|
2024-01-31 20:45:04 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func UpdateFileStat(idx int, fileStat *datapb.ImportFileStats) UpdateAction {
|
|
|
|
return func(task Task) {
|
2024-06-07 14:17:20 +08:00
|
|
|
var t *datapb.PreImportTask
|
|
|
|
switch it := task.(type) {
|
|
|
|
case *PreImportTask:
|
|
|
|
t = it.PreImportTask
|
|
|
|
case *L0PreImportTask:
|
|
|
|
t = it.PreImportTask
|
|
|
|
}
|
|
|
|
if t != nil {
|
|
|
|
t.FileStats[idx].FileSize = fileStat.GetFileSize()
|
|
|
|
t.FileStats[idx].TotalRows = fileStat.GetTotalRows()
|
|
|
|
t.FileStats[idx].TotalMemorySize = fileStat.GetTotalMemorySize()
|
|
|
|
t.FileStats[idx].HashedStats = fileStat.GetHashedStats()
|
2024-01-31 20:45:04 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func UpdateSegmentInfo(info *datapb.ImportSegmentInfo) UpdateAction {
|
|
|
|
mergeFn := func(current []*datapb.FieldBinlog, new []*datapb.FieldBinlog) []*datapb.FieldBinlog {
|
|
|
|
for _, binlog := range new {
|
|
|
|
fieldBinlogs, ok := lo.Find(current, func(log *datapb.FieldBinlog) bool {
|
|
|
|
return log.GetFieldID() == binlog.GetFieldID()
|
|
|
|
})
|
|
|
|
if !ok || fieldBinlogs == nil {
|
|
|
|
current = append(current, binlog)
|
|
|
|
} else {
|
|
|
|
fieldBinlogs.Binlogs = append(fieldBinlogs.Binlogs, binlog.Binlogs...)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return current
|
|
|
|
}
|
|
|
|
return func(task Task) {
|
2024-06-07 14:17:20 +08:00
|
|
|
var segmentsInfo map[int64]*datapb.ImportSegmentInfo
|
|
|
|
switch it := task.(type) {
|
|
|
|
case *ImportTask:
|
|
|
|
segmentsInfo = it.segmentsInfo
|
|
|
|
case *L0ImportTask:
|
|
|
|
segmentsInfo = it.segmentsInfo
|
|
|
|
}
|
|
|
|
if segmentsInfo != nil {
|
2024-01-31 20:45:04 +08:00
|
|
|
segment := info.GetSegmentID()
|
2024-06-07 14:17:20 +08:00
|
|
|
if _, ok := segmentsInfo[segment]; ok {
|
|
|
|
segmentsInfo[segment].ImportedRows = info.GetImportedRows()
|
|
|
|
segmentsInfo[segment].Binlogs = mergeFn(segmentsInfo[segment].Binlogs, info.GetBinlogs())
|
|
|
|
segmentsInfo[segment].Statslogs = mergeFn(segmentsInfo[segment].Statslogs, info.GetStatslogs())
|
|
|
|
segmentsInfo[segment].Deltalogs = mergeFn(segmentsInfo[segment].Deltalogs, info.GetDeltalogs())
|
2024-01-31 20:45:04 +08:00
|
|
|
return
|
|
|
|
}
|
2024-06-07 14:17:20 +08:00
|
|
|
segmentsInfo[segment] = info
|
2024-01-31 20:45:04 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
type Task interface {
|
2024-05-23 11:29:41 +08:00
|
|
|
Execute() []*conc.Future[any]
|
2024-01-31 20:45:04 +08:00
|
|
|
GetJobID() int64
|
|
|
|
GetTaskID() int64
|
|
|
|
GetCollectionID() int64
|
|
|
|
GetPartitionIDs() []int64
|
|
|
|
GetVchannels() []string
|
|
|
|
GetType() TaskType
|
2024-03-01 18:31:02 +08:00
|
|
|
GetState() datapb.ImportTaskStateV2
|
2024-01-31 20:45:04 +08:00
|
|
|
GetReason() string
|
|
|
|
GetSchema() *schemapb.CollectionSchema
|
|
|
|
Cancel()
|
|
|
|
Clone() Task
|
|
|
|
}
|
|
|
|
|
2024-05-23 11:29:41 +08:00
|
|
|
func WrapLogFields(task Task, fields ...zap.Field) []zap.Field {
|
|
|
|
res := []zap.Field{
|
|
|
|
zap.Int64("taskID", task.GetTaskID()),
|
|
|
|
zap.Int64("jobID", task.GetJobID()),
|
|
|
|
zap.Int64("collectionID", task.GetCollectionID()),
|
|
|
|
zap.String("type", task.GetType().String()),
|
2024-01-31 20:45:04 +08:00
|
|
|
}
|
2024-05-23 11:29:41 +08:00
|
|
|
res = append(res, fields...)
|
|
|
|
return res
|
2024-01-31 20:45:04 +08:00
|
|
|
}
|