milvus-io--milvus/internal/proxy/task_upsert.go

// 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 proxy

import (
	"context"
	"fmt"
	"strconv"

	"github.com/samber/lo"
	"go.opentelemetry.io/otel"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
	"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/internal/allocator"
	"github.com/milvus-io/milvus/internal/parser/planparserv2"
	"github.com/milvus-io/milvus/internal/types"
	"github.com/milvus-io/milvus/internal/util/segcore"
	"github.com/milvus-io/milvus/pkg/v3/common"
	"github.com/milvus-io/milvus/pkg/v3/metrics"
	"github.com/milvus-io/milvus/pkg/v3/mlog"
	"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
	"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
	"github.com/milvus-io/milvus/pkg/v3/util/merr"
	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
	"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

type upsertTask struct {
	baseTask
	Condition

	upsertMsg *msgstream.UpsertMsg
	req       *milvuspb.UpsertRequest
	baseMsg   msgstream.BaseMsg

	ctx context.Context

	timestamps       []uint64
	rowIDs           []int64
	result           *milvuspb.MutationResult
	idAllocator      *allocator.IDAllocator
	collectionID     UniqueID
	chMgr            channelsMgr
	chTicker         channelsTimeTicker
	vChannels        []vChan
	pChannels        []pChan
	schema           *schemaInfo
	partitionKeyMode bool
	partitionKeys    *schemapb.FieldData
	// automatic generate pk as new pk wehen autoID == true
	// delete task need use the oldIDs
	oldIDs          *schemapb.IDs
	schemaTimestamp uint64
	schemaVersion   int32

	// write after read, generate write part by queryPreExecute
	node types.ProxyComponent

	deletePKs       *schemapb.IDs
	insertFieldData []*schemapb.FieldData

	storageCost segcore.StorageCost
}

// TraceCtx returns upsertTask context
func (it *upsertTask) TraceCtx() context.Context {
	return it.ctx
}

func (it *upsertTask) ID() UniqueID {
	return it.req.Base.MsgID
}

func (it *upsertTask) SetID(uid UniqueID) {
	it.req.Base.MsgID = uid
}

func (it *upsertTask) Name() string {
	return UpsertTaskName
}

func (it *upsertTask) Type() commonpb.MsgType {
	return it.req.Base.MsgType
}

func (it *upsertTask) BeginTs() Timestamp {
	return it.baseMsg.BeginTimestamp
}

func (it *upsertTask) SetTs(ts Timestamp) {
	it.baseMsg.BeginTimestamp = ts
	it.baseMsg.EndTimestamp = ts
}

func (it *upsertTask) EndTs() Timestamp {
	return it.baseMsg.EndTimestamp
}

func (it *upsertTask) getPChanStats() (map[pChan]pChanStatistics, error) {
	ret := make(map[pChan]pChanStatistics)

	channels := it.getChannels()

	beginTs := it.BeginTs()
	endTs := it.EndTs()

	for _, channel := range channels {
		ret[channel] = pChanStatistics{
			minTs: beginTs,
			maxTs: endTs,
		}
	}
	return ret, nil
}

func (it *upsertTask) setChannels() error {
	collID, err := globalMetaCache.GetCollectionID(it.ctx, it.req.GetDbName(), it.req.CollectionName)
	if err != nil {
		return err
	}
	channels, err := it.chMgr.getChannels(collID)
	if err != nil {
		return err
	}
	it.pChannels = channels
	return nil
}

func (it *upsertTask) getChannels() []pChan {
	return it.pChannels
}

func (it *upsertTask) OnEnqueue() error {
	if it.req.Base == nil {
		it.req.Base = commonpbutil.NewMsgBase()
	}
	it.req.Base.MsgType = commonpb.MsgType_Upsert
	it.req.Base.SourceID = paramtable.GetNodeID()
	return nil
}

func retrieveByPKs(ctx context.Context, t *upsertTask, ids *schemapb.IDs, outputFields []string) (*milvuspb.QueryResults, segcore.StorageCost, error) {
	log := mlog.With(mlog.String("collectionName", t.req.GetCollectionName()))
	var err error
	queryReq := &milvuspb.QueryRequest{
		Base: &commonpb.MsgBase{
			MsgType:   commonpb.MsgType_Retrieve,
			Timestamp: t.BeginTs(),
		},
		DbName:                t.req.GetDbName(),
		CollectionName:        t.req.GetCollectionName(),
		ConsistencyLevel:      commonpb.ConsistencyLevel_Strong,
		NotReturnAllMeta:      false,
		OutputFields:          []string{"*"},
		UseDefaultConsistency: false,
		GuaranteeTimestamp:    t.BeginTs(),
		Namespace:             t.req.Namespace,
	}
	pkField, err := typeutil.GetPrimaryFieldSchema(t.schema.CollectionSchema)
	if err != nil {
		return nil, segcore.StorageCost{}, err
	}

	var partitionIDs []int64
	if t.partitionKeyMode {
		// multi entities with same pk and diff partition keys may be hashed to multi physical partitions
		// if deleteMsg.partitionID = common.InvalidPartition,
		// all segments with this pk under the collection will have the delete record
		partitionIDs = []int64{common.AllPartitionsID}
		queryReq.PartitionNames = []string{}
	} else {
		// partition name could be defaultPartitionName or name specified by sdk
		partName := t.upsertMsg.DeleteMsg.PartitionName
		if err := validatePartitionTag(partName, true); err != nil {
			log.Warn(ctx, "Invalid partition name", mlog.String("partitionName", partName), mlog.Err(err))
			return nil, segcore.StorageCost{}, err
		}
		partID, err := globalMetaCache.GetPartitionID(ctx, t.req.GetDbName(), t.req.GetCollectionName(), partName)
		if err != nil {
			log.Warn(ctx, "Failed to get partition id", mlog.String("partitionName", partName), mlog.Err(err))
			return nil, segcore.StorageCost{}, err
		}
		partitionIDs = []int64{partID}
		queryReq.PartitionNames = []string{partName}
	}

	plan := planparserv2.CreateRequeryPlan(pkField, ids)
	plan.Namespace = namespaceForPlan(t.schema.CollectionSchema, t.req.Namespace)
	qt := &queryTask{
		ctx:       t.ctx,
		Condition: NewTaskCondition(t.ctx),
		RetrieveRequest: &internalpb.RetrieveRequest{
			Base: commonpbutil.NewMsgBase(
				commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
				commonpbutil.WithSourceID(paramtable.GetNodeID()),
			),
			ReqID:            paramtable.GetNodeID(),
			PartitionIDs:     partitionIDs,
			ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
			QueryLabel:       metrics.UpsertQueryLabel,
		},
		request:        queryReq,
		plan:           plan,
		mixCoord:       t.node.(*Proxy).mixCoord,
		lb:             t.node.(*Proxy).lbPolicy,
		shardclientMgr: t.node.(*Proxy).shardMgr,
		chMgr:          t.node.(*Proxy).chMgr,
	}

	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Upsert-retrieveByPKs")
	defer func() {
		sp.End()
	}()
	queryResult, storageCost, err := t.node.(*Proxy).query(ctx, qt, sp)
	if err := merr.CheckRPCCall(queryResult.GetStatus(), err); err != nil {
		return nil, storageCost, err
	}
	return queryResult, storageCost, err
}

func (it *upsertTask) queryPreExecute(ctx context.Context) error {
	log := mlog.With(mlog.String("collectionName", it.req.CollectionName))

	primaryFieldSchema, err := typeutil.GetPrimaryFieldSchema(it.schema.CollectionSchema)
	if err != nil {
		log.Warn(ctx, "get primary field schema failed", mlog.Err(err))
		return err
	}

	primaryFieldData, err := typeutil.GetPrimaryFieldData(it.req.GetFieldsData(), primaryFieldSchema)
	if err != nil {
		log.Error(ctx, "get primary field data failed", mlog.Err(err))
		return merr.WrapErrParameterInvalidMsg("must assign pk when upsert, primary field: %v", primaryFieldSchema.Name)
	}

	upsertIDs, err := parsePrimaryFieldData2IDs(primaryFieldData)
	if err != nil {
		log.Warn(ctx, "parse primary field data to IDs failed", mlog.Err(err))
		return err
	}

	upsertIDSize := typeutil.GetSizeOfIDs(upsertIDs)
	if upsertIDSize == 0 {
		it.deletePKs = &schemapb.IDs{}
		it.insertFieldData = it.req.GetFieldsData()
		log.Info(ctx, "old records not found, just do insert")
		return nil
	}

	tr := timerecord.NewTimeRecorder("Proxy-Upsert-retrieveByPKs")
	// retrieve by primary key to get original field data
	resp, storageCost, err := retrieveByPKs(ctx, it, upsertIDs, []string{"*"})
	if err != nil {
		log.Info(ctx, "retrieve by primary key failed", mlog.Err(err))
		return err
	}
	it.storageCost = storageCost
	if len(resp.GetFieldsData()) == 0 {
		return merr.WrapErrParameterInvalidMsg("retrieve by primary key failed, no data found")
	}

	existFieldData := resp.GetFieldsData()
	pkFieldData, err := typeutil.GetPrimaryFieldData(existFieldData, primaryFieldSchema)
	if err != nil {
		log.Error(ctx, "get primary field data failed", mlog.Err(err))
		return err
	}
	existIDs, err := parsePrimaryFieldData2IDs(pkFieldData)
	if err != nil {
		log.Info(ctx, "parse primary field data to ids failed", mlog.Err(err))
		return err
	}
	log.Info(ctx, "retrieveByPKs cost",
		mlog.Int("resultNum", typeutil.GetSizeOfIDs(existIDs)),
		mlog.Int64("latency", tr.ElapseSpan().Milliseconds()))

	// set field id for user passed field data, prepare for merge logic
	if len(it.upsertMsg.InsertMsg.GetFieldsData()) == 0 {
		return merr.WrapErrParameterInvalidMsg("upsert field data is empty")
	}
	fieldsDataToCheckAligned := make([]*schemapb.FieldData, 0, len(it.upsertMsg.InsertMsg.GetFieldsData()))
	for _, fieldData := range it.upsertMsg.InsertMsg.GetFieldsData() {
		fieldName := fieldData.GetFieldName()
		if fieldData.GetIsDynamic() {
			fieldName = "$meta"
		}
		if typeutil.IsStructSubField(fieldName) {
			return merr.WrapErrParameterInvalidMsg("partial struct update is not supported for struct sub-field '%s'; use the whole struct field instead", fieldName)
		}
		if structSchema := it.schema.schemaHelper.GetStructArrayFieldFromName(fieldName); structSchema != nil {
			fieldData.FieldId = structSchema.GetFieldID()
			fieldData.FieldName = fieldName
			if err := validateWholeStructFieldDataForPartialUpdate(it.schema.schemaHelper, structSchema, fieldData, upsertIDSize); err != nil {
				return err
			}
			for _, subField := range fieldData.GetStructArrays().GetFields() {
				if len(subField.GetValidData()) != 0 && subFieldHasData(subField) {
					subFieldSchema, err := it.schema.schemaHelper.GetFieldFromName(storedStructSubFieldName(structSchema.GetName(), subField.GetFieldName()))
					if err != nil {
						log.Info(ctx, "get struct sub-field schema failed", mlog.Err(err))
						return err
					}
					if subFieldSchema.GetDefaultValue() != nil {
						err = FillWithDefaultValue(subField, subFieldSchema, int(it.upsertMsg.InsertMsg.NRows()))
					} else {
						err = FillWithNullValue(subField, subFieldSchema, int(it.upsertMsg.InsertMsg.NRows()))
					}
					if err != nil {
						log.Info(ctx, "unify struct field data format failed", mlog.Err(err))
						return err
					}
				}
			}
			continue
		}
		fieldSchema, err := it.schema.schemaHelper.GetFieldFromName(fieldName)
		if err != nil {
			log.Info(ctx, "get field schema failed", mlog.Err(err))
			return err
		}
		fieldData.FieldId = fieldSchema.GetFieldID()
		fieldData.FieldName = fieldName

		// Ensure dynamic field has ValidData before merge logic, but only when
		// the field schema actually requires it (nullable or has default value).
		// For 2.5 collections where $meta is non-nullable with no default,
		// ValidData must remain empty — CheckValidData expects len==0 for
		// non-nullable fields.
		if fieldData.GetIsDynamic() && len(fieldData.GetValidData()) == 0 &&
			(fieldSchema.GetNullable() || fieldSchema.GetDefaultValue() != nil) {
			nRows := int(it.upsertMsg.InsertMsg.NRows())
			validData := make([]bool, nRows)
			for i := range validData {
				validData[i] = true
			}
			fieldData.ValidData = validData
		}

		// compatible with different nullable/default_value data format from sdk
		if len(fieldData.GetValidData()) != 0 {
			var err error
			if fieldSchema.GetDefaultValue() != nil {
				err = FillWithDefaultValue(fieldData, fieldSchema, int(it.upsertMsg.InsertMsg.NRows()))
			} else {
				err = FillWithNullValue(fieldData, fieldSchema, int(it.upsertMsg.InsertMsg.NRows()))
			}
			if err != nil {
				log.Info(ctx, "unify field data format failed", mlog.Err(err))
				return err
			}
		}
		fieldsDataToCheckAligned = append(fieldsDataToCheckAligned, fieldData)
	}

	// Validate field data alignment before processing to prevent index out of range panic
	if err := newValidateUtil().checkAligned(fieldsDataToCheckAligned, it.schema.schemaHelper, uint64(upsertIDSize)); err != nil {
		log.Warn(ctx, "check field data aligned failed", mlog.Err(err))
		return err
	}

	// Scalar nullable payloads are expanded before merge. Nullable vector payloads
	// must remain compact: ValidData tracks logical rows, and vector data stores
	// only valid rows.
	insertIdxInUpsert := make([]int, 0)
	updateIdxInUpsert := make([]int, 0)
	// 1. split upsert data into insert and update by query result
	idsChecker, err := typeutil.NewIDsChecker(existIDs)
	if err != nil {
		log.Info(ctx, "create primary key checker failed", mlog.Err(err))
		return err
	}
	for upsertIdx := 0; upsertIdx < upsertIDSize; upsertIdx++ {
		exist, err := idsChecker.Contains(upsertIDs, upsertIdx)
		if err != nil {
			log.Info(ctx, "check primary key exist in query result failed", mlog.Err(err))
			return err
		}
		if exist {
			updateIdxInUpsert = append(updateIdxInUpsert, upsertIdx)
		} else {
			insertIdxInUpsert = append(insertIdxInUpsert, upsertIdx)
		}
	}

	// 2. merge field data on update semantic
	it.deletePKs = &schemapb.IDs{}
	it.insertFieldData = typeutil.PrepareResultFieldData(existFieldData, int64(upsertIDSize))

	if len(updateIdxInUpsert) > 0 {
		upsertFieldMap := lo.SliceToMap(it.upsertMsg.InsertMsg.GetFieldsData(), func(field *schemapb.FieldData) (int64, *schemapb.FieldData) {
			return field.FieldId, field
		})
		// fieldOpMap resolves per-field FieldPartialUpdateOp directives
		// attached to UpsertRequest.field_ops. Empty / missing entries
		// fall back to REPLACE.
		fieldOpMap := buildFieldOpMap(it.req)

		// Build mapping from existing primary keys to their positions in query result
		// This ensures we can correctly locate data even if query results are not in the same order as request
		existIDsLen := typeutil.GetSizeOfIDs(existIDs)
		existPKToIndex := make(map[interface{}]int, existIDsLen)
		for i := 0; i < existIDsLen; i++ {
			pk := typeutil.GetPK(existIDs, int64(i))
			existPKToIndex[pk] = i
		}

		existIndices := make([]int64, len(updateIdxInUpsert))
		for i, upsertIdx := range updateIdxInUpsert {
			typeutil.AppendIDs(it.deletePKs, upsertIDs, upsertIdx)
			oldPK := typeutil.GetPK(upsertIDs, int64(upsertIdx))
			idx, ok := existPKToIndex[oldPK]
			if !ok {
				return merr.WrapErrParameterInvalidMsg("upsert pk %v not found in query result", oldPK)
			}
			existIndices[i] = int64(idx)
		}
		upsertRows := make([]int64, len(updateIdxInUpsert))
		for i, upsertIdx := range updateIdxInUpsert {
			upsertRows[i] = int64(upsertIdx)
		}

		for fieldIdx, existField := range existFieldData {
			dstField := it.insertFieldData[fieldIdx]
			upsertField := upsertFieldMap[existField.FieldId]

			op := schemapb.FieldPartialUpdateOp_REPLACE
			if fieldOpMap != nil {
				if resolved, ok := fieldOpMap[existField.GetFieldName()]; ok {
					op = resolved
				}
			}

			if it.schema.schemaHelper.GetStructArrayFieldFromName(existField.GetFieldName()) != nil {
				if upsertField != nil {
					if op != schemapb.FieldPartialUpdateOp_REPLACE {
						return merr.WrapErrParameterInvalidMsg("op %s is not supported for struct field %q", op.String(), existField.GetFieldName())
					}
					typeutil.AppendFieldDataByColumn(dstField, upsertField, upsertRows)
				} else {
					typeutil.AppendFieldDataByColumn(dstField, existField, existIndices)
				}
				continue
			}

			fieldSchema, err := it.schema.schemaHelper.GetFieldFromName(existField.GetFieldName())
			if err != nil {
				log.Info(ctx, "get field schema failed", mlog.Err(err))
				return err
			}

			// Note: For fields containing default values, default values need to be set according to valid data during insertion,
			// but query results fields do not set valid data when returning default value fields,
			// therefore valid data needs to be manually set to true.
			if fieldSchema.GetDefaultValue() != nil && len(existField.GetValidData()) == 0 {
				existField.ValidData = make([]bool, existIDsLen)
				for i := range existField.ValidData {
					existField.ValidData[i] = true
				}
			}

			isNullableVector := typeutil.IsCompactNullableVectorFieldData(existField)
			existComputer := typeutil.NewFieldDataIdxComputer([]*schemapb.FieldData{existField})
			existSrcIndices := make([]int64, len(updateIdxInUpsert))
			for i, existIdx := range existIndices {
				existSrcIndices[i] = existComputer.Compute(existIdx)[0]
			}
			if upsertField != nil {
				upsertComputer := typeutil.NewFieldDataIdxComputer([]*schemapb.FieldData{upsertField})
				upsertSrcIndices := make([]int64, len(updateIdxInUpsert))
				for i, upsertIdx := range updateIdxInUpsert {
					upsertSrcIndices[i] = upsertComputer.Compute(int64(upsertIdx))[0]
				}
				if isNullableVector {
					// For nullable vector: only copy data for non-null rows
					upsertRowIndices := make([]int64, len(updateIdxInUpsert))
					validDataIndices := make([]int64, 0, len(updateIdxInUpsert))
					for i, upsertIdx := range updateIdxInUpsert {
						upsertRowIndices[i] = int64(upsertIdx)
						if upsertField.ValidData[upsertIdx] {
							validDataIndices = append(validDataIndices, upsertSrcIndices[i])
						}
					}
					typeutil.AppendFieldDataByColumn(dstField, upsertField, validDataIndices, upsertRowIndices)
				} else {
					typeutil.AppendFieldDataByColumn(dstField, existField, existSrcIndices)
					dstIndices := make([]int64, len(updateIdxInUpsert))
					for i := range dstIndices {
						dstIndices[i] = int64(i)
					}
					if op == schemapb.FieldPartialUpdateOp_REPLACE {
						if err := typeutil.UpdateFieldDataByColumn(dstField, upsertField, dstIndices, upsertSrcIndices); err != nil {
							return err
						}
					} else {
						maxCap := readMaxCapacity(fieldSchema)
						if err := typeutil.UpdateArrayFieldByColumnWithOp(
							dstField, upsertField, dstIndices, upsertSrcIndices, op, maxCap,
						); err != nil {
							return err
						}
					}
				}
			} else {
				if isNullableVector {
					validDataIndices := make([]int64, 0, len(existIndices))
					for i, existIdx := range existIndices {
						if existField.ValidData[int(existIdx)] {
							validDataIndices = append(validDataIndices, existSrcIndices[i])
						}
					}
					typeutil.AppendFieldDataByColumn(dstField, existField, validDataIndices, existIndices)
				} else {
					typeutil.AppendFieldDataByColumn(dstField, existField, existSrcIndices)
				}
			}
		}
	}

	// 3. merge field data on insert semantic
	if len(insertIdxInUpsert) > 0 {
		// if necessary field is not exist in upsert request, return error
		lackOfFieldErr := LackOfFieldsDataBySchema(it.schema.CollectionSchema, it.upsertMsg.InsertMsg.GetFieldsData(), false, true)
		if lackOfFieldErr != nil {
			log.Info(ctx, "check fields data by schema failed", mlog.Err(lackOfFieldErr))
			return lackOfFieldErr
		}

		// Build field data for rows that take insert semantics. Request-provided
		// fields are kept as-is; missing nullable/default fields are synthesized
		// so the insert rows still have a complete field set.
		insertFieldsToAppend := make([]*schemapb.FieldData, 0)
		upsertFieldMap := lo.SliceToMap(it.upsertMsg.InsertMsg.GetFieldsData(), func(field *schemapb.FieldData) (string, *schemapb.FieldData) {
			return field.GetFieldName(), field
		})
		for _, fieldSchema := range it.schema.Fields {
			if fieldData, ok := upsertFieldMap[fieldSchema.Name]; !ok {
				if fieldSchema.GetNullable() || fieldSchema.GetDefaultValue() != nil {
					fieldData, err := GenNullableFieldData(fieldSchema, upsertIDSize)
					if err != nil {
						log.Info(ctx, "generate nullable field data failed", mlog.Err(err))
						return err
					}
					insertFieldsToAppend = append(insertFieldsToAppend, fieldData)
				}
			} else {
				insertFieldsToAppend = append(insertFieldsToAppend, fieldData)
			}
		}
		for _, structSchema := range it.schema.GetStructArrayFields() {
			if fieldData, ok := upsertFieldMap[structSchema.GetName()]; ok {
				insertFieldsToAppend = append(insertFieldsToAppend, fieldData)
				continue
			}
			if structSchema.GetNullable() {
				insertFieldsToAppend = append(insertFieldsToAppend, GenNullableStructArrayFieldData(structSchema, upsertIDSize))
			}
		}

		// Build mapping from FieldId to index in it.insertFieldData
		insertFieldDataMap := make(map[int64]int, len(it.insertFieldData))
		for i, fd := range it.insertFieldData {
			insertFieldDataMap[fd.FieldId] = i
		}

		// Process insert data by column (field), similar to update path
		for _, srcField := range insertFieldsToAppend {
			isNullableVector := typeutil.IsCompactNullableVectorFieldData(srcField)
			srcComputer := typeutil.NewFieldDataIdxComputer([]*schemapb.FieldData{srcField})

			// Find or create destination field in it.insertFieldData
			var dstField *schemapb.FieldData
			if idx, ok := insertFieldDataMap[srcField.FieldId]; ok {
				dstField = it.insertFieldData[idx]
			} else {
				// New field not in existFieldData, create empty field data and append
				newFields := typeutil.PrepareResultFieldData([]*schemapb.FieldData{srcField}, int64(len(insertIdxInUpsert)))
				dstField = newFields[0]
				it.insertFieldData = append(it.insertFieldData, dstField)
			}

			if isNullableVector {
				rowIndices := make([]int64, len(insertIdxInUpsert))
				validDataIndices := make([]int64, 0, len(insertIdxInUpsert))
				for i, upsertIdx := range insertIdxInUpsert {
					rowIndices[i] = int64(upsertIdx)
					srcIdx := srcComputer.Compute(int64(upsertIdx))[0]
					if srcField.ValidData[upsertIdx] {
						validDataIndices = append(validDataIndices, srcIdx)
					}
				}
				typeutil.AppendFieldDataByColumn(dstField, srcField, validDataIndices, rowIndices)
			} else {
				srcIndices := make([]int64, len(insertIdxInUpsert))
				for i, upsertIdx := range insertIdxInUpsert {
					srcIndices[i] = srcComputer.Compute(int64(upsertIdx))[0]
				}
				typeutil.AppendFieldDataByColumn(dstField, srcField, srcIndices)
			}
		}
	}

	for _, fieldData := range it.insertFieldData {
		if fieldData.GetType() == schemapb.DataType_ArrayOfStruct {
			if err := ToCompressedFormatNullableStructField(fieldData); err != nil {
				log.Info(ctx, "convert struct field data to compressed format nullable failed", mlog.Err(err))
				return err
			}
			continue
		}
		if len(fieldData.GetValidData()) > 0 {
			err := ToCompressedFormatNullable(fieldData)
			if err != nil {
				log.Info(ctx, "convert to compressed format nullable failed", mlog.Err(err))
				return err
			}
		}
	}
	return nil
}

// ToCompressedFormatNullable converts nullable field data from full format to compressed format.
func ToCompressedFormatNullable(field *schemapb.FieldData) error {
	if getValidNumber(field.GetValidData()) == len(field.GetValidData()) {
		return nil
	}
	switch field.Field.(type) {
	case *schemapb.FieldData_Scalars:
		switch sd := field.GetScalars().GetData().(type) {
		case *schemapb.ScalarField_BoolData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.BoolData.Data = make([]bool, 0)
			} else {
				ret := make([]bool, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.BoolData.Data[i])
					}
				}
				sd.BoolData.Data = ret
			}

		case *schemapb.ScalarField_IntData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.IntData.Data = make([]int32, 0)
			} else {
				ret := make([]int32, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.IntData.Data[i])
					}
				}
				sd.IntData.Data = ret
			}

		case *schemapb.ScalarField_LongData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.LongData.Data = make([]int64, 0)
			} else {
				ret := make([]int64, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.LongData.Data[i])
					}
				}
				sd.LongData.Data = ret
			}

		case *schemapb.ScalarField_FloatData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.FloatData.Data = make([]float32, 0)
			} else {
				ret := make([]float32, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.FloatData.Data[i])
					}
				}
				sd.FloatData.Data = ret
			}

		case *schemapb.ScalarField_DoubleData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.DoubleData.Data = make([]float64, 0)
			} else {
				ret := make([]float64, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.DoubleData.Data[i])
					}
				}
				sd.DoubleData.Data = ret
			}

		case *schemapb.ScalarField_StringData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.StringData.Data = make([]string, 0)
			} else {
				ret := make([]string, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.StringData.Data[i])
					}
				}
				sd.StringData.Data = ret
			}

		case *schemapb.ScalarField_JsonData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.JsonData.Data = make([][]byte, 0)
			} else {
				ret := make([][]byte, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.JsonData.Data[i])
					}
				}
				sd.JsonData.Data = ret
			}

		case *schemapb.ScalarField_ArrayData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.ArrayData.Data = make([]*schemapb.ScalarField, 0)
			} else {
				ret := make([]*schemapb.ScalarField, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.ArrayData.Data[i])
					}
				}
				sd.ArrayData.Data = ret
			}
		case *schemapb.ScalarField_TimestamptzData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.TimestamptzData.Data = make([]int64, 0)
			} else {
				ret := make([]int64, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.TimestamptzData.Data[i])
					}
				}
				sd.TimestamptzData.Data = ret
			}

		case *schemapb.ScalarField_GeometryWktData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.GeometryWktData.Data = make([]string, 0)
			} else {
				ret := make([]string, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.GeometryWktData.Data[i])
					}
				}
				sd.GeometryWktData.Data = ret
			}

		case *schemapb.ScalarField_GeometryData:
			validRowNum := getValidNumber(field.GetValidData())
			if validRowNum == 0 {
				sd.GeometryData.Data = make([][]byte, 0)
			} else {
				ret := make([][]byte, 0, validRowNum)
				for i, valid := range field.GetValidData() {
					if valid {
						ret = append(ret, sd.GeometryData.Data[i])
					}
				}
				sd.GeometryData.Data = ret
			}

		default:
			return merr.WrapErrParameterInvalidMsg("undefined data type:%s", field.Type.String())
		}

	case *schemapb.FieldData_Vectors:
		// Vector data is already in compressed format, skip
		return nil

	default:
		return merr.WrapErrParameterInvalidMsg("undefined data type:%s", field.Type.String())
	}

	return nil
}

