2022-10-18 19:17:27 +08:00
|
|
|
package proxy
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
"fmt"
|
2023-09-12 10:19:17 +08:00
|
|
|
"io"
|
2022-10-18 19:17:27 +08:00
|
|
|
|
2023-03-21 14:06:01 +08:00
|
|
|
"github.com/cockroachdb/errors"
|
2023-09-12 10:19:17 +08:00
|
|
|
"github.com/golang/protobuf/proto"
|
2023-01-12 16:09:39 +08:00
|
|
|
"go.opentelemetry.io/otel"
|
2023-12-26 19:52:48 +08:00
|
|
|
"go.uber.org/atomic"
|
2022-10-18 19:17:27 +08:00
|
|
|
"go.uber.org/zap"
|
|
|
|
|
2023-06-09 01:28:37 +08:00
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
2023-03-21 14:06:01 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/allocator"
|
2023-09-12 10:19:17 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/parser/planparserv2"
|
|
|
|
"github.com/milvus-io/milvus/internal/proto/internalpb"
|
2022-10-18 19:17:27 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/proto/planpb"
|
2023-09-12 10:19:17 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/proto/querypb"
|
|
|
|
"github.com/milvus-io/milvus/internal/types"
|
2024-03-22 13:57:06 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/util/exprutil"
|
2023-04-06 19:14:32 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/common"
|
|
|
|
"github.com/milvus-io/milvus/pkg/log"
|
|
|
|
"github.com/milvus-io/milvus/pkg/mq/msgstream"
|
|
|
|
"github.com/milvus-io/milvus/pkg/util/commonpbutil"
|
2023-09-04 09:57:09 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
2023-04-06 19:14:32 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
|
|
"github.com/milvus-io/milvus/pkg/util/timerecord"
|
|
|
|
"github.com/milvus-io/milvus/pkg/util/typeutil"
|
2022-10-18 19:17:27 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
type BaseDeleteTask = msgstream.DeleteMsg
|
|
|
|
|
|
|
|
type deleteTask struct {
|
2024-02-21 09:52:59 +08:00
|
|
|
baseTask
|
2022-10-18 19:17:27 +08:00
|
|
|
Condition
|
2023-09-12 10:19:17 +08:00
|
|
|
ctx context.Context
|
|
|
|
tr *timerecord.TimeRecorder
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
req *milvuspb.DeleteRequest
|
2023-09-12 10:19:17 +08:00
|
|
|
|
|
|
|
// channel
|
2022-10-18 19:17:27 +08:00
|
|
|
chMgr channelsMgr
|
|
|
|
chTicker channelsTimeTicker
|
|
|
|
pChannels []pChan
|
2023-09-12 10:19:17 +08:00
|
|
|
vChannels []vChan
|
2022-10-18 19:17:27 +08:00
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
idAllocator allocator.Interface
|
2023-09-12 10:19:17 +08:00
|
|
|
|
|
|
|
// delete info
|
2023-12-26 19:52:48 +08:00
|
|
|
primaryKeys *schemapb.IDs
|
2023-11-01 11:26:16 +08:00
|
|
|
collectionID UniqueID
|
|
|
|
partitionID UniqueID
|
|
|
|
partitionKeyMode bool
|
2023-12-26 19:52:48 +08:00
|
|
|
|
|
|
|
// set by scheduler
|
|
|
|
ts Timestamp
|
|
|
|
msgID UniqueID
|
|
|
|
|
|
|
|
// result
|
2024-03-28 06:33:11 +08:00
|
|
|
count int64
|
|
|
|
allQueryCnt int64
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) TraceCtx() context.Context {
|
|
|
|
return dt.ctx
|
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) ID() UniqueID {
|
2023-09-12 10:19:17 +08:00
|
|
|
return dt.msgID
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) SetID(uid UniqueID) {
|
2023-09-12 10:19:17 +08:00
|
|
|
dt.msgID = uid
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) Type() commonpb.MsgType {
|
2023-09-12 10:19:17 +08:00
|
|
|
return commonpb.MsgType_Delete
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) Name() string {
|
2023-01-04 17:21:36 +08:00
|
|
|
return DeleteTaskName
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) BeginTs() Timestamp {
|
2023-09-12 10:19:17 +08:00
|
|
|
return dt.ts
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) EndTs() Timestamp {
|
2023-09-12 10:19:17 +08:00
|
|
|
return dt.ts
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) SetTs(ts Timestamp) {
|
2023-09-12 10:19:17 +08:00
|
|
|
dt.ts = ts
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) OnEnqueue() error {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2023-06-16 16:36:40 +08:00
|
|
|
func (dt *deleteTask) setChannels() error {
|
2023-09-12 10:19:17 +08:00
|
|
|
collID, err := globalMetaCache.GetCollectionID(dt.ctx, dt.req.GetDbName(), dt.req.GetCollectionName())
|
2022-10-18 19:17:27 +08:00
|
|
|
if err != nil {
|
2023-06-16 16:36:40 +08:00
|
|
|
return err
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-04-11 11:00:31 +08:00
|
|
|
channels, err := dt.chMgr.getChannels(collID)
|
|
|
|
if err != nil {
|
2023-06-16 16:36:40 +08:00
|
|
|
return err
|
2023-04-11 11:00:31 +08:00
|
|
|
}
|
|
|
|
dt.pChannels = channels
|
2023-06-16 16:36:40 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) getChannels() []pChan {
|
|
|
|
return dt.pChannels
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func (dt *deleteTask) PreExecute(ctx context.Context) error {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) Execute(ctx context.Context) (err error) {
|
|
|
|
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Delete-Execute")
|
|
|
|
defer sp.End()
|
|
|
|
// log := log.Ctx(ctx)
|
|
|
|
|
|
|
|
if len(dt.req.GetExpr()) == 0 {
|
|
|
|
return merr.WrapErrParameterInvalid("valid expr", "empty expr", "invalid expression")
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
dt.tr = timerecord.NewTimeRecorder(fmt.Sprintf("proxy execute delete %d", dt.ID()))
|
|
|
|
stream, err := dt.chMgr.getOrCreateDmlStream(dt.collectionID)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
2022-12-30 14:17:32 +08:00
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
hashValues := typeutil.HashPK2Channels(dt.primaryKeys, dt.vChannels)
|
|
|
|
// repack delete msg by dmChannel
|
|
|
|
result := make(map[uint32]msgstream.TsMsg)
|
|
|
|
numRows := int64(0)
|
|
|
|
for index, key := range hashValues {
|
|
|
|
vchannel := dt.vChannels[key]
|
|
|
|
_, ok := result[key]
|
|
|
|
if !ok {
|
|
|
|
deleteMsg, err := dt.newDeleteMsg(ctx)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
deleteMsg.ShardName = vchannel
|
|
|
|
result[key] = deleteMsg
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
curMsg := result[key].(*msgstream.DeleteMsg)
|
|
|
|
curMsg.HashValues = append(curMsg.HashValues, hashValues[index])
|
|
|
|
curMsg.Timestamps = append(curMsg.Timestamps, dt.ts)
|
|
|
|
|
|
|
|
typeutil.AppendIDs(curMsg.PrimaryKeys, dt.primaryKeys, index)
|
|
|
|
curMsg.NumRows++
|
|
|
|
numRows++
|
|
|
|
}
|
|
|
|
|
|
|
|
// send delete request to log broker
|
|
|
|
msgPack := &msgstream.MsgPack{
|
|
|
|
BeginTs: dt.BeginTs(),
|
|
|
|
EndTs: dt.EndTs(),
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, msg := range result {
|
|
|
|
if msg != nil {
|
|
|
|
msgPack.Msgs = append(msgPack.Msgs, msg)
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
log.Debug("send delete request to virtual channels",
|
|
|
|
zap.String("collectionName", dt.req.GetCollectionName()),
|
|
|
|
zap.Int64("collectionID", dt.collectionID),
|
|
|
|
zap.Strings("virtual_channels", dt.vChannels),
|
|
|
|
zap.Int64("taskID", dt.ID()),
|
|
|
|
zap.Duration("prepare duration", dt.tr.RecordSpan()))
|
|
|
|
|
|
|
|
err = stream.Produce(msgPack)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
dt.count += numRows
|
|
|
|
return nil
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func (dt *deleteTask) PostExecute(ctx context.Context) error {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (dt *deleteTask) newDeleteMsg(ctx context.Context) (*msgstream.DeleteMsg, error) {
|
|
|
|
msgid, err := dt.idAllocator.AllocOne()
|
|
|
|
if err != nil {
|
|
|
|
return nil, errors.Wrap(err, "failed to allocate MsgID of delete")
|
|
|
|
}
|
|
|
|
sliceRequest := msgpb.DeleteRequest{
|
|
|
|
Base: commonpbutil.NewMsgBase(
|
|
|
|
commonpbutil.WithMsgType(commonpb.MsgType_Delete),
|
|
|
|
// msgid of delete msg must be set
|
|
|
|
// or it will be seen as duplicated msg in mq
|
|
|
|
commonpbutil.WithMsgID(msgid),
|
|
|
|
commonpbutil.WithTimeStamp(dt.ts),
|
|
|
|
commonpbutil.WithSourceID(paramtable.GetNodeID()),
|
|
|
|
),
|
|
|
|
CollectionID: dt.collectionID,
|
|
|
|
PartitionID: dt.partitionID,
|
|
|
|
CollectionName: dt.req.GetCollectionName(),
|
|
|
|
PartitionName: dt.req.GetPartitionName(),
|
|
|
|
PrimaryKeys: &schemapb.IDs{},
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
return &msgstream.DeleteMsg{
|
|
|
|
BaseMsg: msgstream.BaseMsg{
|
|
|
|
Ctx: ctx,
|
|
|
|
},
|
|
|
|
DeleteRequest: sliceRequest,
|
|
|
|
}, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
type deleteRunner struct {
|
|
|
|
req *milvuspb.DeleteRequest
|
|
|
|
result *milvuspb.MutationResult
|
|
|
|
|
|
|
|
// channel
|
|
|
|
chMgr channelsMgr
|
|
|
|
chTicker channelsTimeTicker
|
|
|
|
vChannels []vChan
|
|
|
|
|
|
|
|
idAllocator allocator.Interface
|
|
|
|
tsoAllocatorIns tsoAllocator
|
2022-10-18 19:17:27 +08:00
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
// delete info
|
2024-01-04 17:28:46 +08:00
|
|
|
schema *schemaInfo
|
2023-12-26 19:52:48 +08:00
|
|
|
collectionID UniqueID
|
|
|
|
partitionID UniqueID
|
|
|
|
partitionKeyMode bool
|
|
|
|
|
|
|
|
// for query
|
|
|
|
msgID int64
|
|
|
|
ts uint64
|
|
|
|
lb LBPolicy
|
|
|
|
count atomic.Int64
|
|
|
|
err error
|
|
|
|
|
|
|
|
// task queue
|
|
|
|
queue *dmTaskQueue
|
2024-03-28 06:33:11 +08:00
|
|
|
|
|
|
|
allQueryCnt atomic.Int64
|
2023-12-26 19:52:48 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (dr *deleteRunner) Init(ctx context.Context) error {
|
2023-07-14 16:08:31 +08:00
|
|
|
log := log.Ctx(ctx)
|
2023-12-26 19:52:48 +08:00
|
|
|
var err error
|
|
|
|
|
|
|
|
collName := dr.req.GetCollectionName()
|
2022-10-18 19:17:27 +08:00
|
|
|
if err := validateCollectionName(collName); err != nil {
|
2023-08-30 14:47:00 +08:00
|
|
|
return ErrWithLog(log, "Invalid collection name", err)
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
dr.collectionID, err = globalMetaCache.GetCollectionID(ctx, dr.req.GetDbName(), collName)
|
2022-10-18 19:17:27 +08:00
|
|
|
if err != nil {
|
2023-08-30 14:47:00 +08:00
|
|
|
return ErrWithLog(log, "Failed to get collection id", err)
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
dr.schema, err = globalMetaCache.GetCollectionSchema(ctx, dr.req.GetDbName(), collName)
|
2023-06-06 10:24:34 +08:00
|
|
|
if err != nil {
|
2023-12-26 19:52:48 +08:00
|
|
|
return ErrWithLog(log, "Failed to get collection schema", err)
|
2023-06-06 10:24:34 +08:00
|
|
|
}
|
|
|
|
|
2024-01-04 17:28:46 +08:00
|
|
|
dr.partitionKeyMode = dr.schema.IsPartitionKeyCollection()
|
|
|
|
// get partitionIDs of delete
|
2024-03-20 19:01:05 +08:00
|
|
|
dr.partitionID = common.AllPartitionsID
|
2023-12-26 19:52:48 +08:00
|
|
|
if len(dr.req.PartitionName) > 0 {
|
|
|
|
if dr.partitionKeyMode {
|
|
|
|
return errors.New("not support manually specifying the partition names if partition key mode is used")
|
|
|
|
}
|
|
|
|
|
|
|
|
partName := dr.req.GetPartitionName()
|
2022-10-18 19:17:27 +08:00
|
|
|
if err := validatePartitionTag(partName, true); err != nil {
|
2023-08-30 14:47:00 +08:00
|
|
|
return ErrWithLog(log, "Invalid partition name", err)
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
partID, err := globalMetaCache.GetPartitionID(ctx, dr.req.GetDbName(), collName, partName)
|
2022-10-18 19:17:27 +08:00
|
|
|
if err != nil {
|
2023-08-30 14:47:00 +08:00
|
|
|
return ErrWithLog(log, "Failed to get partition id", err)
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
dr.partitionID = partID
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
2023-09-12 10:19:17 +08:00
|
|
|
// hash primary keys to channels
|
2023-12-26 19:52:48 +08:00
|
|
|
channelNames, err := dr.chMgr.getVChannels(dr.collectionID)
|
2022-10-18 19:17:27 +08:00
|
|
|
if err != nil {
|
2023-08-30 14:47:00 +08:00
|
|
|
return ErrWithLog(log, "Failed to get primary keys from expr", err)
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
dr.vChannels = channelNames
|
2023-07-14 16:08:31 +08:00
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
dr.result = &milvuspb.MutationResult{
|
|
|
|
Status: merr.Success(),
|
|
|
|
IDs: &schemapb.IDs{
|
|
|
|
IdField: nil,
|
|
|
|
},
|
|
|
|
}
|
2022-10-18 19:17:27 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func (dr *deleteRunner) Run(ctx context.Context) error {
|
2024-04-22 11:09:23 +08:00
|
|
|
plan, err := planparserv2.CreateRetrievePlan(dr.schema.schemaHelper, dr.req.GetExpr())
|
2023-09-12 10:19:17 +08:00
|
|
|
if err != nil {
|
2024-02-20 16:02:52 +08:00
|
|
|
return merr.WrapErrParameterInvalidMsg("failed to create delete plan: %v", err)
|
2023-09-12 10:19:17 +08:00
|
|
|
}
|
|
|
|
|
2024-04-22 11:09:23 +08:00
|
|
|
if planparserv2.IsAlwaysTruePlan(plan) {
|
|
|
|
return merr.WrapErrParameterInvalidMsg("delete plan can't be empty or always true : %s", dr.req.GetExpr())
|
|
|
|
}
|
|
|
|
|
2024-01-04 17:28:46 +08:00
|
|
|
isSimple, pk, numRow := getPrimaryKeysFromPlan(dr.schema.CollectionSchema, plan)
|
2023-09-12 10:19:17 +08:00
|
|
|
if isSimple {
|
|
|
|
// if could get delete.primaryKeys from delete expr
|
2023-12-26 19:52:48 +08:00
|
|
|
err := dr.simpleDelete(ctx, pk, numRow)
|
2023-09-12 10:19:17 +08:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
// if get complex delete expr
|
|
|
|
// need query from querynode before delete
|
2023-12-26 19:52:48 +08:00
|
|
|
err = dr.complexDelete(ctx, plan)
|
2023-09-12 10:19:17 +08:00
|
|
|
if err != nil {
|
2023-12-26 19:52:48 +08:00
|
|
|
log.Warn("complex delete failed,but delete some data", zap.Int64("count", dr.result.DeleteCnt), zap.String("expr", dr.req.GetExpr()))
|
2023-09-12 10:19:17 +08:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func (dr *deleteRunner) produce(ctx context.Context, primaryKeys *schemapb.IDs) (*deleteTask, error) {
|
|
|
|
task := &deleteTask{
|
|
|
|
ctx: ctx,
|
|
|
|
Condition: NewTaskCondition(ctx),
|
|
|
|
req: dr.req,
|
|
|
|
idAllocator: dr.idAllocator,
|
|
|
|
chMgr: dr.chMgr,
|
|
|
|
chTicker: dr.chTicker,
|
|
|
|
collectionID: dr.collectionID,
|
|
|
|
partitionID: dr.partitionID,
|
|
|
|
partitionKeyMode: dr.partitionKeyMode,
|
|
|
|
vChannels: dr.vChannels,
|
|
|
|
primaryKeys: primaryKeys,
|
|
|
|
}
|
|
|
|
|
|
|
|
if err := dr.queue.Enqueue(task); err != nil {
|
|
|
|
log.Error("Failed to enqueue delete task: " + err.Error())
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
|
|
|
return task, nil
|
2023-09-12 10:19:17 +08:00
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
// getStreamingQueryAndDelteFunc return query function used by LBPolicy
|
|
|
|
// make sure it concurrent safe
|
|
|
|
func (dr *deleteRunner) getStreamingQueryAndDelteFunc(plan *planpb.PlanNode) executeFunc {
|
2024-01-09 11:38:48 +08:00
|
|
|
return func(ctx context.Context, nodeID int64, qn types.QueryNodeClient, channel string) error {
|
2023-11-28 15:32:31 +08:00
|
|
|
var partitionIDs []int64
|
2023-11-01 11:26:16 +08:00
|
|
|
|
|
|
|
// optimize query when partitionKey on
|
2023-12-26 19:52:48 +08:00
|
|
|
if dr.partitionKeyMode {
|
2024-03-22 13:57:06 +08:00
|
|
|
expr, err := exprutil.ParseExprFromPlan(plan)
|
2023-11-01 11:26:16 +08:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2024-03-22 13:57:06 +08:00
|
|
|
partitionKeys := exprutil.ParseKeys(expr, exprutil.PartitionKey)
|
2023-12-26 19:52:48 +08:00
|
|
|
hashedPartitionNames, err := assignPartitionKeys(ctx, dr.req.GetDbName(), dr.req.GetCollectionName(), partitionKeys)
|
2023-11-01 11:26:16 +08:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
partitionIDs, err = getPartitionIDs(ctx, dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames)
|
2023-11-01 11:26:16 +08:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
} else if dr.partitionID != common.InvalidFieldID {
|
|
|
|
partitionIDs = []int64{dr.partitionID}
|
2023-09-12 10:19:17 +08:00
|
|
|
}
|
|
|
|
|
2023-10-24 14:54:12 +08:00
|
|
|
log := log.Ctx(ctx).With(
|
2023-12-26 19:52:48 +08:00
|
|
|
zap.Int64("collectionID", dr.collectionID),
|
2023-11-28 15:32:31 +08:00
|
|
|
zap.Int64s("partitionIDs", partitionIDs),
|
2024-01-09 11:38:48 +08:00
|
|
|
zap.String("channel", channel),
|
2023-10-24 14:54:12 +08:00
|
|
|
zap.Int64("nodeID", nodeID))
|
2023-12-26 19:52:48 +08:00
|
|
|
|
2023-09-12 10:19:17 +08:00
|
|
|
// set plan
|
2024-01-04 17:28:46 +08:00
|
|
|
_, outputFieldIDs := translatePkOutputFields(dr.schema.CollectionSchema)
|
2023-09-12 10:19:17 +08:00
|
|
|
outputFieldIDs = append(outputFieldIDs, common.TimeStampField)
|
|
|
|
plan.OutputFieldIds = outputFieldIDs
|
|
|
|
|
|
|
|
serializedPlan, err := proto.Marshal(plan)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
queryReq := &querypb.QueryRequest{
|
|
|
|
Req: &internalpb.RetrieveRequest{
|
|
|
|
Base: commonpbutil.NewMsgBase(
|
|
|
|
commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
|
2023-12-26 19:52:48 +08:00
|
|
|
commonpbutil.WithMsgID(dr.msgID),
|
2023-09-12 10:19:17 +08:00
|
|
|
commonpbutil.WithSourceID(paramtable.GetNodeID()),
|
|
|
|
commonpbutil.WithTargetID(nodeID),
|
|
|
|
),
|
2023-12-26 19:52:48 +08:00
|
|
|
MvccTimestamp: dr.ts,
|
2023-09-12 10:19:17 +08:00
|
|
|
ReqID: paramtable.GetNodeID(),
|
|
|
|
DbID: 0, // TODO
|
2023-12-26 19:52:48 +08:00
|
|
|
CollectionID: dr.collectionID,
|
2023-11-28 15:32:31 +08:00
|
|
|
PartitionIDs: partitionIDs,
|
2023-09-12 10:19:17 +08:00
|
|
|
SerializedExprPlan: serializedPlan,
|
|
|
|
OutputFieldsId: outputFieldIDs,
|
2023-12-26 19:52:48 +08:00
|
|
|
GuaranteeTimestamp: parseGuaranteeTsFromConsistency(dr.ts, dr.ts, dr.req.GetConsistencyLevel()),
|
2023-09-12 10:19:17 +08:00
|
|
|
},
|
2024-01-09 11:38:48 +08:00
|
|
|
DmlChannels: []string{channel},
|
2023-09-12 10:19:17 +08:00
|
|
|
Scope: querypb.DataScope_All,
|
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
|
|
defer cancel()
|
|
|
|
log.Debug("start query for delete", zap.Int64("msgID", dr.msgID))
|
2023-09-26 09:57:25 +08:00
|
|
|
client, err := qn.QueryStream(ctx, queryReq)
|
2023-09-12 10:19:17 +08:00
|
|
|
if err != nil {
|
2023-10-24 14:54:12 +08:00
|
|
|
log.Warn("query stream for delete create failed", zap.Error(err))
|
2023-09-12 10:19:17 +08:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
taskCh := make(chan *deleteTask, 256)
|
|
|
|
go dr.receiveQueryResult(ctx, client, taskCh)
|
2024-03-28 06:33:11 +08:00
|
|
|
var allQueryCnt int64
|
2023-12-26 19:52:48 +08:00
|
|
|
// wait all task finish
|
|
|
|
for task := range taskCh {
|
|
|
|
err := task.WaitToFinish()
|
2023-09-12 10:19:17 +08:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
dr.count.Add(task.count)
|
2024-03-28 06:33:11 +08:00
|
|
|
allQueryCnt += task.allQueryCnt
|
2023-12-26 19:52:48 +08:00
|
|
|
}
|
2023-09-12 10:19:17 +08:00
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
// query or produce task failed
|
|
|
|
if dr.err != nil {
|
|
|
|
return dr.err
|
|
|
|
}
|
2024-03-28 06:33:11 +08:00
|
|
|
dr.allQueryCnt.Add(allQueryCnt)
|
2023-12-26 19:52:48 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
}
|
2023-09-12 10:19:17 +08:00
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func (dr *deleteRunner) receiveQueryResult(ctx context.Context, client querypb.QueryNode_QueryStreamClient, taskCh chan *deleteTask) {
|
|
|
|
defer func() {
|
|
|
|
close(taskCh)
|
|
|
|
}()
|
|
|
|
|
|
|
|
for {
|
|
|
|
result, err := client.Recv()
|
|
|
|
if err != nil {
|
|
|
|
if err == io.EOF {
|
|
|
|
log.Debug("query stream for delete finished", zap.Int64("msgID", dr.msgID))
|
|
|
|
return
|
2023-09-12 10:19:17 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
dr.err = err
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
err = merr.Error(result.GetStatus())
|
|
|
|
if err != nil {
|
|
|
|
dr.err = err
|
|
|
|
log.Warn("query stream for delete get error status", zap.Int64("msgID", dr.msgID), zap.Error(err))
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
task, err := dr.produce(ctx, result.GetIds())
|
|
|
|
if err != nil {
|
|
|
|
dr.err = err
|
|
|
|
log.Warn("produce delete task failed", zap.Error(err))
|
|
|
|
return
|
2023-09-12 10:19:17 +08:00
|
|
|
}
|
2024-03-28 06:33:11 +08:00
|
|
|
task.allQueryCnt = result.GetAllRetrieveCount()
|
2023-12-26 19:52:48 +08:00
|
|
|
|
|
|
|
taskCh <- task
|
2023-09-12 10:19:17 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func (dr *deleteRunner) complexDelete(ctx context.Context, plan *planpb.PlanNode) error {
|
|
|
|
rc := timerecord.NewTimeRecorder("QueryStreamDelete")
|
|
|
|
var err error
|
|
|
|
|
|
|
|
dr.msgID, err = dr.idAllocator.AllocOne()
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
dr.ts, err = dr.tsoAllocatorIns.AllocOne(ctx)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
err = dr.lb.Execute(ctx, CollectionWorkLoad{
|
|
|
|
db: dr.req.GetDbName(),
|
|
|
|
collectionName: dr.req.GetCollectionName(),
|
|
|
|
collectionID: dr.collectionID,
|
2023-09-12 10:19:17 +08:00
|
|
|
nq: 1,
|
2023-12-26 19:52:48 +08:00
|
|
|
exec: dr.getStreamingQueryAndDelteFunc(plan),
|
2023-09-12 10:19:17 +08:00
|
|
|
})
|
2023-12-26 19:52:48 +08:00
|
|
|
dr.result.DeleteCnt = dr.count.Load()
|
2022-10-18 19:17:27 +08:00
|
|
|
if err != nil {
|
2023-12-26 19:52:48 +08:00
|
|
|
log.Warn("fail to execute complex delete",
|
|
|
|
zap.Int64("deleteCnt", dr.result.GetDeleteCnt()),
|
|
|
|
zap.Duration("interval", rc.ElapseSpan()),
|
|
|
|
zap.Error(err))
|
2022-10-18 19:17:27 +08:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
log.Info("complex delete finished", zap.Int64("deleteCnt", dr.result.GetDeleteCnt()), zap.Duration("interval", rc.ElapseSpan()))
|
2023-09-12 10:19:17 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func (dr *deleteRunner) simpleDelete(ctx context.Context, pk *schemapb.IDs, numRow int64) error {
|
2023-10-24 14:54:12 +08:00
|
|
|
log.Debug("get primary keys from expr",
|
|
|
|
zap.Int64("len of primary keys", numRow),
|
2023-12-26 19:52:48 +08:00
|
|
|
zap.Int64("collectionID", dr.collectionID),
|
|
|
|
zap.Int64("partitionID", dr.partitionID))
|
|
|
|
|
|
|
|
task, err := dr.produce(ctx, pk)
|
2022-10-18 19:17:27 +08:00
|
|
|
if err != nil {
|
2023-12-26 19:52:48 +08:00
|
|
|
log.Warn("produce delete task failed")
|
2022-10-18 19:17:27 +08:00
|
|
|
return err
|
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
|
|
|
|
err = task.WaitToFinish()
|
|
|
|
if err == nil {
|
|
|
|
dr.result.DeleteCnt = task.count
|
|
|
|
}
|
|
|
|
return err
|
2023-09-12 10:19:17 +08:00
|
|
|
}
|
2022-10-18 19:17:27 +08:00
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func getPrimaryKeysFromPlan(schema *schemapb.CollectionSchema, plan *planpb.PlanNode) (bool, *schemapb.IDs, int64) {
|
|
|
|
// simple delete request need expr with "pk in [a, b]"
|
|
|
|
termExpr, ok := plan.Node.(*planpb.PlanNode_Query).Query.Predicates.Expr.(*planpb.Expr_TermExpr)
|
|
|
|
if ok {
|
|
|
|
if !termExpr.TermExpr.GetColumnInfo().GetIsPrimaryKey() {
|
|
|
|
return false, nil, 0
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-09-12 10:19:17 +08:00
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
ids, rowNum, err := getPrimaryKeysFromTermExpr(schema, termExpr)
|
|
|
|
if err != nil {
|
|
|
|
return false, nil, 0
|
|
|
|
}
|
|
|
|
return true, ids, rowNum
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
// simple delete if expr with "pk == a"
|
|
|
|
unaryRangeExpr, ok := plan.Node.(*planpb.PlanNode_Query).Query.Predicates.Expr.(*planpb.Expr_UnaryRangeExpr)
|
|
|
|
if ok {
|
|
|
|
if unaryRangeExpr.UnaryRangeExpr.GetOp() != planpb.OpType_Equal || !unaryRangeExpr.UnaryRangeExpr.GetColumnInfo().GetIsPrimaryKey() {
|
|
|
|
return false, nil, 0
|
|
|
|
}
|
2023-09-12 10:19:17 +08:00
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
ids, err := getPrimaryKeysFromUnaryRangeExpr(schema, unaryRangeExpr)
|
|
|
|
if err != nil {
|
|
|
|
return false, nil, 0
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
return true, ids, 1
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
return false, nil, 0
|
|
|
|
}
|
2023-03-30 18:56:22 +08:00
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func getPrimaryKeysFromUnaryRangeExpr(schema *schemapb.CollectionSchema, unaryRangeExpr *planpb.Expr_UnaryRangeExpr) (res *schemapb.IDs, err error) {
|
|
|
|
res = &schemapb.IDs{}
|
|
|
|
switch unaryRangeExpr.UnaryRangeExpr.GetColumnInfo().GetDataType() {
|
|
|
|
case schemapb.DataType_Int64:
|
|
|
|
res.IdField = &schemapb.IDs_IntId{
|
|
|
|
IntId: &schemapb.LongArray{
|
|
|
|
Data: []int64{unaryRangeExpr.UnaryRangeExpr.GetValue().GetInt64Val()},
|
|
|
|
},
|
|
|
|
}
|
|
|
|
case schemapb.DataType_VarChar:
|
|
|
|
res.IdField = &schemapb.IDs_StrId{
|
|
|
|
StrId: &schemapb.StringArray{
|
|
|
|
Data: []string{unaryRangeExpr.UnaryRangeExpr.GetValue().GetStringVal()},
|
|
|
|
},
|
|
|
|
}
|
|
|
|
default:
|
|
|
|
return res, fmt.Errorf("invalid field data type specifyed in simple delete expr")
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
|
|
|
|
return res, nil
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|
|
|
|
|
2023-12-26 19:52:48 +08:00
|
|
|
func getPrimaryKeysFromTermExpr(schema *schemapb.CollectionSchema, termExpr *planpb.Expr_TermExpr) (res *schemapb.IDs, rowNum int64, err error) {
|
|
|
|
res = &schemapb.IDs{}
|
|
|
|
rowNum = int64(len(termExpr.TermExpr.Values))
|
|
|
|
switch termExpr.TermExpr.ColumnInfo.GetDataType() {
|
|
|
|
case schemapb.DataType_Int64:
|
|
|
|
ids := make([]int64, 0)
|
|
|
|
for _, v := range termExpr.TermExpr.Values {
|
|
|
|
ids = append(ids, v.GetInt64Val())
|
|
|
|
}
|
|
|
|
res.IdField = &schemapb.IDs_IntId{
|
|
|
|
IntId: &schemapb.LongArray{
|
|
|
|
Data: ids,
|
|
|
|
},
|
|
|
|
}
|
|
|
|
case schemapb.DataType_VarChar:
|
|
|
|
ids := make([]string, 0)
|
|
|
|
for _, v := range termExpr.TermExpr.Values {
|
|
|
|
ids = append(ids, v.GetStringVal())
|
|
|
|
}
|
|
|
|
res.IdField = &schemapb.IDs_StrId{
|
|
|
|
StrId: &schemapb.StringArray{
|
|
|
|
Data: ids,
|
|
|
|
},
|
|
|
|
}
|
|
|
|
default:
|
|
|
|
return res, 0, fmt.Errorf("invalid field data type specifyed in simple delete expr")
|
2023-09-12 10:19:17 +08:00
|
|
|
}
|
2023-12-26 19:52:48 +08:00
|
|
|
|
|
|
|
return res, rowNum, nil
|
2022-10-18 19:17:27 +08:00
|
|
|
}
|