package proxy

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/cockroachdb/errors"
	"github.com/samber/lo"
	"google.golang.org/protobuf/proto"

	"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/schemapb"
	"github.com/milvus-io/milvus/internal/agg"
	"github.com/milvus-io/milvus/internal/parser/planparserv2"
	"github.com/milvus-io/milvus/internal/proxy/accesslog"
	"github.com/milvus-io/milvus/internal/proxy/shardclient"
	"github.com/milvus-io/milvus/internal/types"
	"github.com/milvus-io/milvus/internal/util/exprutil"
	"github.com/milvus-io/milvus/internal/util/reduce"
	"github.com/milvus-io/milvus/internal/util/reduce/orderby"
	"github.com/milvus-io/milvus/internal/util/segcore"
	"github.com/milvus-io/milvus/internal/util/shallowcopy"
	"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/proto/internalpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
	"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
	"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
	"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/retry"
	"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
	"github.com/milvus-io/milvus/pkg/v3/util/timestamptz"
	"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

const (
	WithCache    = true
	WithoutCache = false
)

const (
	RetrieveTaskName = "RetrieveTask"
	QueryTaskName    = "QueryTask"
)

type queryTask struct {
	baseTask
	Condition
	*internalpb.RetrieveRequest

	ctx            context.Context
	result         *milvuspb.QueryResults
	request        *milvuspb.QueryRequest
	mixCoord       types.MixCoordClient
	ids            *schemapb.IDs
	collectionName string
	queryParams    *queryParams
	schema         *schemaInfo

	translatedOutputFields []string
	userOutputFields       []string
	userDynamicFields      []string
	userAggregates         []agg.AggregateBase

	resultBuf *typeutil.ConcurrentSet[*internalpb.RetrieveResults]

	plan             *planpb.PlanNode
	partitionKeyMode bool
	shardclientMgr   shardclient.ShardClientMgr
	lb               shardclient.LBPolicy
	channelsMvcc     map[string]Timestamp
	preferredNodes   map[string]int64
	fastSkip         bool

	reQuery              bool
	allQueryCnt          int64
	totalRelatedDataSize int64
	mustUsePartitionKey  bool
	resolvedTimezoneStr  string
	storageCost          segcore.StorageCost
	aggregationFieldMap  *agg.AggregationFieldMap
	chMgr                channelsMgr
}

func (t *queryTask) getQueryLabel() string {
	if label := t.GetQueryLabel(); label != "" {
		return label
	}
	return metrics.QueryLabel
}

type queryParams struct {
	limit               int64
	offset              int64
	reduceType          reduce.IReduceType
	isIterator          bool
	collectionID        int64
	groupByFields       []string
	orderByFields       []string // NEW: ORDER BY field specifications (e.g., "price:desc")
	timezone            string
	extractTimeFields   []string
	queryIteratorCursor *planpb.QueryIteratorCursor
}

func isSupportedGroupByFieldType(dt schemapb.DataType) bool {
	switch dt {
	case schemapb.DataType_Int8,
		schemapb.DataType_Int16,
		schemapb.DataType_Int32,
		schemapb.DataType_Int64,
		schemapb.DataType_VarChar,
		schemapb.DataType_Timestamptz:
		return true
	default:
		return false
	}
}

func validateGroupByFieldSchema(field *schemapb.FieldSchema) error {
	if field == nil {
		return merr.WrapErrParameterInvalidMsg("group by field schema is nil")
	}
	if !isSupportedGroupByFieldType(field.GetDataType()) {
		return merr.WrapErrParameterInvalidMsg(
			"group by field %s has unsupported data type %s", field.GetName(), field.GetDataType().String(),
		)
	}
	return nil
}

func translateGroupByFieldIds(groupByFieldNames []string, schema *schemapb.CollectionSchema) ([]UniqueID, error) {
	if len(groupByFieldNames) == 0 {
		return nil, nil
	}

	fieldNameToSchema := make(map[string]*schemapb.FieldSchema, len(schema.Fields))
	for _, field := range schema.Fields {
		fieldNameToSchema[field.Name] = field
	}

	groupByFieldIds := make([]UniqueID, 0, len(groupByFieldNames))
	for _, groupByField := range groupByFieldNames {
		groupByField = strings.TrimSpace(groupByField)
		fieldSchema, found := fieldNameToSchema[groupByField]
		if !found {
			return nil, merr.WrapErrParameterInvalidMsg("field %s not exist", groupByField)
		}
		if err := validateGroupByFieldSchema(fieldSchema); err != nil {
			return nil, err
		}
		groupByFieldIds = append(groupByFieldIds, fieldSchema.GetFieldID())
	}
	return groupByFieldIds, nil
}

// validateOrderByFieldsWithGroupBy validates that ORDER BY fields are compatible with GROUP BY.
// When GROUP BY is used, ORDER BY can only reference columns in the GROUP BY clause.
// ORDER BY on aggregate expressions (e.g., count(*)) is not yet supported and is
// explicitly rejected. This restriction may be lifted in a future release.
func validateOrderByFieldsWithGroupBy(
	orderByFieldSpecs []string,
	groupByFields []string,
	aggregates []agg.AggregateBase,
) error {
	if len(orderByFieldSpecs) == 0 {
		return nil
	}

	// If no GROUP BY and no aggregates, any field is valid for ORDER BY
	hasGroupBy := len(groupByFields) > 0 || len(aggregates) > 0
	if !hasGroupBy {
		return nil
	}

	// Build set of valid ORDER BY targets (GROUP BY columns only)
	validTargets := make(map[string]bool)

	// Add GROUP BY fields as valid targets
	for _, field := range groupByFields {
		validTargets[strings.ToLower(strings.TrimSpace(field))] = true
	}

	// Validate each ORDER BY field
	for _, spec := range orderByFieldSpecs {
		spec = strings.TrimSpace(spec)
		if spec == "" {
			continue
		}

		// Extract field name (remove direction suffix like ":desc").
		// Only colon-separated format is supported (e.g., "price:desc").
		// Space-separated format (e.g., "price desc") is NOT supported.
		parts := strings.Split(spec, ":")
		fieldName := strings.ToLower(strings.TrimSpace(parts[0]))

		// Reject aggregate expressions — not yet supported
		if isAgg, _, _ := agg.MatchAggregationExpression(fieldName); isAgg {
			return merr.WrapErrParameterInvalidMsg(
				"ORDER BY on aggregate expression '%s' is not yet supported",
				fieldName,
			)
		}

		if !validTargets[fieldName] {
			return merr.WrapErrParameterInvalidMsg(
				"ORDER BY field '%s' is not valid: when using GROUP BY or aggregates, "+
					"ORDER BY can only reference GROUP BY columns. "+
					"Valid targets are: %v",
				fieldName, getValidTargetList(groupByFields, aggregates),
			)
		}
	}

	return nil
}

