2021-04-19 13:47:10 +08:00
|
|
|
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
|
|
|
//
|
|
|
|
// Licensed 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.
|
|
|
|
|
2021-01-16 10:12:14 +08:00
|
|
|
package querynode
|
2020-11-09 16:27:11 +08:00
|
|
|
|
|
|
|
import (
|
2021-01-11 18:35:54 +08:00
|
|
|
"context"
|
2020-11-09 16:27:11 +08:00
|
|
|
"sync"
|
2020-11-12 12:04:12 +08:00
|
|
|
|
2021-06-09 11:37:55 +08:00
|
|
|
"github.com/opentracing/opentracing-go"
|
|
|
|
"go.uber.org/zap"
|
|
|
|
|
2021-04-22 14:45:57 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/log"
|
|
|
|
"github.com/milvus-io/milvus/internal/proto/commonpb"
|
|
|
|
"github.com/milvus-io/milvus/internal/util/flowgraph"
|
|
|
|
"github.com/milvus-io/milvus/internal/util/trace"
|
2020-11-09 16:27:11 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
type insertNode struct {
|
2021-01-15 15:28:54 +08:00
|
|
|
baseNode
|
2021-06-09 11:37:55 +08:00
|
|
|
replica ReplicaInterface
|
2020-11-09 16:27:11 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
type InsertData struct {
|
2021-01-11 18:35:54 +08:00
|
|
|
insertContext map[int64]context.Context
|
2020-11-09 16:27:11 +08:00
|
|
|
insertIDs map[UniqueID][]UniqueID
|
|
|
|
insertTimestamps map[UniqueID][]Timestamp
|
|
|
|
insertRecords map[UniqueID][]*commonpb.Blob
|
|
|
|
insertOffset map[UniqueID]int64
|
|
|
|
}
|
|
|
|
|
|
|
|
func (iNode *insertNode) Name() string {
|
|
|
|
return "iNode"
|
|
|
|
}
|
|
|
|
|
2021-03-25 14:41:46 +08:00
|
|
|
func (iNode *insertNode) Operate(in []flowgraph.Msg) []flowgraph.Msg {
|
2021-03-05 09:21:35 +08:00
|
|
|
//log.Debug("Do insertNode operation")
|
2020-11-09 16:27:11 +08:00
|
|
|
|
|
|
|
if len(in) != 1 {
|
2021-03-05 09:21:35 +08:00
|
|
|
log.Error("Invalid operate message input in insertNode", zap.Int("input length", len(in)))
|
2020-11-09 16:27:11 +08:00
|
|
|
// TODO: add error handling
|
|
|
|
}
|
|
|
|
|
2021-02-25 17:35:36 +08:00
|
|
|
iMsg, ok := in[0].(*insertMsg)
|
2020-11-09 16:27:11 +08:00
|
|
|
if !ok {
|
2021-03-05 09:21:35 +08:00
|
|
|
log.Error("type assertion failed for insertMsg")
|
2020-11-09 16:27:11 +08:00
|
|
|
// TODO: add error handling
|
|
|
|
}
|
|
|
|
|
|
|
|
insertData := InsertData{
|
|
|
|
insertIDs: make(map[int64][]int64),
|
|
|
|
insertTimestamps: make(map[int64][]uint64),
|
|
|
|
insertRecords: make(map[int64][]*commonpb.Blob),
|
|
|
|
insertOffset: make(map[int64]int64),
|
|
|
|
}
|
|
|
|
|
2021-03-22 16:36:10 +08:00
|
|
|
if iMsg == nil {
|
2021-03-25 14:41:46 +08:00
|
|
|
return []Msg{}
|
|
|
|
}
|
|
|
|
|
|
|
|
var spans []opentracing.Span
|
|
|
|
for _, msg := range iMsg.insertMessages {
|
|
|
|
sp, ctx := trace.StartSpanFromContext(msg.TraceCtx())
|
|
|
|
spans = append(spans, sp)
|
|
|
|
msg.SetTraceCtx(ctx)
|
2021-03-22 16:36:10 +08:00
|
|
|
}
|
|
|
|
|
2020-11-09 16:27:11 +08:00
|
|
|
// 1. hash insertMessages to insertData
|
|
|
|
for _, task := range iMsg.insertMessages {
|
2021-06-15 12:41:40 +08:00
|
|
|
// check if partition exists, if not, create partition
|
|
|
|
if hasPartition := iNode.replica.hasPartition(task.PartitionID); !hasPartition {
|
|
|
|
err := iNode.replica.addPartition(task.CollectionID, task.PartitionID)
|
|
|
|
if err != nil {
|
|
|
|
log.Error(err.Error())
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2020-11-16 17:01:10 +08:00
|
|
|
// check if segment exists, if not, create this segment
|
2020-12-08 14:41:04 +08:00
|
|
|
if !iNode.replica.hasSegment(task.SegmentID) {
|
2021-06-15 12:41:40 +08:00
|
|
|
err := iNode.replica.addSegment(task.SegmentID, task.PartitionID, task.CollectionID, task.ChannelID, segmentTypeGrowing, true)
|
2020-11-16 17:01:10 +08:00
|
|
|
if err != nil {
|
2021-03-05 09:21:35 +08:00
|
|
|
log.Error(err.Error())
|
2020-11-16 17:01:10 +08:00
|
|
|
continue
|
|
|
|
}
|
|
|
|
}
|
2021-03-12 19:23:06 +08:00
|
|
|
|
|
|
|
insertData.insertIDs[task.SegmentID] = append(insertData.insertIDs[task.SegmentID], task.RowIDs...)
|
|
|
|
insertData.insertTimestamps[task.SegmentID] = append(insertData.insertTimestamps[task.SegmentID], task.Timestamps...)
|
|
|
|
insertData.insertRecords[task.SegmentID] = append(insertData.insertRecords[task.SegmentID], task.RowData...)
|
2020-11-09 16:27:11 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 2. do preInsert
|
|
|
|
for segmentID := range insertData.insertRecords {
|
2020-12-08 14:41:04 +08:00
|
|
|
var targetSegment, err = iNode.replica.getSegmentByID(segmentID)
|
2020-11-09 16:27:11 +08:00
|
|
|
if err != nil {
|
2021-03-05 09:21:35 +08:00
|
|
|
log.Error("preInsert failed")
|
2020-11-09 16:27:11 +08:00
|
|
|
// TODO: add error handling
|
|
|
|
}
|
|
|
|
|
|
|
|
var numOfRecords = len(insertData.insertRecords[segmentID])
|
|
|
|
if targetSegment != nil {
|
2021-03-12 19:23:06 +08:00
|
|
|
offset, err := targetSegment.segmentPreInsert(numOfRecords)
|
|
|
|
if err != nil {
|
|
|
|
log.Error(err.Error())
|
|
|
|
}
|
2020-11-09 16:27:11 +08:00
|
|
|
insertData.insertOffset[segmentID] = offset
|
2021-03-31 16:16:58 +08:00
|
|
|
log.Debug("insertNode operator", zap.Int("insert size", numOfRecords), zap.Int64("insert offset", offset), zap.Int64("segment id", segmentID))
|
2020-11-09 16:27:11 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// 3. do insert
|
|
|
|
wg := sync.WaitGroup{}
|
|
|
|
for segmentID := range insertData.insertRecords {
|
|
|
|
wg.Add(1)
|
2021-02-04 19:34:35 +08:00
|
|
|
go iNode.insert(&insertData, segmentID, &wg)
|
2020-11-09 16:27:11 +08:00
|
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
|
|
|
|
var res Msg = &serviceTimeMsg{
|
2021-01-06 16:44:12 +08:00
|
|
|
gcRecord: iMsg.gcRecord,
|
2020-11-09 16:27:11 +08:00
|
|
|
timeRange: iMsg.timeRange,
|
|
|
|
}
|
2021-03-25 14:41:46 +08:00
|
|
|
for _, sp := range spans {
|
|
|
|
sp.Finish()
|
|
|
|
}
|
2021-06-09 11:37:55 +08:00
|
|
|
|
2021-03-25 14:41:46 +08:00
|
|
|
return []Msg{res}
|
2020-11-09 16:27:11 +08:00
|
|
|
}
|
|
|
|
|
2021-02-04 19:34:35 +08:00
|
|
|
func (iNode *insertNode) insert(insertData *InsertData, segmentID int64, wg *sync.WaitGroup) {
|
2021-06-05 16:21:36 +08:00
|
|
|
log.Debug("QueryNode::iNode::insert", zap.Any("SegmentID", segmentID))
|
2020-12-08 14:41:04 +08:00
|
|
|
var targetSegment, err = iNode.replica.getSegmentByID(segmentID)
|
2021-06-09 11:37:55 +08:00
|
|
|
if targetSegment.segmentType != segmentTypeGrowing {
|
2021-03-12 19:23:06 +08:00
|
|
|
wg.Done()
|
|
|
|
return
|
|
|
|
}
|
2020-11-09 16:27:11 +08:00
|
|
|
if err != nil {
|
2021-03-05 09:21:35 +08:00
|
|
|
log.Error("cannot find segment:", zap.Int64("segmentID", segmentID))
|
2020-11-09 16:27:11 +08:00
|
|
|
// TODO: add error handling
|
2020-11-26 16:01:31 +08:00
|
|
|
wg.Done()
|
2020-11-09 16:27:11 +08:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2021-03-22 16:36:10 +08:00
|
|
|
if targetSegment.segmentType != segmentTypeGrowing {
|
|
|
|
wg.Done()
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2020-11-09 16:27:11 +08:00
|
|
|
ids := insertData.insertIDs[segmentID]
|
|
|
|
timestamps := insertData.insertTimestamps[segmentID]
|
|
|
|
records := insertData.insertRecords[segmentID]
|
|
|
|
offsets := insertData.insertOffset[segmentID]
|
|
|
|
|
|
|
|
err = targetSegment.segmentInsert(offsets, &ids, ×tamps, &records)
|
|
|
|
if err != nil {
|
2021-06-05 16:21:36 +08:00
|
|
|
log.Debug("QueryNode: targetSegmentInsert failed", zap.Error(err))
|
2020-11-09 16:27:11 +08:00
|
|
|
// TODO: add error handling
|
2020-11-26 16:01:31 +08:00
|
|
|
wg.Done()
|
2020-11-09 16:27:11 +08:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2021-03-22 16:36:10 +08:00
|
|
|
log.Debug("Do insert done", zap.Int("len", len(insertData.insertIDs[segmentID])),
|
2021-06-09 11:37:55 +08:00
|
|
|
zap.Int64("segmentID", segmentID))
|
2020-11-09 16:27:11 +08:00
|
|
|
wg.Done()
|
|
|
|
}
|
|
|
|
|
2021-06-09 11:37:55 +08:00
|
|
|
func newInsertNode(replica ReplicaInterface) *insertNode {
|
2020-12-10 16:31:09 +08:00
|
|
|
maxQueueLength := Params.FlowGraphMaxQueueLength
|
|
|
|
maxParallelism := Params.FlowGraphMaxParallelism
|
2020-11-18 17:32:52 +08:00
|
|
|
|
2021-01-15 15:28:54 +08:00
|
|
|
baseNode := baseNode{}
|
2020-11-09 16:27:11 +08:00
|
|
|
baseNode.SetMaxQueueLength(maxQueueLength)
|
|
|
|
baseNode.SetMaxParallelism(maxParallelism)
|
|
|
|
|
|
|
|
return &insertNode{
|
2021-06-09 11:37:55 +08:00
|
|
|
baseNode: baseNode,
|
|
|
|
replica: replica,
|
2020-11-09 16:27:11 +08:00
|
|
|
}
|
|
|
|
}
|