// GenNullableFieldData generates all-null nullable field data.
// Scalar fields use expanded zero values; vector fields use compact empty data.
func GenNullableFieldData(field *schemapb.FieldSchema, upsertIDSize int) (*schemapb.FieldData, error) {
	switch field.DataType {
	case schemapb.DataType_Bool:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_BoolData{
						BoolData: &schemapb.BoolArray{
							Data: make([]bool, upsertIDSize),
						},
					},
				},
			},
		}, nil

	case schemapb.DataType_Int32:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_IntData{
						IntData: &schemapb.IntArray{
							Data: make([]int32, upsertIDSize),
						},
					},
				},
			},
		}, nil

	case schemapb.DataType_Int64:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_LongData{
						LongData: &schemapb.LongArray{
							Data: make([]int64, upsertIDSize),
						},
					},
				},
			},
		}, nil

	case schemapb.DataType_Float:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_FloatData{
						FloatData: &schemapb.FloatArray{
							Data: make([]float32, upsertIDSize),
						},
					},
				},
			},
		}, nil

	case schemapb.DataType_Double:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_DoubleData{
						DoubleData: &schemapb.DoubleArray{
							Data: make([]float64, upsertIDSize),
						},
					},
				},
			},
		}, nil

	case schemapb.DataType_VarChar:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_StringData{
						StringData: &schemapb.StringArray{
							Data: make([]string, upsertIDSize),
						},
					},
				},
			},
		}, nil

	case schemapb.DataType_JSON:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_JsonData{
						JsonData: &schemapb.JSONArray{
							Data: make([][]byte, upsertIDSize),
						},
					},
				},
			},
		}, nil

	case schemapb.DataType_Array:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_ArrayData{
						ArrayData: &schemapb.ArrayArray{
							Data:        make([]*schemapb.ScalarField, upsertIDSize),
							ElementType: field.GetElementType(),
						},
					},
				},
			},
		}, nil

	case schemapb.DataType_Timestamptz:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_TimestamptzData{
						TimestamptzData: &schemapb.TimestamptzArray{
							Data: make([]int64, upsertIDSize),
						},
					},
				},
			},
		}, nil

	// the intput data of geometry field is in wkt format
	case schemapb.DataType_Geometry:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			IsDynamic: field.IsDynamic,
			ValidData: make([]bool, upsertIDSize),
			Field: &schemapb.FieldData_Scalars{
				Scalars: &schemapb.ScalarField{
					Data: &schemapb.ScalarField_GeometryWktData{
						GeometryWktData: &schemapb.GeometryWktArray{
							Data: make([]string, upsertIDSize),
						},
					},
				},
			},
		}, nil

	// Nullable vector types
	case schemapb.DataType_FloatVector:
		dim, err := typeutil.GetDim(field)
		if err != nil {
			return nil, err
		}
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			ValidData: make([]bool, upsertIDSize), // all false = all null
			Field: &schemapb.FieldData_Vectors{
				Vectors: &schemapb.VectorField{
					Dim:  dim,
					Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{}}},
				},
			},
		}, nil

	case schemapb.DataType_Float16Vector:
		dim, err := typeutil.GetDim(field)
		if err != nil {
			return nil, err
		}
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			ValidData: make([]bool, upsertIDSize), // all false = all null
			Field: &schemapb.FieldData_Vectors{
				Vectors: &schemapb.VectorField{
					Dim:  dim,
					Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{}},
				},
			},
		}, nil

	case schemapb.DataType_BFloat16Vector:
		dim, err := typeutil.GetDim(field)
		if err != nil {
			return nil, err
		}
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			ValidData: make([]bool, upsertIDSize), // all false = all null
			Field: &schemapb.FieldData_Vectors{
				Vectors: &schemapb.VectorField{
					Dim:  dim,
					Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{}},
				},
			},
		}, nil

	case schemapb.DataType_BinaryVector:
		dim, err := typeutil.GetDim(field)
		if err != nil {
			return nil, err
		}
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			ValidData: make([]bool, upsertIDSize), // all false = all null
			Field: &schemapb.FieldData_Vectors{
				Vectors: &schemapb.VectorField{
					Dim:  dim,
					Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{}},
				},
			},
		}, nil

	case schemapb.DataType_SparseFloatVector:
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			ValidData: make([]bool, upsertIDSize), // all false = all null
			Field: &schemapb.FieldData_Vectors{
				Vectors: &schemapb.VectorField{
					Data: &schemapb.VectorField_SparseFloatVector{SparseFloatVector: &schemapb.SparseFloatArray{
						Contents: [][]byte{},
					}},
				},
			},
		}, nil

	case schemapb.DataType_Int8Vector:
		dim, err := typeutil.GetDim(field)
		if err != nil {
			return nil, err
		}
		return &schemapb.FieldData{
			FieldId:   field.FieldID,
			FieldName: field.Name,
			Type:      field.DataType,
			ValidData: make([]bool, upsertIDSize), // all false = all null
			Field: &schemapb.FieldData_Vectors{
				Vectors: &schemapb.VectorField{
					Dim:  dim,
					Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{}},
				},
			},
		}, nil

	default:
		return nil, merr.WrapErrParameterInvalidMsg("undefined data type:%s", field.DataType.String())
	}
}