// getValidTargetList returns a formatted list of valid ORDER BY targets for error messages.
// The aggregates parameter is currently unused because ORDER BY on aggregate expressions
// (e.g., count(*)) is not yet supported. When enabled in the future, aggregate original
// names should be appended to the target list.
func getValidTargetList(groupByFields []string, _ []agg.AggregateBase) []string {
	targets := make([]string, 0, len(groupByFields))
	targets = append(targets, groupByFields...)
	return targets
}

// translateOrderByFields converts ORDER BY field specifications to planpb.OrderByField messages.
// Delegates parsing to orderby.ParseOrderByFields to ensure consistent behavior
// (direction validation, nullsFirst defaults) between C++ segcore and Go proxy pipeline.
func translateOrderByFields(orderByFieldSpecs []string, schema *schemapb.CollectionSchema) ([]*planpb.OrderByField, error) {
	parsed, err := orderby.ParseOrderByFields(orderByFieldSpecs, schema)
	if err != nil {
		return nil, merr.WrapErrParameterInvalidMsg(err.Error())
	}
	if len(parsed) == 0 {
		return nil, nil
	}

	result := make([]*planpb.OrderByField, len(parsed))
	for i, f := range parsed {
		result[i] = &planpb.OrderByField{
			FieldId:    f.FieldID,
			Ascending:  f.Ascending,
			NullsFirst: f.NullsFirst,
		}
	}
	return result, nil
}

// translateToOutputFieldIDs translates output fields name to output fields id.
// If no output fields specified, return only pk field
func translateToOutputFieldIDs(outputFields []string, schema *schemapb.CollectionSchema) ([]UniqueID, error) {
	outputFieldIDs := make([]UniqueID, 0, len(outputFields)+1)
	if len(outputFields) == 0 {
		for _, field := range schema.Fields {
			if field.IsPrimaryKey {
				outputFieldIDs = append(outputFieldIDs, field.FieldID)
			}
		}
	} else {
		var pkFieldID UniqueID
		for _, field := range schema.Fields {
			if field.IsPrimaryKey {
				pkFieldID = field.FieldID
			}
		}
		for _, reqField := range outputFields {
			var fieldFound bool
			for _, field := range schema.Fields {
				if reqField == field.Name {
					outputFieldIDs = append(outputFieldIDs, field.FieldID)
					fieldFound = true
					break
				}
			}

			if !fieldFound {
			structFieldLoop:
				for _, structField := range schema.StructArrayFields {
					for _, field := range structField.Fields {
						if reqField == field.Name {
							outputFieldIDs = append(outputFieldIDs, field.FieldID)
							fieldFound = true
							break structFieldLoop
						}
					}
				}
			}

			if !fieldFound {
				return nil, merr.WrapErrParameterInvalidMsg("field %s not exist", reqField)
			}
		}

		// pk field needs to be in output field list
		var pkFound bool
		for _, outputField := range outputFieldIDs {
			if outputField == pkFieldID {
				pkFound = true
				break
			}
		}

		if !pkFound {
			outputFieldIDs = append(outputFieldIDs, pkFieldID)
		}
	}
	return outputFieldIDs, nil
}

func filterSystemFields(outputFieldIDs []UniqueID) []UniqueID {
	filtered := make([]UniqueID, 0, len(outputFieldIDs))
	for _, outputFieldID := range outputFieldIDs {
		if !common.IsSystemField(outputFieldID) {
			filtered = append(filtered, outputFieldID)
		}
	}
	return filtered
}

// parseQueryParams get limit and offset from queryParamsPair, both are optional.
func parseQueryParams(queryParamsPair []*commonpb.KeyValuePair, largeTopKEnabled bool, pkDataType schemapb.DataType) (*queryParams, error) {
	var (
		limit             int64
		offset            int64
		reduceStopForBest bool
		isIterator        bool
		err               error
		collectionID      int64
		timezone          string
		extractTimeFields []string
	)
	reduceStopForBestStr, err := funcutil.GetAttrByKeyFromRepeatedKV(ReduceStopForBestKey, queryParamsPair)
	// if reduce_stop_for_best is provided
	if err == nil {
		reduceStopForBest, err = strconv.ParseBool(reduceStopForBestStr)
		if err != nil {
			return nil, merr.WrapErrParameterInvalid("true or false", reduceStopForBestStr,
				"value for reduce_stop_for_best is invalid")
		}
	}

	isIteratorStr, err := funcutil.GetAttrByKeyFromRepeatedKV(IteratorField, queryParamsPair)
	// if reduce_stop_for_best is provided
	if err == nil {
		isIterator, err = strconv.ParseBool(isIteratorStr)
		if err != nil {
			return nil, merr.WrapErrParameterInvalid("true or false", isIteratorStr,
				"value for iterator field is invalid")
		}
	}

	collectionIdStr, err := funcutil.GetAttrByKeyFromRepeatedKV(CollectionID, queryParamsPair)
	if err == nil {
		collectionID, err = strconv.ParseInt(collectionIdStr, 0, 64)
		if err != nil {
			return nil, merr.WrapErrParameterInvalid("int value for collection_id", CollectionID,
				"value for collection id is invalid")
		}
	}

	reduceType := reduce.IReduceNoOrder
	if isIterator {
		if reduceStopForBest {
			reduceType = reduce.IReduceInOrderForBest
		} else {
			reduceType = reduce.IReduceInOrder
		}
	}

	limit = typeutil.Unlimited
	isLimitProvided := false
	limitStr, err := funcutil.GetAttrByKeyFromRepeatedKV(LimitKey, queryParamsPair)
	// if limit is provided
	if err == nil {
		isLimitProvided = true
		limit, err = strconv.ParseInt(limitStr, 0, 64)
		if err != nil {
			return nil, merr.WrapErrParameterInvalidMsg("%s [%s] is invalid", LimitKey, limitStr)
		}
	}
	if isLimitProvided {
		offsetStr, err := funcutil.GetAttrByKeyFromRepeatedKV(OffsetKey, queryParamsPair)
		// if offset is provided
		if err == nil {
			offset, err = strconv.ParseInt(offsetStr, 0, 64)
			if err != nil {
				return nil, merr.WrapErrParameterInvalidMsg("%s [%s] is invalid", OffsetKey, offsetStr)
			}
		}
		// validate max result window.
		if err = validateMaxQueryResultWindow(offset, limit, largeTopKEnabled); err != nil {
			return nil, merr.WrapErrParameterInvalidMsg("invalid max query result window, %v", err)
		}
	}

	timezone, _ = funcutil.TryGetAttrByKeyFromRepeatedKV(common.TimezoneKey, queryParamsPair)
	if (timezone != "") && !timestamptz.IsTimezoneValid(timezone) {
		return nil, merr.WrapErrParameterInvalidMsg("unknown or invalid IANA Time Zone ID: %s", timezone)
	}

	extractTimeFieldsStr, err := funcutil.GetAttrByKeyFromRepeatedKV(TimefieldsKey, queryParamsPair)
	if err == nil {
		extractTimeFields = strings.FieldsFunc(extractTimeFieldsStr, func(r rune) bool {
			return r == ',' || r == ' '
		})
	}

	// parse group by fields
	groupByFieldsStr, err := funcutil.GetAttrByKeyFromRepeatedKV(GroupByFieldsKey, queryParamsPair)
	var groupByFields []string
	if err == nil {
		splitFields := strings.Split(groupByFieldsStr, ",")
		for _, field := range splitFields {
			trimmed := strings.TrimSpace(field)
			if trimmed != "" {
				groupByFields = append(groupByFields, trimmed)
			}
		}
	}

	// parse order by fields (e.g., "price:desc,rating:asc"). Only colon-separated format is supported.
	orderByFieldsStr, err := funcutil.GetAttrByKeyFromRepeatedKV(OrderByFieldsKey, queryParamsPair)
	var orderByFields []string
	if err == nil {
		splitFields := strings.Split(orderByFieldsStr, ",")
		for _, field := range splitFields {
			trimmed := strings.TrimSpace(field)
			if trimmed != "" {
				orderByFields = append(orderByFields, trimmed)
			}
		}
	}

	queryIteratorCursor, err := parseQueryIteratorCursor(queryParamsPair, isIterator, pkDataType)
	if err != nil {
		return nil, err
	}

	return &queryParams{
		limit:               limit,
		offset:              offset,
		reduceType:          reduceType,
		isIterator:          isIterator,
		collectionID:        collectionID,
		groupByFields:       groupByFields,
		orderByFields:       orderByFields,
		queryIteratorCursor: queryIteratorCursor,
		timezone:            timezone,
		extractTimeFields:   extractTimeFields,
	}, nil
}

func parseQueryIteratorCursor(queryParamsPair []*commonpb.KeyValuePair, isIterator bool, pkDataType schemapb.DataType) (*planpb.QueryIteratorCursor, error) {
	lastPK, hasLastPK := funcutil.TryGetAttrByKeyFromRepeatedKV(QueryIterLastPKKey, queryParamsPair)
	lastOffsetStr, hasLastOffset := funcutil.TryGetAttrByKeyFromRepeatedKV(QueryIterLastOffsetKey, queryParamsPair)
	if !hasLastPK && !hasLastOffset {
		return nil, nil
	}
	if !isIterator {
		return nil, merr.WrapErrParameterInvalidMsg(
			"invalid query iterator cursor params: %s and %s can only be used when iterator=true",
			QueryIterLastPKKey, QueryIterLastOffsetKey)
	}
	if !hasLastPK || !hasLastOffset {
		return nil, merr.WrapErrParameterInvalidMsg(
			"incomplete query iterator cursor params: %s and %s must be provided together, has_last_pk=%t, has_last_element_offset=%t",
			QueryIterLastPKKey, QueryIterLastOffsetKey, hasLastPK, hasLastOffset)
	}

	lastOffset, err := strconv.ParseInt(lastOffsetStr, 0, 64)
	if err != nil || lastOffset < 0 {
		return nil, merr.WrapErrParameterInvalid("non-negative int value", lastOffsetStr,
			"value for query iterator last element offset is invalid")
	}

	cursor := &planpb.QueryIteratorCursor{
		LastElementOffset: lastOffset,
	}
	switch pkDataType {
	case schemapb.DataType_Int64:
		lastIntPK, err := strconv.ParseInt(lastPK, 0, 64)
		if err != nil {
			return nil, merr.WrapErrParameterInvalid("int64 primary key", lastPK,
				"value for query iterator last primary key is invalid")
		}
		cursor.LastIntPk = &lastIntPK
	case schemapb.DataType_VarChar:
		cursor.LastStrPk = &lastPK
	default:
		return nil, merr.WrapErrParameterInvalidMsg("unsupported primary key type %s for query iterator cursor", pkDataType.String())
	}
	return cursor, nil
}

func getPrimaryKeyDataType(schema *schemapb.CollectionSchema) schemapb.DataType {
	for _, field := range schema.GetFields() {
		if field.GetIsPrimaryKey() {
			return field.GetDataType()
		}
	}
	return schemapb.DataType_None
}

func matchCountRule(outputs []string) bool {
	return len(outputs) == 1 && strings.ToLower(strings.TrimSpace(outputs[0])) == "count(*)"
}

func createCntPlan(expr string, schemaHelper *typeutil.SchemaHelper, exprTemplateValues map[string]*schemapb.TemplateValue) (*planpb.PlanNode, error) {
	if expr == "" {
		return &planpb.PlanNode{
			Node: &planpb.PlanNode_Query{
				Query: &planpb.QueryPlanNode{
					Predicates: nil,
					IsCount:    true,
				},
			},
		}, nil
	}
	start := time.Now()
	plan, err := planparserv2.CreateRetrievePlan(schemaHelper, expr, exprTemplateValues)
	if err != nil {
		metrics.ProxyParseExpressionLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.QueryLabel, metrics.FailLabel).Observe(float64(time.Since(start).Microseconds()) / 1000.0)
		return nil, merr.WrapErrParameterInvalidMsg("failed to create query plan: %v", err)
	}
	metrics.ProxyParseExpressionLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.QueryLabel, metrics.SuccessLabel).Observe(float64(time.Since(start).Microseconds()) / 1000.0)
	plan.Node.(*planpb.PlanNode_Query).Query.IsCount = true

	return plan, nil
}

func (t *queryTask) createPlan(ctx context.Context) error {
	return t.createPlanArgs(ctx, &planparserv2.ParserVisitorArgs{})
}