func GenNullableStructArrayFieldData(structField *schemapb.StructArrayFieldSchema, rowCount int) *schemapb.FieldData {
	fields := make([]*schemapb.FieldData, 0, len(structField.GetFields()))
	for _, subField := range structField.GetFields() {
		fieldName := subField.GetName()
		if typeutil.IsStructSubField(fieldName) {
			if originalName, err := typeutil.ExtractStructFieldName(fieldName); err == nil {
				fieldName = originalName
			}
		}
		fields = append(fields, &schemapb.FieldData{
			FieldName: fieldName,
			FieldId:   subField.GetFieldID(),
			Type:      subField.GetDataType(),
			// The source field is indexed by original upsert row IDs before
			// appending the insert subset, so ValidData uses the same row domain.
			ValidData: make([]bool, rowCount),
		})
	}
	return &schemapb.FieldData{
		FieldId:   structField.GetFieldID(),
		FieldName: structField.GetName(),
		Type:      schemapb.DataType_ArrayOfStruct,
		Field: &schemapb.FieldData_StructArrays{
			StructArrays: &schemapb.StructArrayField{Fields: fields},
		},
	}
}

// ToCompressedFormatNullableStructField converts nullable struct sub-fields from full format to compressed format.
func ToCompressedFormatNullableStructField(fieldData *schemapb.FieldData) error {
	structArrays := fieldData.GetStructArrays()
	if structArrays == nil {
		return nil
	}
	for _, subField := range structArrays.GetFields() {
		if len(subField.GetValidData()) == 0 {
			continue
		}

		validData := subField.GetValidData()
		validRows := getValidNumber(validData)
		if validRows == 0 && !subFieldHasData(subField) {
			continue
		}
		payloadRows, ok := structSubFieldPayloadRowsForPartialUpdate(subField)
		if !ok {
			return merr.WrapErrParameterInvalidMsg("struct sub-field '%s' has invalid field data", subField.GetFieldName())
		}
		if payloadRows == validRows {
			continue
		}
		if payloadRows != len(validData) {
			return merr.WrapErrParameterInvalidMsg("struct sub-field '%s' payload row count %d does not match valid rows %d or logical rows %d",
				subField.GetFieldName(), payloadRows, validRows, len(validData))
		}
		if subField.GetType() == schemapb.DataType_ArrayOfVector {
			vectorArray := subField.GetVectors().GetVectorArray()
			compactRows := make([]*schemapb.VectorField, 0, validRows)
			for row, valid := range validData {
				if valid {
					compactRows = append(compactRows, vectorArray.GetData()[row])
				}
			}
			vectorArray.Data = compactRows
			continue
		}
		if err := ToCompressedFormatNullable(subField); err != nil {
			return err
		}
	}
	return nil
}