func (t *queryTask) createPlanArgs(ctx context.Context, visitorArgs *planparserv2.ParserVisitorArgs) error {
	schema := t.schema
	var err error
	if t.plan == nil {
		start := time.Now()
		t.plan, err = planparserv2.CreateRetrievePlanArgs(schema.schemaHelper, t.request.Expr, t.request.GetExprTemplateValues(), visitorArgs)
		if err != nil {
			metrics.ProxyParseExpressionLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.QueryLabel, metrics.FailLabel).Observe(float64(time.Since(start).Microseconds()) / 1000.0)
			return merr.WrapErrParameterInvalidMsg("failed to create query plan: %v", err)
		}
		metrics.ProxyParseExpressionLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.QueryLabel, metrics.SuccessLabel).Observe(float64(time.Since(start).Microseconds()) / 1000.0)
	}
	// parse output fields names
	originalOuputFields := t.request.GetOutputFields()
	t.translatedOutputFields, t.userOutputFields, t.userDynamicFields, t.userAggregates, _, err = translateOutputFields(t.request.GetOutputFields(), t.schema, false)
	if err != nil {
		return err
	}

	// parse aggregates
	t.plan.GetQuery().Aggregates = agg.AggregatesToPB(t.userAggregates)
	t.Aggregates = t.plan.GetQuery().GetAggregates()
	// parse group by field ids
	groupByFieldsIDs, err := translateGroupByFieldIds(t.queryParams.groupByFields, t.schema.CollectionSchema)
	if err != nil {
		return err
	}
	t.plan.GetQuery().GroupByFieldIds = groupByFieldsIDs
	t.GroupByFieldIds = groupByFieldsIDs

	// Validate ORDER BY fields compatibility with GROUP BY
	// When GROUP BY is used, ORDER BY can only reference groupBy columns or aggregate results
	if err := validateOrderByFieldsWithGroupBy(
		t.queryParams.orderByFields,
		t.queryParams.groupByFields,
		t.userAggregates,
	); err != nil {
		return err
	}

	// parse order by fields
	orderByFields, err := translateOrderByFields(t.queryParams.orderByFields, t.schema.CollectionSchema)
	if err != nil {
		return err
	}
	t.plan.GetQuery().OrderByFields = orderByFields
	// Also populate on RetrieveRequest so QN/Delegator can read directly
	// without re-parsing serialized_expr_plan.
	t.OrderByFields = orderByFields

	hasAgg := len(t.GroupByFieldIds) > 0 || len(t.Aggregates) > 0
	// parse output field ids
	if hasAgg {
		emptyOutputFields := make([]UniqueID, 0)
		t.OutputFieldsId = emptyOutputFields
		t.plan.OutputFieldIds = emptyOutputFields
		aggFieldMap, err := agg.NewAggregationFieldMap(originalOuputFields, t.queryParams.groupByFields, t.userAggregates)
		if err != nil {
			return merr.WrapErrParameterInvalidMsg(err.Error())
		}
		t.aggregationFieldMap = aggFieldMap
	} else {
		outputFieldIDs, err := translateToOutputFieldIDs(t.translatedOutputFields, schema.CollectionSchema)
		if err != nil {
			return err
		}
		outputFieldIDs = append(outputFieldIDs, common.TimeStampField)
		t.OutputFieldsId = outputFieldIDs
		t.plan.OutputFieldIds = outputFieldIDs
		t.plan.DynamicFields = t.userDynamicFields
	}

	mlog.Debug(ctx, "translate output fields to field ids",
		mlog.Int64s("OutputFieldsID", t.OutputFieldsId),
		mlog.String("requestType", t.getQueryLabel()))
	return nil
}

func (t *queryTask) hasCountStar() bool {
	for _, agg := range t.userAggregates {
		if agg.Name() == "count" && agg.FieldID() == 0 {
			return true
		}
	}
	return false
}

func (t *queryTask) CanSkipAllocTimestamp() bool {
	var consistencyLevel commonpb.ConsistencyLevel
	useDefaultConsistency := t.request.GetUseDefaultConsistency()
	if !useDefaultConsistency {
		// legacy SDK & resultful behavior
		if t.request.GetConsistencyLevel() == commonpb.ConsistencyLevel_Strong && t.request.GetGuaranteeTimestamp() > 0 {
			return true
		}
		consistencyLevel = t.request.GetConsistencyLevel()
	} else {
		collID, err := globalMetaCache.GetCollectionID(context.Background(), t.request.GetDbName(), t.request.GetCollectionName())
		if err != nil { // err is not nil if collection not exists
			mlog.Warn(t.ctx, "query task get collectionID failed, can't skip alloc timestamp",
				mlog.String("collectionName", t.request.GetCollectionName()), mlog.Err(err))
			return false
		}

		collectionInfo, err2 := globalMetaCache.GetCollectionInfo(context.Background(), t.request.GetDbName(), t.request.GetCollectionName(), collID)
		if err2 != nil {
			mlog.Warn(t.ctx, "query task get collection info failed, can't skip alloc timestamp",
				mlog.String("collectionName", t.request.GetCollectionName()), mlog.Err(err))
			return false
		}
		consistencyLevel = collectionInfo.consistencyLevel
	}
	return consistencyLevel != commonpb.ConsistencyLevel_Strong
}

func (t *queryTask) PreExecute(ctx context.Context) error {
	t.Base.MsgType = commonpb.MsgType_Retrieve
	t.Base.SourceID = paramtable.GetNodeID()

	collectionName := t.request.CollectionName
	t.collectionName = collectionName

	log := mlog.With(mlog.String("collectionName", collectionName),
		mlog.Strings("partitionNames", t.request.GetPartitionNames()),
		mlog.String("requestType", t.getQueryLabel()))

	if err := validateCollectionName(collectionName); err != nil {
		log.Warn(ctx, "Invalid collectionName.")
		return err
	}
	log.Debug(ctx, "Validate collectionName.")

	collID, err := globalMetaCache.GetCollectionID(ctx, t.request.GetDbName(), collectionName)
	if err != nil {
		log.Warn(ctx, "Failed to get collection id.", mlog.String("collectionName", collectionName), mlog.Err(err))
		return err
	}
	t.CollectionID = collID

	colInfo, err := globalMetaCache.GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
	if err != nil {
		log.Warn(ctx, "Failed to get collection info.", mlog.String("collectionName", collectionName),
			mlog.Int64("collectionID", t.CollectionID), mlog.Err(err))
		// The name was already resolved above (GetCollectionID succeeded), so a
		// not-found here means the collection was concurrently dropped between the
		// two lookups — a TOCTOU race, not the caller's input error. Leave it as the
		// default SystemError; do not stamp InputError.
		return err
	}
	log.Debug(ctx, "Get collection ID by name", mlog.Int64("collectionID", t.CollectionID))

	schema, err := globalMetaCache.GetCollectionSchema(ctx, t.request.GetDbName(), t.collectionName)
	if err != nil {
		log.Warn(ctx, "get collection schema failed", mlog.Err(err))
		return err
	}
	t.schema = schema
	if err := validateTextStorageV3Enabled(t.schema.CollectionSchema); err != nil {
		return err
	}
	partitionNames, namespaceAsPartition, err := resolveNamespacePartitionNames(t.schema.CollectionSchema, t.request.Namespace, t.request.GetPartitionNames())
	if err != nil {
		return err
	}
	if namespaceAsPartition {
		t.request.PartitionNames = partitionNames
	}

	t.partitionKeyMode, err = isPartitionKeyMode(ctx, t.request.GetDbName(), collectionName)
	if err != nil {
		log.Warn(ctx, "check partition key mode failed", mlog.Int64("collectionID", t.CollectionID), mlog.Err(err))
		return err
	}
	if t.partitionKeyMode && len(t.request.GetPartitionNames()) != 0 {
		return merr.WrapErrParameterInvalidMsg("not support manually specifying the partition names if partition key mode is used")
	}
	if t.mustUsePartitionKey && !t.partitionKeyMode {
		return merr.WrapErrParameterInvalidMsg("must use partition key in the query request " +
			"because the mustUsePartitionKey config is true")
	}

	for _, tag := range t.request.PartitionNames {
		if err := validatePartitionTag(tag, false); err != nil {
			log.Warn(ctx, "invalid partition name", mlog.String("partition name", tag))
			return err
		}
	}
	log.Debug(ctx, "Validate partition names.")

	// fetch search_growing from query param
	if t.IgnoreGrowing, err = isIgnoreGrowing(t.request.GetQueryParams()); err != nil {
		return err
	}
	queryParams, err := parseQueryParams(t.request.GetQueryParams(), colInfo.queryMode == common.QueryModeLargeTopK, getPrimaryKeyDataType(schema.CollectionSchema))
	if err != nil {
		return err
	}
	if queryParams.collectionID > 0 && queryParams.collectionID != t.GetCollectionID() {
		return merr.WrapErrParameterInvalidMsg("Input collection id is not consistent to collectionID in the context," +
			"alias or database may have changed")
	}
	if queryParams.reduceType == reduce.IReduceInOrderForBest {
		t.ReduceStopForBest = true
	}
	t.ReduceType = int32(queryParams.reduceType)

	t.queryParams = queryParams

	// ORDER BY requires explicit limit to prevent segment-level OOM.
	// SortBuffer loads all matching rows into memory for sorting;
	// MaxOutputSize only guards proxy reduce, not segment sorting.
	if len(queryParams.orderByFields) > 0 && queryParams.limit == typeutil.Unlimited {
		return merr.WrapErrParameterInvalidMsg("ORDER BY requires explicit limit")
	}

	// ORDER BY with iterator is not yet supported. Current iterator relies on
	// PK-ordered pagination (plain query pipeline), while ORDER BY uses a
	// different pipeline that sorts by arbitrary fields. Future work will
	// enable iterator with user-specified ORDER BY fields.
	if len(queryParams.orderByFields) > 0 && queryParams.isIterator {
		return merr.WrapErrParameterInvalidMsg("ORDER BY with iterator is not supported")
	}

	t.Limit = queryParams.limit + queryParams.offset

	if t.ids != nil {
		pkField := ""
		for _, field := range schema.Fields {
			if field.IsPrimaryKey {
				pkField = field.Name
			}
		}
		t.request.Expr = IDs2Expr(pkField, t.ids)
	}

	if t.queryParams.timezone != "" {
		// validated in queryParams, no need to validate again
		t.resolvedTimezoneStr = t.queryParams.timezone
		log.Debug(ctx, "determine timezone from request", mlog.String("user defined timezone", t.resolvedTimezoneStr))
	} else {
		t.resolvedTimezoneStr = getColTimezone(colInfo)
		log.Debug(ctx, "determine timezone from collection", mlog.Any("collection timezone", t.resolvedTimezoneStr))
	}

	if err := t.createPlanArgs(ctx, &planparserv2.ParserVisitorArgs{Timezone: t.resolvedTimezoneStr}); err != nil {
		return err
	}
	t.plan.GetQuery().Limit = t.Limit
	if t.queryParams.queryIteratorCursor != nil {
		t.plan.GetQuery().QueryIteratorCursor = t.queryParams.queryIteratorCursor
	}

	// Aggregation queries have bounded result sizes:
	// - global aggregation (no GROUP BY) returns exactly one row
	// - GROUP BY aggregation returns at most one row per distinct group value
	// Both are safe without a limit, so exempt them from the limit requirement.
	hasAgg := len(t.userAggregates) > 0

	if planparserv2.IsAlwaysTruePlan(t.plan) && t.Limit == typeutil.Unlimited && !hasAgg {
		return merr.WrapErrParameterInvalidMsg("empty expression should be used with limit")
	}

	// convert partition names only when requery is false
	if !t.reQuery {
		partitionNames := t.request.GetPartitionNames()
		if namespacePartitionKeyMode(t.schema.CollectionSchema) && t.request.Namespace != nil {
			hashedPartitionNames, err := assignNamespacePartitionKey(ctx, t.request.GetDbName(), t.request.CollectionName, t.request.Namespace)
			if err != nil {
				return err
			}

			partitionNames = append(partitionNames, hashedPartitionNames...)
		} else if t.partitionKeyMode {
			expr, err := exprutil.ParseExprFromPlan(t.plan)
			if err != nil {
				return err
			}
			partitionKeys := exprutil.ParseKeys(expr, exprutil.PartitionKey)
			hashedPartitionNames, err := assignPartitionKeys(ctx, t.request.GetDbName(), t.request.CollectionName, partitionKeys)
			if err != nil {
				return err
			}

			partitionNames = append(partitionNames, hashedPartitionNames...)
		}
		t.PartitionIDs, err = getPartitionIDs(ctx, t.request.GetDbName(), t.request.CollectionName, partitionNames)
		if err != nil {
			return err
		}
	}

	// count(*) without GROUP BY is a single-value result, pagination is meaningless.
	// But count(*) with GROUP BY + limit is valid (limits the number of groups returned).
	if t.hasCountStar() && t.queryParams.limit != typeutil.Unlimited && len(t.GetGroupByFieldIds()) == 0 {
		return merr.WrapErrParameterInvalidMsg("count entities with pagination is not allowed")
	}
	t.plan.Namespace = namespaceForPlan(t.schema.CollectionSchema, t.request.Namespace)

	t.SerializedExprPlan, err = proto.Marshal(t.plan)
	if err != nil {
		return err
	}
	t.PkFilter = checkSegmentFilter(t.plan)
	// Set username for this query request,
	if username, _ := GetCurUserFromContext(ctx); username != "" {
		t.Username = username
	}

	collectionInfo, err2 := globalMetaCache.GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
	if err2 != nil {
		log.Warn(ctx, "Proxy::queryTask::PreExecute failed to GetCollectionInfo from cache",
			mlog.String("collectionName", collectionName), mlog.Int64("collectionID", t.CollectionID),
			mlog.Err(err2))
		return err2
	}

	guaranteeTs := t.request.GetGuaranteeTimestamp()
	var consistencyLevel commonpb.ConsistencyLevel
	useDefaultConsistency := t.request.GetUseDefaultConsistency()
	t.ConsistencyLevel = t.request.GetConsistencyLevel()
	if useDefaultConsistency {
		consistencyLevel = collectionInfo.consistencyLevel
		guaranteeTs = parseGuaranteeTsFromConsistency(guaranteeTs, t.BeginTs(), consistencyLevel)
	} else {
		consistencyLevel = t.request.GetConsistencyLevel()
		// Compatibility logic, parse guarantee timestamp
		if consistencyLevel == 0 && guaranteeTs > 0 {
			guaranteeTs = parseGuaranteeTs(guaranteeTs, t.BeginTs())
		} else {
			// parse from guarantee timestamp and user input consistency level
			guaranteeTs = parseGuaranteeTsFromConsistency(guaranteeTs, t.BeginTs(), consistencyLevel)
		}
	}

	// update actual consistency level
	accesslog.SetActualConsistencyLevel(ctx, consistencyLevel)

	// use collection schema updated timestamp if it's greater than calculate guarantee timestamp
	// this make query view updated happens before new read request happens
	// see also schema change design
	if collectionInfo.updateTimestamp > guaranteeTs {
		guaranteeTs = collectionInfo.updateTimestamp
	}

	t.GuaranteeTimestamp = guaranteeTs
	// Extract physical time for entity-level TTL (issue #47413)
	physicalTimeMs, _ := tsoutil.ParseHybridTs(guaranteeTs)
	t.EntityTtlPhysicalTime = uint64(physicalTimeMs * 1000)
	// need modify mvccTs and guaranteeTs for iterator specially
	if t.queryParams.isIterator && t.request.GetGuaranteeTimestamp() > 0 {
		t.MvccTimestamp = t.request.GetGuaranteeTimestamp()
		t.GuaranteeTimestamp = t.request.GetGuaranteeTimestamp()
	}
	t.IsIterator = queryParams.isIterator

	if collectionInfo.collectionTTL != 0 {
		physicalTime := tsoutil.PhysicalTime(t.GetBase().GetTimestamp())
		expireTime := physicalTime.Add(-time.Duration(collectionInfo.collectionTTL))
		t.CollectionTtlTimestamps = tsoutil.ComposeTSByTime(expireTime)
		// preventing overflow, abort
		if t.CollectionTtlTimestamps > t.GetBase().GetTimestamp() {
			return merr.WrapErrServiceInternalMsg("ttl timestamp overflow, base timestamp: %d, ttl duration %v", t.GetBase().GetTimestamp(), collectionInfo.collectionTTL)
		}
	}
	deadline, ok := t.TraceCtx().Deadline()
	if ok {
		t.TimeoutTimestamp = tsoutil.ComposeTSByTime(deadline)
	}

	t.DbID = 0 // TODO
	log.Debug(ctx, "Query PreExecute done.",
		mlog.Uint64("guarantee_ts", guaranteeTs),
		mlog.Uint64("mvcc_ts", t.GetMvccTimestamp()),
		mlog.Uint64("timeout_ts", t.GetTimeoutTimestamp()),
		mlog.Uint64("collection_ttl_timestamps", t.CollectionTtlTimestamps))
	return nil
}