// validateWholeStructFieldDataForPartialUpdate validates only the logical
// top-level struct payload used by partial-update whole struct REPLACE.
// It checks:
// 1. ArrayOfStruct wrapper and sub-field count/name/type.
// 2. All-or-none sub-field payload presence and nullable struct null-mask consistency.
// 3. Payload row counts needed for safe row append.
// It does not flatten struct data or merge sub-fields; final insert
// pre-execute still runs checkAndFlattenStructFieldData before storage writes.
func validateWholeStructFieldDataForPartialUpdate(schemaHelper *typeutil.SchemaHelper, structSchema *schemapb.StructArrayFieldSchema, fieldData *schemapb.FieldData, rowCount int) error {
	if fieldData.GetType() != schemapb.DataType_ArrayOfStruct {
		return merr.WrapErrParameterInvalidMsg("field %q expects ArrayOfStruct payload, got %s", fieldData.GetFieldName(), fieldData.GetType().String())
	}
	structArrays := fieldData.GetStructArrays()
	if structArrays == nil {
		return merr.WrapErrParameterInvalidMsg("field %q expects non-nil StructArrays payload", fieldData.GetFieldName())
	}
	if len(structArrays.GetFields()) != len(structSchema.GetFields()) {
		return merr.WrapErrParameterInvalidMsg("length of fields of struct field mismatch length of the fields in schema, fieldName: %s, fieldData fields length:%d, schema fields length:%d",
			fieldData.GetFieldName(), len(structArrays.GetFields()), len(structSchema.GetFields()))
	}

	seenSubFieldIDs := make(map[int64]struct{}, len(structArrays.GetFields()))
	for _, subField := range structArrays.GetFields() {
		subFieldSchema, err := schemaHelper.GetFieldFromName(storedStructSubFieldName(structSchema.GetName(), subField.GetFieldName()))
		if err != nil {
			return err
		}
		subField.FieldId = subFieldSchema.GetFieldID()
		if _, ok := seenSubFieldIDs[subFieldSchema.GetFieldID()]; ok {
			return merr.WrapErrParameterInvalidMsg("duplicated sub-field '%s' in struct '%s'", subField.GetFieldName(), fieldData.GetFieldName())
		}
		seenSubFieldIDs[subFieldSchema.GetFieldID()] = struct{}{}
		if subField.GetType() != subFieldSchema.GetDataType() {
			return merr.WrapErrParameterInvalidMsg("sub-field '%s' in struct '%s' expects type %s, got %s",
				subField.GetFieldName(), fieldData.GetFieldName(), subFieldSchema.GetDataType().String(), subField.GetType().String())
		}
		if len(subField.GetValidData()) > 0 && len(subField.GetValidData()) != rowCount {
			return merr.WrapErrParameterInvalidMsg("sub-field ValidData length mismatch in struct '%s': '%s' has %d, expected %d",
				fieldData.GetFieldName(), subField.GetFieldName(), len(subField.GetValidData()), rowCount)
		}
	}

	hasDataCount := 0
	for _, subField := range structArrays.GetFields() {
		if subFieldHasData(subField) {
			hasDataCount++
		}
	}
	// All sub-fields must be either present together or absent together.
	// A partial payload would mean partial struct replace, which is not
	// supported by this path.
	if hasDataCount == 0 {
		for _, subField := range structArrays.GetFields() {
			if len(subField.GetValidData()) != 0 && len(subField.GetValidData()) != rowCount {
				return merr.WrapErrParameterInvalidMsg("sub-field ValidData length mismatch in struct '%s': '%s' has %d, expected %d",
					fieldData.GetFieldName(), subField.GetFieldName(), len(subField.GetValidData()), rowCount)
			}
			for j, valid := range subField.GetValidData() {
				if valid {
					return merr.WrapErrParameterInvalidMsg("sub-field '%s' in struct '%s' claims row %d is valid but no payload is provided",
						subField.GetFieldName(), fieldData.GetFieldName(), j)
				}
			}
		}
		return nil
	}
	if hasDataCount != len(structArrays.GetFields()) {
		return merr.WrapErrParameterInvalidMsg("inconsistent sub-field data in struct '%s': %d of %d sub-fields have data, all must be present or all absent",
			fieldData.GetFieldName(), hasDataCount, len(structArrays.GetFields()))
	}

	var refValidData []bool
	var refFieldName string
	refInitialized := false
	for _, subField := range structArrays.GetFields() {
		// Nullable is a property of the whole struct row. The sub-fields
		// therefore must share the same null mask.
		if structSchema.GetNullable() {
			if !refInitialized {
				refValidData = subField.GetValidData()
				refFieldName = subField.GetFieldName()
				refInitialized = true
			} else {
				if len(subField.GetValidData()) != len(refValidData) {
					return merr.WrapErrParameterInvalidMsg("sub-field ValidData length mismatch in struct '%s': '%s' has %d, '%s' has %d",
						fieldData.GetFieldName(), refFieldName, len(refValidData), subField.GetFieldName(), len(subField.GetValidData()))
				}
				for j := range refValidData {
					if subField.GetValidData()[j] != refValidData[j] {
						return merr.WrapErrParameterInvalidMsg("sub-field ValidData mismatch in struct '%s' at row %d: '%s'=%v, '%s'=%v",
							fieldData.GetFieldName(), j, refFieldName, refValidData[j], subField.GetFieldName(), subField.GetValidData()[j])
					}
				}
			}
		}

		// Struct sub-fields without ValidData use dense logical rows. With
		// ValidData, the payload must be compact so FillWithNullValue can expand
		// it consistently before merge.
		payloadRows, ok := structSubFieldPayloadRowsForPartialUpdate(subField)
		if !ok {
			return merr.WrapErrParameterInvalidMsg("struct sub-field '%s' has invalid field data", subField.GetFieldName())
		}
		if len(subField.GetValidData()) == 0 {
			if payloadRows != rowCount {
				return merr.WrapErrParameterInvalidMsg("struct sub-field '%s' payload row count %d does not match logical rows %d",
					subField.GetFieldName(), payloadRows, rowCount)
			}
			continue
		}
		validRows := getValidNumber(subField.GetValidData())
		if payloadRows != validRows {
			return merr.WrapErrParameterInvalidMsg("nullable struct sub-field '%s' payload must be compact: payload row count %d does not match valid rows %d",
				subField.GetFieldName(), payloadRows, validRows)
		}
	}
	return nil
}