func (t *queryTask) Execute(ctx context.Context) error {
	tr := timerecord.NewTimeRecorder(fmt.Sprintf("proxy execute query %d", t.ID()))
	defer tr.CtxElapse(ctx, "done")
	log := mlog.With(mlog.Int64("collection", t.GetCollectionID()),
		mlog.Int64s("partitionIDs", t.GetPartitionIDs()),
		mlog.String("requestType", t.getQueryLabel()))

	t.resultBuf = typeutil.NewConcurrentSet[*internalpb.RetrieveResults]()
	if namespacePartitionKeyModeEnabled(t.schema.CollectionSchema) && t.request.Namespace != nil {
		channelNames, err := t.chMgr.getVChannels(t.CollectionID)
		if err != nil {
			log.Warn(ctx, "get vChannels failed", mlog.Int64("collectionID", t.CollectionID), mlog.Err(err))
			return err
		}
		channelName, ok, err := namespaceShardingChannel(t.schema.CollectionSchema, t.request.Namespace, channelNames)
		if err != nil {
			return err
		}
		if ok {
			if err := t.lb.ExecuteWithRetry(ctx, shardclient.ChannelWorkload{
				Db:              t.request.GetDbName(),
				CollectionName:  t.collectionName,
				CollectionID:    t.CollectionID,
				Channel:         channelName,
				Nq:              1,
				Exec:            t.queryShard,
				PreferredNodeID: preferredNodeForChannel(t.preferredNodes, channelName),
			}); err != nil {
				log.Warn(ctx, "fail to execute query", mlog.Err(err))
				return errors.Wrap(err, "failed to query")
			}

			log.Debug(ctx, "Query Execute done.")
			return nil
		}
	}
	err := t.lb.Execute(ctx, shardclient.CollectionWorkLoad{
		Db:             t.request.GetDbName(),
		CollectionID:   t.CollectionID,
		CollectionName: t.collectionName,
		Nq:             1,
		Exec:           t.queryShard,
		PreferredNodes: t.preferredNodes,
	})
	if err != nil {
		log.Warn(ctx, "fail to execute query", mlog.Err(err))
		return errors.Wrap(err, "failed to query")
	}

	log.Debug(ctx, "Query Execute done.")
	return nil
}

func (t *queryTask) PostExecute(ctx context.Context) error {
	tr := timerecord.NewTimeRecorder("queryTask PostExecute")
	defer func() {
		tr.CtxElapse(ctx, "done")
	}()

	log := mlog.With(mlog.Int64("collection", t.GetCollectionID()),
		mlog.Int64s("partitionIDs", t.GetPartitionIDs()),
		mlog.String("requestType", t.getQueryLabel()))

	var err error

	toReduceResults := make([]*internalpb.RetrieveResults, 0)
	t.allQueryCnt = 0
	t.totalRelatedDataSize = 0
	t.storageCost = segcore.StorageCost{}
	select {
	case <-t.TraceCtx().Done():
		log.Warn(ctx, "proxy", mlog.Int64("Query: wait to finish failed, timeout!, msgID:", t.ID()))
		return merr.Wrapf(t.TraceCtx().Err(), "Query wait to finish timeout, msgID=%d", t.ID())
	default:
		log.Debug(ctx, "all queries are finished or canceled")
		t.resultBuf.Range(func(res *internalpb.RetrieveResults) bool {
			toReduceResults = append(toReduceResults, res)
			t.allQueryCnt += res.GetAllRetrieveCount()
			t.storageCost.ScannedRemoteBytes += res.GetScannedRemoteBytes()
			t.storageCost.ScannedTotalBytes += res.GetScannedTotalBytes()
			t.totalRelatedDataSize += res.GetCostAggregation().GetTotalRelatedDataSize()
			log.Debug(ctx, "proxy receives one query result", mlog.Int64("sourceID", res.GetBase().GetSourceID()))
			return true
		})
	}

	metrics.ProxyDecodeResultLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), t.getQueryLabel()).Observe(0.0)
	tr.CtxRecord(ctx, "reduceResultStart")

	// Parse ORDER BY fields if present
	var orderByFields []*orderby.OrderByField
	if len(t.queryParams.orderByFields) > 0 {
		orderByFields, err = orderby.ParseOrderByFields(t.queryParams.orderByFields, t.schema.CollectionSchema)
		if err != nil {
			log.Warn(ctx, "fail to parse order by fields", mlog.Err(err))
			return err
		}
	}

	primaryFieldSchema, err := t.schema.GetPkField()
	if err != nil {
		log.Warn(ctx, "failed to get primary field schema", mlog.Err(err))
		return err
	}

	pipeline, err := NewQueryPipeline(
		t.schema.CollectionSchema,
		t.queryParams.limit,
		t.queryParams.offset,
		t.queryParams.reduceType,
		orderByFields,
		t.GetGroupByFieldIds(),
		t.GetAggregates(),
		t.aggregationFieldMap,
		filterSystemFields(t.GetOutputFieldsId()),
	)
	if err != nil {
		log.Warn(ctx, "fail to create query pipeline", mlog.Err(err))
		return err
	}
	t.result, err = pipeline.Execute(ctx, toReduceResults)
	if err != nil {
		log.Warn(ctx, "fail to reduce query result", mlog.Err(err))
		return err
	}

	// FieldName/Type/IsDynamic setting and timestamp column removal are now
	// handled by complementFieldOperator in the pipeline (for non-aggregation queries).
	// Only geometry WKB→WKT conversion still needs to happen here.
	for i, fieldData := range t.result.FieldsData {
		if fieldData.Type == schemapb.DataType_Geometry {
			if err := validateGeometryFieldSearchResult(&t.result.FieldsData[i]); err != nil {
				log.Warn(ctx, "fail to validate geometry field search result", mlog.Err(err))
				return err
			}
		}
	}
	t.result.OutputFields = t.userOutputFields
	if !t.reQuery {
		reconstructStructFieldDataForQuery(t.result, t.schema.CollectionSchema)
	}

	t.result.CollectionName = t.collectionName
	t.result.PrimaryFieldName = primaryFieldSchema.GetName()
	metrics.ProxyReduceResultLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), t.getQueryLabel()).Observe(float64(tr.RecordSpan().Microseconds()) / 1000.0)

	if t.queryParams.isIterator && t.request.GetGuaranteeTimestamp() == 0 {
		// first page for iteration, need to set up sessionTs for iterator
		t.result.SessionTs = getMaxMvccTsFromChannels(t.channelsMvcc, t.BeginTs())
	}
	if !t.reQuery {
		if len(t.queryParams.extractTimeFields) > 0 {
			log.Debug(ctx, "extracting fields for timestamptz", mlog.Strings("fields", t.queryParams.extractTimeFields))
			err = extractFieldsFromResults(t.result.GetFieldsData(), t.resolvedTimezoneStr, t.queryParams.extractTimeFields)
			if err != nil {
				log.Warn(ctx, "fail to extract fields for timestamptz", mlog.Err(err))
				return err
			}
		} else {
			log.Debug(ctx, "translate timestamp to ISO string", mlog.String("user define timezone", t.queryParams.timezone))
			err = timestamptzUTC2IsoStr(t.result.GetFieldsData(), t.resolvedTimezoneStr)
			if err != nil {
				log.Warn(ctx, "fail to translate timestamp", mlog.Err(err))
				return err
			}
		}
	}
	log.Debug(ctx, "Query PostExecute done")
	return nil
}

func (t *queryTask) IsSubTask() bool {
	return t.reQuery
}

func (t *queryTask) queryShard(ctx context.Context, nodeID int64, qn types.QueryNodeClient, channel string) error {
	ctx = retry.WithMaxAttemptsContext(ctx, 1)
	needOverrideMvcc := false
	mvccTs := t.MvccTimestamp
	if len(t.channelsMvcc) > 0 {
		mvccTs, needOverrideMvcc = t.channelsMvcc[channel]
		// In fast mode, if there is no corresponding channel in channelsMvcc, quickly skip this query.
		if !needOverrideMvcc && t.fastSkip {
			return nil
		}
	}

	retrieveReq := shallowcopy.ShallowCopyRetrieveRequest(t.RetrieveRequest, nodeID)
	if needOverrideMvcc && mvccTs > 0 {
		retrieveReq.MvccTimestamp = mvccTs
		retrieveReq.GuaranteeTimestamp = mvccTs
	}
	retrieveReq.ConsistencyLevel = t.ConsistencyLevel
	req := &querypb.QueryRequest{
		Req:         retrieveReq,
		DmlChannels: []string{channel},
		Scope:       querypb.DataScope_All,
	}

	log := mlog.With(mlog.Int64("collection", t.GetCollectionID()),
		mlog.Int64s("partitionIDs", t.GetPartitionIDs()),
		mlog.Int64("nodeID", nodeID),
		mlog.String("channel", channel))

	result, err := qn.Query(ctx, req)
	if err != nil {
		log.Warn(ctx, "QueryNode query return error", mlog.Err(err))
		t.shardclientMgr.DeprecateShardCache(t.request.GetDbName(), t.collectionName)
		return err
	}
	if result.GetStatus().GetErrorCode() == commonpb.ErrorCode_NotShardLeader {
		log.Warn(ctx, "QueryNode is not shardLeader")
		t.shardclientMgr.DeprecateShardCache(t.request.GetDbName(), t.collectionName)
		return merr.Error(result.GetStatus())
	}
	if result.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success {
		log.Warn(ctx, "QueryNode query result error", mlog.Any("errorCode", result.GetStatus().GetErrorCode()), mlog.String("reason", result.GetStatus().GetReason()))
		return errors.Wrapf(merr.Error(result.GetStatus()), "fail to Query on QueryNode %d", nodeID)
	}

	log.Debug(ctx, "get query result")
	t.resultBuf.Insert(result)
	t.lb.UpdateCostMetrics(nodeID, result.CostAggregation)
	return nil
}

// IDs2Expr converts ids slices to bool expresion with specified field name
func IDs2Expr(fieldName string, ids *schemapb.IDs) string {
	var idsStr string
	switch ids.GetIdField().(type) {
	case *schemapb.IDs_IntId:
		idsStr = strings.Trim(strings.Join(strings.Fields(fmt.Sprint(ids.GetIntId().GetData())), ", "), "[]")
	case *schemapb.IDs_StrId:
		strs := lo.Map(ids.GetStrId().GetData(), func(str string, _ int) string {
			return fmt.Sprintf("\"%s\"", str)
		})
		idsStr = strings.Trim(strings.Join(strs, ", "), "[]")
	}

	return fieldName + " in [ " + idsStr + " ]"
}