func structSubFieldPayloadRowsForPartialUpdate(fieldData *schemapb.FieldData) (int, bool) {
	switch fieldData.GetType() {
	case schemapb.DataType_Array:
		scalars := fieldData.GetScalars()
		if scalars == nil || scalars.GetArrayData() == nil {
			return 0, false
		}
		return len(scalars.GetArrayData().GetData()), true
	case schemapb.DataType_ArrayOfVector:
		vectors := fieldData.GetVectors()
		if vectors == nil || vectors.GetVectorArray() == nil {
			return 0, false
		}
		return len(vectors.GetVectorArray().GetData()), true
	default:
		return 0, false
	}
}

func (it *upsertTask) insertPreExecute(ctx context.Context) error {
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Upsert-insertPreExecute")
	defer sp.End()
	collectionName := it.upsertMsg.InsertMsg.CollectionName
	log := mlog.With(mlog.String("collectionName", collectionName))
	if err := validateCollectionName(collectionName); err != nil {
		log.Error(ctx, "valid collection name failed", mlog.String("collectionName", collectionName), mlog.Err(err))
		return err
	}

	bm25Fields := GetBM25FunctionOutputFields(it.schema.CollectionSchema)
	minHashFields := GetMinHashFunctionOutputFields(it.schema.CollectionSchema)
	fieldsToFilter := typeutil.NewSet[string](append(bm25Fields, minHashFields...)...)
	if it.req.PartialUpdate {
		// remove old BM25 and MinHash function output fields, which will be regenerated downstream
		// Note: do NOT filter other function outputs (e.g. text embedding) as they are generated
		// by genFunctionFields and expected by validateFieldDataColumns
		ret := make([]*schemapb.FieldData, 0)
		for _, fieldData := range it.upsertMsg.InsertMsg.GetFieldsData() {
			if fieldsToFilter.Contain(fieldData.GetFieldName()) {
				continue
			}
			ret = append(ret, fieldData)
		}
		it.upsertMsg.InsertMsg.FieldsData = ret
	}

	rowNums := uint32(it.upsertMsg.InsertMsg.NRows())
	// set upsertTask.insertRequest.rowIDs
	tr := timerecord.NewTimeRecorder("applyPK")
	clusterID := Params.CommonCfg.ClusterID.GetAsUint64()
	rowIDBegin, rowIDEnd, allocateErr := common.AllocAutoID(it.idAllocator.Alloc, rowNums, clusterID)
	metrics.ProxyApplyPrimaryKeyLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10)).Observe(float64(tr.ElapseSpan().Microseconds()) / 1000.0)
	if allocateErr != nil {
		log.Warn(ctx, "failed to allocate auto id for upsert",
			mlog.String("collectionName", collectionName),
			mlog.Int64("collectionID", it.upsertMsg.InsertMsg.CollectionID),
			mlog.Uint32("rowNums", rowNums),
			mlog.Err(allocateErr))
		return allocateErr
	}

	it.upsertMsg.InsertMsg.RowIDs = make([]UniqueID, rowNums)
	it.rowIDs = make([]UniqueID, rowNums)
	for i := rowIDBegin; i < rowIDEnd; i++ {
		offset := i - rowIDBegin
		it.upsertMsg.InsertMsg.RowIDs[offset] = i
		it.rowIDs[offset] = i
	}
	// set upsertTask.insertRequest.timeStamps
	rowNum := it.upsertMsg.InsertMsg.NRows()
	it.upsertMsg.InsertMsg.Timestamps = make([]uint64, rowNum)
	it.timestamps = make([]uint64, rowNum)
	for index := range it.timestamps {
		it.upsertMsg.InsertMsg.Timestamps[index] = it.BeginTs()
		it.timestamps[index] = it.BeginTs()
	}
	// set result.SuccIndex
	sliceIndex := make([]uint32, rowNums)
	for i := uint32(0); i < rowNums; i++ {
		sliceIndex[i] = i
	}
	it.result.SuccIndex = sliceIndex

	var err error
	if it.schema.EnableDynamicField {
		var err error
		if it.req.GetPartialUpdate() {
			err = checkDynamicFieldDataForPartialUpdate(it.schema.CollectionSchema, it.upsertMsg.InsertMsg)
		} else {
			err = checkDynamicFieldData(it.schema.CollectionSchema, it.upsertMsg.InsertMsg)
		}
		if err != nil {
			return err
		}
	}

	err = addNamespaceData(it.schema.CollectionSchema, it.upsertMsg.InsertMsg)
	if err != nil {
		return err
	}

	if err = checkAndFlattenStructFieldData(it.schema.CollectionSchema, it.upsertMsg.InsertMsg); err != nil {
		return err
	}

	allFields := typeutil.GetAllFieldSchemas(it.schema.CollectionSchema)

	// use the passed pk as new pk when autoID == false
	// automatic generate pk as new pk wehen autoID == true
	it.result.IDs, it.oldIDs, err = checkUpsertPrimaryFieldData(ctx, allFields, it.schema.CollectionSchema, it.upsertMsg.InsertMsg)
	if err != nil {
		log.Warn(ctx, "check primary field data and hash primary key failed when upsert",
			mlog.Err(err))
		return merr.WrapErrAsInputErrorWhen(err, merr.ErrParameterInvalid)
	}

	// check varchar/text with analyzer was utf-8 format
	err = checkInputUtf8Compatiable(allFields, it.upsertMsg.InsertMsg)
	if err != nil {
		log.Warn(ctx, "check varchar/text format failed", mlog.Err(err))
		return err
	}

	// Validate and set field ID to insert field data
	err = validateFieldDataColumns(it.upsertMsg.InsertMsg.GetFieldsData(), it.schema)
	if err != nil {
		log.Warn(ctx, "validate field data columns failed when upsert", mlog.Err(err))
		return merr.WrapErrAsInputErrorWhen(err, merr.ErrParameterInvalid)
	}
	err = fillFieldPropertiesOnly(it.upsertMsg.InsertMsg.GetFieldsData(), it.schema)
	if err != nil {
		log.Warn(ctx, "fill field properties failed when upsert", mlog.Err(err))
		return merr.WrapErrAsInputErrorWhen(err, merr.ErrParameterInvalid)
	}
	err = normalizeFP32ToFP16BF16VectorFieldData(it.upsertMsg.InsertMsg.GetFieldsData(), it.schema)
	if err != nil {
		log.Warn(ctx, "normalize fp32 to fp16/bf16 vector field data failed when upsert", mlog.Err(err))
		return merr.WrapErrAsInputErrorWhen(err, merr.ErrParameterInvalid)
	}

	if it.partitionKeyMode {
		fieldSchema, _ := typeutil.GetPartitionKeyFieldSchema(it.schema.CollectionSchema)
		it.partitionKeys, err = getPartitionKeyFieldData(fieldSchema, it.upsertMsg.InsertMsg)
		if err != nil {
			log.Warn(ctx, "get partition keys from insert request failed",
				mlog.String("collectionName", collectionName),
				mlog.Err(err))
			return err
		}
	} else {
		partitionTag := it.upsertMsg.InsertMsg.PartitionName
		if err = validatePartitionTag(partitionTag, true); err != nil {
			log.Warn(ctx, "valid partition name failed", mlog.String("partition name", partitionTag), mlog.Err(err))
			return err
		}
	}

	if err := newValidateUtil(withNANCheck(), withOverflowCheck(), withMaxLenCheck(), withMaxCapCheck()).
		Validate(it.upsertMsg.InsertMsg.GetFieldsData(), it.schema.schemaHelper, it.upsertMsg.InsertMsg.NRows()); err != nil {
		return err
	}

	log.Debug(ctx, "Proxy Upsert insertPreExecute done")

	return nil
}