func reduceRetrieveResults(ctx context.Context, retrieveResults []*internalpb.RetrieveResults, queryParams *queryParams) (*milvuspb.QueryResults, error) {
	mlog.Debug(ctx, "reduceInternalRetrieveResults", mlog.Int("len(retrieveResults)", len(retrieveResults)))
	var (
		ret     = &milvuspb.QueryResults{}
		loopEnd int
	)

	// Detect if this is an element-level query
	isElementLevel := len(retrieveResults) > 0 && retrieveResults[0].GetElementLevel()

	validRetrieveResults := []*internalpb.RetrieveResults{}
	for _, r := range retrieveResults {
		size := typeutil.GetSizeOfIDs(r.GetIds())
		if r == nil || len(r.GetFieldsData()) == 0 || size == 0 {
			continue
		}
		// Validate element-level consistency: if any result is element-level, all
		// must be. The flag and element_indices come from querynode results, never
		// from the request, so a mismatch is an internal contract violation.
		if isElementLevel && !r.GetElementLevel() {
			return nil, merr.WrapErrServiceInternalMsg("inconsistent element-level flag: expected all results to be element-level")
		}
		// Validate element_indices length matches ids length for element-level
		if isElementLevel && len(r.GetElementIndices()) != size {
			return nil, merr.WrapErrServiceInternalMsg("element_indices length (%d) does not match ids length (%d)", len(r.GetElementIndices()), size)
		}
		validRetrieveResults = append(validRetrieveResults, r)
		loopEnd += size
	}

	if len(validRetrieveResults) == 0 {
		return ret, nil
	}

	cursors := make([]int64, len(validRetrieveResults))
	idxComputers := make([]*typeutil.FieldDataIdxComputer, len(validRetrieveResults))
	for i, vr := range validRetrieveResults {
		idxComputers[i] = typeutil.NewFieldDataIdxComputer(vr.GetFieldsData())
	}

	// Used in element-level query to limit the number of elements returned
	elementLimit := -1
	if queryParams != nil && queryParams.limit != typeutil.Unlimited {
		// IReduceInOrderForBest will try to get as many results as possible
		// so loopEnd in this case will be set to the sum of all results' size
		// to get as many qualified results as possible
		if reduce.ShouldUseInputLimit(queryParams.reduceType) {
			if !isElementLevel {
				loopEnd = int(queryParams.limit)
			}
			elementLimit = int(queryParams.limit)
		}
	}

	// handle offset
	if queryParams != nil && queryParams.offset > 0 {
		var skipped int64
		for skipped < queryParams.offset {
			sel, drainOneResult := typeutil.SelectMinPK(validRetrieveResults, cursors)
			if sel == -1 || (reduce.ShouldStopWhenDrained(queryParams.reduceType) && drainOneResult) {
				return ret, nil
			}
			if isElementLevel {
				elemIndices := validRetrieveResults[sel].GetElementIndices()[cursors[sel]]
				indicesCount := int64(len(elemIndices.GetIndices()))
				if skipped+indicesCount > queryParams.offset {
					elemIndices.Indices = elemIndices.Indices[queryParams.offset-skipped:]
					break
				} else {
					skipped += indicesCount
				}
			} else {
				skipped++
			}
			cursors[sel]++
		}
	}

	ret.FieldsData = typeutil.PrepareResultFieldData(validRetrieveResults[0].GetFieldsData(), int64(loopEnd))
	var retSize int64
	var availableCount int // for element-level: element count; for doc-level: doc count
	maxOutputSize := paramtable.Get().QuotaConfig.MaxOutputSize.GetAsInt64()
	for j := 0; j < loopEnd && (elementLimit == -1 || availableCount < elementLimit); j++ {
		sel, drainOneResult := typeutil.SelectMinPK(validRetrieveResults, cursors)
		if sel == -1 || (reduce.ShouldStopWhenDrained(queryParams.reduceType) && drainOneResult) {
			break
		}

		// Get element indices for element-level query
		elemCount := 1 // default for doc-level
		if isElementLevel {
			elemIndices := validRetrieveResults[sel].GetElementIndices()[cursors[sel]]
			elemCount = len(elemIndices.GetIndices())
			ret.ElementIndices = append(ret.ElementIndices, convertInternalElementIndicesToMilvus(elemIndices))
		}

		fieldIdxs := idxComputers[sel].Compute(cursors[sel])
		retSize += typeutil.AppendFieldData(ret.FieldsData, validRetrieveResults[sel].GetFieldsData(), cursors[sel], fieldIdxs...)
		availableCount += elemCount

		// limit retrieve result to avoid oom
		if retSize > maxOutputSize {
			return nil, merr.WrapErrParameterInvalidMsg("query results exceed the maxOutputSize Limit %d", maxOutputSize)
		}

		cursors[sel]++
	}

	return ret, nil
}

// convertInternalElementIndicesToMilvus converts internalpb.ElementIndices (int32) to milvuspb.ElementIndices (int64)
func convertInternalElementIndicesToMilvus(src *internalpb.ElementIndices) *milvuspb.ElementIndices {
	if src == nil {
		return nil
	}
	indices := src.GetIndices()
	data := make([]int64, len(indices))
	for i, v := range indices {
		data[i] = int64(v)
	}
	return &milvuspb.ElementIndices{
		Indices: &schemapb.LongArray{Data: data},
	}
}

func (t *queryTask) TraceCtx() context.Context {
	return t.ctx
}

func (t *queryTask) ID() UniqueID {
	return t.Base.MsgID
}

func (t *queryTask) SetID(uid UniqueID) {
	t.Base.MsgID = uid
}

func (t *queryTask) Name() string {
	return RetrieveTaskName
}

func (t *queryTask) Type() commonpb.MsgType {
	return t.Base.MsgType
}

func (t *queryTask) BeginTs() Timestamp {
	return t.Base.Timestamp
}

func (t *queryTask) EndTs() Timestamp {
	return t.Base.Timestamp
}

func (t *queryTask) SetTs(ts Timestamp) {
	if t.reQuery && t.Base.Timestamp != 0 {
		return
	}
	t.Base.Timestamp = ts
}

func (t *queryTask) OnEnqueue() error {
	if t.Base == nil {
		t.Base = commonpbutil.NewMsgBase()
	}
	t.Base.MsgType = commonpb.MsgType_Retrieve
	t.Base.SourceID = paramtable.GetNodeID()
	return nil
}