func (it *upsertTask) deletePreExecute(ctx context.Context) error {
	collName := it.upsertMsg.DeleteMsg.CollectionName
	log := mlog.With(
		mlog.String("collectionName", collName))

	if it.upsertMsg.DeleteMsg.PrimaryKeys == nil {
		// if primary keys are not set by queryPreExecute, use oldIDs to delete all given records
		it.upsertMsg.DeleteMsg.PrimaryKeys = it.oldIDs
	}
	if typeutil.GetSizeOfIDs(it.upsertMsg.DeleteMsg.PrimaryKeys) == 0 {
		log.Info(ctx, "deletePKs is empty, skip deleteExecute")
		return nil
	}

	if err := validateCollectionName(collName); err != nil {
		log.Info(ctx, "Invalid collectionName", mlog.Err(err))
		return err
	}

	if it.partitionKeyMode {
		// multi entities with same pk and diff partition keys may be hashed to multi physical partitions
		// if deleteMsg.partitionID = common.InvalidPartition,
		// all segments with this pk under the collection will have the delete record
		it.upsertMsg.DeleteMsg.PartitionID = common.AllPartitionsID
	} else {
		// partition name could be defaultPartitionName or name specified by sdk
		partName := it.upsertMsg.DeleteMsg.PartitionName
		if err := validatePartitionTag(partName, true); err != nil {
			log.Warn(ctx, "Invalid partition name", mlog.String("partitionName", partName), mlog.Err(err))
			return err
		}
		partID, err := globalMetaCache.GetPartitionID(ctx, it.req.GetDbName(), collName, partName)
		if err != nil {
			log.Warn(ctx, "Failed to get partition id", mlog.String("collectionName", collName), mlog.String("partitionName", partName), mlog.Err(err))
			return err
		}
		it.upsertMsg.DeleteMsg.PartitionID = partID
	}

	it.upsertMsg.DeleteMsg.Timestamps = make([]uint64, it.upsertMsg.DeleteMsg.NumRows)
	for index := range it.upsertMsg.DeleteMsg.Timestamps {
		it.upsertMsg.DeleteMsg.Timestamps[index] = it.BeginTs()
	}
	log.Debug(ctx, "Proxy Upsert deletePreExecute done")
	return nil
}

func (it *upsertTask) PreExecute(ctx context.Context) error {
	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Upsert-PreExecute")
	defer sp.End()

	collectionName := it.req.CollectionName
	log := mlog.With(mlog.String("collectionName", collectionName))

	it.result = &milvuspb.MutationResult{
		Status: merr.Success(),
		IDs: &schemapb.IDs{
			IdField: nil,
		},
		Timestamp: it.EndTs(),
	}

	// check collection exists
	collID, err := globalMetaCache.GetCollectionID(context.Background(), it.req.GetDbName(), collectionName)
	if err != nil {
		log.Warn(ctx, "fail to get collection id", mlog.Err(err))
		return err
	}
	it.collectionID = collID

	colInfo, err := globalMetaCache.GetCollectionInfo(ctx, it.req.GetDbName(), collectionName, collID)
	if err != nil {
		log.Warn(ctx, "fail to get collection info", mlog.Err(err))
		return err
	}

	if it.schemaTimestamp != 0 {
		if it.schemaTimestamp != colInfo.updateTimestamp {
			err := merr.WrapErrCollectionSchemaMisMatch(collectionName)
			log.Info(ctx, "collection schema mismatch", mlog.String("collectionName", collectionName),
				mlog.Uint64("requestSchemaTs", it.schemaTimestamp),
				mlog.Uint64("collectionSchemaTs", colInfo.updateTimestamp),
				mlog.Err(err))
			return err
		}
	}

	schema, err := globalMetaCache.GetCollectionSchema(ctx, it.req.GetDbName(), collectionName)
	if err != nil {
		log.Warn(ctx, "Failed to get collection schema",
			mlog.String("collectionName", collectionName),
			mlog.Err(err))
		return err
	}
	it.schema = schema
	it.schemaVersion = schema.Version
	if err := validateTextStorageV3Enabled(schema.CollectionSchema); err != nil {
		return err
	}

	// Validate any FieldPartialUpdateOp directives attached to FieldData.
	// A non-REPLACE op implicitly promotes the request to partial_update=true
	// so users do not need to set both fields explicitly.
	nonReplaceSeen, err := validateFieldPartialUpdateOps(it.req, schema.CollectionSchema)
	if err != nil {
		log.Warn(ctx, "validate field partial update ops failed", mlog.Err(err))
		return err
	}
	if nonReplaceSeen && !it.req.GetPartialUpdate() {
		it.req.PartialUpdate = true
	}

	partitionName, namespaceAsPartition, err := resolveNamespacePartitionName(schema.CollectionSchema, it.req.Namespace, it.req.GetPartitionName())
	if err != nil {
		return err
	}
	if namespaceAsPartition {
		it.req.PartitionName = partitionName
	}

	it.partitionKeyMode, err = isPartitionKeyMode(ctx, it.req.GetDbName(), collectionName)
	if err != nil {
		log.Warn(ctx, "check partition key mode failed",
			mlog.String("collectionName", collectionName),
			mlog.Err(err))
		return err
	}
	if it.partitionKeyMode {
		if len(it.req.GetPartitionName()) > 0 {
			return merr.WrapErrParameterInvalidMsg("not support manually specifying the partition names if partition key mode is used")
		}
	} else {
		// set default partition name if not use partition key
		// insert to _default partition
		partitionTag := it.req.GetPartitionName()
		if len(partitionTag) <= 0 {
			pinfo, err := globalMetaCache.GetPartitionInfo(ctx, it.req.GetDbName(), collectionName, "")
			if err != nil {
				log.Warn(ctx, "get partition info failed", mlog.String("collectionName", collectionName), mlog.Err(err))
				return err
			}
			it.req.PartitionName = pinfo.name
		}
	}

	// check for duplicate primary keys in the same batch
	primaryFieldSchema, err := typeutil.GetPrimaryFieldSchema(schema.CollectionSchema)
	if err != nil {
		log.Warn(ctx, "fail to get primary field schema", mlog.Err(err))
		return err
	}
	duplicate, err := CheckDuplicatePkExist(primaryFieldSchema, it.req.GetFieldsData())
	if err != nil {
		log.Warn(ctx, "fail to check duplicate primary keys", mlog.Err(err))
		return err
	}
	if duplicate {
		return merr.WrapErrParameterInvalidMsg("duplicate primary keys are not allowed in the same batch")
	}

	it.upsertMsg = &msgstream.UpsertMsg{
		InsertMsg: &msgstream.InsertMsg{
			InsertRequest: &msgpb.InsertRequest{
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Insert),
					commonpbutil.WithSourceID(paramtable.GetNodeID()),
				),
				CollectionName: it.req.CollectionName,
				CollectionID:   it.collectionID,
				PartitionName:  it.req.PartitionName,
				FieldsData:     it.req.FieldsData,
				NumRows:        uint64(it.req.NumRows),
				Version:        msgpb.InsertDataVersion_ColumnBased,
				DbName:         it.req.DbName,
				Namespace:      it.req.Namespace,
			},
		},
		DeleteMsg: &msgstream.DeleteMsg{
			DeleteRequest: &msgpb.DeleteRequest{
				Base: commonpbutil.NewMsgBase(
					commonpbutil.WithMsgType(commonpb.MsgType_Delete),
					commonpbutil.WithSourceID(paramtable.GetNodeID()),
				),
				DbName:         it.req.DbName,
				CollectionName: it.req.CollectionName,
				CollectionID:   it.collectionID,
				NumRows:        int64(it.req.NumRows),
				PartitionName:  it.req.PartitionName,
			},
		},
	}

	// check if num_rows is valid
	if it.req.NumRows <= 0 {
		return merr.WrapErrParameterInvalid("invalid num_rows", fmt.Sprint(it.req.NumRows), "num_rows should be greater than 0")
	}

	if err := genFunctionFields(ctx, it.upsertMsg.InsertMsg, it.schema, it.req.GetPartialUpdate()); err != nil {
		return err
	}

	if it.req.GetPartialUpdate() {
		err = it.queryPreExecute(ctx)
		if err != nil {
			log.Warn(ctx, "Fail to queryPreExecute", mlog.Err(err))
			return err
		}
		// reconstruct upsert msg after queryPreExecute
		it.upsertMsg.InsertMsg.FieldsData = it.insertFieldData
		it.upsertMsg.DeleteMsg.PrimaryKeys = it.deletePKs
		it.upsertMsg.DeleteMsg.NumRows = int64(typeutil.GetSizeOfIDs(it.deletePKs))
	}

	err = it.insertPreExecute(ctx)
	if err != nil {
		log.Warn(ctx, "Fail to insertPreExecute", mlog.Err(err))
		return err
	}

	err = it.deletePreExecute(ctx)
	if err != nil {
		log.Warn(ctx, "Fail to deletePreExecute", mlog.Err(err))
		return err
	}

	it.result.DeleteCnt = it.upsertMsg.DeleteMsg.NumRows
	it.result.InsertCnt = int64(it.upsertMsg.InsertMsg.NumRows)
	if it.result.DeleteCnt != it.result.InsertCnt {
		log.Info(ctx, "DeleteCnt and InsertCnt are not the same when upsert",
			mlog.Int64("DeleteCnt", it.result.DeleteCnt),
			mlog.Int64("InsertCnt", it.result.InsertCnt))
	}
	it.result.UpsertCnt = it.result.InsertCnt
	log.Debug(ctx, "Proxy Upsert PreExecute done")
	return nil
}

func (it *upsertTask) PostExecute(ctx context.Context) error {
	return nil
}