enhance: extract proxy DQL tasks into dql package (#53685)

Extract the proxy's DQL (search/query/statistics) task cluster out of
the root proxy package into internal/proxy/dql, following the proxy
task-package split (issue #44761).

What moved:
- 16 task/pipeline files (search/query/statistics tasks, search & query
pipelines, highlighter, rerank meta, vector-type convert, function-chain
validator, plan-size & segment-filter helpers).
- The query-param *Key constants, HighlightTask and isIgnoreGrowing from
task.go, and the DQL util closure from util.go.
- The matching unit tests, rewritten with mockey instead of hand-rolled
testify mocks so global patches are cleaned up per test.

Why this design:
- The TASKS<->PIPELINE cycle stays inside dql, so it is resolved by
containment instead of surgery.
- Tasks keep private fields: root builds them via exported constructors
(NewSearchTask/NewQueryTask/NewGetStatisticsTask) that derive everything
from the taskmodel.TaskNode contract and paramtable, and reads results
through getters. No field-export or callback-injection churn.
- Root references the package only through tasks_alias.go, the
sanctioned shim, so proxy->dql->taskmodel stays acyclic (verified with
go list).
- The upsert requery keeps the accepted dml->dql edge: it builds a
dql.QueryTask and runs it via the new taskmodel.QueryRunner
(Proxy.ExecuteQuery).

The requery integration tests that spun up a real *Proxy + scheduler
were rewritten with a mock taskmodel.TaskNode, and the statistic suite's
shard dispatch is stubbed at the lb level.

Verification: go build ./internal/proxy/...; go vet -tags dynamic,test
./internal/proxy/...; proxy + dql test suites green; go list shows no
proxy-root import in dql's dependency closure.

issue: #44761

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

internal/proxy/dql/OWNERS

@@ -0,0 +1,10 @@
+reviewers:
+  - MrPresent-Han
+  - junjiejiangjjj
+  - aoiasd
+  - XuanYang-cn
+  - czs007
+  - congqixia
+
+approvers:
+  - maintainers

internal/proxy/dql/README.md

@@ -0,0 +1,132 @@
+# DQL Package
+
+The `dql` package owns the proxy's **DQL (data query language)** task
+implementations: search, query, and statistics, together with the search/query
+pipelines that execute them. It was extracted from the proxy root package
+(issue #44761) as part of the proxy task-package split.
+
+The package keeps the search/query TASKS and their PIPELINE together: the
+`searchTask`/`queryTask`/statistics tasks, the pipeline operators, the
+highlighter, plan-size checks, rerank metadata, vector-type conversion, and the
+util closure they share all live here. The only cross-group edge is
+`dml -> dql`: the upsert requery (currently in the proxy root package, to be
+extracted into `internal/proxy/dml`) builds a `QueryTask` and executes it
+through `QueryRunner`. The graph stays acyclic.
+
+## Overview
+
+DQL requests arrive at the proxy's gRPC/REST handlers, which construct a task
+via the exported constructors, enqueue it on the scheduler's `DqQueue`/`DdQueue`,
+and wait for the result. Execution flows through the search/query pipelines
+(`PreExecute -> Execute -> PostExecute`), which fan out to query nodes via the
+shard client, reduce per-shard results, and reconstruct the final response.
+
+### Tasks
+
+- **`SearchTask`** — a search request (`*milvuspb.SearchRequest`). `NewSearchTask`
+  wires the host node (`taskmodel.TaskNode`) and scheduler, derives
+  `MetaCache`/LB policy/shard manager/channel manager from the node, and stores
+  request-specific inputs only.
+- **`QueryTask`** — a query-by-PK / expression request. `NewQueryTask` is also
+  used by the upsert requery path (the accepted `dml -> dql` edge).
+- **`GetStatisticsTask` / `GetCollectionStatisticsTask` /
+  `GetPartitionStatisticsTask`** — collection/partition statistics.
+- **`HighlightTask`** — the post-search lexical highlight task, enqueued on the
+  scheduler's `DqQueue` by the highlighter operator.
+
+### Pipelines
+
+The search pipeline is a chain of operators built from the request and the
+schema: query-plan generation, partition-key resolution, query execution,
+reduce, rerank, highlight, and result organization. The requery operator
+rebuilds a `QueryTask` from the search result IDs and runs it back through the
+host node's `taskmodel.QueryRunner` (`Proxy.ExecuteQuery` in the root package).
+
+## Responsibilities
+
+1. **Task construction** — exported constructors that take only request-specific
+   inputs; everything derived from the host node comes from the
+   `taskmodel.TaskNode` contract and `paramtable`, so the root package never
+   reaches into private task fields.
+2. **Search/query execution** — plan building (`planparserv2`), placeholder
+   conversion, output-field translation, aggregation, iteration, group-by and
+   hybrid search, and multi-vector-field handling.
+3. **Reduce & rerank** — per-shard result merging, top-k selection, group-by
+   reduce, function-chain rerank metadata, and rank/score merging.
+4. **Query-param constants** — the search/query `*Key` constants (in `keys.go`)
+   consumed across the proxy.
+5. **Shared util closure** — `util_dql.go` holds the dql-only helpers plus copies
+   of small shared helpers (name/partition-tag validation, guarantee-ts parsing,
+   namespace routing), so DQL does not depend on the root package.
+
+## Architecture
+
+```
+┌──────────────────────────────────────────────────────────────┐
+│                            dql                                │
+│                                                               │
+│   SearchTask ──► searchPipeline ──► operator chain             │
+│      │              │  plan · partition · requery · reduce    │
+│      │              │  rerank · highlight · organize          │
+│      └──► highlighter ──► HighlightTask ──► sched.DqQueue     │
+│                                                               │
+│   QueryTask ──► queryPipeline ──► shard read ──► reduce        │
+│      │                                                         │
+│      └──► (requery from dml/upsert via QueryRunner)            │
+│                                                               │
+│   GetStatisticsTask / GetCollection / GetPartitionStatistics  │
+│   HighlightTask                                               │
+└──────────────────────────────────────────────────────────────┘
+```
+
+### Key constructors
+
+```go
+func NewSearchTask(node taskmodel.TaskNode, sched *scheduler.TaskScheduler,
+    ctx context.Context, request *milvuspb.SearchRequest,
+    optimizedSearch bool, isRecallEvaluation bool,
+    tr *timerecord.TimeRecorder) *SearchTask
+
+func NewQueryTask(node taskmodel.TaskNode, ctx context.Context,
+    request *milvuspb.QueryRequest, plan *planpb.PlanNode,
+    retrieveReq *internalpb.RetrieveRequest, cache metacache.Cache) *QueryTask
+
+func NewGetStatisticsTask(node taskmodel.TaskNode, ctx context.Context,
+    request *milvuspb.GetStatisticsRequest, tr *timerecord.TimeRecorder) *GetStatisticsTask
+
+func ConvertHybridSearchToSearch(req *milvuspb.HybridSearchRequest) *milvuspb.SearchRequest
+func PickFieldData(ids *schemapb.IDs, pkOffset map[any]int,
+    fields []*schemapb.FieldData, schema *schemapb.CollectionSchema,
+    collectionID int64) ([]*schemapb.FieldData, error)
+```
+
+## Host-node contract
+
+Tasks hold a `taskmodel.TaskNode` (implemented by `*Proxy` in the root package)
+and read everything they need through it: `GetMetaCache`, `LBPolicy`, `ShardMgr`,
+`ChMgr`, and `TsoAllocator`. Requery execution goes through the separate
+`taskmodel.QueryRunner` interface (`Proxy.ExecuteQuery`), whose concrete task
+type is asserted inside the root implementation.
+
+## Dependency rule
+
+`dql` imports `taskmodel`, `scheduler`, `metacache`, `channelmgr`, `shardclient`,
+`fieldvalidator`, `search_agg`, and `accesslog` (all leaf/earlier-extracted
+proxy sub-packages), plus `internal/types`, `agg`, `planparserv2`, `reduce`,
+`segcore`, the function-chain packages, and `pkg/v3`. It never imports the
+`internal/proxy` root package — verified by `go list -deps` — so the edges
+`root -> dql -> taskmodel/...` stay acyclic.
+
+## Related components
+
+- **taskmodel** (`internal/proxy/taskmodel/`): the `Task`/`DMLTask` contracts,
+  `TaskNode`, `QueryRunner`, and the shared value types.
+- **scheduler** (`internal/proxy/scheduler/`): the `DqQueue`/`DdQueue` the DQL
+  tasks are enqueued on; the highlighter injects its `*TaskScheduler`.
+- **shardclient / channelmgr / metacache** — shard-leader selection, channel
+  management, and schema/metadata cache used during execution.
+- **proxy root** (`internal/proxy/`): constructs the tasks, implements
+  `TaskNode`/`QueryRunner`, and owns `tasks_alias.go` (the only place allowed to
+  reference this package). The upsert requery (`task_upsert.go`) consumes
+  `QueryTask` — the single `dml -> dql` edge, to be extracted with the DML
+  package.

internal/proxy/dql/aliases.go

@@ -0,0 +1,45 @@
+// 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 dql
+
+import (
+	"context"
+
+	"github.com/milvus-io/milvus/internal/proxy/metacache"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
+	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
+)
+
+// Local aliases for the shared model types, declared here so the concrete task
+// structs keep their original (embedded) names without importing the root
+// proxy package. Follows the pChan/vChan precedent in DEPENDENCIES.md.
+type (
+	baseTask        = taskmodel.BaseTask
+	Condition       = taskmodel.Condition
+	schemaInfo      = metacache.SchemaInfo
+	Timestamp       = typeutil.Timestamp
+	UniqueID        = typeutil.UniqueID
+	pChan           = taskmodel.PChan
+	vChan           = taskmodel.VChan
+	pChanStatistics = taskmodel.PChanStatistics
+	Cache           = metacache.Cache
+	collectionInfo  = metacache.CollectionInfo
+)
+
+func NewTaskCondition(ctx context.Context) *taskmodel.TaskCondition {
+	return taskmodel.NewTaskCondition(ctx)
+}

internal/proxy/dql/function_chain_validator.go

@@ -13,7 +13,7 @@
 // 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
+package dql
 
 import (
 	"strings"

internal/proxy/dql/function_chain_validator_test.go

@@ -13,7 +13,7 @@
 // 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
+package dql
 
 import (
 	"strings"
@@ -25,6 +25,7 @@ import (
 	"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/proxy/metacache"
 	chainexpr "github.com/milvus-io/milvus/internal/util/function/chain/expr"
 	"github.com/milvus-io/milvus/internal/util/function/chain/types"
 	"github.com/milvus-io/milvus/pkg/v3/common"
@@ -758,3 +759,13 @@ func mapOp(output string, exprName string, args ...*schemapb.FunctionChainExprAr
 func columnArg(name string) *schemapb.FunctionChainExprArg {
 	return &schemapb.FunctionChainExprArg{Arg: &schemapb.FunctionChainExprArg_Column{Column: &schemapb.FunctionChainColumnArg{Name: name}}}
 }
+
+// mustNewSchemaInfo builds a schemaInfo for tests, mirroring the helper that
+// lives in the root proxy test suite.
+func mustNewSchemaInfo(schema *schemapb.CollectionSchema) *schemaInfo {
+	si, err := metacache.NewSchemaInfo(schema)
+	if err != nil {
+		panic(err)
+	}
+	return si
+}

internal/proxy/dql/highlight_task.go

@@ -0,0 +1,139 @@
+// 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 dql
+
+import (
+	"context"
+	"strconv"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
+	"github.com/milvus-io/milvus/internal/proxy/shardclient"
+	"github.com/milvus-io/milvus/internal/types"
+	"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/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/typeutil"
+)
+
+// HighlightTaskName is the task name for the post-search highlight task.
+const HighlightTaskName = "Highlight"
+
+// git highlight after search
+type HighlightTask struct {
+	baseTask
+	Condition
+	*querypb.GetHighlightRequest
+	ctx            context.Context
+	collectionName string
+	collectionID   typeutil.UniqueID
+	dbName         string
+	lb             shardclient.LBPolicy
+
+	result *querypb.GetHighlightResponse
+}
+
+func (t *HighlightTask) TraceCtx() context.Context {
+	return t.ctx
+}
+
+func (t *HighlightTask) ID() UniqueID {
+	return t.Base.MsgID
+}
+
+func (t *HighlightTask) SetID(uid UniqueID) {
+	t.Base.MsgID = uid
+}
+
+func (t *HighlightTask) Name() string {
+	return HighlightTaskName
+}
+
+func (t *HighlightTask) Type() commonpb.MsgType {
+	return t.Base.MsgType
+}
+
+func (t *HighlightTask) BeginTs() Timestamp {
+	return t.Base.Timestamp
+}
+
+func (t *HighlightTask) EndTs() Timestamp {
+	return t.Base.Timestamp
+}
+
+func (t *HighlightTask) SetTs(ts Timestamp) {
+	t.Base.Timestamp = ts
+}
+
+func (t *HighlightTask) OnEnqueue() error {
+	if t.Base == nil {
+		t.Base = commonpbutil.NewMsgBase()
+	}
+	t.Base.MsgType = commonpb.MsgType_Undefined
+	t.Base.SourceID = paramtable.GetNodeID()
+	return nil
+}
+
+func (t *HighlightTask) PreExecute(ctx context.Context) error {
+	return nil
+}
+
+func (t *HighlightTask) getHighlightOnShardleader(ctx context.Context, nodeID int64, qn types.QueryNodeClient, channel string) error {
+	ctx = retry.WithMaxAttemptsContext(ctx, 1)
+	t.Channel = channel
+	resp, err := qn.GetHighlight(ctx, t.GetHighlightRequest)
+	if err != nil {
+		return err
+	}
+
+	if err := merr.Error(resp.GetStatus()); err != nil {
+		return err
+	}
+	t.result = resp
+	return nil
+}
+
+func (t *HighlightTask) Execute(ctx context.Context) error {
+	err := t.lb.ExecuteOneChannel(ctx, shardclient.CollectionWorkLoad{
+		Db:             t.dbName,
+		CollectionName: t.collectionName,
+		CollectionID:   t.collectionID,
+		Nq:             int64(len(t.GetTopks()) * len(t.GetTasks())),
+		Exec:           t.getHighlightOnShardleader,
+	})
+
+	return err
+}
+
+func (t *HighlightTask) PostExecute(ctx context.Context) error {
+	return nil
+}
+
+// isIgnoreGrowing is used to check if the request should ignore growing
+func isIgnoreGrowing(params []*commonpb.KeyValuePair) (bool, error) {
+	for _, kv := range params {
+		if kv.GetKey() == IgnoreGrowingKey {
+			ignoreGrowing, err := strconv.ParseBool(kv.GetValue())
+			if err != nil {
+				return false, merr.WrapErrParameterInvalidMsg("parse ignore growing field failed")
+			}
+			return ignoreGrowing, nil
+		}
+	}
+	return false, nil
+}

internal/proxy/dql/highlight_task_test.go

@@ -0,0 +1,270 @@
+// 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 dql
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
+	"github.com/milvus-io/milvus/internal/mocks"
+	"github.com/milvus-io/milvus/internal/proxy/shardclient"
+	"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/merr"
+	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/v3/util/uniquegenerator"
+)
+
+func TestHighlightTask(t *testing.T) {
+	paramtable.Init()
+	ctx := context.Background()
+
+	t.Run("basic methods", func(t *testing.T) {
+		task := &HighlightTask{
+			ctx: ctx,
+			GetHighlightRequest: &querypb.GetHighlightRequest{
+				Base: commonpbutil.NewMsgBase(),
+			},
+			Condition: NewTaskCondition(ctx),
+		}
+
+		t.Run("traceCtx", func(t *testing.T) {
+			traceCtx := task.TraceCtx()
+			assert.NotNil(t, traceCtx)
+			assert.Equal(t, ctx, traceCtx)
+		})
+
+		t.Run("id", func(t *testing.T) {
+			id := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
+			task.SetID(id)
+			assert.Equal(t, id, task.ID())
+		})
+
+		t.Run("name", func(t *testing.T) {
+			assert.Equal(t, HighlightTaskName, task.Name())
+		})
+
+		t.Run("type", func(t *testing.T) {
+			assert.Equal(t, commonpb.MsgType_Undefined, task.Type())
+		})
+
+		t.Run("ts", func(t *testing.T) {
+			ts := Timestamp(time.Now().UnixNano())
+			task.SetTs(ts)
+			assert.Equal(t, ts, task.BeginTs())
+			assert.Equal(t, ts, task.EndTs())
+		})
+	})
+
+	t.Run("OnEnqueue", func(t *testing.T) {
+		task := &HighlightTask{
+			ctx:                 ctx,
+			GetHighlightRequest: &querypb.GetHighlightRequest{},
+			Condition:           NewTaskCondition(ctx),
+		}
+
+		err := task.OnEnqueue()
+		assert.NoError(t, err)
+		assert.NotNil(t, task.Base)
+		assert.Equal(t, commonpb.MsgType_Undefined, task.Base.MsgType)
+		assert.Equal(t, paramtable.GetNodeID(), task.Base.SourceID)
+	})
+
+	t.Run("PreExecute", func(t *testing.T) {
+		task := &HighlightTask{
+			ctx: ctx,
+			GetHighlightRequest: &querypb.GetHighlightRequest{
+				Base: commonpbutil.NewMsgBase(),
+			},
+			Condition: NewTaskCondition(ctx),
+		}
+
+		err := task.PreExecute(ctx)
+		assert.NoError(t, err)
+	})
+
+	t.Run("getHighlightOnShardleader success", func(t *testing.T) {
+		task := &HighlightTask{
+			ctx: ctx,
+			GetHighlightRequest: &querypb.GetHighlightRequest{
+				Base:  commonpbutil.NewMsgBase(),
+				Topks: []int64{10},
+				Tasks: []*querypb.HighlightTask{
+					{},
+				},
+			},
+			Condition: NewTaskCondition(ctx),
+		}
+
+		mockQN := mocks.NewMockQueryNodeClient(t)
+		expectedResp := &querypb.GetHighlightResponse{
+			Status: merr.Success(),
+			Results: []*querypb.HighlightResult{
+				{
+					Fragments: []*querypb.HighlightFragment{
+						{
+							StartOffset: 0,
+							EndOffset:   10,
+							Offsets:     []int64{0, 5, 5, 10},
+						},
+					},
+				},
+			},
+		}
+
+		mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(expectedResp, nil)
+
+		err := task.getHighlightOnShardleader(ctx, 1, mockQN, "test_channel")
+		assert.NoError(t, err)
+		assert.NotNil(t, task.result)
+		assert.Equal(t, expectedResp, task.result)
+		assert.Equal(t, "test_channel", task.Channel)
+	})
+
+	t.Run("getHighlightOnShardleader rpc error", func(t *testing.T) {
+		task := &HighlightTask{
+			ctx: ctx,
+			GetHighlightRequest: &querypb.GetHighlightRequest{
+				Base: commonpbutil.NewMsgBase(),
+			},
+			Condition: NewTaskCondition(ctx),
+		}
+
+		mockQN := mocks.NewMockQueryNodeClient(t)
+		mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(nil, errors.New("rpc error"))
+
+		err := task.getHighlightOnShardleader(ctx, 1, mockQN, "test_channel")
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "rpc error")
+	})
+
+	t.Run("getHighlightOnShardleader status error", func(t *testing.T) {
+		task := &HighlightTask{
+			ctx: ctx,
+			GetHighlightRequest: &querypb.GetHighlightRequest{
+				Base: commonpbutil.NewMsgBase(),
+			},
+			Condition: NewTaskCondition(ctx),
+		}
+
+		mockQN := mocks.NewMockQueryNodeClient(t)
+		expectedResp := &querypb.GetHighlightResponse{
+			Status: merr.Status(errors.New("status error")),
+		}
+
+		mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(expectedResp, nil)
+
+		err := task.getHighlightOnShardleader(ctx, 1, mockQN, "test_channel")
+		assert.Error(t, err)
+	})
+
+	t.Run("Execute success", func(t *testing.T) {
+		task := &HighlightTask{
+			ctx:            ctx,
+			collectionName: "test_collection",
+			collectionID:   100,
+			dbName:         "default",
+			GetHighlightRequest: &querypb.GetHighlightRequest{
+				Base:  commonpbutil.NewMsgBase(),
+				Topks: []int64{10, 20},
+				Tasks: []*querypb.HighlightTask{
+					{},
+					{},
+				},
+			},
+			Condition: NewTaskCondition(ctx),
+		}
+
+		mockLB := shardclient.NewMockLBPolicy(t)
+		mockQN := mocks.NewMockQueryNodeClient(t)
+
+		expectedResp := &querypb.GetHighlightResponse{
+			Status: merr.Success(),
+			Results: []*querypb.HighlightResult{
+				{
+					Fragments: []*querypb.HighlightFragment{
+						{
+							StartOffset: 0,
+							EndOffset:   10,
+							Offsets:     []int64{0, 5, 5, 10},
+						},
+					},
+				},
+			},
+		}
+
+		mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(expectedResp, nil)
+
+		mockLB.EXPECT().ExecuteOneChannel(mock.Anything, mock.Anything).Run(
+			func(ctx context.Context, workload shardclient.CollectionWorkLoad) {
+				assert.Equal(t, "default", workload.Db)
+				assert.Equal(t, "test_collection", workload.CollectionName)
+				assert.Equal(t, int64(100), workload.CollectionID)
+				assert.Equal(t, int64(4), workload.Nq) // len(topks) * len(tasks) = 2 * 2 = 4
+				err := workload.Exec(ctx, 1, mockQN, "test_channel")
+				assert.NoError(t, err)
+			},
+		).Return(nil)
+
+		task.lb = mockLB
+
+		err := task.Execute(ctx)
+		assert.NoError(t, err)
+		assert.NotNil(t, task.result)
+	})
+
+	t.Run("Execute lb error", func(t *testing.T) {
+		task := &HighlightTask{
+			ctx:            ctx,
+			collectionName: "test_collection",
+			collectionID:   100,
+			dbName:         "default",
+			GetHighlightRequest: &querypb.GetHighlightRequest{
+				Base: commonpbutil.NewMsgBase(),
+			},
+			Condition: NewTaskCondition(ctx),
+		}
+
+		mockLB := shardclient.NewMockLBPolicy(t)
+		mockLB.EXPECT().ExecuteOneChannel(mock.Anything, mock.Anything).Return(errors.New("lb error"))
+
+		task.lb = mockLB
+
+		err := task.Execute(ctx)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "lb error")
+	})
+
+	t.Run("PostExecute", func(t *testing.T) {
+		task := &HighlightTask{
+			ctx: ctx,
+			GetHighlightRequest: &querypb.GetHighlightRequest{
+				Base: commonpbutil.NewMsgBase(),
+			},
+			Condition: NewTaskCondition(ctx),
+		}
+
+		err := task.PostExecute(ctx)
+		assert.NoError(t, err)
+	})
+}

internal/proxy/dql/highlighter.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"
@@ -40,7 +40,7 @@ const (
 )
 
 type Highlighter interface {
-	AsSearchPipelineOperator(t *searchTask) (operator, error)
+	AsSearchPipelineOperator(t *SearchTask) (operator, error)
 	FieldIDs() []int64
 	RequiredFieldIDs() []int64
 	DynamicFieldNames() []string
@@ -151,7 +151,7 @@ func (h *LexicalHighlighter) addTaskWithQuery(collInfo *schemaInfo, fieldID int6
 	return nil
 }
 
-func (h *LexicalHighlighter) initHighlightQueries(t *searchTask, analyzerName string) error {
+func (h *LexicalHighlighter) initHighlightQueries(t *SearchTask, analyzerName string) error {
 	// add query to highlight tasks
 	for _, query := range h.queries {
 		fieldID, ok := t.schema.MapFieldID(query.fieldName)
@@ -165,7 +165,7 @@ func (h *LexicalHighlighter) initHighlightQueries(t *searchTask, analyzerName st
 	return nil
 }
 
-func (h *LexicalHighlighter) AsSearchPipelineOperator(t *searchTask) (operator, error) {
+func (h *LexicalHighlighter) AsSearchPipelineOperator(t *SearchTask) (operator, error) {
 	return newLexicalHighlightOperator(t, lo.Values(h.tasks))
 }
 
@@ -328,11 +328,11 @@ type lexicalHighlightOperator struct {
 	dbName         string
 }
 
-func newLexicalHighlightOperator(t *searchTask, tasks []*highlightTask) (operator, error) {
+func newLexicalHighlightOperator(t *SearchTask, tasks []*highlightTask) (operator, error) {
 	return &lexicalHighlightOperator{
 		tasks:          tasks,
 		lbPolicy:       t.lb,
-		scheduler:      t.node.(*Proxy).sched,
+		scheduler:      t.sched,
 		schema:         t.schema,
 		collectionName: t.request.CollectionName,
 		collectionID:   t.CollectionID,
@@ -528,7 +528,7 @@ type SemanticHighlighter struct {
 	highlight *highlight.SemanticHighlight
 }
 
-func newSemanticHighlighter(t *searchTask, extraInfo *models.ModelExtraInfo) (*SemanticHighlighter, error) {
+func newSemanticHighlighter(t *SearchTask, extraInfo *models.ModelExtraInfo) (*SemanticHighlighter, error) {
 	conf := paramtable.Get().FunctionCfg.ZillizProviders.GetValue()
 	highlight, err := highlight.NewSemanticHighlight(t.schema.CollectionSchema, t.request.GetHighlighter().GetParams(), conf, extraInfo)
 	if err != nil {
@@ -549,7 +549,7 @@ func (h *SemanticHighlighter) DynamicFieldNames() []string {
 	return h.highlight.DynamicFieldNames()
 }
 
-func (h *SemanticHighlighter) AsSearchPipelineOperator(t *searchTask) (operator, error) {
+func (h *SemanticHighlighter) AsSearchPipelineOperator(t *SearchTask) (operator, error) {
 	return &semanticHighlightOperator{highlight: h.highlight}, nil
 }
 

internal/proxy/dql/highlighter_test.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"encoding/json"
@@ -78,7 +78,7 @@ func (s *HighlighterSuite) TestLexicalHighlighter_AsSearchPipelineOperator() {
 	}
 
 	// Mock newLexicalHighlightOperator to avoid nil pointer dereference
-	mockOp := mockey.Mock(newLexicalHighlightOperator).To(func(t *searchTask, tasks []*highlightTask) (operator, error) {
+	mockOp := mockey.Mock(newLexicalHighlightOperator).To(func(t *SearchTask, tasks []*highlightTask) (operator, error) {
 		return &lexicalHighlightOperator{tasks: tasks}, nil
 	}).Build()
 	defer mockOp.UnPatch()
@@ -93,9 +93,9 @@ func (s *HighlighterSuite) TestLexicalHighlighter_AsSearchPipelineOperator() {
 // ============== SemanticHighlighter Tests ==============
 
 func (s *HighlighterSuite) TestSemanticHighlighter_RequiredFieldIDs_SchemaFieldsOnly() {
-	mock1 := mockey.Mock(zilliz.NewZilliClient).To(func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
+	mock1 := mockTestTo(s.T(), zilliz.NewZilliClient, func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
 		return &zilliz.ZillizClient{}, nil
-	}).Build()
+	})
 	defer mock1.UnPatch()
 
 	schema := &schemapb.CollectionSchema{
@@ -136,9 +136,9 @@ func (s *HighlighterSuite) TestSemanticHighlighter_RequiredFieldIDs_SchemaFields
 }
 
 func (s *HighlighterSuite) TestSemanticHighlighter_RequiredFieldIDs_WithDynamicFields() {
-	mock1 := mockey.Mock(zilliz.NewZilliClient).To(func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
+	mock1 := mockTestTo(s.T(), zilliz.NewZilliClient, func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
 		return &zilliz.ZillizClient{}, nil
-	}).Build()
+	})
 	defer mock1.UnPatch()
 
 	schema := &schemapb.CollectionSchema{
@@ -182,9 +182,9 @@ func (s *HighlighterSuite) TestSemanticHighlighter_RequiredFieldIDs_WithDynamicF
 }
 
 func (s *HighlighterSuite) TestSemanticHighlighter_RequiredFieldIDs_DoesNotMutateFieldIDs() {
-	mock1 := mockey.Mock(zilliz.NewZilliClient).To(func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
+	mock1 := mockTestTo(s.T(), zilliz.NewZilliClient, func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
 		return &zilliz.ZillizClient{}, nil
-	}).Build()
+	})
 	defer mock1.UnPatch()
 
 	schema := &schemapb.CollectionSchema{
@@ -243,9 +243,9 @@ func (s *HighlighterSuite) TestSemanticHighlighter_RequiredFieldIDs_DoesNotMutat
 }
 
 func (s *HighlighterSuite) TestSemanticHighlighter_DynamicFieldNames_Empty() {
-	mock1 := mockey.Mock(zilliz.NewZilliClient).To(func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
+	mock1 := mockTestTo(s.T(), zilliz.NewZilliClient, func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
 		return &zilliz.ZillizClient{}, nil
-	}).Build()
+	})
 	defer mock1.UnPatch()
 
 	schema := &schemapb.CollectionSchema{
@@ -282,9 +282,9 @@ func (s *HighlighterSuite) TestSemanticHighlighter_DynamicFieldNames_Empty() {
 }
 
 func (s *HighlighterSuite) TestSemanticHighlighter_DynamicFieldNames_WithDynamicFields() {
-	mock1 := mockey.Mock(zilliz.NewZilliClient).To(func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
+	mock1 := mockTestTo(s.T(), zilliz.NewZilliClient, func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
 		return &zilliz.ZillizClient{}, nil
-	}).Build()
+	})
 	defer mock1.UnPatch()
 
 	schema := &schemapb.CollectionSchema{
@@ -323,9 +323,9 @@ func (s *HighlighterSuite) TestSemanticHighlighter_DynamicFieldNames_WithDynamic
 }
 
 func (s *HighlighterSuite) TestSemanticHighlighter_AsSearchPipelineOperator() {
-	mock1 := mockey.Mock(zilliz.NewZilliClient).To(func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
+	mock1 := mockTestTo(s.T(), zilliz.NewZilliClient, func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
 		return &zilliz.ZillizClient{}, nil
-	}).Build()
+	})
 	defer mock1.UnPatch()
 
 	schema := &schemapb.CollectionSchema{
@@ -378,7 +378,7 @@ func (s *HighlighterSuite) TestHighlighterInterface_LexicalHighlighter() {
 	s.NotNil(h.DynamicFieldNames())
 
 	// Mock newLexicalHighlightOperator to avoid nil pointer dereference
-	mockOp := mockey.Mock(newLexicalHighlightOperator).To(func(t *searchTask, tasks []*highlightTask) (operator, error) {
+	mockOp := mockey.Mock(newLexicalHighlightOperator).To(func(t *SearchTask, tasks []*highlightTask) (operator, error) {
 		return &lexicalHighlightOperator{}, nil
 	}).Build()
 	defer mockOp.UnPatch()
@@ -389,9 +389,9 @@ func (s *HighlighterSuite) TestHighlighterInterface_LexicalHighlighter() {
 }
 
 func (s *HighlighterSuite) TestHighlighterInterface_SemanticHighlighter() {
-	mock1 := mockey.Mock(zilliz.NewZilliClient).To(func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
+	mock1 := mockTestTo(s.T(), zilliz.NewZilliClient, func(_ string, _ string, _ string, _ map[string]string, _ int64) (*zilliz.ZillizClient, error) {
 		return &zilliz.ZillizClient{}, nil
-	}).Build()
+	})
 	defer mock1.UnPatch()
 
 	schema := &schemapb.CollectionSchema{

internal/proxy/dql/keys.go

@@ -0,0 +1,71 @@
+// 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 dql
+
+import (
+	"math"
+
+	"github.com/milvus-io/milvus/pkg/v3/common"
+)
+
+const (
+	SumScorer string = "sum"
+	MaxScorer string = "max"
+	AvgScorer string = "avg"
+
+	IgnoreGrowingKey     = "ignore_growing"
+	ReduceStopForBestKey = "reduce_stop_for_best"
+	IteratorField        = "iterator"
+	CollectionID         = "collection_id"
+	GroupByFieldKey      = "group_by_field"
+	GroupSizeKey         = "group_size"
+	StrictGroupSize      = "strict_group_size"
+	JSONPath             = "json_path"
+	JSONType             = "json_type"
+	StrictCastKey        = "strict_cast"
+	RankGroupScorer      = "rank_group_scorer"
+	AnnsFieldKey         = "anns_field"
+	AnalyzerKey          = "analyzer_name"
+	TopKKey              = "topk"
+	NQKey                = "nq"
+	MetricTypeKey        = common.MetricTypeKey
+	ParamsKey            = common.ParamsKey
+	ExprParamsKey        = "expr_params"
+	RoundDecimalKey      = "round_decimal"
+	OffsetKey            = "offset"
+	LimitKey             = "limit"
+	// key for timestamptz translation
+	TimefieldsKey = "time_fields"
+
+	SearchIterV2Key        = "search_iter_v2"
+	SearchIterBatchSizeKey = "search_iter_batch_size"
+	SearchIterLastBoundKey = "search_iter_last_bound"
+	SearchIterIdKey        = "search_iter_id"
+	QueryIterLastPKKey     = "query_iter_last_pk"
+	QueryIterLastOffsetKey = "query_iter_last_element_offset"
+	GroupByFieldsKey       = "group_by_fields"
+	OrderByFieldsKey       = "order_by_fields"
+	PipelineTraceKey       = "pipeline_trace"
+
+	// minFloat32 minimum float.
+	minFloat32 = -1 * float32(math.MaxFloat32)
+
+	RankTypeKey      = "strategy"
+	RRFParamsKey     = "k"
+	WeightsParamsKey = "weights"
+	NormScoreKey     = "norm_score"
+)

internal/proxy/dql/membership_filter_plan_size.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"fmt"
@@ -28,7 +28,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
 )
 
-// wrapPlanCreationError adds request context to a plan-creation failure and
+// WrapPlanCreationError adds request context to a plan-creation failure and
 // keeps the Proxy's existing ParameterInvalid (1100) projection for expression
 // parsing errors, while preserving the cause in the chain -- errors.Is still
 // matches merr.ErrQueryPlan.
@@ -42,7 +42,7 @@ import (
 // this boundary to 2201 is a worthwhile cross-cutting change, but it needs its
 // own PR: REST, Go SDK and Delete expectations updated together, plus a
 // release note.
-func wrapPlanCreationError(err error, context string) error {
+func WrapPlanCreationError(err error, context string) error {
 	// Membership-filter size checks are a new, more specific public error. Do
 	// not collapse them into the legacy ParameterInvalid projection below.
 	if errors.Is(err, merr.ErrParameterTooLarge) {
@@ -51,7 +51,7 @@ func wrapPlanCreationError(err error, context string) error {
 	return merr.Combine(merr.Wrap(err, context), merr.ErrParameterInvalid)
 }
 
-// marshalPlanWithMembershipFilterSizeLimit serializes a plan after accounting
+// MarshalPlanWithMembershipFilterSizeLimit serializes a plan after accounting
 // for its size against the request-wide membership-filter plan budget. Plans
 // carrying no client pre-built membership blob do not consume it. For
 // HybridSearch, callers pass the returned accumulated size into the next
@@ -67,7 +67,7 @@ func wrapPlanCreationError(err error, context string) error {
 // proto.Size is intentionally evaluated before proto.Marshal: repeated uses of
 // one client blob create independent blob fields in the compiled plan, and
 // rejecting here avoids allocating the amplified serialized buffer.
-func marshalPlanWithMembershipFilterSizeLimit(plan *planpb.PlanNode, accumulatedSize int64) ([]byte, int64, error) {
+func MarshalPlanWithMembershipFilterSizeLimit(plan *planpb.PlanNode, accumulatedSize int64) ([]byte, int64, error) {
 	nextSize := accumulatedSize
 	if planparserv2.PlanContainsMembershipFilter(plan) {
 		planSize := int64(proto.Size(plan))

internal/proxy/dql/membership_filter_plan_size_test.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"context"
@@ -22,7 +22,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
 
@@ -30,6 +29,7 @@ import (
 	"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/client/v3/membership/sbbf"
+	"github.com/milvus-io/milvus/internal/proxy/metacache"
 	"github.com/milvus-io/milvus/pkg/v3/common"
 	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
 	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
@@ -47,7 +47,7 @@ func TestWrapPlanCreationError(t *testing.T) {
 	// the Go SDK geometry cases.
 	t.Run("keeps ordinary parser failures parameter invalid", func(t *testing.T) {
 		cause := merr.WrapErrQueryPlanMsg("invalid expression")
-		err := wrapPlanCreationError(cause, "failed to create query plan")
+		err := WrapPlanCreationError(cause, "failed to create query plan")
 		require.Error(t, err)
 		assert.ErrorIs(t, err, cause)
 		// The cause survives in the chain even though the projected code is 1100.
@@ -61,7 +61,7 @@ func TestWrapPlanCreationError(t *testing.T) {
 
 	t.Run("delete keeps its InputError classification", func(t *testing.T) {
 		cause := merr.WrapErrQueryPlanMsg("invalid expression")
-		err := merr.WrapErrAsInputError(wrapPlanCreationError(cause, "failed to create delete plan"))
+		err := merr.WrapErrAsInputError(WrapPlanCreationError(cause, "failed to create delete plan"))
 		require.Error(t, err)
 		assert.ErrorIs(t, err, merr.ErrQueryPlan)
 		assert.Equal(t, int32(1100), merr.Code(err))
@@ -69,7 +69,7 @@ func TestWrapPlanCreationError(t *testing.T) {
 
 	t.Run("preserves parameter too large", func(t *testing.T) {
 		cause := merr.WrapErrParameterTooLarge("membership-filter plan")
-		err := wrapPlanCreationError(cause, "failed to create query plan")
+		err := WrapPlanCreationError(cause, "failed to create query plan")
 		require.Error(t, err)
 		assert.ErrorIs(t, err, merr.ErrParameterTooLarge)
 		assert.Equal(t, int32(1102), merr.Code(err))
@@ -119,7 +119,7 @@ func TestMarshalPlanWithMembershipFilterSizeLimit(t *testing.T) {
 
 	t.Run("exact limit succeeds", func(t *testing.T) {
 		params.Save(key, strconv.Itoa(repeatedPlanSize))
-		serialized, accumulated, err := marshalPlanWithMembershipFilterSizeLimit(repeatedPlan, 0)
+		serialized, accumulated, err := MarshalPlanWithMembershipFilterSizeLimit(repeatedPlan, 0)
 		require.NoError(t, err)
 		assert.Len(t, serialized, repeatedPlanSize)
 		assert.Equal(t, int64(repeatedPlanSize), accumulated)
@@ -127,7 +127,7 @@ func TestMarshalPlanWithMembershipFilterSizeLimit(t *testing.T) {
 
 	t.Run("repeated reference over limit is input error 1102", func(t *testing.T) {
 		params.Save(key, strconv.Itoa(repeatedPlanSize-1))
-		serialized, accumulated, err := marshalPlanWithMembershipFilterSizeLimit(repeatedPlan, 0)
+		serialized, accumulated, err := MarshalPlanWithMembershipFilterSizeLimit(repeatedPlan, 0)
 		require.Error(t, err)
 		assert.Nil(t, serialized, "the amplified plan must be rejected before proto.Marshal allocates its output")
 		assert.Zero(t, accumulated)
@@ -149,12 +149,12 @@ func TestMarshalPlanWithMembershipFilterSizeLimit(t *testing.T) {
 		singlePlanSize := proto.Size(singlePlan)
 		params.Save(key, strconv.Itoa(singlePlanSize*2-1))
 
-		first, accumulated, err := marshalPlanWithMembershipFilterSizeLimit(singlePlan, 0)
+		first, accumulated, err := MarshalPlanWithMembershipFilterSizeLimit(singlePlan, 0)
 		require.NoError(t, err)
 		assert.Len(t, first, singlePlanSize)
 		assert.Equal(t, int64(singlePlanSize), accumulated)
 
-		second, unchanged, err := marshalPlanWithMembershipFilterSizeLimit(singlePlan, accumulated)
+		second, unchanged, err := MarshalPlanWithMembershipFilterSizeLimit(singlePlan, accumulated)
 		require.Error(t, err)
 		assert.Nil(t, second)
 		assert.Equal(t, accumulated, unchanged)
@@ -166,7 +166,7 @@ func TestMarshalPlanWithMembershipFilterSizeLimit(t *testing.T) {
 			Query: &planpb.QueryPlanNode{Predicates: roaringPlanSizeTestLeaf(blob)},
 		}}
 		params.Save(key, strconv.Itoa(proto.Size(roaringPlan)-1))
-		serialized, accumulated, err := marshalPlanWithMembershipFilterSizeLimit(roaringPlan, 0)
+		serialized, accumulated, err := MarshalPlanWithMembershipFilterSizeLimit(roaringPlan, 0)
 		require.Error(t, err)
 		assert.Nil(t, serialized)
 		assert.Zero(t, accumulated)
@@ -186,7 +186,7 @@ func TestMarshalPlanWithMembershipFilterSizeLimit(t *testing.T) {
 					}},
 				}
 				params.Save(key, strconv.Itoa(proto.Size(plan)-1))
-				serialized, accumulated, err := marshalPlanWithMembershipFilterSizeLimit(plan, 0)
+				serialized, accumulated, err := MarshalPlanWithMembershipFilterSizeLimit(plan, 0)
 				require.Error(t, err)
 				assert.Nil(t, serialized)
 				assert.Zero(t, accumulated)
@@ -201,7 +201,7 @@ func TestMarshalPlanWithMembershipFilterSizeLimit(t *testing.T) {
 			OutputFieldIds: make([]int64, 4096),
 		}
 		params.Save(key, "1")
-		serialized, accumulated, err := marshalPlanWithMembershipFilterSizeLimit(nonMembershipPlan, 0)
+		serialized, accumulated, err := MarshalPlanWithMembershipFilterSizeLimit(nonMembershipPlan, 0)
 		require.NoError(t, err)
 		assert.NotEmpty(t, serialized)
 		assert.Zero(t, accumulated)
@@ -226,8 +226,7 @@ func bloomPlanSizeIntegrationFixture(t *testing.T) (*schemaInfo, map[string]*sch
 			},
 		},
 	}
-	schemaInfo, err := newSchemaInfo(schema)
-	require.NoError(t, err)
+	schemaInfo := mustNewSchemaInfo(schema)
 	return schemaInfo, map[string]*schemapb.TemplateValue{
 		"bf": {Val: &schemapb.TemplateValue_BytesVal{BytesVal: blob}},
 	}
@@ -252,7 +251,7 @@ func TestSearchTaskMembershipFilterPlanSizeLimit(t *testing.T) {
 	schema, templateValues := bloomPlanSizeIntegrationFixture(t)
 
 	t.Run("normal search", func(t *testing.T) {
-		task := &searchTask{
+		task := &SearchTask{
 			ctx:            ctx,
 			collectionName: schema.GetName(),
 			SearchRequest: &internalpb.SearchRequest{
@@ -276,7 +275,7 @@ func TestSearchTaskMembershipFilterPlanSizeLimit(t *testing.T) {
 	})
 
 	t.Run("hybrid search", func(t *testing.T) {
-		task := &searchTask{
+		task := &SearchTask{
 			ctx:            ctx,
 			collectionName: schema.GetName(),
 			SearchRequest: &internalpb.SearchRequest{
@@ -328,7 +327,7 @@ func TestSearchTaskMembershipFilterPlanSizeLimit(t *testing.T) {
 				SearchParams:       bloomPlanSizeSearchParams(),
 			}
 		}
-		task := &searchTask{
+		task := &SearchTask{
 			ctx:            ctx,
 			collectionName: schema.GetName(),
 			SearchRequest: &internalpb.SearchRequest{
@@ -362,13 +361,11 @@ func TestQueryTaskMembershipFilterPlanSizeLimit(t *testing.T) {
 
 	ctx := context.Background()
 	schema, templateValues := bloomPlanSizeIntegrationFixture(t)
-	cache := NewMockCache(t)
-	cache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(int64(1), nil)
-	cache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
-		Return(&collectionInfo{}, nil).Maybe()
-	cache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).
-		Return(schema, nil).Maybe()
-	task := &queryTask{
+	cache := newTestCache()
+	mockTest(t, (*metacache.MetaCache).GetCollectionID, int64(1), nil)
+	mockTest(t, (*metacache.MetaCache).GetCollectionInfo, &collectionInfo{}, nil)
+	mockTest(t, (*metacache.MetaCache).GetCollectionSchema, schema, nil)
+	task := &QueryTask{
 		baseTask: baseTask{
 			MetaCache: cache,
 		},

internal/proxy/dql/pipeline_trace.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"context"

internal/proxy/dql/pipeline_trace_test.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"context"

internal/proxy/dql/query_pipeline.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"context"

internal/proxy/dql/query_pipeline_test.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"context"

internal/proxy/dql/requery_integration_test.go

@@ -0,0 +1,311 @@
+// 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 dql
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace"
+
+	"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/proxy/metacache"
+	"github.com/milvus-io/milvus/internal/proxy/shardclient"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
+	"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/proto/internalpb"
+	"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/timerecord"
+	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
+)
+
+// TestSearchTask_Requery exercises the requeryOperator: it builds a requery
+// QueryTask from the search results and runs it through the node's
+// QueryRunner, then verifies the reconstructed fields. The host node is a mock
+// taskmodel.TaskNode whose QueryRunner is patched by mockey.
+func TestSearchTask_Requery(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	const (
+		dim        = 128
+		rows       = 5
+		collection = "test-requery"
+
+		pkField  = "pk"
+		vecField = "vec"
+	)
+
+	ids := make([]int64, rows)
+	for i := range ids {
+		ids[i] = int64(i)
+	}
+
+	collectionName := "col"
+	cache := newTestCache()
+	collSchema := constructCollectionSchema(pkField, vecField, dim, collection)
+	schema := mustNewSchemaInfo(collSchema)
+	mockTest(t, (*metacache.MetaCache).GetCollectionID, UniqueID(0), nil)
+	mockTest(t, (*metacache.MetaCache).GetCollectionSchema, schema, nil)
+	mockTest(t, (*metacache.MetaCache).GetPartitions, map[string]int64{"_default": UniqueID(1)}, nil)
+	mockTest(t, (*metacache.MetaCache).GetCollectionInfo, &collectionInfo{}, nil)
+
+	node := &namespaceRequeryMockNode{}
+
+	t.Run("Test normal", func(t *testing.T) {
+		resultIDs := &schemapb.IDs{
+			IdField: &schemapb.IDs_IntId{
+				IntId: &schemapb.LongArray{
+					Data: ids,
+				},
+			},
+		}
+
+		outputFields := []string{pkField, vecField}
+		qt := &SearchTask{
+			baseTask: baseTask{MetaCache: cache},
+			ctx:      ctx,
+			SearchRequest: &internalpb.SearchRequest{
+				Base: &commonpb.MsgBase{
+					MsgType:  commonpb.MsgType_Search,
+					SourceID: paramtable.GetNodeID(),
+				},
+			},
+			request: &milvuspb.SearchRequest{
+				CollectionName: collectionName,
+				OutputFields:   outputFields,
+			},
+			result: &milvuspb.SearchResults{
+				Results: &schemapb.SearchResultData{
+					Ids: resultIDs,
+				},
+			},
+			schema:                 schema,
+			tr:                     timerecord.NewTimeRecorder("search"),
+			node:                   node,
+			translatedOutputFields: outputFields,
+			queryChannelsNode:      typeutil.NewConcurrentMap[string, int64](),
+		}
+		qt.queryChannelsNode.Insert("mock_qn", 1)
+
+		// The requery query task is executed on the host node's QueryRunner.
+		// Return a result carrying pk + vec fields.
+		mockTestTo(t, (*namespaceRequeryMockNode).ExecuteQuery, func(_ *namespaceRequeryMockNode, _ context.Context, qt taskmodel.Task, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
+			queryTask := qt.(*QueryTask)
+			require.Equal(t, metrics.ReQueryLabel, queryTask.GetQueryLabel())
+			require.True(t, queryTask.ReQuery())
+			return &milvuspb.QueryResults{
+				Status: merr.Success(),
+				FieldsData: []*schemapb.FieldData{
+					{FieldName: pkField, FieldId: 100, Type: schemapb.DataType_Int64},
+					{FieldName: vecField, FieldId: 101, Type: schemapb.DataType_FloatVector},
+				},
+			}, segcore.StorageCost{}, nil
+		})
+
+		op, err := newRequeryOperator(qt, nil)
+		assert.NoError(t, err)
+		queryResult, storageCost, err := op.(*requeryOperator).requery(ctx, nil, qt.result.Results.Ids, outputFields)
+		assert.NoError(t, err)
+		assert.Equal(t, int64(0), storageCost.ScannedRemoteBytes)
+		assert.Equal(t, int64(0), storageCost.ScannedTotalBytes)
+		assert.Len(t, queryResult.FieldsData, 2)
+		for _, field := range queryResult.FieldsData {
+			assert.Contains(t, []string{pkField, vecField}, field.GetFieldName())
+		}
+	})
+
+	t.Run("Test no primary key", func(t *testing.T) {
+		collSchema := &schemapb.CollectionSchema{}
+		schema := mustNewSchemaInfo(collSchema)
+
+		qt := &SearchTask{
+			baseTask: baseTask{MetaCache: cache},
+			ctx:      ctx,
+			SearchRequest: &internalpb.SearchRequest{
+				Base: &commonpb.MsgBase{
+					MsgType:  commonpb.MsgType_Search,
+					SourceID: paramtable.GetNodeID(),
+				},
+			},
+			request: &milvuspb.SearchRequest{},
+			schema:  schema,
+			tr:      timerecord.NewTimeRecorder("search"),
+			node:    node,
+		}
+
+		_, err := newRequeryOperator(qt, nil)
+		assert.Error(t, err)
+	})
+
+	t.Run("Test requery failed", func(t *testing.T) {
+		resultIDs := &schemapb.IDs{
+			IdField: &schemapb.IDs_IntId{
+				IntId: &schemapb.LongArray{
+					Data: ids,
+				},
+			},
+		}
+		outputFields := []string{pkField, vecField}
+		qt := &SearchTask{
+			baseTask: baseTask{MetaCache: cache},
+			ctx:      ctx,
+			SearchRequest: &internalpb.SearchRequest{
+				Base: &commonpb.MsgBase{
+					MsgType:  commonpb.MsgType_Search,
+					SourceID: paramtable.GetNodeID(),
+				},
+			},
+			request: &milvuspb.SearchRequest{
+				CollectionName: collectionName,
+				OutputFields:   outputFields,
+			},
+			result: &milvuspb.SearchResults{
+				Results: &schemapb.SearchResultData{
+					Ids: resultIDs,
+				},
+			},
+			schema:                 schema,
+			tr:                     timerecord.NewTimeRecorder("search"),
+			node:                   node,
+			translatedOutputFields: outputFields,
+			queryChannelsNode:      typeutil.NewConcurrentMap[string, int64](),
+		}
+		qt.queryChannelsNode.Insert("mock_qn", 1)
+
+		mockTestTo(t, (*namespaceRequeryMockNode).ExecuteQuery, func(_ *namespaceRequeryMockNode, _ context.Context, _ taskmodel.Task, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
+			return nil, segcore.StorageCost{}, errors.New("mock requery failure")
+		})
+
+		op, err := newRequeryOperator(qt, nil)
+		assert.NoError(t, err)
+		_, _, err = op.(*requeryOperator).requery(ctx, nil, &schemapb.IDs{}, []string{})
+		assert.Error(t, err)
+	})
+}
+
+// TestSearchTask_ErrExecute exercises the error propagation of SearchTask
+// execution. The shard dispatch is stubbed at the lb level: Execute runs the
+// workload, whose Exec invokes the query node client; each scenario returns a
+// different failure and verifies the surfaced error.
+func TestSearchTask_ErrExecute(t *testing.T) {
+	ctx := context.TODO()
+	collectionName := t.Name() + funcutil.GenRandomStr()
+
+	cache := newTestCache()
+	fieldName2Types := map[string]schemapb.DataType{
+		testBoolField:     schemapb.DataType_Bool,
+		testInt32Field:    schemapb.DataType_Int32,
+		testInt64Field:    schemapb.DataType_Int64,
+		testFloatField:    schemapb.DataType_Float,
+		testDoubleField:   schemapb.DataType_Double,
+		testFloatVecField: schemapb.DataType_FloatVector,
+	}
+	if enableMultipleVectorFields {
+		fieldName2Types[testBinaryVecField] = schemapb.DataType_BinaryVector
+	}
+	schema := constructCollectionSchemaByDataType(collectionName, fieldName2Types, testInt64Field, false)
+	schemaInfo := mustNewSchemaInfo(schema)
+	collectionID := int64(1000)
+	mockTest(t, (*metacache.MetaCache).GetCollectionID, collectionID, nil)
+	mockTest(t, (*metacache.MetaCache).GetCollectionSchema, schemaInfo, nil)
+	mockTest(t, (*metacache.MetaCache).GetCollectionInfo, &collectionInfo{CollID: collectionID}, nil)
+	mockTest(t, (*metacache.MetaCache).GetPartitions, map[string]int64{"_default": 1}, nil)
+
+	qn := getQueryNodeClient()
+	mgr := shardclient.NewMockShardClientManager(t)
+	mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(qn, nil).Maybe()
+	mgr.EXPECT().GetShardLeaderList(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]string{"mock_qn"}, nil).Maybe()
+	mgr.EXPECT().GetShard(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]shardclient.NodeInfo{
+		{NodeID: 1, Address: "mock_qn", Serviceable: true},
+	}, nil).Maybe()
+	mgr.EXPECT().InvalidateShardLeaderCache(mock.Anything).Return().Maybe()
+	lb := shardclient.NewMockLBPolicy(t)
+
+	task := &SearchTask{
+		baseTask:  baseTask{MetaCache: cache},
+		Condition: NewTaskCondition(ctx),
+		SearchRequest: &internalpb.SearchRequest{
+			Base: &commonpb.MsgBase{
+				MsgType:  commonpb.MsgType_Search,
+				SourceID: paramtable.GetNodeID(),
+			},
+			CollectionID:   collectionID,
+			OutputFieldsId: make([]int64, len(fieldName2Types)),
+		},
+		ctx: ctx,
+		result: &milvuspb.SearchResults{
+			Status: merr.Success(),
+		},
+		request: &milvuspb.SearchRequest{
+			Base: &commonpb.MsgBase{
+				MsgType:  commonpb.MsgType_Search,
+				SourceID: paramtable.GetNodeID(),
+			},
+			CollectionName: collectionName,
+			Nq:             1,
+			DslType:        commonpb.DslType_BoolExprV1,
+		},
+		lb:             lb,
+		shardClientMgr: mgr,
+	}
+	for i := 0; i < len(fieldName2Types); i++ {
+		task.OutputFieldsId[i] = int64(common.StartOfUserFieldID + i)
+	}
+
+	assert.NoError(t, task.OnEnqueue())
+	task.ctx = ctx
+	if enableMultipleVectorFields {
+		err := task.PreExecute(ctx)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "multiple anns_fields exist, please specify a anns_field in search_params")
+	} else {
+		assert.NoError(t, task.PreExecute(ctx))
+	}
+
+	// stub shard dispatch: Execute runs the workload against the mock qn and
+	// propagates the query node error back through lb.
+	lb.EXPECT().Execute(mock.Anything, mock.Anything).RunAndReturn(
+		func(_ context.Context, workload shardclient.CollectionWorkLoad) error {
+			return workload.Exec(ctx, 1, qn, "mock_qn")
+		}).Maybe()
+	lb.EXPECT().UpdateCostMetrics(mock.Anything, mock.Anything).Return().Maybe()
+
+	qn.EXPECT().Search(mock.Anything, mock.Anything).Return(nil, errors.New("mock error")).Once()
+	assert.Error(t, task.Execute(ctx))
+
+	qn.EXPECT().Search(mock.Anything, mock.Anything).Return(&internalpb.SearchResults{
+		Status: merr.Status(merr.ErrChannelNotAvailable),
+	}, nil).Once()
+	err := task.Execute(ctx)
+	assert.ErrorIs(t, err, merr.ErrChannelNotAvailable)
+
+	qn.EXPECT().Search(mock.Anything, mock.Anything).Return(&internalpb.SearchResults{
+		Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
+	}, nil).Once()
+	assert.Error(t, task.Execute(ctx))
+}

internal/proxy/dql/rerank_meta.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"

internal/proxy/dql/rerank_meta_test.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"testing"

internal/proxy/dql/search_pipeline.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"bytes"
@@ -36,7 +36,7 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
 	"github.com/milvus-io/milvus/internal/parser/planparserv2"
 	"github.com/milvus-io/milvus/internal/proxy/search_agg"
-	"github.com/milvus-io/milvus/internal/types"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
 	"github.com/milvus-io/milvus/internal/util/function/chain"
 	chaintypes "github.com/milvus-io/milvus/internal/util/function/chain/types"
 	"github.com/milvus-io/milvus/internal/util/function/models"
@@ -137,7 +137,7 @@ const (
 	pipelineStorageCost = "storage_cost"
 )
 
-var opFactory = map[string]func(t *searchTask, params map[string]any) (operator, error){
+var opFactory = map[string]func(t *SearchTask, params map[string]any) (operator, error){
 	searchReduceOp:        newSearchReduceOperator,
 	hybridSearchReduceOp:  newHybridSearchReduceOperator,
 	aggOp:                 newAggregateOperator,
@@ -152,7 +152,7 @@ var opFactory = map[string]func(t *searchTask, params map[string]any) (operator,
 	orderByOp:             newOrderByOperator,
 }
 
-func NewNode(info *nodeDef, t *searchTask) (*Node, error) {
+func NewNode(info *nodeDef, t *SearchTask) (*Node, error) {
 	n := Node{
 		name:    info.name,
 		opName:  info.opName,
@@ -182,7 +182,7 @@ type searchReduceOperator struct {
 
 const reduceOffsetParamKey = "reduce_offset"
 
-func newSearchReduceOperator(t *searchTask, params map[string]any) (operator, error) {
+func newSearchReduceOperator(t *SearchTask, params map[string]any) (operator, error) {
 	pkField, err := t.schema.GetPkField()
 	if err != nil {
 		return nil, err
@@ -230,7 +230,7 @@ type hybridSearchReduceOperator struct {
 	collSchema         *schemapb.CollectionSchema
 }
 
-func newHybridSearchReduceOperator(t *searchTask, _ map[string]any) (operator, error) {
+func newHybridSearchReduceOperator(t *SearchTask, _ map[string]any) (operator, error) {
 	pkField, err := t.schema.GetPkField()
 	if err != nil {
 		return nil, err
@@ -299,14 +299,14 @@ type elementBestCollapseOperator struct {
 	elementLevelHybrid bool
 }
 
-func newElementBestCollapseOperator(t *searchTask, _ map[string]any) (operator, error) {
+func newElementBestCollapseOperator(t *SearchTask, _ map[string]any) (operator, error) {
 	return &elementBestCollapseOperator{
 		configs:            t.hybridCollapseConfigs(),
 		elementLevelHybrid: t.hybridElementLevel,
 	}, nil
 }
 
-func (t *searchTask) hybridCollapseConfigs() []elementCollapseConfig {
+func (t *SearchTask) hybridCollapseConfigs() []elementCollapseConfig {
 	if len(t.hybridSubSearchInfos) == 0 {
 		return nil
 	}
@@ -667,7 +667,7 @@ type aggregateOperator struct {
 	roundDecimal int64
 }
 
-func newAggregateOperator(t *searchTask, _ map[string]any) (operator, error) {
+func newAggregateOperator(t *SearchTask, _ map[string]any) (operator, error) {
 	if t.aggCtx == nil {
 		return nil, merr.WrapErrServiceInternal("aggregate operator requires non-nil aggCtx")
 	}
@@ -1006,7 +1006,7 @@ func fillFieldNames(schema *schemapb.CollectionSchema, resultData *schemapb.Sear
 	}
 }
 
-func newRerankOperator(t *searchTask, params map[string]any) (operator, error) {
+func newRerankOperator(t *SearchTask, params map[string]any) (operator, error) {
 	if t.GetIsAdvanced() {
 		useRequestLimit := false
 		if meta, ok := t.rerankMeta.(*functionChainRerankMeta); ok {
@@ -1234,10 +1234,10 @@ type requeryOperator struct {
 	namespace          *string
 	planNamespace      *string
 
-	node types.ProxyComponent
+	node taskmodel.TaskNode
 }
 
-func newRequeryOperator(t *searchTask, _ map[string]any) (operator, error) {
+func newRequeryOperator(t *SearchTask, _ map[string]any) (operator, error) {
 	pkField, err := t.schema.GetPkField()
 	if err != nil {
 		return nil, err
@@ -1331,9 +1331,9 @@ func (op *requeryOperator) requery(ctx context.Context, span trace.Span, ids *sc
 	for k, v := range op.queryChannelsNode {
 		preferredNodes[k] = v
 	}
-	qt := &queryTask{
+	qt := &QueryTask{
 		baseTask: baseTask{
-			MetaCache: op.node.(*Proxy).GetMetaCache(),
+			MetaCache: op.node.GetMetaCache(),
 		},
 		ctx:       op.traceCtx,
 		Condition: NewTaskCondition(op.traceCtx),
@@ -1349,16 +1349,19 @@ func (op *requeryOperator) requery(ctx context.Context, span trace.Span, ids *sc
 		},
 		request:        queryReq,
 		plan:           plan,
-		mixCoord:       op.node.(*Proxy).mixCoord,
-		lb:             op.node.(*Proxy).lbPolicy,
-		shardclientMgr: op.node.(*Proxy).shardMgr,
+		lb:             op.node.LBPolicy(),
+		shardclientMgr: op.node.ShardMgr(),
 		channelsMvcc:   channelsMvcc,
 		preferredNodes: preferredNodes,
 		fastSkip:       true,
 		reQuery:        true,
-		chMgr:          op.node.(*Proxy).chMgr,
+		chMgr:          op.node.ChMgr(),
 	}
-	queryResult, storageCost, err := op.node.(*Proxy).query(op.traceCtx, qt, span)
+	queryRunner, ok := op.node.(taskmodel.QueryRunner)
+	if !ok {
+		return nil, segcore.StorageCost{}, merr.WrapErrServiceInternalMsg("host node does not implement QueryRunner")
+	}
+	queryResult, storageCost, err := queryRunner.ExecuteQuery(op.traceCtx, qt, span)
 	if err != nil {
 		return nil, segcore.StorageCost{}, err
 	}
@@ -1376,7 +1379,7 @@ type organizeOperator struct {
 	collectionID       int64
 }
 
-func newOrganizeOperator(t *searchTask, _ map[string]any) (operator, error) {
+func newOrganizeOperator(t *SearchTask, _ map[string]any) (operator, error) {
 	pkField, err := t.schema.GetPkField()
 	if err != nil {
 		return nil, err
@@ -1449,7 +1452,7 @@ func (op *organizeOperator) run(ctx context.Context, span trace.Span, inputs ...
 			allFieldData[idx] = emptyFields
 			continue
 		}
-		if fieldData, err := pickFieldData(ids, offsets, fields, op.schema, op.collectionID); err != nil {
+		if fieldData, err := PickFieldData(ids, offsets, fields, op.schema, op.collectionID); err != nil {
 			return nil, err
 		} else {
 			allFieldData[idx] = fieldData
@@ -1458,7 +1461,7 @@ func (op *organizeOperator) run(ctx context.Context, span trace.Span, inputs ...
 	return []any{allFieldData}, nil
 }
 
-func pickFieldData(ids *schemapb.IDs, pkOffset map[any]int, fields []*schemapb.FieldData, schema *schemapb.CollectionSchema, collectionID int64) ([]*schemapb.FieldData, error) {
+func PickFieldData(ids *schemapb.IDs, pkOffset map[any]int, fields []*schemapb.FieldData, schema *schemapb.CollectionSchema, collectionID int64) ([]*schemapb.FieldData, error) {
 	// Reorganize Results. The order of query result ids will be altered and differ from queried ids.
 	// We should reorganize query results to keep the order of original queried ids. For example:
 	// ===========================================
@@ -1522,7 +1525,7 @@ type hybridElementCandidateKey struct {
 	elementIndex int64
 }
 
-func newHybridAssembleOperator(t *searchTask, _ map[string]any) (operator, error) {
+func newHybridAssembleOperator(t *SearchTask, _ map[string]any) (operator, error) {
 	return &hybridAssembleOperator{
 		collectionID:       t.GetCollectionID(),
 		elementLevelHybrid: t.hybridElementLevel,
@@ -1671,7 +1674,7 @@ type lambdaOperator struct {
 	f func(ctx context.Context, span trace.Span, inputs ...any) ([]any, error)
 }
 
-func newLambdaOperator(_ *searchTask, params map[string]any) (operator, error) {
+func newLambdaOperator(_ *SearchTask, params map[string]any) (operator, error) {
 	return &lambdaOperator{
 		f: params[lambdaParamKey].(func(ctx context.Context, span trace.Span, inputs ...any) ([]any, error)),
 	}, nil
@@ -1687,7 +1690,7 @@ type endOperator struct {
 	roundDecimal     int64
 }
 
-func newEndOperator(t *searchTask, _ map[string]any) (operator, error) {
+func newEndOperator(t *SearchTask, _ map[string]any) (operator, error) {
 	roundDecimal := int64(-1)
 	if !t.GetIsAdvanced() && len(t.queryInfos) > 0 && t.queryInfos[0] != nil {
 		roundDecimal = t.queryInfos[0].GetRoundDecimal()
@@ -1756,7 +1759,7 @@ func roundAggHitScores(buckets []*search_agg.AggBucketResult, roundDecimal int64
 	}
 }
 
-func newHighlightOperator(t *searchTask, _ map[string]any) (operator, error) {
+func newHighlightOperator(t *SearchTask, _ map[string]any) (operator, error) {
 	return t.highlighter.AsSearchPipelineOperator(t)
 }
 
@@ -1768,7 +1771,7 @@ type orderByOperator struct {
 	offset         int64
 }
 
-func newOrderByOperator(t *searchTask, _ map[string]any) (operator, error) {
+func newOrderByOperator(t *SearchTask, _ map[string]any) (operator, error) {
 	var groupByFieldId int64 = -1
 	var groupSize int64 = 1
 	if len(t.queryInfos) > 0 && t.queryInfos[0] != nil {
@@ -2945,7 +2948,7 @@ type pipeline struct {
 	traceEnabled bool
 }
 
-func newPipeline(pipeDef *pipelineDef, t *searchTask) (*pipeline, error) {
+func newPipeline(pipeDef *pipelineDef, t *SearchTask) (*pipeline, error) {
 	nodes := make([]*Node, len(pipeDef.nodes))
 	for i, def := range pipeDef.nodes {
 		node, err := NewNode(def, t)
@@ -2957,7 +2960,7 @@ func newPipeline(pipeDef *pipelineDef, t *searchTask) (*pipeline, error) {
 	return &pipeline{name: pipeDef.name, nodes: nodes, traceEnabled: t.traceEnabled}, nil
 }
 
-func (p *pipeline) AddNodes(t *searchTask, nodes ...*nodeDef) error {
+func (p *pipeline) AddNodes(t *SearchTask, nodes ...*nodeDef) error {
 	for _, def := range nodes {
 		node, err := NewNode(def, t)
 		if err != nil {
@@ -3507,7 +3510,7 @@ var searchWithOrderByPipe = &pipelineDef{
 	},
 }
 
-func newBuiltInPipeline(t *searchTask) (*pipeline, error) {
+func newBuiltInPipeline(t *SearchTask) (*pipeline, error) {
 	if t.aggCtx != nil {
 		return newPipeline(searchWithAggPipe, t)
 	}
@@ -3551,7 +3554,7 @@ func newBuiltInPipeline(t *searchTask) (*pipeline, error) {
 	return nil, merr.WrapErrServiceInternal("unsupported pipeline")
 }
 
-func newSearchPipeline(t *searchTask) (*pipeline, error) {
+func newSearchPipeline(t *SearchTask) (*pipeline, error) {
 	p, err := newBuiltInPipeline(t)
 	if err != nil {
 		return nil, err

internal/proxy/dql/search_pipeline_hybrid_function_chain_test.go

@@ -14,14 +14,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"context"
 	"math"
 	"strconv"
 
-	"github.com/bytedance/mockey"
 	"go.opentelemetry.io/otel/trace"
 
 	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
@@ -187,7 +186,7 @@ func (s *SearchPipelineSuite) TestHybridFunctionChainRequestLimitFallback() {
 					makeResult(nil, nil, []int64{0, 0}),
 				}
 			}
-			opFactory[hybridSearchReduceOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+			opFactory[hybridSearchReduceOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 				return searchPipelineTestOperator(func(context.Context, trace.Span, ...any) ([]any, error) {
 					return []any{queryResults, []string{metric.IP, metric.IP}}, nil
 				}), nil
@@ -202,7 +201,7 @@ func (s *SearchPipelineSuite) TestHybridFunctionChainRequestLimitFallback() {
 			task := getHybridSearchTask("test_collection", [][]string{{"1", "2"}, {"3", "4"}}, nil)
 			// Exercise the wire conversion and request parser used for SDK-supplied limits.
 			// A nested L1 Limit must not suppress the top-level L2 fallback.
-			task.request = convertHybridSearchToSearch(&milvuspb.HybridSearchRequest{
+			task.request = ConvertHybridSearchToSearch(&milvuspb.HybridSearchRequest{
 				CollectionName: "test_collection",
 				Requests: []*milvuspb.SearchRequest{
 					{FunctionChains: []*schemapb.FunctionChain{l1LimitFunctionChain(1)}},
@@ -276,9 +275,9 @@ func (s *SearchPipelineSuite) TestHybridFunctionChainRequeryInputDoesNotLeak() {
 	outputField.FieldId = 102
 	pkField := testutils.GenerateScalarFieldData(schemapb.DataType_Int64, "int64", 20)
 	pkField.FieldId = 100
-	requeryMock := mockey.Mock((*requeryOperator).requery).Return(&milvuspb.QueryResults{
+	requeryMock := mockTest(s.T(), (*requeryOperator).requery, &milvuspb.QueryResults{
 		FieldsData: []*schemapb.FieldData{intField, outputField, pkField},
-	}, segcore.StorageCost{}, nil).Build()
+	}, segcore.StorageCost{}, nil)
 	defer requeryMock.UnPatch()
 
 	pipeline, err := newSearchPipeline(task)
@@ -364,7 +363,7 @@ func (s *SearchPipelineSuite) TestHybridFunctionChainElementLevelPreservesNative
 				makeResult(test.ids2, []int64{1, 0, 2}, []float32{0.95, 0.75, 0.60}, []int64{110, 200, 220}),
 			}
 
-			opFactory[hybridSearchReduceOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+			opFactory[hybridSearchReduceOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 				return searchPipelineTestOperator(func(context.Context, trace.Span, ...any) ([]any, error) {
 					return []any{reduced, []string{"IP", "IP"}}, nil
 				}), nil
@@ -422,7 +421,7 @@ func (s *SearchPipelineSuite) TestHybridFunctionChainUsesRuntimeMetricsInSubRequ
 		makeResult([]int64{1, 2}, []float32{10, 0}),
 		makeResult([]int64{3, 4}, []float32{0, 10}),
 	}
-	opFactory[hybridSearchReduceOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[hybridSearchReduceOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(context.Context, trace.Span, ...any) ([]any, error) {
 			return []any{reduced, []string{metric.IP, metric.L2}}, nil
 		}), nil
@@ -504,7 +503,7 @@ func (s *SearchPipelineSuite) TestHybridFunctionChainOwnsSortOffsetAndRounding()
 			},
 		},
 	}
-	opFactory[hybridSearchReduceOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[hybridSearchReduceOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(context.Context, trace.Span, ...any) ([]any, error) {
 			return []any{reduced, []string{metric.IP, metric.IP}}, nil
 		}), nil
@@ -594,7 +593,7 @@ func (s *SearchPipelineSuite) TestHybridFunctionChainGroupByDoesNotOverrideFinal
 		makeResult([]int64{1, 2}, []float32{0.9, 0.7}, []int64{100, 200}),
 		makeResult([]int64{3, 4}, []float32{0.8, 0.6}, []int64{300, 400}),
 	}
-	opFactory[hybridSearchReduceOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[hybridSearchReduceOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(context.Context, trace.Span, ...any) ([]any, error) {
 			return []any{reduced, []string{metric.IP, metric.IP}}, nil
 		}), nil
@@ -767,7 +766,7 @@ func (s *SearchPipelineSuite) TestHybridFunctionChainEmptyResultValidatesDownstr
 			},
 		}
 	}
-	opFactory[hybridSearchReduceOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[hybridSearchReduceOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(context.Context, trace.Span, ...any) ([]any, error) {
 			return []any{[]*milvuspb.SearchResults{emptyResult(), emptyResult()}, []string{metric.IP, metric.IP}}, nil
 		}), nil

internal/proxy/dql/search_pipeline_test.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"context"
@@ -1195,14 +1195,14 @@ func (s *SearchPipelineSuite) TestHybridSearchPipe_RerankReceivesCollapsedRowLev
 		opFactory[hybridAssembleOp] = originalAssembleFactory
 	}()
 
-	opFactory[hybridSearchReduceOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[hybridSearchReduceOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(ctx context.Context, span trace.Span, inputs ...any) ([]any, error) {
 			return []any{[]*milvuspb.SearchResults{input}, []string{"IP"}}, nil
 		}), nil
 	}
 
 	rerankCalled := false
-	opFactory[rerankOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[rerankOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(ctx context.Context, span trace.Span, inputs ...any) ([]any, error) {
 			rerankCalled = true
 
@@ -1231,7 +1231,7 @@ func (s *SearchPipelineSuite) TestHybridSearchPipe_RerankReceivesCollapsedRowLev
 		}), nil
 	}
 
-	opFactory[hybridAssembleOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[hybridAssembleOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(ctx context.Context, span trace.Span, inputs ...any) ([]any, error) {
 			collapsed, ok := inputs[0].([]*milvuspb.SearchResults)
 			s.Require().True(ok)
@@ -1244,9 +1244,9 @@ func (s *SearchPipelineSuite) TestHybridSearchPipe_RerankReceivesCollapsedRowLev
 		}), nil
 	}
 
-	pipeline, err := newPipeline(hybridSearchPipe, &searchTask{})
+	pipeline, err := newPipeline(hybridSearchPipe, &SearchTask{})
 	s.Require().NoError(err)
-	err = pipeline.AddNodes(&searchTask{}, &nodeDef{
+	err = pipeline.AddNodes(&SearchTask{}, &nodeDef{
 		name:    "finish",
 		inputs:  []string{"result"},
 		outputs: []string{pipelineOutput},
@@ -1893,9 +1893,9 @@ func (s *SearchPipelineSuite) TestRequeryOp() {
 	f1 := testutils.GenerateScalarFieldData(schemapb.DataType_Int64, "int64", 20)
 	f1.FieldId = 101
 
-	mocker := mockey.Mock((*requeryOperator).requery).Return(&milvuspb.QueryResults{
+	mocker := mockTest(s.T(), (*requeryOperator).requery, &milvuspb.QueryResults{
 		FieldsData: []*schemapb.FieldData{f1},
-	}, segcore.StorageCost{}, nil).Build()
+	}, segcore.StorageCost{}, nil)
 	defer mocker.UnPatch()
 
 	op := requeryOperator{
@@ -1972,17 +1972,12 @@ func (s *SearchPipelineSuite) TestHighlightOp() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	proxy := &Proxy{}
-	proxy.tsoAllocator = &timestampAllocator{
-		tso: newMockTimestampAllocatorInterface(),
-	}
-	sched, err := scheduler.NewTaskScheduler(ctx, proxy.tsoAllocator)
+	sched, err := scheduler.NewTaskScheduler(ctx, &mockTsoAllocator{})
 	s.Require().NoError(err)
 
 	err = sched.Start()
 	s.Require().NoError(err)
 	defer sched.Close()
-	proxy.sched = sched
 
 	collName := "test_coll_highlight"
 	fieldName2Types := map[string]schemapb.DataType{
@@ -2008,8 +2003,8 @@ func (s *SearchPipelineSuite) TestHighlightOp() {
 	}
 
 	mockLb := shardclient.NewMockLBPolicy(s.T())
-	searchTask := &searchTask{
-		node: proxy,
+	SearchTask := &SearchTask{
+		sched: sched,
 		highlighter: &LexicalHighlighter{
 			tasks: highlightTasks,
 		},
@@ -2022,7 +2017,7 @@ func (s *SearchPipelineSuite) TestHighlightOp() {
 		},
 	}
 
-	op, err := opFactory[highlightOp](searchTask, map[string]any{})
+	op, err := opFactory[highlightOp](SearchTask, map[string]any{})
 	s.Require().NoError(err)
 
 	// mockery
@@ -2095,17 +2090,12 @@ func (s *SearchPipelineSuite) TestLexicalHighlightOpNullableStringKeepsEmptyHigh
 			ctx, cancel := context.WithCancel(context.Background())
 			defer cancel()
 
-			proxy := &Proxy{}
-			proxy.tsoAllocator = &timestampAllocator{
-				tso: newMockTimestampAllocatorInterface(),
-			}
-			sched, err := scheduler.NewTaskScheduler(ctx, proxy.tsoAllocator)
+			sched, err := scheduler.NewTaskScheduler(ctx, &mockTsoAllocator{})
 			s.Require().NoError(err)
 
 			err = sched.Start()
 			s.Require().NoError(err)
 			defer sched.Close()
-			proxy.sched = sched
 
 			collName := "test_coll_highlight_nullable"
 			fieldName2Types := map[string]schemapb.DataType{
@@ -2127,8 +2117,8 @@ func (s *SearchPipelineSuite) TestLexicalHighlightOpNullableStringKeepsEmptyHigh
 			}
 
 			mockLb := shardclient.NewMockLBPolicy(s.T())
-			searchTask := &searchTask{
-				node: proxy,
+			SearchTask := &SearchTask{
+				sched: sched,
 				highlighter: &LexicalHighlighter{
 					tasks: highlightTasks,
 				},
@@ -2139,7 +2129,7 @@ func (s *SearchPipelineSuite) TestLexicalHighlightOpNullableStringKeepsEmptyHigh
 				SearchRequest:  &internalpb.SearchRequest{CollectionID: 0},
 			}
 
-			op, err := opFactory[highlightOp](searchTask, map[string]any{})
+			op, err := opFactory[highlightOp](SearchTask, map[string]any{})
 			s.Require().NoError(err)
 
 			queryNodeResults := make([]*querypb.HighlightResult, tc.rowNum)
@@ -2284,14 +2274,14 @@ func (s *SearchPipelineSuite) TestSemanticHighlightOp() {
 		}).Build()
 	defer mockProcess.UnPatch()
 
-	mockFieldIDs := mockey.Mock((*highlight.SemanticHighlight).FieldIDs).To(func(h *highlight.SemanticHighlight) []int64 {
+	mockFieldIDs := mockTestTo(s.T(), (*highlight.SemanticHighlight).FieldIDs, func(h *highlight.SemanticHighlight) []int64 {
 		return []int64{101}
-	}).Build()
+	})
 	defer mockFieldIDs.UnPatch()
 
-	mockGetFieldName := mockey.Mock((*highlight.SemanticHighlight).GetFieldName).To(func(h *highlight.SemanticHighlight, id int64) string {
+	mockGetFieldName := mockTestTo(s.T(), (*highlight.SemanticHighlight).GetFieldName, func(h *highlight.SemanticHighlight, id int64) string {
 		return testVarCharField
-	}).Build()
+	})
 	defer mockGetFieldName.UnPatch()
 
 	// Create operator
@@ -2363,10 +2353,10 @@ func (s *SearchPipelineSuite) TestSemanticHighlightOpNullableStringAlignsRows()
 		}).Build()
 	defer mockProcess.UnPatch()
 
-	mockFieldIDs := mockey.Mock((*highlight.SemanticHighlight).FieldIDs).Return([]int64{101}).Build()
+	mockFieldIDs := mockTest(s.T(), (*highlight.SemanticHighlight).FieldIDs, []int64{101})
 	defer mockFieldIDs.UnPatch()
 
-	mockGetFieldName := mockey.Mock((*highlight.SemanticHighlight).GetFieldName).Return(testVarCharField).Build()
+	mockGetFieldName := mockTest(s.T(), (*highlight.SemanticHighlight).GetFieldName, testVarCharField)
 	defer mockGetFieldName.UnPatch()
 
 	op := &semanticHighlightOperator{
@@ -2421,7 +2411,7 @@ func (s *SearchPipelineSuite) TestSemanticHighlightOpMissingField() {
 	ctx := context.Background()
 
 	// Mock FieldIDs to return field 999 (not in results)
-	mockFieldIDs := mockey.Mock((*highlight.SemanticHighlight).FieldIDs).Return([]int64{999}).Build()
+	mockFieldIDs := mockTest(s.T(), (*highlight.SemanticHighlight).FieldIDs, []int64{999})
 	defer mockFieldIDs.UnPatch()
 
 	op := &semanticHighlightOperator{
@@ -2478,15 +2468,15 @@ func (s *SearchPipelineSuite) TestSemanticHighlightOpMultipleFields() {
 		}).Build()
 	defer mockProcess.UnPatch()
 
-	mockFieldIDs := mockey.Mock((*highlight.SemanticHighlight).FieldIDs).Return([]int64{101, 102}).Build()
+	mockFieldIDs := mockTest(s.T(), (*highlight.SemanticHighlight).FieldIDs, []int64{101, 102})
 	defer mockFieldIDs.UnPatch()
 
-	mockGetFieldName := mockey.Mock((*highlight.SemanticHighlight).GetFieldName).To(func(h *highlight.SemanticHighlight, id int64) string {
+	mockGetFieldName := mockTestTo(s.T(), (*highlight.SemanticHighlight).GetFieldName, func(h *highlight.SemanticHighlight, id int64) string {
 		if id == 101 {
 			return "field1"
 		}
 		return "field2"
-	}).Build()
+	})
 	defer mockGetFieldName.UnPatch()
 
 	op := &semanticHighlightOperator{
@@ -2562,12 +2552,12 @@ func (s *SearchPipelineSuite) TestSemanticHighlightOpEmptyResults() {
 		}).Build()
 	defer mockProcess.UnPatch()
 
-	mockFieldIDs := mockey.Mock((*highlight.SemanticHighlight).FieldIDs).Return([]int64{101}).Build()
+	mockFieldIDs := mockTest(s.T(), (*highlight.SemanticHighlight).FieldIDs, []int64{101})
 	defer mockFieldIDs.UnPatch()
 
-	mockGetFieldName := mockey.Mock((*highlight.SemanticHighlight).GetFieldName).To(func(h *highlight.SemanticHighlight, id int64) string {
+	mockGetFieldName := mockTestTo(s.T(), (*highlight.SemanticHighlight).GetFieldName, func(h *highlight.SemanticHighlight, id int64) string {
 		return testVarCharField
-	}).Build()
+	})
 	defer mockGetFieldName.UnPatch()
 
 	op := &semanticHighlightOperator{
@@ -2612,7 +2602,7 @@ func (s *SearchPipelineSuite) TestSemanticHighlightOpEmptyResults() {
 func (s *SearchPipelineSuite) TestSemanticHighlightOpZeroHitWithNonEmptyFieldsData() {
 	ctx := context.Background()
 
-	mockFieldIDs := mockey.Mock((*highlight.SemanticHighlight).FieldIDs).Return([]int64{999}).Build()
+	mockFieldIDs := mockTest(s.T(), (*highlight.SemanticHighlight).FieldIDs, []int64{999})
 	defer mockFieldIDs.UnPatch()
 
 	op := &semanticHighlightOperator{
@@ -2682,16 +2672,16 @@ func (s *SearchPipelineSuite) TestSemanticHighlightOpDynamicField() {
 		}).Build()
 	defer mockProcess.UnPatch()
 
-	mockFieldIDs := mockey.Mock((*highlight.SemanticHighlight).FieldIDs).Return([]int64{}).Build()
+	mockFieldIDs := mockTest(s.T(), (*highlight.SemanticHighlight).FieldIDs, []int64{})
 	defer mockFieldIDs.UnPatch()
 
-	mockHasDynamicFields := mockey.Mock((*highlight.SemanticHighlight).HasDynamicFields).Return(true).Build()
+	mockHasDynamicFields := mockTest(s.T(), (*highlight.SemanticHighlight).HasDynamicFields, true)
 	defer mockHasDynamicFields.UnPatch()
 
-	mockDynamicFieldNames := mockey.Mock((*highlight.SemanticHighlight).DynamicFieldNames).Return([]string{"dyn_content"}).Build()
+	mockDynamicFieldNames := mockTest(s.T(), (*highlight.SemanticHighlight).DynamicFieldNames, []string{"dyn_content"})
 	defer mockDynamicFieldNames.UnPatch()
 
-	mockDynamicFieldID := mockey.Mock((*highlight.SemanticHighlight).DynamicFieldID).Return(int64(102)).Build()
+	mockDynamicFieldID := mockTest(s.T(), (*highlight.SemanticHighlight).DynamicFieldID, int64(102))
 	defer mockDynamicFieldID.UnPatch()
 
 	// Create operator
@@ -2763,19 +2753,19 @@ func (s *SearchPipelineSuite) TestSemanticHighlightOpMixedFields() {
 		}).Build()
 	defer mockProcess.UnPatch()
 
-	mockFieldIDs := mockey.Mock((*highlight.SemanticHighlight).FieldIDs).Return([]int64{101}).Build()
+	mockFieldIDs := mockTest(s.T(), (*highlight.SemanticHighlight).FieldIDs, []int64{101})
 	defer mockFieldIDs.UnPatch()
 
-	mockGetFieldName := mockey.Mock((*highlight.SemanticHighlight).GetFieldName).Return(testVarCharField).Build()
+	mockGetFieldName := mockTest(s.T(), (*highlight.SemanticHighlight).GetFieldName, testVarCharField)
 	defer mockGetFieldName.UnPatch()
 
-	mockHasDynamicFields := mockey.Mock((*highlight.SemanticHighlight).HasDynamicFields).Return(true).Build()
+	mockHasDynamicFields := mockTest(s.T(), (*highlight.SemanticHighlight).HasDynamicFields, true)
 	defer mockHasDynamicFields.UnPatch()
 
-	mockDynamicFieldNames := mockey.Mock((*highlight.SemanticHighlight).DynamicFieldNames).Return([]string{"dyn_content"}).Build()
+	mockDynamicFieldNames := mockTest(s.T(), (*highlight.SemanticHighlight).DynamicFieldNames, []string{"dyn_content"})
 	defer mockDynamicFieldNames.UnPatch()
 
-	mockDynamicFieldID := mockey.Mock((*highlight.SemanticHighlight).DynamicFieldID).Return(int64(102)).Build()
+	mockDynamicFieldID := mockTest(s.T(), (*highlight.SemanticHighlight).DynamicFieldID, int64(102))
 	defer mockDynamicFieldID.UnPatch()
 
 	// Create operator
@@ -2859,16 +2849,16 @@ func (s *SearchPipelineSuite) TestSemanticHighlightOpDynamicFieldRowCountMismatc
 		}).Build()
 	defer mockProcess.UnPatch()
 
-	mockFieldIDs := mockey.Mock((*highlight.SemanticHighlight).FieldIDs).Return([]int64{}).Build()
+	mockFieldIDs := mockTest(s.T(), (*highlight.SemanticHighlight).FieldIDs, []int64{})
 	defer mockFieldIDs.UnPatch()
 
-	mockHasDynamicFields := mockey.Mock((*highlight.SemanticHighlight).HasDynamicFields).Return(true).Build()
+	mockHasDynamicFields := mockTest(s.T(), (*highlight.SemanticHighlight).HasDynamicFields, true)
 	defer mockHasDynamicFields.UnPatch()
 
-	mockDynamicFieldNames := mockey.Mock((*highlight.SemanticHighlight).DynamicFieldNames).Return([]string{"dyn_content"}).Build()
+	mockDynamicFieldNames := mockTest(s.T(), (*highlight.SemanticHighlight).DynamicFieldNames, []string{"dyn_content"})
 	defer mockDynamicFieldNames.UnPatch()
 
-	mockDynamicFieldID := mockey.Mock((*highlight.SemanticHighlight).DynamicFieldID).Return(int64(102)).Build()
+	mockDynamicFieldID := mockTest(s.T(), (*highlight.SemanticHighlight).DynamicFieldID, int64(102))
 	defer mockDynamicFieldID.UnPatch()
 
 	op := &semanticHighlightOperator{
@@ -2975,7 +2965,7 @@ func (s *SearchPipelineSuite) TestExtractMultipleDynamicFieldTexts() {
 
 func (s *SearchPipelineSuite) TestSearchPipeline() {
 	collectionName := "test"
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: collectionName,
 		SearchRequest: &internalpb.SearchRequest{
@@ -3032,7 +3022,7 @@ func (s *SearchPipelineSuite) TestSearchPipeline() {
 
 func (s *SearchPipelineSuite) TestSearchPipelineWithRequery() {
 	collectionName := "test_collection"
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: collectionName,
 		SearchRequest: &internalpb.SearchRequest{
@@ -3067,9 +3057,9 @@ func (s *SearchPipelineSuite) TestSearchPipelineWithRequery() {
 	f1.FieldId = 101
 	f2 := testutils.GenerateScalarFieldData(schemapb.DataType_Int64, "int64", 20)
 	f2.FieldId = 100
-	mocker := mockey.Mock((*requeryOperator).requery).Return(&milvuspb.QueryResults{
+	mocker := mockTest(s.T(), (*requeryOperator).requery, &milvuspb.QueryResults{
 		FieldsData: []*schemapb.FieldData{f1, f2},
-	}, segcore.StorageCost{ScannedRemoteBytes: 100, ScannedTotalBytes: 200}, nil).Build()
+	}, segcore.StorageCost{ScannedRemoteBytes: 100, ScannedTotalBytes: 200}, nil)
 	defer mocker.UnPatch()
 
 	pipeline, err := newPipeline(searchWithRequeryPipe, task)
@@ -3122,7 +3112,7 @@ func (s *SearchPipelineSuite) TestSearchWithRerankPipe() {
 	}
 	funcScoreSchema := &schemapb.FunctionScore{Functions: []*schemapb.FunctionSchema{functionSchema}}
 
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: "test_collection",
 		SearchRequest: &internalpb.SearchRequest{
@@ -3191,7 +3181,7 @@ func (s *SearchPipelineSuite) TestSearchWithRerankRequeryPipe() {
 	}
 	funcScoreSchema := &schemapb.FunctionScore{Functions: []*schemapb.FunctionSchema{functionSchema}}
 
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: "test_collection",
 		SearchRequest: &internalpb.SearchRequest{
@@ -3220,9 +3210,9 @@ func (s *SearchPipelineSuite) TestSearchWithRerankRequeryPipe() {
 	f1.FieldId = 101
 	f2 := testutils.GenerateScalarFieldData(schemapb.DataType_Int64, "int64", 20)
 	f2.FieldId = 100
-	mocker := mockey.Mock((*requeryOperator).requery).Return(&milvuspb.QueryResults{
+	mocker := mockTest(s.T(), (*requeryOperator).requery, &milvuspb.QueryResults{
 		FieldsData: []*schemapb.FieldData{f1, f2},
-	}, segcore.StorageCost{}, nil).Build()
+	}, segcore.StorageCost{}, nil)
 	defer mocker.UnPatch()
 
 	pipeline, err := newPipeline(searchWithRerankRequeryPipe, task)
@@ -3302,7 +3292,7 @@ func (s *SearchPipelineSuite) TestFilterFieldOperatorWithStructArrayFields() {
 		},
 	}
 
-	task := &searchTask{
+	task := &SearchTask{
 		schema: &schemaInfo{
 			CollectionSchema: schema,
 		},
@@ -3358,7 +3348,7 @@ func (s *SearchPipelineSuite) TestFilterFieldOperatorWithStructArrayFields() {
 }
 
 func (s *SearchPipelineSuite) TestEndOperatorRoundsScores() {
-	task := &searchTask{
+	task := &SearchTask{
 		queryInfos: []*planpb.QueryInfo{{RoundDecimal: 0}},
 		schema: &schemaInfo{
 			CollectionSchema: &schemapb.CollectionSchema{},
@@ -3413,7 +3403,7 @@ func (s *SearchPipelineSuite) TestRoundAggHitScoresDisabled() {
 }
 
 func (s *SearchPipelineSuite) TestEndOperatorKeepsScoresWhenRoundDecimalDisabled() {
-	task := &searchTask{
+	task := &SearchTask{
 		queryInfos: []*planpb.QueryInfo{{RoundDecimal: -1}},
 		schema: &schemaInfo{
 			CollectionSchema: &schemapb.CollectionSchema{},
@@ -3437,7 +3427,7 @@ func (s *SearchPipelineSuite) TestEndOperatorKeepsScoresWhenRoundDecimalDisabled
 }
 
 func (s *SearchPipelineSuite) TestEndOperatorKeepsAdvancedRerankScores() {
-	task := &searchTask{
+	task := &SearchTask{
 		SearchRequest: &internalpb.SearchRequest{IsAdvanced: true},
 		rankParams:    &rankParams{roundDecimal: 0},
 		schema: &schemaInfo{
@@ -3473,9 +3463,9 @@ func (s *SearchPipelineSuite) TestHybridSearchWithRequeryAndRerankByDataPipe() {
 	f1.FieldId = 101
 	f2 := testutils.GenerateScalarFieldData(schemapb.DataType_Int64, "int64", 20)
 	f2.FieldId = 100
-	mocker := mockey.Mock((*requeryOperator).requery).Return(&milvuspb.QueryResults{
+	mocker := mockTest(s.T(), (*requeryOperator).requery, &milvuspb.QueryResults{
 		FieldsData: []*schemapb.FieldData{f1, f2},
-	}, segcore.StorageCost{}, nil).Build()
+	}, segcore.StorageCost{}, nil)
 	defer mocker.UnPatch()
 
 	pipeline, err := newPipeline(hybridSearchWithRequeryAndRerankByFieldDataPipe, task)
@@ -3560,14 +3550,14 @@ func (s *SearchPipelineSuite) TestHybridSearchWithRequeryAndRerankByDataPipe_Ele
 		opFactory[rerankOp] = originalRerankFactory
 	}()
 
-	opFactory[hybridSearchReduceOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[hybridSearchReduceOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(ctx context.Context, span trace.Span, inputs ...any) ([]any, error) {
 			return []any{[]*milvuspb.SearchResults{input}, []string{"IP"}}, nil
 		}), nil
 	}
 
 	requeryCalled := false
-	opFactory[requeryOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[requeryOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(ctx context.Context, span trace.Span, inputs ...any) ([]any, error) {
 			requeryCalled = true
 
@@ -3582,7 +3572,7 @@ func (s *SearchPipelineSuite) TestHybridSearchWithRequeryAndRerankByDataPipe_Ele
 	}
 
 	rerankCalled := false
-	opFactory[rerankOp] = func(_ *searchTask, _ map[string]any) (operator, error) {
+	opFactory[rerankOp] = func(_ *SearchTask, _ map[string]any) (operator, error) {
 		return searchPipelineTestOperator(func(ctx context.Context, span trace.Span, inputs ...any) ([]any, error) {
 			rerankCalled = true
 
@@ -3641,9 +3631,9 @@ func (s *SearchPipelineSuite) TestHybridSearchWithRequeryPipe() {
 	f1.FieldId = 101
 	f2 := testutils.GenerateScalarFieldData(schemapb.DataType_Int64, "int64", 20)
 	f2.FieldId = 100
-	mocker := mockey.Mock((*requeryOperator).requery).Return(&milvuspb.QueryResults{
+	mocker := mockTest(s.T(), (*requeryOperator).requery, &milvuspb.QueryResults{
 		FieldsData: []*schemapb.FieldData{f1, f2},
-	}, segcore.StorageCost{}, nil).Build()
+	}, segcore.StorageCost{}, nil)
 	defer mocker.UnPatch()
 
 	pipeline, err := newPipeline(hybridSearchWithRequeryPipe, task)
@@ -3672,7 +3662,7 @@ func (s *SearchPipelineSuite) TestHybridSearchWithRequeryPipe() {
 	s.Equal(int64(2*2*10), results.GetResults().AllSearchCount)
 }
 
-func getHybridSearchTask(collName string, data [][]string, outputFields []string) *searchTask {
+func getHybridSearchTask(collName string, data [][]string, outputFields []string) *SearchTask {
 	subReqs := []*milvuspb.SubSearchRequest{}
 	for _, item := range data {
 		subReq := &milvuspb.SubSearchRequest{
@@ -3704,7 +3694,7 @@ func getHybridSearchTask(collName string, data [][]string, outputFields []string
 	if err != nil {
 		panic(err)
 	}
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: collName,
 		SearchRequest: &internalpb.SearchRequest{
@@ -3739,8 +3729,8 @@ func getHybridSearchTask(collName string, data [][]string, outputFields []string
 			CollectionSchema: schema,
 			PkField:          &schemapb.FieldSchema{FieldID: 100, Name: "int64", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
 		},
-		mixCoord: nil,
-		tr:       timerecord.NewTimeRecorder("test-search"),
+
+		tr: timerecord.NewTimeRecorder("test-search"),
 		rankParams: &rankParams{
 			limit:        10,
 			offset:       0,
@@ -4269,7 +4259,7 @@ func (s *SearchPipelineSuite) TestIsSameGroupByValue() {
 }
 
 func (s *SearchPipelineSuite) TestNewSearchReduceOperatorUsesPipelineOffsetParam() {
-	task := &searchTask{
+	task := &SearchTask{
 		ctx: context.Background(),
 		SearchRequest: &internalpb.SearchRequest{
 			Nq:           1,
@@ -4293,7 +4283,7 @@ func (s *SearchPipelineSuite) TestNewSearchReduceOperatorUsesPipelineOffsetParam
 }
 
 func (s *SearchPipelineSuite) TestNewOrderByOperatorUsesPluralGroupByFieldIDs() {
-	task := &searchTask{
+	task := &SearchTask{
 		orderByFields: []OrderByField{{FieldName: "price", Ascending: true}},
 		queryInfos: []*planpb.QueryInfo{{
 			GroupByFieldId:  -1,
@@ -7388,12 +7378,12 @@ func (s *SearchPipelineSuite) TestNewSearchPipelineWithOrderBy() {
 	f1.FieldId = 101
 	f2 := testutils.GenerateScalarFieldData(schemapb.DataType_Int64, "int64", 20)
 	f2.FieldId = 100
-	mocker := mockey.Mock((*requeryOperator).requery).Return(&milvuspb.QueryResults{
+	mocker := mockTest(s.T(), (*requeryOperator).requery, &milvuspb.QueryResults{
 		FieldsData: []*schemapb.FieldData{f1, f2},
-	}, segcore.StorageCost{ScannedRemoteBytes: 100, ScannedTotalBytes: 200}, nil).Build()
+	}, segcore.StorageCost{ScannedRemoteBytes: 100, ScannedTotalBytes: 200}, nil)
 	defer mocker.UnPatch()
 
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: collectionName,
 		SearchRequest: &internalpb.SearchRequest{
@@ -7455,12 +7445,12 @@ func (s *SearchPipelineSuite) TestNewRequeryOperator_WithHighlightDynamicFields(
 
 	// Mock highlighter with dynamic field names
 	mockHighlighter := &SemanticHighlighter{highlight: &highlight.SemanticHighlight{}}
-	mockDynFields := mockey.Mock((*SemanticHighlighter).DynamicFieldNames).To(func(h *SemanticHighlighter) []string {
+	mockDynFields := mockTestTo(s.T(), (*SemanticHighlighter).DynamicFieldNames, func(h *SemanticHighlighter) []string {
 		return []string{"dyn_field1", "dyn_field2"}
-	}).Build()
+	})
 	defer mockDynFields.UnPatch()
 
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: "test_collection",
 		SearchRequest: &internalpb.SearchRequest{
@@ -7503,7 +7493,7 @@ func (s *SearchPipelineSuite) TestNewRequeryOperatorIncludesOrderByOutputFieldNa
 		},
 		PkField: &schemapb.FieldSchema{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
 	}
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: "test_collection",
 		SearchRequest:  &internalpb.SearchRequest{Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Search}},
@@ -7538,7 +7528,7 @@ func (s *SearchPipelineSuite) TestNewRequeryOperator_WithoutHighlighter() {
 		PkField: &schemapb.FieldSchema{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
 	}
 
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: "test_collection",
 		SearchRequest: &internalpb.SearchRequest{
@@ -7573,7 +7563,7 @@ func (s *SearchPipelineSuite) TestNewBuiltInPipelineWithAggCtx() {
 	pkField := &schemapb.FieldSchema{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}
 	aggCtx, err := search_agg.NewContext(1, nil, nil, nil)
 	s.Require().NoError(err)
-	task := &searchTask{
+	task := &SearchTask{
 		SearchRequest: &internalpb.SearchRequest{IsAdvanced: false},
 		aggCtx:        aggCtx,
 		schema: &schemaInfo{
@@ -7593,7 +7583,7 @@ func (s *SearchPipelineSuite) TestNewSearchPipelineWithAggCtxSkipsEndNode() {
 	pkField := &schemapb.FieldSchema{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}
 	aggCtx, err := search_agg.NewContext(1, nil, nil, nil)
 	s.Require().NoError(err)
-	task := &searchTask{
+	task := &SearchTask{
 		SearchRequest: &internalpb.SearchRequest{IsAdvanced: false},
 		aggCtx:        aggCtx,
 		schema: &schemaInfo{
@@ -7672,12 +7662,12 @@ func (s *SearchPipelineSuite) TestNewRequeryOperator_WithHighlighterNoDynamicFie
 
 	// Mock highlighter with empty dynamic field names
 	mockHighlighter := &SemanticHighlighter{highlight: &highlight.SemanticHighlight{}}
-	mockDynFields := mockey.Mock((*SemanticHighlighter).DynamicFieldNames).To(func(h *SemanticHighlighter) []string {
+	mockDynFields := mockTestTo(s.T(), (*SemanticHighlighter).DynamicFieldNames, func(h *SemanticHighlighter) []string {
 		return []string{} // Empty dynamic fields
-	}).Build()
+	})
 	defer mockDynFields.UnPatch()
 
-	task := &searchTask{
+	task := &SearchTask{
 		ctx:            context.Background(),
 		collectionName: "test_collection",
 		SearchRequest: &internalpb.SearchRequest{
@@ -7706,13 +7696,13 @@ func (s *SearchPipelineSuite) TestNewRequeryOperator_WithHighlighterNoDynamicFie
 	s.Equal([]string{"title"}, reqOp.outputFieldNames)
 }
 
-// TestPickFieldDataWithNullableSparseVector tests pickFieldData with nullable sparse vector fields.
+// TestPickFieldDataWithNullableSparseVector tests PickFieldData with nullable sparse vector fields.
 // This test ensures that when reordering query results (which may return in different order than requested),
 // nullable sparse vectors are correctly handled using FieldDataIdxComputer to map row indices to data indices.
 func (s *SearchPipelineSuite) TestPickFieldDataWithNullableSparseVector() {
 	// Scenario: Search returns IDs [3, 1, 2], but Query returns them in order [1, 2, 3]
 	// with a nullable sparse vector where row 1 (middle) is null.
-	// pickFieldData should correctly reorder the data to match search result order.
+	// PickFieldData should correctly reorder the data to match search result order.
 
 	// Search result IDs (the order we want)
 	searchIDs := &schemapb.IDs{
@@ -7771,8 +7761,8 @@ func (s *SearchPipelineSuite) TestPickFieldDataWithNullableSparseVector() {
 		},
 	}
 
-	// Call pickFieldData - this should NOT panic
-	result, err := pickFieldData(searchIDs, pkOffset, queryFields, nil, 12345)
+	// Call PickFieldData - this should NOT panic
+	result, err := PickFieldData(searchIDs, pkOffset, queryFields, nil, 12345)
 	s.NoError(err)
 	s.NotNil(result)
 	s.Len(result, 2)
@@ -7792,7 +7782,7 @@ func (s *SearchPipelineSuite) TestPickFieldDataWithNullableSparseVector() {
 	s.Len(sparseContents, 2)
 }
 
-// TestPickFieldDataWithAllNullSparseVector tests pickFieldData when all sparse vector values are null.
+// TestPickFieldDataWithAllNullSparseVector tests PickFieldData when all sparse vector values are null.
 func (s *SearchPipelineSuite) TestPickFieldDataWithAllNullSparseVector() {
 	searchIDs := &schemapb.IDs{
 		IdField: &schemapb.IDs_IntId{
@@ -7841,7 +7831,7 @@ func (s *SearchPipelineSuite) TestPickFieldDataWithAllNullSparseVector() {
 		},
 	}
 
-	result, err := pickFieldData(searchIDs, pkOffset, queryFields, nil, 12345)
+	result, err := PickFieldData(searchIDs, pkOffset, queryFields, nil, 12345)
 	s.NoError(err)
 	s.NotNil(result)
 
@@ -7904,7 +7894,7 @@ func (s *SearchPipelineSuite) TestPickFieldDataWithNullableSparseVectorMissingVa
 		},
 	}
 
-	result, err := pickFieldData(searchIDs, pkOffset, queryFields, schema, 12345)
+	result, err := PickFieldData(searchIDs, pkOffset, queryFields, schema, 12345)
 	s.NoError(err)
 	s.Equal([]int64{2, 1}, result[0].GetScalars().GetLongData().GetData())
 	s.Equal([]bool{false, false}, typeutil.GetFieldDataValidData(result[1]))

internal/proxy/dql/search_reduce_multi_groupby_test.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"

internal/proxy/dql/search_reduce_util.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"

internal/proxy/dql/search_reduce_util_test.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"

internal/proxy/dql/search_shard_test.go

@@ -0,0 +1,77 @@
+// 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 dql
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
+	"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/internal/mocks"
+	"github.com/milvus-io/milvus/internal/proxy/shardclient"
+	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
+	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
+	"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/typeutil"
+)
+
+func TestSearchTask_SearchShardRecordsNodeHint(t *testing.T) {
+	ctx := context.Background()
+	const channel = "by-dev-rootcoord-dml_0_100v0"
+	const nodeID int64 = 101
+
+	qn := mocks.NewMockQueryNodeClient(t)
+	qn.EXPECT().Search(mock.Anything, mock.MatchedBy(func(req *querypb.SearchRequest) bool {
+		return len(req.GetDmlChannels()) == 1 && req.GetDmlChannels()[0] == channel
+	})).Return(&internalpb.SearchResults{
+		Status:          merr.Success(),
+		CostAggregation: &internalpb.CostAggregation{},
+	}, nil)
+
+	lb := shardclient.NewMockLBPolicy(t)
+	lb.EXPECT().UpdateCostMetrics(nodeID, mock.Anything).Return()
+
+	task := &SearchTask{
+		ctx: ctx,
+		SearchRequest: &internalpb.SearchRequest{
+			Base: &commonpb.MsgBase{
+				MsgType:  commonpb.MsgType_Search,
+				SourceID: paramtable.GetNodeID(),
+			},
+		},
+		request: &milvuspb.SearchRequest{
+			DbName:         "default",
+			CollectionName: "test_search_shard_records_node_hint",
+		},
+		Condition:         NewTaskCondition(ctx),
+		lb:                lb,
+		resultBuf:         typeutil.NewConcurrentSet[*internalpb.SearchResults](),
+		queryChannelsNode: typeutil.NewConcurrentMap[string, int64](),
+	}
+
+	err := task.searchShard(ctx, nodeID, qn, channel)
+	require.NoError(t, err)
+	recordedNodeID, ok := task.queryChannelsNode.Get(channel)
+	require.True(t, ok)
+	assert.Equal(t, nodeID, recordedNodeID)
+}

internal/proxy/dql/search_util.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"
@@ -20,6 +20,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
 	"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/typeutil"
 )
 
@@ -517,9 +518,9 @@ func parseSearchInfo(searchParamsPair []*commonpb.KeyValuePair, schema *schemapb
 			// 1. if the request is from iterator, we set topK to QuotaLimit as the iterator can resolve too large topK problem
 			// 2. GetAsInt64 has cached inside, no need to worry about cpu cost for parsing here
 			if largeTopKEnabled {
-				topK = Params.QuotaConfig.LargeTopKLimit.GetAsInt64()
+				topK = paramtable.Get().QuotaConfig.LargeTopKLimit.GetAsInt64()
 			} else {
-				topK = Params.QuotaConfig.TopKLimit.GetAsInt64()
+				topK = paramtable.Get().QuotaConfig.TopKLimit.GetAsInt64()
 			}
 		} else {
 			return nil, merr.WrapErrParameterInvalidMsg("%s [%d] is invalid, %v", TopKKey, topK, err)
@@ -716,7 +717,7 @@ func getNq(req *milvuspb.SearchRequest) (int64, error) {
 	return req.GetNq(), nil
 }
 
-func getPartitionIDs(ctx context.Context, metaCache Cache, dbName string, collectionName string, partitionNames []string) (partitionIDs []UniqueID, err error) {
+func GetPartitionIDs(ctx context.Context, metaCache Cache, dbName string, collectionName string, partitionNames []string) (partitionIDs []UniqueID, err error) {
 	for _, tag := range partitionNames {
 		if err := validatePartitionTag(tag, false); err != nil {
 			return nil, err
@@ -728,7 +729,7 @@ func getPartitionIDs(ctx context.Context, metaCache Cache, dbName string, collec
 		return nil, err
 	}
 
-	useRegexp := Params.ProxyCfg.PartitionNameRegexp.GetAsBool()
+	useRegexp := paramtable.Get().ProxyCfg.PartitionNameRegexp.GetAsBool()
 
 	partitionsSet := typeutil.NewUniqueSet()
 	for _, partitionName := range partitionNames {
@@ -1015,9 +1016,9 @@ func parseGroupByInfo(searchParamsPair []*commonpb.KeyValuePair, schema *schemap
 				fmt.Sprintf("input group size:%d is negative, failed to do search_groupby", groupSize))
 		}
 	}
-	if groupSize > Params.QuotaConfig.MaxGroupSize.GetAsInt64() {
+	if groupSize > paramtable.Get().QuotaConfig.MaxGroupSize.GetAsInt64() {
 		return nil, merr.WrapErrParameterInvalidMsg(
-			fmt.Sprintf("input group size:%d exceeds configured max group size:%d", groupSize, Params.QuotaConfig.MaxGroupSize.GetAsInt64()))
+			fmt.Sprintf("input group size:%d exceeds configured max group size:%d", groupSize, paramtable.Get().QuotaConfig.MaxGroupSize.GetAsInt64()))
 	}
 	ret.groupSize = groupSize
 
@@ -1155,7 +1156,7 @@ func getGroupScorerStr(params []*commonpb.KeyValuePair) string {
 	return groupScorerStr
 }
 
-func convertHybridSearchToSearch(req *milvuspb.HybridSearchRequest) *milvuspb.SearchRequest {
+func ConvertHybridSearchToSearch(req *milvuspb.HybridSearchRequest) *milvuspb.SearchRequest {
 	ret := &milvuspb.SearchRequest{
 		Base:                  req.GetBase(),
 		DbName:                req.GetDbName(),

internal/proxy/dql/search_util_test.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"testing"
@@ -26,13 +26,15 @@ import (
 	"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/pkg/v3/proto/internalpb"
+	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
 	"github.com/milvus-io/milvus/pkg/v3/util/merr"
 )
 
 func TestConvertHybridSearchToSearchKeepsNamespace(t *testing.T) {
 	namespace := "tenant_a"
 
-	searchReq := convertHybridSearchToSearch(&milvuspb.HybridSearchRequest{
+	searchReq := ConvertHybridSearchToSearch(&milvuspb.HybridSearchRequest{
 		CollectionName: "coll",
 		Namespace:      &namespace,
 		Requests: []*milvuspb.SearchRequest{
@@ -55,7 +57,7 @@ func TestConvertHybridSearchToSearchKeepsFunctionChains(t *testing.T) {
 		},
 	}
 
-	searchReq := convertHybridSearchToSearch(&milvuspb.HybridSearchRequest{
+	searchReq := ConvertHybridSearchToSearch(&milvuspb.HybridSearchRequest{
 		CollectionName: "coll",
 		FunctionChains: functionChains,
 		Requests: []*milvuspb.SearchRequest{
@@ -76,7 +78,7 @@ func TestConvertHybridSearchToSearchKeepsSubSearchFunctionChainsByIndex(t *testi
 		Ops:   []*schemapb.FunctionChainOp{{Op: "limit"}},
 	}
 
-	searchReq := convertHybridSearchToSearch(&milvuspb.HybridSearchRequest{
+	searchReq := ConvertHybridSearchToSearch(&milvuspb.HybridSearchRequest{
 		CollectionName: "coll",
 		Requests: []*milvuspb.SearchRequest{
 			{Nq: 1, FunctionChains: []*schemapb.FunctionChain{l0Chain}},
@@ -215,3 +217,36 @@ func TestRankParamsCarryJSONGroupAttributes(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "/brand", searchInfo.planInfo.GetJsonPath())
 }
+
+func TestSearchInfoDetermineSearchTypeWithPluralGroupByFieldIDs(t *testing.T) {
+	info := &SearchInfo{
+		planInfo: &planpb.QueryInfo{
+			GroupByFieldIds: []int64{101},
+		},
+	}
+
+	assert.Equal(t, internalpb.SearchType_DEFAULT, info.DetermineSearchType(false))
+}
+
+func TestParseGroupByInfoLegacyFieldPrecedence(t *testing.T) {
+	schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
+		{FieldID: 101, Name: "brand", DataType: schemapb.DataType_VarChar},
+		{FieldID: 102, Name: "category", DataType: schemapb.DataType_VarChar},
+	}}
+
+	info, err := parseGroupByInfo([]*commonpb.KeyValuePair{
+		{Key: GroupByFieldKey, Value: "brand"},
+		{Key: GroupByFieldsKey, Value: "category"},
+	}, schema)
+	assert.NoError(t, err)
+	assert.Equal(t, []int64{101}, info.groupByFieldIds)
+	assert.Equal(t, []string{"brand"}, info.groupByFieldNames)
+
+	info, err = parseGroupByInfo([]*commonpb.KeyValuePair{
+		{Key: GroupByFieldKey, Value: " "},
+		{Key: GroupByFieldsKey, Value: "brand, category"},
+	}, schema)
+	assert.NoError(t, err)
+	assert.Equal(t, []int64{101, 102}, info.groupByFieldIds)
+	assert.Equal(t, []string{"brand", "category"}, info.groupByFieldNames)
+}

internal/proxy/dql/segment_filter_helper.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"context"

internal/proxy/dql/struct_hybrid_search.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"encoding/json"

internal/proxy/dql/task_query.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"
@@ -18,7 +18,9 @@ import (
 	"github.com/milvus-io/milvus/internal/proxy/accesslog"
 	"github.com/milvus-io/milvus/internal/proxy/channelmgr"
 	"github.com/milvus-io/milvus/internal/proxy/fieldvalidator"
+	"github.com/milvus-io/milvus/internal/proxy/metacache"
 	"github.com/milvus-io/milvus/internal/proxy/shardclient"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/exprutil"
 	"github.com/milvus-io/milvus/internal/util/reduce"
@@ -50,7 +52,7 @@ const (
 	RetrieveTaskName = "RetrieveTask"
 )
 
-type queryTask struct {
+type QueryTask struct {
 	baseTask
 	Condition
 	*internalpb.RetrieveRequest
@@ -58,7 +60,6 @@ type queryTask struct {
 	ctx            context.Context
 	result         *milvuspb.QueryResults
 	request        *milvuspb.QueryRequest
-	mixCoord       types.MixCoordClient
 	ids            *schemapb.IDs
 	collectionName string
 	queryParams    *queryParams
@@ -92,7 +93,86 @@ type queryTask struct {
 	chMgr                channelmgr.ChannelsMgr
 }
 
-func (t *queryTask) getQueryLabel() string {
+// NewQueryTask constructs a query task. The retrieve request carries the
+// per-invocation differences (partition IDs, consistency level, query label);
+// everything derived from the host node comes from the taskmodel.TaskNode
+// contract and paramtable. cache is the meta cache the task reads from, and
+// mustUsePartitionKey carries the proxy.mustUsePartitionKey policy. The caller
+// decides its value: external Query requests pass the configured flag, while
+// internal reads (e.g. the upsert requery) pass false to keep the previous
+// behavior where such reads are not gated by the partition-key check.
+func NewQueryTask(ctx context.Context, node taskmodel.TaskNode, request *milvuspb.QueryRequest, plan *planpb.PlanNode, retrieveReq *internalpb.RetrieveRequest, cache metacache.Cache, mustUsePartitionKey bool) *QueryTask {
+	return &QueryTask{
+		baseTask: baseTask{
+			MetaCache: cache,
+		},
+		ctx:                 ctx,
+		Condition:           NewTaskCondition(ctx),
+		request:             request,
+		plan:                plan,
+		lb:                  node.LBPolicy(),
+		shardclientMgr:      node.ShardMgr(),
+		chMgr:               node.ChMgr(),
+		RetrieveRequest:     retrieveReq,
+		mustUsePartitionKey: mustUsePartitionKey,
+	}
+}
+
+// Result returns the query result after execution.
+func (t *QueryTask) Result() *milvuspb.QueryResults {
+	return t.result
+}
+
+// SetActualChannelsMvcc installs the output snapshot for internal partial-update
+// reads, independent of the channelsMvcc input overrides.
+func (t *QueryTask) SetActualChannelsMvcc(channelReadTs *typeutil.ConcurrentMap[string, uint64]) {
+	t.actualChannelsMvcc = channelReadTs
+}
+
+// RecordChannelMvcc records the MVCC timestamp observed on a channel during a
+// shard read. It is a no-op when the output snapshot has not been installed.
+func (t *QueryTask) RecordChannelMvcc(channel string, ts uint64) {
+	if t.actualChannelsMvcc != nil {
+		t.actualChannelsMvcc.Insert(channel, ts)
+	}
+}
+
+// HasRecordedChannelMvcc reports whether an output MVCC snapshot is installed.
+func (t *QueryTask) HasRecordedChannelMvcc() bool {
+	return t.actualChannelsMvcc != nil
+}
+
+// Request returns the original query request.
+func (t *QueryTask) Request() *milvuspb.QueryRequest {
+	return t.request
+}
+
+// StorageCost returns the storage cost reported by the last execution.
+func (t *QueryTask) StorageCost() segcore.StorageCost {
+	return t.storageCost
+}
+
+// ReQuery reports whether this task was built as a requery of an earlier search.
+func (t *QueryTask) ReQuery() bool {
+	return t.reQuery
+}
+
+// AllQueryCnt returns the total number of queries counted across shards.
+func (t *QueryTask) AllQueryCnt() int64 {
+	return t.allQueryCnt
+}
+
+// TotalRelatedDataSize returns the accumulated related data size.
+func (t *QueryTask) TotalRelatedDataSize() int64 {
+	return t.totalRelatedDataSize
+}
+
+// GetMetricLabel returns the metric label of this query.
+func (t *QueryTask) GetMetricLabel() string {
+	return t.getQueryLabel()
+}
+
+func (t *QueryTask) getQueryLabel() string {
 	if label := t.GetQueryLabel(); label != "" {
 		return label
 	}
@@ -504,7 +584,7 @@ func getPrimaryKeyDataType(schema *schemapb.CollectionSchema) schemapb.DataType
 	return schemapb.DataType_None
 }
 
-func matchCountRule(outputs []string) bool {
+func MatchCountRule(outputs []string) bool {
 	return len(outputs) == 1 && strings.ToLower(strings.TrimSpace(outputs[0])) == "count(*)"
 }
 
@@ -523,7 +603,7 @@ func createCntPlan(expr string, schemaHelper *typeutil.SchemaHelper, exprTemplat
 	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, wrapPlanCreationError(err, "failed to create query plan")
+		return nil, WrapPlanCreationError(err, "failed to create query plan")
 	}
 	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
@@ -531,11 +611,11 @@ func createCntPlan(expr string, schemaHelper *typeutil.SchemaHelper, exprTemplat
 	return plan, nil
 }
 
-func (t *queryTask) createPlan(ctx context.Context) error {
+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 {
+func (t *QueryTask) createPlanArgs(ctx context.Context, visitorArgs *planparserv2.ParserVisitorArgs) error {
 	schema := t.schema
 	var err error
 	if t.plan == nil {
@@ -543,7 +623,7 @@ func (t *queryTask) createPlanArgs(ctx context.Context, visitorArgs *planparserv
 		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 wrapPlanCreationError(err, "failed to create query plan")
+			return WrapPlanCreationError(err, "failed to create query plan")
 		}
 		metrics.ProxyParseExpressionLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), metrics.QueryLabel, metrics.SuccessLabel).Observe(float64(time.Since(start).Microseconds()) / 1000.0)
 	}
@@ -613,7 +693,7 @@ func (t *queryTask) createPlanArgs(ctx context.Context, visitorArgs *planparserv
 	return nil
 }
 
-func (t *queryTask) hasCountStar() bool {
+func (t *QueryTask) hasCountStar() bool {
 	for _, agg := range t.userAggregates {
 		if agg.Name() == "count" && agg.FieldID() == 0 {
 			return true
@@ -622,7 +702,7 @@ func (t *queryTask) hasCountStar() bool {
 	return false
 }
 
-func (t *queryTask) CanSkipAllocTimestamp() bool {
+func (t *QueryTask) CanSkipAllocTimestamp() bool {
 	var consistencyLevel commonpb.ConsistencyLevel
 	useDefaultConsistency := t.request.GetUseDefaultConsistency()
 	if !useDefaultConsistency {
@@ -650,7 +730,7 @@ func (t *queryTask) CanSkipAllocTimestamp() bool {
 	return consistencyLevel != commonpb.ConsistencyLevel_Strong
 }
 
-func (t *queryTask) PreExecute(ctx context.Context) error {
+func (t *QueryTask) PreExecute(ctx context.Context) error {
 	t.Base.MsgType = commonpb.MsgType_Retrieve
 	t.Base.SourceID = paramtable.GetNodeID()
 
@@ -817,7 +897,7 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
 
 			partitionNames = append(partitionNames, hashedPartitionNames...)
 		}
-		t.PartitionIDs, err = getPartitionIDs(ctx, t.GetMetaCache(), t.request.GetDbName(), t.request.CollectionName, partitionNames)
+		t.PartitionIDs, err = GetPartitionIDs(ctx, t.GetMetaCache(), t.request.GetDbName(), t.request.CollectionName, partitionNames)
 		if err != nil {
 			return err
 		}
@@ -830,7 +910,7 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
 	}
 	t.plan.Namespace = namespaceForPlan(t.schema.CollectionSchema, t.request.Namespace)
 
-	t.SerializedExprPlan, _, err = marshalPlanWithMembershipFilterSizeLimit(t.plan, 0)
+	t.SerializedExprPlan, _, err = MarshalPlanWithMembershipFilterSizeLimit(t.plan, 0)
 	if err != nil {
 		return err
 	}
@@ -842,7 +922,7 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
 
 	collectionInfo, err2 := t.GetMetaCache().GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
 	if err2 != nil {
-		log.Warn(ctx, "Proxy::queryTask::PreExecute failed to GetCollectionInfo from cache",
+		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
@@ -910,7 +990,7 @@ func (t *queryTask) PreExecute(ctx context.Context) error {
 	return nil
 }
 
-func (t *queryTask) Execute(ctx context.Context) error {
+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()),
@@ -960,8 +1040,8 @@ func (t *queryTask) Execute(ctx context.Context) error {
 	return nil
 }
 
-func (t *queryTask) PostExecute(ctx context.Context) error {
-	tr := timerecord.NewTimeRecorder("queryTask PostExecute")
+func (t *QueryTask) PostExecute(ctx context.Context) error {
+	tr := timerecord.NewTimeRecorder("QueryTask PostExecute")
 	defer func() {
 		tr.CtxElapse(ctx, "done")
 	}()
@@ -1078,11 +1158,11 @@ func (t *queryTask) PostExecute(ctx context.Context) error {
 	return nil
 }
 
-func (t *queryTask) IsSubTask() bool {
+func (t *QueryTask) IsSubTask() bool {
 	return t.reQuery
 }
 
-func (t *queryTask) queryShard(ctx context.Context, nodeID int64, qn types.QueryNodeClient, channel string) error {
+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
@@ -1278,42 +1358,42 @@ func convertInternalElementIndicesToMilvus(src *internalpb.ElementIndices) *milv
 	}
 }
 
-func (t *queryTask) TraceCtx() context.Context {
+func (t *QueryTask) TraceCtx() context.Context {
 	return t.ctx
 }
 
-func (t *queryTask) ID() UniqueID {
+func (t *QueryTask) ID() UniqueID {
 	return t.Base.MsgID
 }
 
-func (t *queryTask) SetID(uid UniqueID) {
+func (t *QueryTask) SetID(uid UniqueID) {
 	t.Base.MsgID = uid
 }
 
-func (t *queryTask) Name() string {
+func (t *QueryTask) Name() string {
 	return RetrieveTaskName
 }
 
-func (t *queryTask) Type() commonpb.MsgType {
+func (t *QueryTask) Type() commonpb.MsgType {
 	return t.Base.MsgType
 }
 
-func (t *queryTask) BeginTs() Timestamp {
+func (t *QueryTask) BeginTs() Timestamp {
 	return t.Base.Timestamp
 }
 
-func (t *queryTask) EndTs() Timestamp {
+func (t *QueryTask) EndTs() Timestamp {
 	return t.Base.Timestamp
 }
 
-func (t *queryTask) SetTs(ts Timestamp) {
+func (t *QueryTask) SetTs(ts Timestamp) {
 	if t.reQuery && t.Base.Timestamp != 0 {
 		return
 	}
 	t.Base.Timestamp = ts
 }
 
-func (t *queryTask) OnEnqueue() error {
+func (t *QueryTask) OnEnqueue() error {
 	if t.Base == nil {
 		t.Base = commonpbutil.NewMsgBase()
 	}

internal/proxy/dql/task_query_namespace_test.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"
@@ -11,6 +11,7 @@ import (
 	"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/parser/planparserv2"
+	"github.com/milvus-io/milvus/internal/proxy/metacache"
 	"github.com/milvus-io/milvus/pkg/v3/common"
 	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
 	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
@@ -53,15 +54,15 @@ func expectedNamespacePartitionID(namespace string, partitionNames []string, par
 
 func TestQueryTask_PlanNamespace_AfterPreExecute(t *testing.T) {
 	mockey.PatchConvey("TestQueryTask_PlanNamespace_AfterPreExecute", t, func() {
-		cache := &MetaCache{}
-		mockey.Mock((*MetaCache).GetCollectionID).Return(int64(1001), nil).Build()
-		mockey.Mock((*MetaCache).GetCollectionInfo).Return(&collectionInfo{UpdateTimestamp: 12345, ConsistencyLevel: commonpb.ConsistencyLevel_Strong}, nil).Build()
-		mockey.Mock(isPartitionKeyMode).Return(false, nil).Build()
-		mockey.Mock(validatePartitionTag).Return(nil).Build()
-		mockey.Mock(isIgnoreGrowing).Return(false, nil).Build()
+		cache := newTestCache()
+		mockTest(t, (*metacache.MetaCache).GetCollectionID, int64(1001), nil)
+		mockTest(t, (*metacache.MetaCache).GetCollectionInfo, &collectionInfo{UpdateTimestamp: 12345, ConsistencyLevel: commonpb.ConsistencyLevel_Strong}, nil)
+		mockTest(t, isPartitionKeyMode, false, nil)
+		mockTest(t, validatePartitionTag, nil)
+		mockTest(t, isIgnoreGrowing, false, nil)
 
 		// Schema with namespace enabled
-		mockey.Mock((*MetaCache).GetCollectionSchema).To(func(_ *MetaCache, _ context.Context, _ string, _ string) (*schemaInfo, error) {
+		mockTestTo(t, (*metacache.MetaCache).GetCollectionSchema, func(_ *metacache.MetaCache, _ context.Context, _ string, _ string) (*schemaInfo, error) {
 			schema := &schemapb.CollectionSchema{
 				Name: "test_collection",
 				Fields: []*schemapb.FieldSchema{
@@ -71,19 +72,19 @@ func TestQueryTask_PlanNamespace_AfterPreExecute(t *testing.T) {
 				EnableNamespace: true,
 			}
 			return mustNewSchemaInfo(schema), nil
-		}).Build()
+		})
 
 		// Capture plan to verify namespace by mocking plan creation inside createPlan
 		var capturedPlan *planpb.PlanNode
-		mockey.Mock((*queryTask).createPlanArgs).To(func(q *queryTask, ctx context.Context, visitorArgs *planparserv2.ParserVisitorArgs) error {
+		mockTestTo(t, (*QueryTask).createPlanArgs, func(q *QueryTask, ctx context.Context, visitorArgs *planparserv2.ParserVisitorArgs) error {
 			capturedPlan = &planpb.PlanNode{Node: &planpb.PlanNode_Query{Query: &planpb.QueryPlanNode{}}}
 			q.plan = capturedPlan
 			q.translatedOutputFields = []string{"id"}
 			q.userOutputFields = []string{"id"}
 			return nil
-		}).Build()
+		})
 
-		task := &queryTask{
+		task := &QueryTask{
 			baseTask:        baseTask{MetaCache: cache},
 			Condition:       NewTaskCondition(context.Background()),
 			RetrieveRequest: &internalpb.RetrieveRequest{QueryLabel: "query", Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Retrieve}},
@@ -115,7 +116,7 @@ func TestQueryTask_PlanNamespace_AfterPreExecute(t *testing.T) {
 
 func TestQueryTask_NamespaceSetsPartitionIDs(t *testing.T) {
 	mockey.PatchConvey("TestQueryTask_NamespaceSetsPartitionIDs", t, func() {
-		cache := &MetaCache{}
+		cache := newTestCache()
 
 		partitionNames := []string{"_default_0", "_default_1"}
 		partitionIDs := map[string]int64{"_default_0": 101, "_default_1": 102}
@@ -125,15 +126,15 @@ func TestQueryTask_NamespaceSetsPartitionIDs(t *testing.T) {
 			&schemapb.FieldSchema{FieldID: 101, Name: "value", DataType: schemapb.DataType_Int64},
 		)
 
-		mockey.Mock((*MetaCache).GetCollectionID).Return(int64(1001), nil).Build()
-		mockey.Mock((*MetaCache).GetCollectionInfo).Return(&collectionInfo{UpdateTimestamp: 12345, ConsistencyLevel: commonpb.ConsistencyLevel_Strong}, nil).Build()
-		mockey.Mock((*MetaCache).GetCollectionSchema).Return(mustNewSchemaInfo(schema), nil).Build()
-		mockey.Mock((*MetaCache).GetPartitionsIndex).Return(partitionNames, nil).Build()
-		mockey.Mock((*MetaCache).GetPartitions).Return(partitionIDs, nil).Build()
-		mockey.Mock(validatePartitionTag).Return(nil).Build()
-		mockey.Mock(isIgnoreGrowing).Return(false, nil).Build()
+		mockTest(t, (*metacache.MetaCache).GetCollectionID, int64(1001), nil)
+		mockTest(t, (*metacache.MetaCache).GetCollectionInfo, &collectionInfo{UpdateTimestamp: 12345, ConsistencyLevel: commonpb.ConsistencyLevel_Strong}, nil)
+		mockTest(t, (*metacache.MetaCache).GetCollectionSchema, mustNewSchemaInfo(schema), nil)
+		mockTest(t, (*metacache.MetaCache).GetPartitionsIndex, partitionNames, nil)
+		mockTest(t, (*metacache.MetaCache).GetPartitions, partitionIDs, nil)
+		mockTest(t, validatePartitionTag, nil)
+		mockTest(t, isIgnoreGrowing, false, nil)
 
-		mockey.Mock((*queryTask).createPlanArgs).To(func(q *queryTask, ctx context.Context, visitorArgs *planparserv2.ParserVisitorArgs) error {
+		mockTestTo(t, (*QueryTask).createPlanArgs, func(q *QueryTask, ctx context.Context, visitorArgs *planparserv2.ParserVisitorArgs) error {
 			q.plan = &planpb.PlanNode{
 				Node: &planpb.PlanNode_Query{
 					Query: &planpb.QueryPlanNode{
@@ -144,11 +145,11 @@ func TestQueryTask_NamespaceSetsPartitionIDs(t *testing.T) {
 			q.translatedOutputFields = []string{"id"}
 			q.userOutputFields = []string{"id"}
 			return nil
-		}).Build()
+		})
 
 		for _, ns := range namespaces {
 			namespace := ns
-			task := &queryTask{
+			task := &QueryTask{
 				baseTask:        baseTask{MetaCache: cache},
 				Condition:       NewTaskCondition(context.Background()),
 				RetrieveRequest: &internalpb.RetrieveRequest{QueryLabel: "query", Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Retrieve}},

internal/proxy/dql/task_query_test.go

@@ -13,7 +13,7 @@
 // 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
+package dql
 
 import (
 	"context"
@@ -29,11 +29,11 @@ import (
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
-	"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/proxy/metacache"
 	"github.com/milvus-io/milvus/internal/proxy/shardclient"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/reduce"
@@ -65,17 +65,17 @@ func TestQueryTaskPreExecuteSnapshotFences(t *testing.T) {
 		t.Run(scenario.name, func(t *testing.T) {
 			schema := createTestSchema()
 			patch := func(target any, values ...any) {
-				m := mockey.Mock(target).Return(values...).Build()
+				m := mockTest(t, target, values...)
 				t.Cleanup(func() { m.UnPatch() })
 			}
-			patch((*MetaCache).GetCollectionID, int64(1001), nil)
-			patch((*MetaCache).GetCollectionInfo, &collectionInfo{Schema: schema, UpdateTimestamp: 300}, nil)
-			patch((*MetaCache).GetCollectionSchema, schema, nil)
+			patch((*metacache.MetaCache).GetCollectionID, int64(1001), nil)
+			patch((*metacache.MetaCache).GetCollectionInfo, &collectionInfo{Schema: schema, UpdateTimestamp: 300}, nil)
+			patch((*metacache.MetaCache).GetCollectionSchema, schema, nil)
 			patch(isPartitionKeyMode, false, nil)
-			patch(getPartitionIDs, []int64{1}, nil)
+			patch(GetPartitionIDs, []int64{1}, nil)
 			ctx := context.Background()
-			task := &queryTask{
-				baseTask:        baseTask{MetaCache: &MetaCache{}},
+			task := &QueryTask{
+				baseTask:        baseTask{MetaCache: newTestCache()},
 				ctx:             ctx,
 				RetrieveRequest: &internalpb.RetrieveRequest{Base: &commonpb.MsgBase{Timestamp: 200}},
 				request:         &milvuspb.QueryRequest{CollectionName: "test_collection", Expr: "id in [1]", OutputFields: []string{"id"}, GuaranteeTimestamp: scenario.guarantee},
@@ -124,13 +124,13 @@ func TestQueryShardCollectsOnlySuccessfulSnapshots(t *testing.T) {
 	defer patch.UnPatch()
 	type lbTarget struct{ shardclient.LBPolicy }
 	lb := &lbTarget{}
-	patchCost := mockey.Mock((*lbTarget).UpdateCostMetrics).Return().Build()
+	patchCost := mockTest(t, (*lbTarget).UpdateCostMetrics)
 	defer patchCost.UnPatch()
 	type managerTarget struct{ shardclient.ShardClientMgr }
 	manager := &managerTarget{}
-	patchInvalidate := mockey.Mock((*managerTarget).InvalidateShardLeaderCache).Return().Build()
+	patchInvalidate := mockTest(t, (*managerTarget).InvalidateShardLeaderCache)
 	defer patchInvalidate.UnPatch()
-	task := &queryTask{
+	task := &QueryTask{
 		shardclientMgr:     manager,
 		RetrieveRequest:    &internalpb.RetrieveRequest{Base: &commonpb.MsgBase{}, ConsistencyLevel: commonpb.ConsistencyLevel_Strong},
 		resultBuf:          typeutil.NewConcurrentSet[*internalpb.RetrieveResults](),
@@ -196,15 +196,15 @@ func TestQueryTaskPreExecuteUsesStrongVisibility(t *testing.T) {
 		Schema:        schema,
 		CollectionTTL: uint64(collectionTTL),
 	}
-	metaCache := &MetaCache{}
-	collectionIDMocker := mockey.Mock((*MetaCache).GetCollectionID).Return(collectionID, nil).Build()
+	metaCache := newTestCache()
+	collectionIDMocker := mockTest(t, (*metacache.MetaCache).GetCollectionID, collectionID, nil)
 	defer collectionIDMocker.UnPatch()
-	collectionInfoMocker := mockey.Mock((*MetaCache).GetCollectionInfo).Return(collectionInfo, nil).Build()
+	collectionInfoMocker := mockTest(t, (*metacache.MetaCache).GetCollectionInfo, collectionInfo, nil)
 	defer collectionInfoMocker.UnPatch()
-	collectionSchemaMocker := mockey.Mock((*MetaCache).GetCollectionSchema).Return(schema, nil).Build()
+	collectionSchemaMocker := mockTest(t, (*metacache.MetaCache).GetCollectionSchema, schema, nil)
 	defer collectionSchemaMocker.UnPatch()
-	partitionsMocker := mockey.Mock((*MetaCache).GetPartitions).Return(map[string]UniqueID{
-		Params.CommonCfg.DefaultPartitionName.GetValue(): partitionID,
+	partitionsMocker := mockey.Mock((*metacache.MetaCache).GetPartitions).Return(map[string]UniqueID{
+		paramtable.Get().CommonCfg.DefaultPartitionName.GetValue(): partitionID,
 	}, nil).Build()
 	defer partitionsMocker.UnPatch()
 
@@ -219,7 +219,7 @@ func TestQueryTaskPreExecuteUsesStrongVisibility(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			collectionInfo.UpdateTimestamp = tsoutil.ComposeTSByTime(tc.updateTime)
 			ctx := context.Background()
-			task := &queryTask{
+			task := &QueryTask{
 				baseTask:  baseTask{MetaCache: metaCache},
 				Condition: NewTaskCondition(ctx),
 				RetrieveRequest: &internalpb.RetrieveRequest{
@@ -254,13 +254,10 @@ func TestQueryTaskPreExecuteUsesStrongVisibility(t *testing.T) {
 
 func TestQueryTask_all(t *testing.T) {
 	var (
-		err error
 		ctx = context.TODO()
 
-		qc = NewMixCoordMock()
 		qn = getQueryNodeClient()
 
-		shardsNum      = common.DefaultShardsNum
 		collectionName = t.Name() + funcutil.GenRandomStr()
 
 		expr   = fmt.Sprintf("%s > 0", testInt64Field)
@@ -278,8 +275,7 @@ func TestQueryTask_all(t *testing.T) {
 	mgr.EXPECT().InvalidateShardLeaderCache(mock.Anything).Return().Maybe()
 	lb := shardclient.NewLBPolicyImpl(mgr)
 
-	cache, err := initMetaCache(ctx, qc)
-	assert.NoError(t, err)
+	cache := newTestCache()
 
 	fieldName2Types := map[string]schemapb.DataType{
 		testBoolField:     schemapb.DataType_Bool,
@@ -294,39 +290,19 @@ func TestQueryTask_all(t *testing.T) {
 	}
 
 	schema := constructCollectionSchemaByDataType(collectionName, fieldName2Types, testInt64Field, false)
-	marshaledSchema, err := proto.Marshal(schema)
-	assert.NoError(t, err)
+	schemaInfo := mustNewSchemaInfo(schema)
+	collectionID := int64(1000)
 
-	createColT := &createCollectionTask{
-		Condition: NewTaskCondition(ctx),
-		CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
-			CollectionName: collectionName,
-			Schema:         marshaledSchema,
-			ShardsNum:      shardsNum,
-		},
-		ctx:      ctx,
-		mixCoord: qc,
-	}
-
-	require.NoError(t, createColT.OnEnqueue())
-	require.NoError(t, createColT.PreExecute(ctx))
-	require.NoError(t, createColT.Execute(ctx))
-	require.NoError(t, createColT.PostExecute(ctx))
-	collectionID, err := cache.GetCollectionID(ctx, GetCurDBNameFromContextOrDefault(ctx), collectionName)
-	assert.NoError(t, err)
-
-	status, err := qc.LoadCollection(ctx, &querypb.LoadCollectionRequest{
-		Base: &commonpb.MsgBase{
-			MsgType:  commonpb.MsgType_LoadCollection,
-			SourceID: paramtable.GetNodeID(),
-		},
-		CollectionID: collectionID,
-	})
-	require.NoError(t, err)
-	require.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
+	mockTest(t, (*metacache.MetaCache).GetCollectionID, collectionID, nil)
+	mockTest(t, (*metacache.MetaCache).GetCollectionSchema, schemaInfo, nil)
+	mockTest(t, (*metacache.MetaCache).GetCollectionInfo, &collectionInfo{
+		CollID: collectionID,
+		Schema: schemaInfo,
+	}, nil)
+	mockTest(t, (*metacache.MetaCache).GetPartitions, map[string]int64{"_default": 1}, nil)
 
 	t.Run("test query task parameters", func(t *testing.T) {
-		task := &queryTask{
+		task := &QueryTask{
 			baseTask:  baseTask{MetaCache: cache},
 			Condition: NewTaskCondition(ctx),
 			RetrieveRequest: &internalpb.RetrieveRequest{
@@ -357,7 +333,7 @@ func TestQueryTask_all(t *testing.T) {
 					},
 				},
 			},
-			mixCoord:       qc,
+
 			lb:             lb,
 			shardclientMgr: mgr,
 		}
@@ -474,7 +450,7 @@ func TestQueryTask_all(t *testing.T) {
 	})
 
 	t.Run("test query for iterator", func(t *testing.T) {
-		qt := &queryTask{
+		qt := &QueryTask{
 			baseTask:  baseTask{MetaCache: cache},
 			Condition: NewTaskCondition(ctx),
 			RetrieveRequest: &internalpb.RetrieveRequest{
@@ -509,7 +485,7 @@ func TestQueryTask_all(t *testing.T) {
 					},
 				},
 			},
-			mixCoord:       qc,
+
 			lb:             lb,
 			shardclientMgr: mgr,
 			resultBuf:      &typeutil.ConcurrentSet[*internalpb.RetrieveResults]{},
@@ -527,7 +503,7 @@ func TestQueryTask_all(t *testing.T) {
 		assert.True(t, qt.result.GetSessionTs() > 0)
 
 		// next page query task
-		qt = &queryTask{
+		qt = &QueryTask{
 			baseTask:  baseTask{MetaCache: cache},
 			Condition: NewTaskCondition(ctx),
 			RetrieveRequest: &internalpb.RetrieveRequest{
@@ -563,7 +539,7 @@ func TestQueryTask_all(t *testing.T) {
 				},
 				GuaranteeTimestamp: enqueTs,
 			},
-			mixCoord:       qc,
+
 			lb:             lb,
 			shardclientMgr: mgr,
 			resultBuf:      &typeutil.ConcurrentSet[*internalpb.RetrieveResults]{},
@@ -580,7 +556,7 @@ func TestQueryTask_all(t *testing.T) {
 		// 1. Aggregation queries exempt from "empty expression needs limit" check (#47326)
 		// 2. count(*) pagination only blocked when no GROUP BY (#47326)
 
-		makeTask := func(exprStr string, outputFields []string, limitValue int64, groupByField string) *queryTask {
+		makeTask := func(exprStr string, outputFields []string, limitValue int64, groupByField string) *QueryTask {
 			queryParams := []*commonpb.KeyValuePair{
 				{Key: IgnoreGrowingKey, Value: "false"},
 			}
@@ -596,7 +572,7 @@ func TestQueryTask_all(t *testing.T) {
 					Value: groupByField,
 				})
 			}
-			task := &queryTask{
+			task := &QueryTask{
 				baseTask:  baseTask{MetaCache: cache},
 				Condition: NewTaskCondition(ctx),
 				RetrieveRequest: &internalpb.RetrieveRequest{
@@ -623,7 +599,7 @@ func TestQueryTask_all(t *testing.T) {
 					OutputFields:   outputFields,
 					QueryParams:    queryParams,
 				},
-				mixCoord:       qc,
+
 				lb:             lb,
 				shardclientMgr: mgr,
 			}
@@ -696,7 +672,7 @@ func TestQueryTask_all(t *testing.T) {
 
 	t.Run("test order by without limit", func(t *testing.T) {
 		// ORDER BY without explicit limit should fail with an error
-		task := &queryTask{
+		task := &QueryTask{
 			baseTask:  baseTask{MetaCache: cache},
 			Condition: NewTaskCondition(ctx),
 			RetrieveRequest: &internalpb.RetrieveRequest{
@@ -727,7 +703,7 @@ func TestQueryTask_all(t *testing.T) {
 					// No limit specified
 				},
 			},
-			mixCoord:       qc,
+
 			lb:             lb,
 			shardclientMgr: mgr,
 		}
@@ -739,7 +715,7 @@ func TestQueryTask_all(t *testing.T) {
 
 	t.Run("test order by with limit succeeds", func(t *testing.T) {
 		// ORDER BY with explicit limit should pass the validation
-		task := &queryTask{
+		task := &QueryTask{
 			baseTask:  baseTask{MetaCache: cache},
 			Condition: NewTaskCondition(ctx),
 			RetrieveRequest: &internalpb.RetrieveRequest{
@@ -773,7 +749,7 @@ func TestQueryTask_all(t *testing.T) {
 					},
 				},
 			},
-			mixCoord:       qc,
+
 			lb:             lb,
 			shardclientMgr: mgr,
 		}
@@ -1905,7 +1881,7 @@ func Test_queryTask_createPlan(t *testing.T) {
 	collSchema := newTestSchema()
 	t.Run("query without expression", func(t *testing.T) {
 		schema := mustNewSchemaInfo(collSchema)
-		tsk := &queryTask{
+		tsk := &QueryTask{
 			request: &milvuspb.QueryRequest{
 				OutputFields: []string{"Int64"},
 			},
@@ -1918,7 +1894,7 @@ func Test_queryTask_createPlan(t *testing.T) {
 	t.Run("invalid expression", func(t *testing.T) {
 		schema := mustNewSchemaInfo(collSchema)
 
-		tsk := &queryTask{
+		tsk := &QueryTask{
 			schema: schema,
 			request: &milvuspb.QueryRequest{
 				OutputFields: []string{"a"},
@@ -1932,7 +1908,7 @@ func Test_queryTask_createPlan(t *testing.T) {
 	t.Run("invalid output fields", func(t *testing.T) {
 		schema := mustNewSchemaInfo(collSchema)
 
-		tsk := &queryTask{
+		tsk := &QueryTask{
 			schema: schema,
 			request: &milvuspb.QueryRequest{
 				OutputFields: []string{"b"},
@@ -1973,10 +1949,10 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
 	dbName := "test_query"
 	collName := "test_skip_alloc_timestamp"
 	collID := UniqueID(111)
-	mockMetaCache := NewMockCache(t)
+	mockMetaCache := newTestCache()
 
 	t.Run("default consistency level", func(t *testing.T) {
-		qt := &queryTask{
+		qt := &QueryTask{
 			baseTask: baseTask{MetaCache: mockMetaCache},
 			request: &milvuspb.QueryRequest{
 				Base:                  nil,
@@ -1985,41 +1961,38 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
 				UseDefaultConsistency: true,
 			},
 		}
-		mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, nil)
-		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionInfo{
-				CollID:           collID,
-				ConsistencyLevel: commonpb.ConsistencyLevel_Eventually,
-			}, nil).Once()
+		mockTest(t, (*metacache.MetaCache).GetCollectionID, collID, nil)
+		// Sequential consistency-level responses: Eventually, Bounded, Strong.
+		levels := []commonpb.ConsistencyLevel{
+			commonpb.ConsistencyLevel_Eventually,
+			commonpb.ConsistencyLevel_Bounded,
+			commonpb.ConsistencyLevel_Strong,
+		}
+		idx := 0
+		mockTestTo(t, (*metacache.MetaCache).GetCollectionInfo, func(_ *metacache.MetaCache, _ context.Context, _, _ string, _ int64) (*collectionInfo, error) {
+			level := levels[idx]
+			idx++
+			return &collectionInfo{CollID: collID, ConsistencyLevel: level}, nil
+		})
 
 		skip := qt.CanSkipAllocTimestamp()
 		assert.True(t, skip)
 
-		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionInfo{
-				CollID:           collID,
-				ConsistencyLevel: commonpb.ConsistencyLevel_Bounded,
-			}, nil).Once()
 		skip = qt.CanSkipAllocTimestamp()
 		assert.True(t, skip)
 
-		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionInfo{
-				CollID:           collID,
-				ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
-			}, nil).Once()
 		skip = qt.CanSkipAllocTimestamp()
 		assert.False(t, skip)
 	})
 
 	t.Run("request consistency level", func(t *testing.T) {
-		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
+		mockTest(t, (*metacache.MetaCache).GetCollectionInfo,
 			&collectionInfo{
 				CollID:           collID,
 				ConsistencyLevel: commonpb.ConsistencyLevel_Eventually,
-			}, nil).Times(3)
+			}, nil)
 
-		qt := &queryTask{
+		qt := &QueryTask{
 			baseTask: baseTask{MetaCache: mockMetaCache},
 			request: &milvuspb.QueryRequest{
 				Base:                  nil,
@@ -2043,7 +2016,7 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
 	})
 
 	t.Run("legacy_guarantee_ts", func(t *testing.T) {
-		qt := &queryTask{
+		qt := &QueryTask{
 			baseTask: baseTask{MetaCache: mockMetaCache},
 			request: &milvuspb.QueryRequest{
 				Base:                  nil,
@@ -2067,12 +2040,26 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
 	})
 
 	t.Run("failed", func(t *testing.T) {
-		mockMetaCache.ExpectedCalls = nil
-		mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, nil)
-		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			nil, errors.New("mock error")).Once()
+		// GetCollectionInfo returns an error on the first call, then succeeds;
+		// GetCollectionID errors on the second call.
+		infoCalls := 0
+		idCalls := 0
+		mockTestTo(t, (*metacache.MetaCache).GetCollectionID, func(_ *metacache.MetaCache, _ context.Context, _, _ string) (int64, error) {
+			idCalls++
+			if idCalls >= 2 {
+				return 0, errors.New("mock error")
+			}
+			return collID, nil
+		})
+		mockTestTo(t, (*metacache.MetaCache).GetCollectionInfo, func(_ *metacache.MetaCache, _ context.Context, _, _ string, _ int64) (*collectionInfo, error) {
+			infoCalls++
+			if infoCalls == 1 {
+				return nil, errors.New("mock error")
+			}
+			return &collectionInfo{CollID: collID, ConsistencyLevel: commonpb.ConsistencyLevel_Eventually}, nil
+		})
 
-		qt := &queryTask{
+		qt := &QueryTask{
 			baseTask: baseTask{MetaCache: mockMetaCache},
 			request: &milvuspb.QueryRequest{
 				Base:                  nil,
@@ -2086,18 +2073,10 @@ func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
 		skip := qt.CanSkipAllocTimestamp()
 		assert.False(t, skip)
 
-		mockMetaCache.ExpectedCalls = nil
-		mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, errors.New("mock error"))
-		mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
-			&collectionInfo{
-				CollID:           collID,
-				ConsistencyLevel: commonpb.ConsistencyLevel_Eventually,
-			}, nil)
-
 		skip = qt.CanSkipAllocTimestamp()
 		assert.False(t, skip)
 
-		qt2 := &queryTask{
+		qt2 := &QueryTask{
 			baseTask: baseTask{MetaCache: mockMetaCache},
 			request: &milvuspb.QueryRequest{
 				Base:                  nil,

internal/proxy/dql/task_search.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"
@@ -22,6 +22,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proxy/accesslog"
 	"github.com/milvus-io/milvus/internal/proxy/channelmgr"
 	"github.com/milvus-io/milvus/internal/proxy/fieldvalidator"
+	"github.com/milvus-io/milvus/internal/proxy/scheduler"
 	"github.com/milvus-io/milvus/internal/proxy/search_agg"
 	"github.com/milvus-io/milvus/internal/proxy/shardclient"
 	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
@@ -63,7 +64,7 @@ const (
 
 // type requery func(span trace.Span, ids *schemapb.IDs, outputFields []string) (*milvuspb.QueryResults, error)
 
-type searchTask struct {
+type SearchTask struct {
 	baseTask
 	Condition
 	ctx context.Context
@@ -93,8 +94,8 @@ type searchTask struct {
 
 	partitionIDsSet *typeutil.ConcurrentSet[UniqueID]
 
-	mixCoord          types.MixCoordClient
-	node              types.ProxyComponent
+	node              taskmodel.TaskNode
+	sched             *scheduler.TaskScheduler
 	lb                shardclient.LBPolicy
 	shardClientMgr    shardclient.ShardClientMgr
 	queryChannelsTs   map[string]Timestamp
@@ -130,7 +131,70 @@ type searchTask struct {
 	chMgr channelmgr.ChannelsMgr
 }
 
-func (t *searchTask) CanSkipAllocTimestamp() bool {
+// NewSearchTask constructs a search task. Only request-specific inputs are
+// taken from the caller; the host node and its dependencies are derived from
+// the taskmodel.TaskNode contract and paramtable, so the root package does not
+// need to reach into the task's private fields.
+func NewSearchTask(ctx context.Context, node taskmodel.TaskNode, sched *scheduler.TaskScheduler, request *milvuspb.SearchRequest, optimizedSearch bool, isRecallEvaluation bool, tr *timerecord.TimeRecorder) *SearchTask {
+	return &SearchTask{
+		baseTask: baseTask{
+			MetaCache: node.GetMetaCache(),
+		},
+		ctx:       ctx,
+		Condition: NewTaskCondition(ctx),
+		SearchRequest: &internalpb.SearchRequest{
+			Base: commonpbutil.NewMsgBase(
+				commonpbutil.WithMsgType(commonpb.MsgType_Search),
+				commonpbutil.WithSourceID(paramtable.GetNodeID()),
+			),
+			ReqID:              paramtable.GetNodeID(),
+			IsTopkReduce:       optimizedSearch,
+			IsRecallEvaluation: isRecallEvaluation,
+		},
+		request:                request,
+		tr:                     tr,
+		node:                   node,
+		sched:                  sched,
+		lb:                     node.LBPolicy(),
+		shardClientMgr:         node.ShardMgr(),
+		enableMaterializedView: paramtable.Get().CommonCfg.EnableMaterializedView.GetAsBool(),
+		mustUsePartitionKey:    paramtable.Get().ProxyCfg.MustUsePartitionKey.GetAsBool(),
+		chMgr:                  node.ChMgr(),
+	}
+}
+
+// Result returns the search result after execution.
+func (t *SearchTask) Result() *milvuspb.SearchResults {
+	return t.result
+}
+
+// StorageCost returns the storage cost reported by the last execution.
+func (t *SearchTask) StorageCost() segcore.StorageCost {
+	return t.storageCost
+}
+
+// RelatedDataSize returns the related data size reported by the last execution.
+func (t *SearchTask) RelatedDataSize() int64 {
+	return t.relatedDataSize
+}
+
+// ResultSizeInsufficient reports whether the result was truncated because the
+// requested top-k could not fit in the configured result size window.
+func (t *SearchTask) ResultSizeInsufficient() bool {
+	return t.resultSizeInsufficient
+}
+
+// IsTopkReduce reports whether the search ran in top-k reduce mode.
+func (t *SearchTask) IsTopkReduce() bool {
+	return t.isTopkReduce
+}
+
+// IsRecallEvaluation reports whether the search ran in recall-evaluation mode.
+func (t *SearchTask) IsRecallEvaluation() bool {
+	return t.isRecallEvaluation
+}
+
+func (t *SearchTask) CanSkipAllocTimestamp() bool {
 	var consistencyLevel commonpb.ConsistencyLevel
 	useDefaultConsistency := t.request.GetUseDefaultConsistency()
 	if !useDefaultConsistency {
@@ -158,7 +222,7 @@ func (t *searchTask) CanSkipAllocTimestamp() bool {
 	return consistencyLevel != commonpb.ConsistencyLevel_Strong
 }
 
-func (t *searchTask) PreExecute(ctx context.Context) error {
+func (t *SearchTask) PreExecute(ctx context.Context) error {
 	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Search-PreExecute")
 	defer sp.End()
 
@@ -187,7 +251,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
 
 	collectionInfo, err2 := t.GetMetaCache().GetCollectionInfo(ctx, t.request.GetDbName(), collectionName, t.CollectionID)
 	if err2 != nil {
-		log.Warn(ctx, "Proxy::searchTask::PreExecute failed to GetCollectionInfo from cache",
+		log.Warn(ctx, "Proxy::SearchTask::PreExecute failed to GetCollectionInfo from cache",
 			mlog.String("collectionName", collectionName), mlog.Int64("collectionID", t.CollectionID), mlog.Err(err2))
 		return err2
 	}
@@ -216,7 +280,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
 
 	if !t.partitionKeyMode && len(t.request.GetPartitionNames()) > 0 {
 		// translate partition name to partition ids. Use regex-pattern to match partition name.
-		t.PartitionIDs, err = getPartitionIDs(ctx, t.GetMetaCache(), t.request.GetDbName(), collectionName, t.request.GetPartitionNames())
+		t.PartitionIDs, err = GetPartitionIDs(ctx, t.GetMetaCache(), t.request.GetDbName(), collectionName, t.request.GetPartitionNames())
 		if err != nil {
 			log.Warn(ctx, "failed to get partition ids", mlog.Err(err))
 			return err
@@ -361,7 +425,7 @@ func (t *searchTask) PreExecute(ctx context.Context) error {
 	return nil
 }
 
-func (t *searchTask) checkNq(ctx context.Context) (int64, error) {
+func (t *SearchTask) checkNq(ctx context.Context) (int64, error) {
 	var nq int64
 	if t.GetIsAdvanced() {
 		// In the context of Advanced Search, it is essential to verify that the number of vectors
@@ -400,7 +464,7 @@ func (t *searchTask) checkNq(ctx context.Context) (int64, error) {
 	return nq, nil
 }
 
-func (t *searchTask) initSearchAggregation() error {
+func (t *SearchTask) initSearchAggregation() error {
 	spec := t.request.GetSearchAggregation()
 	if spec == nil {
 		t.aggCtx = nil
@@ -448,7 +512,7 @@ func (t *searchTask) initSearchAggregation() error {
 	return nil
 }
 
-func (t *searchTask) validatePartitionKeyIsolation(plan *planpb.PlanNode) error {
+func (t *SearchTask) validatePartitionKeyIsolation(plan *planpb.PlanNode) error {
 	if !t.partitionKeyIsolation {
 		return nil
 	}
@@ -460,7 +524,7 @@ func (t *searchTask) validatePartitionKeyIsolation(plan *planpb.PlanNode) error
 	return exprutil.ValidatePartitionKeyIsolation(expr)
 }
 
-func setQueryInfoIfMvEnable(queryInfo *planpb.QueryInfo, t *searchTask) error {
+func setQueryInfoIfMvEnable(queryInfo *planpb.QueryInfo, t *SearchTask) error {
 	if t.enableMaterializedView {
 		partitionKeyFieldSchema, err := typeutil.GetPartitionKeyFieldSchema(t.schema.CollectionSchema)
 		if err != nil {
@@ -479,7 +543,7 @@ func setQueryInfoIfMvEnable(queryInfo *planpb.QueryInfo, t *searchTask) error {
 	return nil
 }
 
-func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error {
+func (t *SearchTask) initAdvancedSearchRequest(ctx context.Context) error {
 	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "init advanced search request")
 	defer sp.End()
 	t.partitionIDsSet = typeutil.NewConcurrentSet[UniqueID]()
@@ -702,7 +766,7 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error {
 		plan.Namespace = namespaceForPlan(t.schema.CollectionSchema, t.request.Namespace)
 		plan.QuerynodeFunctionChains = querynodeFunctionChains
 
-		internalSubReq.SerializedExprPlan, membershipFilterPlanSize, err = marshalPlanWithMembershipFilterSizeLimit(plan, membershipFilterPlanSize)
+		internalSubReq.SerializedExprPlan, membershipFilterPlanSize, err = MarshalPlanWithMembershipFilterSizeLimit(plan, membershipFilterPlanSize)
 		if err != nil {
 			return err
 		}
@@ -754,7 +818,7 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error {
 	return nil
 }
 
-func (t *searchTask) validateHybridArrayOfVectorGroupBy() error {
+func (t *SearchTask) validateHybridArrayOfVectorGroupBy() error {
 	groupByFieldIDs := t.rankParams.GetGroupByFieldIds()
 	if len(groupByFieldIDs) == 0 {
 		return nil
@@ -798,7 +862,7 @@ func (t *searchTask) validateHybridArrayOfVectorGroupBy() error {
 	return nil
 }
 
-func (t *searchTask) fillResult() {
+func (t *SearchTask) fillResult() {
 	limit := t.GetTopk() - t.GetOffset()
 	resultSizeInsufficient := false
 	if t.aggCtx == nil {
@@ -813,7 +877,7 @@ func (t *searchTask) fillResult() {
 	t.result.CollectionName = t.collectionName
 }
 
-func (t *searchTask) getBM25SearchTexts(placeholder []byte) ([]string, error) {
+func (t *SearchTask) getBM25SearchTexts(placeholder []byte) ([]string, error) {
 	pb := &commonpb.PlaceholderGroup{}
 	if err := proto.Unmarshal(placeholder, pb); err != nil {
 		return nil, merr.WrapErrParameterInvalidMsg("failed to unmarshal BM25 search placeholder group: %v", err)
@@ -832,7 +896,7 @@ func (t *searchTask) getBM25SearchTexts(placeholder []byte) ([]string, error) {
 	return texts, nil
 }
 
-func (t *searchTask) createLexicalHighlighter(highlighter *commonpb.Highlighter, metricType string, annsField int64, placeholder []byte, analyzerName string) error {
+func (t *SearchTask) createLexicalHighlighter(highlighter *commonpb.Highlighter, metricType string, annsField int64, placeholder []byte, analyzerName string) error {
 	h, err := NewLexicalHighlighter(highlighter)
 	if err != nil {
 		return err
@@ -861,7 +925,7 @@ func (t *searchTask) createLexicalHighlighter(highlighter *commonpb.Highlighter,
 	return h.initHighlightQueries(t, analyzerName)
 }
 
-func (t *searchTask) addHighlightTask(highlighter *commonpb.Highlighter, metricType string, annsField int64, placeholder []byte, analyzerName string) error {
+func (t *SearchTask) addHighlightTask(highlighter *commonpb.Highlighter, metricType string, annsField int64, placeholder []byte, analyzerName string) error {
 	if highlighter == nil {
 		return nil
 	}
@@ -881,7 +945,7 @@ func (t *searchTask) addHighlightTask(highlighter *commonpb.Highlighter, metricT
 	}
 }
 
-func (t *searchTask) initSearchRequest(ctx context.Context) error {
+func (t *SearchTask) initSearchRequest(ctx context.Context) error {
 	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "init search request")
 	defer sp.End()
 
@@ -1064,7 +1128,7 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error {
 		queryInfo.StrictGroupSize = true
 	}
 
-	t.SerializedExprPlan, _, err = marshalPlanWithMembershipFilterSizeLimit(plan, 0)
+	t.SerializedExprPlan, _, err = MarshalPlanWithMembershipFilterSizeLimit(plan, 0)
 	if err != nil {
 		return err
 	}
@@ -1155,7 +1219,7 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error {
 // retainRerankDynamicFields keeps hidden L2 inputs in the materialized dynamic
 // root for ordinary Search and each Hybrid sub-search. Empty DynamicFields
 // already fetches the whole root; userOutputFields still owns client projection.
-func (t *searchTask) retainRerankDynamicFields(plan *planpb.PlanNode) {
+func (t *SearchTask) retainRerankDynamicFields(plan *planpb.PlanNode) {
 	if len(plan.DynamicFields) == 0 || t.rerankMeta == nil {
 		return
 	}
@@ -1172,7 +1236,7 @@ func (t *searchTask) retainRerankDynamicFields(plan *planpb.PlanNode) {
 	plan.DynamicFields = fields.Collect()
 }
 
-func (t *searchTask) skipRequeryByNamespacePartitionMode() bool {
+func (t *SearchTask) skipRequeryByNamespacePartitionMode() bool {
 	return t.schema != nil &&
 		t.schema.CollectionSchema != nil &&
 		namespacePartitionModeEnabled(t.schema.CollectionSchema)
@@ -1180,7 +1244,7 @@ func (t *searchTask) skipRequeryByNamespacePartitionMode() bool {
 
 // convertPlaceholderIfNeeded converts fp32 vectors to fp16/bf16 if the target field uses lower precision.
 // Returns converted bytes and the original placeholder type (before any conversion).
-func (t *searchTask) convertPlaceholderIfNeeded(phgBytes []byte, fieldID int64) ([]byte, commonpb.PlaceholderType, error) {
+func (t *SearchTask) convertPlaceholderIfNeeded(phgBytes []byte, fieldID int64) ([]byte, commonpb.PlaceholderType, error) {
 	field := typeutil.GetFieldByID(t.schema.CollectionSchema, fieldID)
 	if field == nil {
 		return phgBytes, 0, nil
@@ -1188,7 +1252,7 @@ func (t *searchTask) convertPlaceholderIfNeeded(phgBytes []byte, fieldID int64)
 	return ConvertPlaceholderGroup(phgBytes, field)
 }
 
-func (t *searchTask) tryGeneratePlan(
+func (t *SearchTask) tryGeneratePlan(
 	params []*commonpb.KeyValuePair,
 	dsl string,
 	exprTemplateValues map[string]*schemapb.TemplateValue,
@@ -1247,7 +1311,7 @@ func (t *searchTask) tryGeneratePlan(
 			mlog.Int("dsl_bytes", len(dsl)),
 			mlog.String("anns field", annsFieldName), mlog.Any("query info", searchInfo.planInfo))
 		metrics.ProxyParseExpressionLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), "search", metrics.FailLabel).Observe(float64(time.Since(start).Microseconds()) / 1000.0)
-		return nil, nil, 0, false, nil, internalpb.SearchType_DEFAULT, wrapPlanCreationError(planErr, "failed to create query plan")
+		return nil, nil, 0, false, nil, internalpb.SearchType_DEFAULT, WrapPlanCreationError(planErr, "failed to create query plan")
 	}
 	metrics.ProxyParseExpressionLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), "search", metrics.SuccessLabel).Observe(float64(time.Since(start).Microseconds()) / 1000.0)
 	mlog.Debug(t.ctx, "create query plan",
@@ -1256,7 +1320,7 @@ func (t *searchTask) tryGeneratePlan(
 	return plan, searchInfo.planInfo, searchInfo.offset, searchInfo.isIterator, searchInfo.orderByFields, searchType, nil
 }
 
-func (t *searchTask) tryParsePartitionIDsFromPlan(plan *planpb.PlanNode) ([]int64, error) {
+func (t *SearchTask) tryParsePartitionIDsFromPlan(plan *planpb.PlanNode) ([]int64, error) {
 	if namespacePartitionKeyMode(t.schema.CollectionSchema) && t.request.Namespace != nil {
 		hashedPartitionNames, err := assignNamespacePartitionKey(t.ctx, t.GetMetaCache(), t.request.GetDbName(), t.collectionName, t.request.Namespace)
 		if err != nil {
@@ -1264,7 +1328,7 @@ func (t *searchTask) tryParsePartitionIDsFromPlan(plan *planpb.PlanNode) ([]int6
 			return nil, err
 		}
 		if len(hashedPartitionNames) > 0 {
-			PartitionIDs, err2 := getPartitionIDs(t.ctx, t.GetMetaCache(), t.request.GetDbName(), t.collectionName, hashedPartitionNames)
+			PartitionIDs, err2 := GetPartitionIDs(t.ctx, t.GetMetaCache(), t.request.GetDbName(), t.collectionName, hashedPartitionNames)
 			if err2 != nil {
 				mlog.Warn(t.ctx, "failed to get namespace partition ids", mlog.Err(err2))
 				return nil, err2
@@ -1288,7 +1352,7 @@ func (t *searchTask) tryParsePartitionIDsFromPlan(plan *planpb.PlanNode) ([]int6
 
 	if len(hashedPartitionNames) > 0 {
 		// translate partition name to partition ids. Use regex-pattern to match partition name.
-		PartitionIDs, err2 := getPartitionIDs(t.ctx, t.GetMetaCache(), t.request.GetDbName(), t.collectionName, hashedPartitionNames)
+		PartitionIDs, err2 := GetPartitionIDs(t.ctx, t.GetMetaCache(), t.request.GetDbName(), t.collectionName, hashedPartitionNames)
 		if err2 != nil {
 			mlog.Warn(t.ctx, "failed to get partition ids", mlog.Err(err2))
 			return nil, err2
@@ -1298,7 +1362,7 @@ func (t *searchTask) tryParsePartitionIDsFromPlan(plan *planpb.PlanNode) ([]int6
 	return nil, nil
 }
 
-func (t *searchTask) Execute(ctx context.Context) error {
+func (t *SearchTask) Execute(ctx context.Context) error {
 	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Search-Execute")
 	defer sp.End()
 	log := mlog.WithLazy(mlog.Int64("nq", t.GetNq()))
@@ -1408,11 +1472,11 @@ func validateElementFilterVectorSearch(plan *planpb.PlanNode, schema *schemapb.C
 		"element_filter is only supported for element-level search on vector sub-fields of the same struct array; use MATCH_ANY/MATCH_* for row-level vector search")
 }
 
-func (t *searchTask) PostExecute(ctx context.Context) error {
+func (t *SearchTask) PostExecute(ctx context.Context) error {
 	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Search-PostExecute")
 	defer sp.End()
 
-	tr := timerecord.NewTimeRecorder("searchTask PostExecute")
+	tr := timerecord.NewTimeRecorder("SearchTask PostExecute")
 	defer func() {
 		tr.CtxElapse(ctx, "done")
 	}()
@@ -1560,7 +1624,7 @@ func (t *searchTask) PostExecute(ctx context.Context) error {
 	return nil
 }
 
-func (t *searchTask) searchShard(ctx context.Context, nodeID int64, qn types.QueryNodeClient, channel string) error {
+func (t *SearchTask) searchShard(ctx context.Context, nodeID int64, qn types.QueryNodeClient, channel string) error {
 	ctx = retry.WithMaxAttemptsContext(ctx, 1)
 	searchReq := shallowcopy.ShallowCopySearchRequest(t.SearchRequest, nodeID)
 	req := &querypb.SearchRequest{
@@ -1605,7 +1669,7 @@ func (t *searchTask) searchShard(ctx context.Context, nodeID int64, qn types.Que
 	return nil
 }
 
-func (t *searchTask) estimateResultSize(nq int64, topK int64) (int64, error) {
+func (t *SearchTask) estimateResultSize(nq int64, topK int64) (int64, error) {
 	vectorOutputFields := lo.Filter(t.schema.GetFields(), func(field *schemapb.FieldSchema, _ int) bool {
 		return lo.Contains(t.translatedOutputFields, field.GetName()) && typeutil.IsVectorType(field.GetDataType())
 	})
@@ -1638,7 +1702,7 @@ func (t *searchTask) estimateResultSize(nq int64, topK int64) (int64, error) {
 	//return int64(sizePerRecord) * nq * topK, nil
 }
 
-func (t *searchTask) collectSearchResults(ctx context.Context) ([]*internalpb.SearchResults, error) {
+func (t *SearchTask) collectSearchResults(ctx context.Context) ([]*internalpb.SearchResults, error) {
 	select {
 	case <-t.TraceCtx().Done():
 		mlog.Warn(ctx, "search task wait to finish timeout!")
@@ -1656,39 +1720,39 @@ func (t *searchTask) collectSearchResults(ctx context.Context) ([]*internalpb.Se
 	}
 }
 
-func (t *searchTask) TraceCtx() context.Context {
+func (t *SearchTask) TraceCtx() context.Context {
 	return t.ctx
 }
 
-func (t *searchTask) ID() UniqueID {
+func (t *SearchTask) ID() UniqueID {
 	return t.Base.MsgID
 }
 
-func (t *searchTask) SetID(uid UniqueID) {
+func (t *SearchTask) SetID(uid UniqueID) {
 	t.Base.MsgID = uid
 }
 
-func (t *searchTask) Name() string {
+func (t *SearchTask) Name() string {
 	return taskmodel.SearchTaskName
 }
 
-func (t *searchTask) Type() commonpb.MsgType {
+func (t *SearchTask) Type() commonpb.MsgType {
 	return t.Base.MsgType
 }
 
-func (t *searchTask) BeginTs() Timestamp {
+func (t *SearchTask) BeginTs() Timestamp {
 	return t.Base.Timestamp
 }
 
-func (t *searchTask) EndTs() Timestamp {
+func (t *SearchTask) EndTs() Timestamp {
 	return t.Base.Timestamp
 }
 
-func (t *searchTask) SetTs(ts Timestamp) {
+func (t *SearchTask) SetTs(ts Timestamp) {
 	t.Base.Timestamp = ts
 }
 
-func (t *searchTask) OnEnqueue() error {
+func (t *SearchTask) OnEnqueue() error {
 	t.Base = commonpbutil.NewMsgBase()
 	t.Base.MsgType = commonpb.MsgType_Search
 	t.Base.SourceID = paramtable.GetNodeID()

internal/proxy/dql/task_search_hybrid_dynamic_test.go

@@ -14,14 +14,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"context"
 	"fmt"
 	"testing"
 
-	"github.com/bytedance/mockey"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel/trace"
@@ -41,7 +40,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
 )
 
-func newHybridDynamicSearchTask(t *testing.T, partitionMode bool, outputs []string, topChain *schemapb.FunctionChain) *searchTask {
+func newHybridDynamicSearchTask(t *testing.T, partitionMode bool, outputs []string, topChain *schemapb.FunctionChain) *SearchTask {
 	t.Helper()
 	schema := newFunctionChainJSONTestSchema().CollectionSchema
 	schema.Name = "hybrid_dynamic"
@@ -88,10 +87,10 @@ func newHybridDynamicSearchTask(t *testing.T, partitionMode bool, outputs []stri
 	l1.Params = map[string]*schemapb.FunctionParamValue{types.InputDataTypesParam: chainDataTypesParam(schemapb.DataType_Int64)}
 	request.Requests[0].FunctionChains = []*schemapb.FunctionChain{l0FunctionChain(l0)}
 	request.Requests[1].FunctionChains = []*schemapb.FunctionChain{l1FunctionChain(l1)}
-	task := &searchTask{
+	task := &SearchTask{
 		ctx: context.Background(), collectionName: schema.Name,
 		SearchRequest: &internalpb.SearchRequest{Base: &commonpb.MsgBase{Timestamp: 1}, CollectionID: 1, PartitionIDs: []int64{1}, IsAdvanced: true, Nq: 2},
-		request:       convertHybridSearchToSearch(request), schema: mustNewSchemaInfo(schema),
+		request:       ConvertHybridSearchToSearch(request), schema: mustNewSchemaInfo(schema),
 		tr: timerecord.NewTimeRecorder("hybrid-dynamic-test"),
 	}
 	task.translatedOutputFields, task.userOutputFields, task.userDynamicFields, _, task.userRequestedPkFieldExplicitly, err = translateOutputFields(outputs, task.schema, true)
@@ -157,7 +156,7 @@ func TestHybridDynamicPipeline(t *testing.T) {
 					}
 					original := opFactory[hybridSearchReduceOp]
 					t.Cleanup(func() { opFactory[hybridSearchReduceOp] = original })
-					opFactory[hybridSearchReduceOp] = func(*searchTask, map[string]any) (operator, error) {
+					opFactory[hybridSearchReduceOp] = func(*SearchTask, map[string]any) (operator, error) {
 						return searchPipelineTestOperator(func(context.Context, trace.Span, ...any) ([]any, error) {
 							return []any{[]*milvuspb.SearchResults{
 								makeResult([]int64{1, 2}, []int64{2, 0}),
@@ -165,7 +164,7 @@ func TestHybridDynamicPipeline(t *testing.T) {
 							}, []string{metric.IP, metric.IP}}, nil
 						}), nil
 					}
-					queryMock := mockey.Mock((*requeryOperator).requery).To(func(_ *requeryOperator, _ context.Context, _ trace.Span, ids *schemapb.IDs, fields []string) (*milvuspb.QueryResults, segcore.StorageCost, error) {
+					queryMock := mockTestTo(t, (*requeryOperator).requery, func(_ *requeryOperator, _ context.Context, _ trace.Span, ids *schemapb.IDs, fields []string) (*milvuspb.QueryResults, segcore.StorageCost, error) {
 						assert.ElementsMatch(t, []int64{1, 2, 3}, ids.GetIntId().GetData())
 						for _, name := range task.rerankMeta.GetInputFieldNames() {
 							assert.Contains(t, fields, name)
@@ -173,7 +172,7 @@ func TestHybridDynamicPipeline(t *testing.T) {
 						assert.NotContains(t, fields, path)
 						// Requery order differs from every sub-result's order.
 						return &milvuspb.QueryResults{FieldsData: hybridDynamicFields([]int64{3, 1, 2})}, segcore.StorageCost{}, nil
-					}).Build()
+					})
 					t.Cleanup(func() { queryMock.UnPatch() })
 					pipeline, err := newSearchPipeline(task)
 					require.NoError(t, err)

internal/proxy/dql/task_search_namespace_test.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"
@@ -12,6 +12,11 @@ import (
 	"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/parser/planparserv2"
+	"github.com/milvus-io/milvus/internal/proxy/channelmgr"
+	"github.com/milvus-io/milvus/internal/proxy/metacache"
+	"github.com/milvus-io/milvus/internal/proxy/shardclient"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
+	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/segcore"
 	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
 	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
@@ -20,14 +25,14 @@ import (
 
 func TestSearchTask_PlanNamespace_AfterPreExecute(t *testing.T) {
 	mockey.PatchConvey("TestSearchTask_PlanNamespace_AfterPreExecute", t, func() {
-		cache := &MetaCache{}
-		mockey.Mock((*MetaCache).GetCollectionID).Return(int64(1001), nil).Build()
-		mockey.Mock((*MetaCache).GetCollectionInfo).Return(&collectionInfo{UpdateTimestamp: 12345, ConsistencyLevel: commonpb.ConsistencyLevel_Strong}, nil).Build()
-		mockey.Mock(isPartitionKeyMode).Return(false, nil).Build()
-		mockey.Mock(isIgnoreGrowing).Return(false, nil).Build()
+		cache := newTestCache()
+		mockTest(t, (*metacache.MetaCache).GetCollectionID, int64(1001), nil)
+		mockTest(t, (*metacache.MetaCache).GetCollectionInfo, &collectionInfo{UpdateTimestamp: 12345, ConsistencyLevel: commonpb.ConsistencyLevel_Strong}, nil)
+		mockTest(t, isPartitionKeyMode, false, nil)
+		mockTest(t, isIgnoreGrowing, false, nil)
 
 		// Schema with namespace enabled and a vector field
-		mockey.Mock((*MetaCache).GetCollectionSchema).To(func(_ *MetaCache, _ context.Context, _ string, _ string) (*schemaInfo, error) {
+		mockTestTo(t, (*metacache.MetaCache).GetCollectionSchema, func(_ *metacache.MetaCache, _ context.Context, _ string, _ string) (*schemaInfo, error) {
 			schema := &schemapb.CollectionSchema{
 				Name: "test_collection",
 				Fields: []*schemapb.FieldSchema{
@@ -37,21 +42,21 @@ func TestSearchTask_PlanNamespace_AfterPreExecute(t *testing.T) {
 				EnableNamespace: true,
 			}
 			return mustNewSchemaInfo(schema), nil
-		}).Build()
+		})
 
 		// Patch checkNq to bypass placeholder parsing
-		mockey.Mock((*searchTask).checkNq).Return(int64(1), nil).Build()
+		mockTest(t, (*SearchTask).checkNq, int64(1), nil)
 
 		// Capture plan to verify namespace by mocking tryGeneratePlan
 		var capturedPlan *planpb.PlanNode
-		mockey.Mock((*searchTask).tryGeneratePlan).To(func(_ *searchTask, _ []*commonpb.KeyValuePair, _ string, _ map[string]*schemapb.TemplateValue, _ *planparserv2.MembershipPreflightBudget, _ bool) (*planpb.PlanNode, *planpb.QueryInfo, int64, bool, []OrderByField, internalpb.SearchType, error) {
+		mockTestTo(t, (*SearchTask).tryGeneratePlan, func(_ *SearchTask, _ []*commonpb.KeyValuePair, _ string, _ map[string]*schemapb.TemplateValue, _ *planparserv2.MembershipPreflightBudget, _ bool) (*planpb.PlanNode, *planpb.QueryInfo, int64, bool, []OrderByField, internalpb.SearchType, error) {
 			capturedPlan = &planpb.PlanNode{}
 			qi := &planpb.QueryInfo{Topk: 10, MetricType: "L2", QueryFieldId: 101, GroupByFieldId: -1}
 			return capturedPlan, qi, 0, false, nil, internalpb.SearchType_DEFAULT, nil
-		}).Build()
+		})
 
 		// Build task
-		task := &searchTask{
+		task := &SearchTask{
 			baseTask:      baseTask{MetaCache: cache},
 			Condition:     NewTaskCondition(context.Background()),
 			SearchRequest: &internalpb.SearchRequest{Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Search}},
@@ -72,7 +77,7 @@ func TestSearchTask_PlanNamespace_AfterPreExecute(t *testing.T) {
 
 func TestSearchTask_NamespaceSetsPartitionIDs(t *testing.T) {
 	mockey.PatchConvey("TestSearchTask_NamespaceSetsPartitionIDs", t, func() {
-		cache := &MetaCache{}
+		cache := newTestCache()
 
 		partitionNames := []string{"_default_0", "_default_1"}
 		partitionIDs := map[string]int64{"_default_0": 101, "_default_1": 102}
@@ -83,14 +88,14 @@ func TestSearchTask_NamespaceSetsPartitionIDs(t *testing.T) {
 			&schemapb.FieldSchema{FieldID: 102, Name: "value", DataType: schemapb.DataType_Int64},
 		)
 
-		mockey.Mock((*MetaCache).GetCollectionID).Return(int64(1001), nil).Build()
-		mockey.Mock((*MetaCache).GetCollectionInfo).Return(&collectionInfo{UpdateTimestamp: 12345, ConsistencyLevel: commonpb.ConsistencyLevel_Strong}, nil).Build()
-		mockey.Mock((*MetaCache).GetCollectionSchema).Return(mustNewSchemaInfo(schema), nil).Build()
-		mockey.Mock((*MetaCache).GetPartitionsIndex).Return(partitionNames, nil).Build()
-		mockey.Mock((*MetaCache).GetPartitions).Return(partitionIDs, nil).Build()
-		mockey.Mock(isIgnoreGrowing).Return(false, nil).Build()
-		mockey.Mock((*searchTask).checkNq).Return(int64(1), nil).Build()
-		mockey.Mock((*searchTask).tryGeneratePlan).To(func(_ *searchTask, _ []*commonpb.KeyValuePair, _ string, _ map[string]*schemapb.TemplateValue, _ *planparserv2.MembershipPreflightBudget, _ bool) (*planpb.PlanNode, *planpb.QueryInfo, int64, bool, []OrderByField, internalpb.SearchType, error) {
+		mockTest(t, (*metacache.MetaCache).GetCollectionID, int64(1001), nil)
+		mockTest(t, (*metacache.MetaCache).GetCollectionInfo, &collectionInfo{UpdateTimestamp: 12345, ConsistencyLevel: commonpb.ConsistencyLevel_Strong}, nil)
+		mockTest(t, (*metacache.MetaCache).GetCollectionSchema, mustNewSchemaInfo(schema), nil)
+		mockTest(t, (*metacache.MetaCache).GetPartitionsIndex, partitionNames, nil)
+		mockTest(t, (*metacache.MetaCache).GetPartitions, partitionIDs, nil)
+		mockTest(t, isIgnoreGrowing, false, nil)
+		mockTest(t, (*SearchTask).checkNq, int64(1), nil)
+		mockTestTo(t, (*SearchTask).tryGeneratePlan, func(_ *SearchTask, _ []*commonpb.KeyValuePair, _ string, _ map[string]*schemapb.TemplateValue, _ *planparserv2.MembershipPreflightBudget, _ bool) (*planpb.PlanNode, *planpb.QueryInfo, int64, bool, []OrderByField, internalpb.SearchType, error) {
 			plan := &planpb.PlanNode{
 				Node: &planpb.PlanNode_VectorAnns{
 					VectorAnns: &planpb.VectorANNS{
@@ -100,11 +105,11 @@ func TestSearchTask_NamespaceSetsPartitionIDs(t *testing.T) {
 			}
 			qi := &planpb.QueryInfo{Topk: 10, MetricType: "L2", QueryFieldId: 101, GroupByFieldId: -1}
 			return plan, qi, 0, false, nil, internalpb.SearchType_DEFAULT, nil
-		}).Build()
+		})
 
 		for _, ns := range namespaces {
 			namespace := ns
-			task := &searchTask{
+			task := &SearchTask{
 				baseTask:      baseTask{MetaCache: cache},
 				Condition:     NewTaskCondition(context.Background()),
 				SearchRequest: &internalpb.SearchRequest{Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Search}},
@@ -122,7 +127,7 @@ func TestSearchTask_NamespaceSetsPartitionIDs(t *testing.T) {
 
 func TestSearchTask_RequeryPlanNamespace(t *testing.T) {
 	mockey.PatchConvey("TestSearchTask_RequeryPlanNamespace", t, func() {
-		// Minimal searchTask with schema and namespace
+		// Minimal SearchTask with schema and namespace
 		schema := &schemapb.CollectionSchema{
 			Name: "test_collection",
 			Fields: []*schemapb.FieldSchema{
@@ -131,31 +136,32 @@ func TestSearchTask_RequeryPlanNamespace(t *testing.T) {
 			},
 			EnableNamespace: true,
 		}
-		tsk := &searchTask{
+		tsk := &SearchTask{
 			Condition:     NewTaskCondition(context.Background()),
 			ctx:           context.Background(),
 			schema:        mustNewSchemaInfo(schema),
 			request:       &milvuspb.SearchRequest{CollectionName: "test_collection"},
 			SearchRequest: &internalpb.SearchRequest{Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Search}},
-			node:          &Proxy{},
+			node:          &namespaceRequeryMockNode{},
 		}
 		ns := "ns-1"
 		tsk.request.Namespace = &ns
 
 		// Capture plan created in requery
 		var capturedPlan *planpb.PlanNode
-		mockey.Mock(planparserv2.CreateRequeryPlan).To(func(_ *schemapb.FieldSchema, _ *schemapb.IDs) *planpb.PlanNode {
+		mockTestTo(t, planparserv2.CreateRequeryPlan, func(_ *schemapb.FieldSchema, _ *schemapb.IDs) *planpb.PlanNode {
 			capturedPlan = &planpb.PlanNode{}
 			return capturedPlan
-		}).Build()
+		})
 
-		// Capture qt.request as well to ensure request namespace is wired
-		mockey.Mock((*Proxy).query).To(func(_ *Proxy, _ context.Context, qt *queryTask, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
-			if qt.plan == nil || qt.plan.Namespace == nil || *qt.plan.Namespace != *tsk.request.Namespace {
-				t.Fatalf("requery plan namespace mismatch, got=%v want=%v", qt.plan.Namespace, *tsk.request.Namespace)
+		// Capture qt.plan to ensure request namespace is wired
+		mockTestTo(t, (*namespaceRequeryMockNode).ExecuteQuery, func(_ *namespaceRequeryMockNode, _ context.Context, qt taskmodel.Task, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
+			queryTask := qt.(*QueryTask)
+			if queryTask.plan == nil || queryTask.plan.Namespace == nil || *queryTask.plan.Namespace != *tsk.request.Namespace {
+				t.Fatalf("requery plan namespace mismatch, got=%v want=%v", queryTask.plan.Namespace, *tsk.request.Namespace)
 			}
 			return &milvuspb.QueryResults{Status: merr.Success()}, segcore.StorageCost{}, nil
-		}).Build()
+		})
 
 		op, err := newRequeryOperator(tsk, nil)
 		assert.NoError(t, err)
@@ -168,3 +174,22 @@ func TestSearchTask_RequeryPlanNamespace(t *testing.T) {
 		assert.Equal(t, *tsk.request.Namespace, *capturedPlan.Namespace)
 	})
 }
+
+// namespaceRequeryMockNode is a taskmodel.TaskNode + QueryRunner stub used to
+// capture the requery plan namespace in tests.
+type namespaceRequeryMockNode struct{}
+
+func (n *namespaceRequeryMockNode) GetMetaCache() metacache.Cache        { return newTestCache() }
+func (n *namespaceRequeryMockNode) MixCoord() types.MixCoordClient       { return nil }
+func (n *namespaceRequeryMockNode) LBPolicy() shardclient.LBPolicy       { return nil }
+func (n *namespaceRequeryMockNode) ShardMgr() shardclient.ShardClientMgr { return nil }
+func (n *namespaceRequeryMockNode) ChMgr() channelmgr.ChannelsMgr        { return nil }
+func (n *namespaceRequeryMockNode) TsoAllocator() taskmodel.TsoAllocator { return &mockTsoAllocator{} }
+func (n *namespaceRequeryMockNode) ExecuteQuery(ctx context.Context, qt taskmodel.Task, sp trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
+	panic("ExecuteQuery must be patched by mockey")
+}
+
+var (
+	_ taskmodel.TaskNode    = (*namespaceRequeryMockNode)(nil)
+	_ taskmodel.QueryRunner = (*namespaceRequeryMockNode)(nil)
+)

internal/proxy/dql/task_search_pk_hint_test.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"testing"

internal/proxy/dql/task_statistic.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"context"
@@ -12,6 +12,7 @@ import (
 	"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/internal/proxy/shardclient"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/pkg/v3/mlog"
 	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
@@ -31,7 +32,7 @@ const (
 	GetPartitionStatisticsTaskName  = "GetPartitionStatisticsTask"
 )
 
-type getStatisticsTask struct {
+type GetStatisticsTask struct {
 	request *milvuspb.GetStatisticsRequest
 	result  *milvuspb.GetStatisticsResponse
 	baseTask
@@ -58,39 +59,59 @@ type getStatisticsTask struct {
 	lb             shardclient.LBPolicy
 }
 
-func (g *getStatisticsTask) TraceCtx() context.Context {
+// NewGetStatisticsTask constructs a statistics task. Host-node dependencies are
+// derived from the taskmodel.TaskNode contract.
+func NewGetStatisticsTask(ctx context.Context, node taskmodel.TaskNode, request *milvuspb.GetStatisticsRequest, tr *timerecord.TimeRecorder) *GetStatisticsTask {
+	return &GetStatisticsTask{
+		request:        request,
+		baseTask:       baseTask{MetaCache: node.GetMetaCache()},
+		Condition:      NewTaskCondition(ctx),
+		ctx:            ctx,
+		tr:             tr,
+		mixc:           node.MixCoord(),
+		lb:             node.LBPolicy(),
+		shardclientMgr: node.ShardMgr(),
+	}
+}
+
+// Result returns the statistics result after execution.
+func (g *GetStatisticsTask) Result() *milvuspb.GetStatisticsResponse {
+	return g.result
+}
+
+func (g *GetStatisticsTask) TraceCtx() context.Context {
 	return g.ctx
 }
 
-func (g *getStatisticsTask) ID() UniqueID {
+func (g *GetStatisticsTask) ID() UniqueID {
 	return g.Base.MsgID
 }
 
-func (g *getStatisticsTask) SetID(uid UniqueID) {
+func (g *GetStatisticsTask) SetID(uid UniqueID) {
 	g.Base.MsgID = uid
 }
 
-func (g *getStatisticsTask) Name() string {
+func (g *GetStatisticsTask) Name() string {
 	return GetPartitionStatisticsTaskName
 }
 
-func (g *getStatisticsTask) Type() commonpb.MsgType {
+func (g *GetStatisticsTask) Type() commonpb.MsgType {
 	return g.Base.MsgType
 }
 
-func (g *getStatisticsTask) BeginTs() Timestamp {
+func (g *GetStatisticsTask) BeginTs() Timestamp {
 	return g.Base.Timestamp
 }
 
-func (g *getStatisticsTask) EndTs() Timestamp {
+func (g *GetStatisticsTask) EndTs() Timestamp {
 	return g.Base.Timestamp
 }
 
-func (g *getStatisticsTask) SetTs(ts Timestamp) {
+func (g *GetStatisticsTask) SetTs(ts Timestamp) {
 	g.Base.Timestamp = ts
 }
 
-func (g *getStatisticsTask) OnEnqueue() error {
+func (g *GetStatisticsTask) OnEnqueue() error {
 	g.GetStatisticsRequest = &internalpb.GetStatisticsRequest{
 		Base: commonpbutil.NewMsgBase(),
 	}
@@ -100,7 +121,7 @@ func (g *getStatisticsTask) OnEnqueue() error {
 	return nil
 }
 
-func (g *getStatisticsTask) PreExecute(ctx context.Context) error {
+func (g *GetStatisticsTask) PreExecute(ctx context.Context) error {
 	g.DbID = 0
 	g.collectionName = g.request.GetCollectionName()
 	g.partitionNames = g.request.GetPartitionNames()
@@ -114,7 +135,7 @@ func (g *getStatisticsTask) PreExecute(ctx context.Context) error {
 	if err != nil { // err is not nil if collection not exists
 		return err
 	}
-	partIDs, err := getPartitionIDs(ctx, g.GetMetaCache(), g.request.GetDbName(), g.collectionName, g.partitionNames)
+	partIDs, err := GetPartitionIDs(ctx, g.GetMetaCache(), g.request.GetDbName(), g.collectionName, g.partitionNames)
 	if err != nil { // err is not nil if partition not exists
 		return err
 	}
@@ -158,7 +179,7 @@ func (g *getStatisticsTask) PreExecute(ctx context.Context) error {
 	return nil
 }
 
-func (g *getStatisticsTask) Execute(ctx context.Context) error {
+func (g *GetStatisticsTask) Execute(ctx context.Context) error {
 	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetStatistics-Execute")
 	defer sp.End()
 	if g.fromQueryNode {
@@ -183,7 +204,7 @@ func (g *getStatisticsTask) Execute(ctx context.Context) error {
 	return nil
 }
 
-func (g *getStatisticsTask) PostExecute(ctx context.Context) error {
+func (g *GetStatisticsTask) PostExecute(ctx context.Context) error {
 	_, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetStatistic-PostExecute")
 	defer sp.End()
 	tr := timerecord.NewTimeRecorder("getStatisticTask PostExecute")
@@ -224,7 +245,7 @@ func (g *getStatisticsTask) PostExecute(ctx context.Context) error {
 	return nil
 }
 
-func (g *getStatisticsTask) getStatisticsFromDataCoord(ctx context.Context) error {
+func (g *GetStatisticsTask) getStatisticsFromDataCoord(ctx context.Context) error {
 	collID := g.CollectionID
 	partIDs := g.unloadedPartitionIDs
 
@@ -254,7 +275,7 @@ func (g *getStatisticsTask) getStatisticsFromDataCoord(ctx context.Context) erro
 	return nil
 }
 
-func (g *getStatisticsTask) getStatisticsFromQueryNode(ctx context.Context) error {
+func (g *GetStatisticsTask) getStatisticsFromQueryNode(ctx context.Context) error {
 	g.PartitionIDs = g.loadedPartitionIDs
 	if g.resultBuf == nil {
 		g.resultBuf = typeutil.NewConcurrentSet[*internalpb.GetStatisticsResponse]()
@@ -273,7 +294,7 @@ func (g *getStatisticsTask) getStatisticsFromQueryNode(ctx context.Context) erro
 	return nil
 }
 
-func (g *getStatisticsTask) getStatisticsShard(ctx context.Context, nodeID int64, qn types.QueryNodeClient, channel string) error {
+func (g *GetStatisticsTask) getStatisticsShard(ctx context.Context, nodeID int64, qn types.QueryNodeClient, channel string) error {
 	nodeReq := proto.Clone(g.GetStatisticsRequest).(*internalpb.GetStatisticsRequest)
 	nodeReq.Base.TargetID = nodeID
 	req := &querypb.GetStatisticsRequest{
@@ -423,24 +444,24 @@ func reduceStatisticResponse(results []map[string]string) ([]*commonpb.KeyValueP
 // implement Task
 // try to compatible with old API (getCollectionStatistics & getPartitionStatistics)
 
-//type getPartitionStatisticsTask struct {
-//	getStatisticsTask
+//type GetPartitionStatisticsTask struct {
+//	GetStatisticsTask
 //	request *milvuspb.GetPartitionStatisticsRequest
 //	result  *milvuspb.GetPartitionStatisticsResponse
 //}
 //
-//func (g *getPartitionStatisticsTask) PreExecute(ctx context.Context) error {
-//	g.getStatisticsTask.DbID = 0
-//	g.getStatisticsTask.collectionName = g.request.GetCollectionName()
-//	g.getStatisticsTask.partitionNames = []string{g.request.GetPartitionName()}
+//func (g *GetPartitionStatisticsTask) PreExecute(ctx context.Context) error {
+//	g.GetStatisticsTask.DbID = 0
+//	g.GetStatisticsTask.collectionName = g.request.GetCollectionName()
+//	g.GetStatisticsTask.partitionNames = []string{g.request.GetPartitionName()}
 //	// g.TravelTimestamp = g.request.GetTravelTimestamp()
 //	// g.GuaranteeTimestamp = g.request.GetGuaranteeTimestamp()
-//	return g.getStatisticsTask.PreExecute(ctx)
+//	return g.GetStatisticsTask.PreExecute(ctx)
 //}
 //
-//func (g *getPartitionStatisticsTask) Execute(ctx context.Context) error {
+//func (g *GetPartitionStatisticsTask) Execute(ctx context.Context) error {
 //	if g.fromQueryNode {
-//		err := g.getStatisticsTask.Execute(ctx)
+//		err := g.GetStatisticsTask.Execute(ctx)
 //		if err != nil {
 //			return err
 //		}
@@ -478,8 +499,8 @@ func reduceStatisticResponse(results []map[string]string) ([]*commonpb.KeyValueP
 //	return nil
 //}
 //
-//func (g *getPartitionStatisticsTask) PostExecute(ctx context.Context) error {
-//	err := g.getStatisticsTask.PostExecute(ctx)
+//func (g *GetPartitionStatisticsTask) PostExecute(ctx context.Context) error {
+//	err := g.GetStatisticsTask.PostExecute(ctx)
 //	if err != nil {
 //		return err
 //	}
@@ -490,28 +511,28 @@ func reduceStatisticResponse(results []map[string]string) ([]*commonpb.KeyValueP
 //	return nil
 //}
 //
-//type getCollectionStatisticsTask struct {
-//	getStatisticsTask
+//type GetCollectionStatisticsTask struct {
+//	GetStatisticsTask
 //	request *milvuspb.GetCollectionStatisticsRequest
 //	result  *milvuspb.GetCollectionStatisticsResponse
 //}
 //
-//func (g *getCollectionStatisticsTask) PreExecute(ctx context.Context) error {
-//	g.getStatisticsTask.DbID = 0
-//	g.getStatisticsTask.collectionName = g.request.GetCollectionName()
-//	g.getStatisticsTask.partitionNames = []string{}
+//func (g *GetCollectionStatisticsTask) PreExecute(ctx context.Context) error {
+//	g.GetStatisticsTask.DbID = 0
+//	g.GetStatisticsTask.collectionName = g.request.GetCollectionName()
+//	g.GetStatisticsTask.partitionNames = []string{}
 //	// g.TravelTimestamp = g.request.GetTravelTimestamp()
 //	// g.GuaranteeTimestamp = g.request.GetGuaranteeTimestamp()
-//	return g.getStatisticsTask.PreExecute(ctx)
+//	return g.GetStatisticsTask.PreExecute(ctx)
 //}
 //
-//func (g *getCollectionStatisticsTask) Execute(ctx context.Context) error {
+//func (g *GetCollectionStatisticsTask) Execute(ctx context.Context) error {
 //	if g.fromQueryNode {
 //		// if you get entire collection, we need to pass partition ids param.
 //		if len(g.unloadedPartitionIDs) == 0 {
 //			g.GetStatisticsRequest.PartitionIDs = nil
 //		}
-//		err := g.getStatisticsTask.Execute(ctx)
+//		err := g.GetStatisticsTask.Execute(ctx)
 //		if err != nil {
 //			return err
 //		}
@@ -574,8 +595,8 @@ func reduceStatisticResponse(results []map[string]string) ([]*commonpb.KeyValueP
 //	return nil
 //}
 //
-//func (g *getCollectionStatisticsTask) PostExecute(ctx context.Context) error {
-//	err := g.getStatisticsTask.PostExecute(ctx)
+//func (g *GetCollectionStatisticsTask) PostExecute(ctx context.Context) error {
+//	err := g.GetStatisticsTask.PostExecute(ctx)
 //	if err != nil {
 //		return err
 //	}
@@ -587,8 +608,8 @@ func reduceStatisticResponse(results []map[string]string) ([]*commonpb.KeyValueP
 //}
 
 // old version of get statistics
-// please remove it after getStatisticsTask below is stable
-type getCollectionStatisticsTask struct {
+// please remove it after GetStatisticsTask below is stable
+type GetCollectionStatisticsTask struct {
 	baseTask
 	Condition
 	*milvuspb.GetCollectionStatisticsRequest
@@ -599,50 +620,67 @@ type getCollectionStatisticsTask struct {
 	collectionID UniqueID
 }
 
-func (g *getCollectionStatisticsTask) TraceCtx() context.Context {
+// NewGetCollectionStatisticsTask constructs a collection-statistics task.
+// Host-node dependencies are derived from the taskmodel.TaskNode contract.
+func NewGetCollectionStatisticsTask(ctx context.Context, node taskmodel.TaskNode, request *milvuspb.GetCollectionStatisticsRequest) *GetCollectionStatisticsTask {
+	return &GetCollectionStatisticsTask{
+		baseTask:                       baseTask{MetaCache: node.GetMetaCache()},
+		ctx:                            ctx,
+		Condition:                      NewTaskCondition(ctx),
+		GetCollectionStatisticsRequest: request,
+		mixCoord:                       node.MixCoord(),
+	}
+}
+
+// Result returns the statistics result after execution.
+func (g *GetCollectionStatisticsTask) Result() *milvuspb.GetCollectionStatisticsResponse {
+	return g.result
+}
+
+func (g *GetCollectionStatisticsTask) TraceCtx() context.Context {
 	return g.ctx
 }
 
-func (g *getCollectionStatisticsTask) ID() UniqueID {
+func (g *GetCollectionStatisticsTask) ID() UniqueID {
 	return g.Base.MsgID
 }
 
-func (g *getCollectionStatisticsTask) SetID(uid UniqueID) {
+func (g *GetCollectionStatisticsTask) SetID(uid UniqueID) {
 	g.Base.MsgID = uid
 }
 
-func (g *getCollectionStatisticsTask) Name() string {
+func (g *GetCollectionStatisticsTask) Name() string {
 	return GetCollectionStatisticsTaskName
 }
 
-func (g *getCollectionStatisticsTask) Type() commonpb.MsgType {
+func (g *GetCollectionStatisticsTask) Type() commonpb.MsgType {
 	return g.Base.MsgType
 }
 
-func (g *getCollectionStatisticsTask) BeginTs() Timestamp {
+func (g *GetCollectionStatisticsTask) BeginTs() Timestamp {
 	return g.Base.Timestamp
 }
 
-func (g *getCollectionStatisticsTask) EndTs() Timestamp {
+func (g *GetCollectionStatisticsTask) EndTs() Timestamp {
 	return g.Base.Timestamp
 }
 
-func (g *getCollectionStatisticsTask) SetTs(ts Timestamp) {
+func (g *GetCollectionStatisticsTask) SetTs(ts Timestamp) {
 	g.Base.Timestamp = ts
 }
 
-func (g *getCollectionStatisticsTask) OnEnqueue() error {
+func (g *GetCollectionStatisticsTask) OnEnqueue() error {
 	g.Base = commonpbutil.NewMsgBase()
 	g.Base.MsgType = commonpb.MsgType_GetCollectionStatistics
 	g.Base.SourceID = paramtable.GetNodeID()
 	return nil
 }
 
-func (g *getCollectionStatisticsTask) PreExecute(ctx context.Context) error {
+func (g *GetCollectionStatisticsTask) PreExecute(ctx context.Context) error {
 	return nil
 }
 
-func (g *getCollectionStatisticsTask) Execute(ctx context.Context) error {
+func (g *GetCollectionStatisticsTask) Execute(ctx context.Context) error {
 	collID, err := g.GetMetaCache().GetCollectionID(ctx, g.GetDbName(), g.CollectionName)
 	if err != nil {
 		return err
@@ -667,11 +705,11 @@ func (g *getCollectionStatisticsTask) Execute(ctx context.Context) error {
 	return nil
 }
 
-func (g *getCollectionStatisticsTask) PostExecute(ctx context.Context) error {
+func (g *GetCollectionStatisticsTask) PostExecute(ctx context.Context) error {
 	return nil
 }
 
-type getPartitionStatisticsTask struct {
+type GetPartitionStatisticsTask struct {
 	baseTask
 	Condition
 	*milvuspb.GetPartitionStatisticsRequest
@@ -682,50 +720,67 @@ type getPartitionStatisticsTask struct {
 	collectionID UniqueID
 }
 
-func (g *getPartitionStatisticsTask) TraceCtx() context.Context {
+// NewGetPartitionStatisticsTask constructs a partition-statistics task.
+// Host-node dependencies are derived from the taskmodel.TaskNode contract.
+func NewGetPartitionStatisticsTask(ctx context.Context, node taskmodel.TaskNode, request *milvuspb.GetPartitionStatisticsRequest) *GetPartitionStatisticsTask {
+	return &GetPartitionStatisticsTask{
+		baseTask:                      baseTask{MetaCache: node.GetMetaCache()},
+		ctx:                           ctx,
+		Condition:                     NewTaskCondition(ctx),
+		GetPartitionStatisticsRequest: request,
+		mixCoord:                      node.MixCoord(),
+	}
+}
+
+// Result returns the statistics result after execution.
+func (g *GetPartitionStatisticsTask) Result() *milvuspb.GetPartitionStatisticsResponse {
+	return g.result
+}
+
+func (g *GetPartitionStatisticsTask) TraceCtx() context.Context {
 	return g.ctx
 }
 
-func (g *getPartitionStatisticsTask) ID() UniqueID {
+func (g *GetPartitionStatisticsTask) ID() UniqueID {
 	return g.Base.MsgID
 }
 
-func (g *getPartitionStatisticsTask) SetID(uid UniqueID) {
+func (g *GetPartitionStatisticsTask) SetID(uid UniqueID) {
 	g.Base.MsgID = uid
 }
 
-func (g *getPartitionStatisticsTask) Name() string {
+func (g *GetPartitionStatisticsTask) Name() string {
 	return GetPartitionStatisticsTaskName
 }
 
-func (g *getPartitionStatisticsTask) Type() commonpb.MsgType {
+func (g *GetPartitionStatisticsTask) Type() commonpb.MsgType {
 	return g.Base.MsgType
 }
 
-func (g *getPartitionStatisticsTask) BeginTs() Timestamp {
+func (g *GetPartitionStatisticsTask) BeginTs() Timestamp {
 	return g.Base.Timestamp
 }
 
-func (g *getPartitionStatisticsTask) EndTs() Timestamp {
+func (g *GetPartitionStatisticsTask) EndTs() Timestamp {
 	return g.Base.Timestamp
 }
 
-func (g *getPartitionStatisticsTask) SetTs(ts Timestamp) {
+func (g *GetPartitionStatisticsTask) SetTs(ts Timestamp) {
 	g.Base.Timestamp = ts
 }
 
-func (g *getPartitionStatisticsTask) OnEnqueue() error {
+func (g *GetPartitionStatisticsTask) OnEnqueue() error {
 	g.Base = commonpbutil.NewMsgBase()
 	g.Base.MsgType = commonpb.MsgType_GetPartitionStatistics
 	g.Base.SourceID = paramtable.GetNodeID()
 	return nil
 }
 
-func (g *getPartitionStatisticsTask) PreExecute(ctx context.Context) error {
+func (g *GetPartitionStatisticsTask) PreExecute(ctx context.Context) error {
 	return nil
 }
 
-func (g *getPartitionStatisticsTask) Execute(ctx context.Context) error {
+func (g *GetPartitionStatisticsTask) Execute(ctx context.Context) error {
 	collID, err := g.GetMetaCache().GetCollectionID(ctx, g.GetDbName(), g.CollectionName)
 	if err != nil {
 		return err
@@ -758,6 +813,6 @@ func (g *getPartitionStatisticsTask) Execute(ctx context.Context) error {
 	return nil
 }
 
-func (g *getPartitionStatisticsTask) PostExecute(ctx context.Context) error {
+func (g *GetPartitionStatisticsTask) PostExecute(ctx context.Context) error {
 	return nil
 }

internal/proxy/dql/task_statistic_integration_test.go

@@ -13,7 +13,7 @@
 // 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
+package dql
 
 import (
 	"context"
@@ -22,16 +22,14 @@ import (
 
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
-	"google.golang.org/grpc"
-	"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/mocks"
+	"github.com/milvus-io/milvus/internal/proxy/metacache"
 	"github.com/milvus-io/milvus/internal/proxy/shardclient"
 	"github.com/milvus-io/milvus/internal/types"
-	"github.com/milvus-io/milvus/pkg/v3/common"
 	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
 	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
 	"github.com/milvus-io/milvus/pkg/v3/util/merr"
@@ -46,7 +44,7 @@ type StatisticTaskSuite struct {
 
 	lb        shardclient.LBPolicy
 	mgr       shardclient.ShardClientMgr
-	metaCache Cache
+	metaCache *metacache.MetaCache
 
 	collectionName string
 	collectionID   int64
@@ -58,30 +56,27 @@ func (s *StatisticTaskSuite) SetupSuite() {
 
 func (s *StatisticTaskSuite) SetupTest() {
 	successStatus := commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
-	mixc := NewMixCoordMock()
+	mixc := mocks.NewMockMixCoordClient(s.T())
 
-	mixc.GetShardLeadersFunc = func(ctx context.Context, req *querypb.GetShardLeadersRequest, opts ...grpc.CallOption) (*querypb.GetShardLeadersResponse, error) {
-		return &querypb.GetShardLeadersResponse{
-			Status: &successStatus,
-			Shards: []*querypb.ShardLeadersList{
-				{
-					ChannelName: "channel-1",
-					NodeIds:     []int64{1, 2, 3},
-					NodeAddrs:   []string{"localhost:9000", "localhost:9001", "localhost:9002"},
-					Serviceable: []bool{true, true, true},
-				},
+	mixc.EXPECT().GetShardLeaders(mock.Anything, mock.Anything, mock.Anything).Return(&querypb.GetShardLeadersResponse{
+		Status: &successStatus,
+		Shards: []*querypb.ShardLeadersList{
+			{
+				ChannelName: "channel-1",
+				NodeIds:     []int64{1, 2, 3},
+				NodeAddrs:   []string{"localhost:9000", "localhost:9001", "localhost:9002"},
+				Serviceable: []bool{true, true, true},
 			},
-		}, nil
-	}
+		},
+	}, nil).Maybe()
 
-	mixc.ShowLoadPartitionsFunc = func(ctx context.Context, req *querypb.ShowPartitionsRequest, opts ...grpc.CallOption) (*querypb.ShowPartitionsResponse, error) {
-		return &querypb.ShowPartitionsResponse{
-			Status:       &successStatus,
-			PartitionIDs: []int64{1, 2, 3},
-		}, nil
-	}
+	mixc.EXPECT().ShowLoadPartitions(mock.Anything, mock.Anything, mock.Anything).Return(&querypb.ShowPartitionsResponse{
+		Status:       &successStatus,
+		PartitionIDs: []int64{1, 2, 3},
+	}, nil).Maybe()
 
 	s.mixc = mixc
+	mixc.EXPECT().Close().Return(nil).Maybe()
 	s.qn = mocks.NewMockQueryNodeClient(s.T())
 
 	s.qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
@@ -93,10 +88,16 @@ func (s *StatisticTaskSuite) SetupTest() {
 	}, nil).Maybe()
 	mgr.EXPECT().InvalidateShardLeaderCache(mock.Anything).Return().Maybe()
 	s.mgr = mgr
-	s.lb = shardclient.NewLBPolicyImpl(mgr)
+	// Stub shard dispatch: Execute runs the workload directly against the mock
+	// query node, bypassing shard-leader resolution against the coordinator.
+	lb := shardclient.NewMockLBPolicy(s.T())
+	lb.EXPECT().Execute(mock.Anything, mock.Anything).RunAndReturn(
+		func(_ context.Context, workload shardclient.CollectionWorkLoad) error {
+			return workload.Exec(context.Background(), 1, s.qn, "mock_qn")
+		}).Maybe()
+	s.lb = lb
 
-	cache, err := initMetaCache(context.Background(), s.mixc)
-	s.NoError(err)
+	cache := newTestCache()
 	s.metaCache = cache
 
 	s.collectionName = "test_statistics_task"
@@ -117,38 +118,13 @@ func (s *StatisticTaskSuite) loadCollection() {
 	}
 
 	schema := constructCollectionSchemaByDataType(s.collectionName, fieldName2Types, testInt64Field, false)
-	marshaledSchema, err := proto.Marshal(schema)
-	s.NoError(err)
+	schemaInfo := mustNewSchemaInfo(schema)
 
-	ctx := context.Background()
-	createColT := &createCollectionTask{
-		Condition: NewTaskCondition(ctx),
-		CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
-			CollectionName: s.collectionName,
-			Schema:         marshaledSchema,
-			ShardsNum:      common.DefaultShardsNum,
-		},
-		ctx:      ctx,
-		mixCoord: s.mixc,
-	}
-
-	s.NoError(createColT.OnEnqueue())
-	s.NoError(createColT.PreExecute(ctx))
-	s.NoError(createColT.Execute(ctx))
-	s.NoError(createColT.PostExecute(ctx))
-
-	collectionID, err := s.metaCache.GetCollectionID(ctx, GetCurDBNameFromContextOrDefault(ctx), s.collectionName)
-	s.NoError(err)
-
-	status, err := s.mixc.LoadCollection(ctx, &querypb.LoadCollectionRequest{
-		Base: &commonpb.MsgBase{
-			MsgType:  commonpb.MsgType_LoadCollection,
-			SourceID: paramtable.GetNodeID(),
-		},
-		CollectionID: collectionID,
-	})
-	s.NoError(err)
-	s.Equal(commonpb.ErrorCode_Success, status.ErrorCode)
+	collectionID := int64(1000)
+	mockTest(s.T(), (*metacache.MetaCache).GetCollectionID, collectionID, nil)
+	mockTest(s.T(), (*metacache.MetaCache).GetCollectionSchema, schemaInfo, nil)
+	mockTest(s.T(), (*metacache.MetaCache).GetCollectionInfo, &collectionInfo{CollID: collectionID}, nil)
+	mockTest(s.T(), (*metacache.MetaCache).GetPartitions, map[string]int64{"_default": 1}, nil)
 	s.collectionID = collectionID
 }
 
@@ -158,7 +134,7 @@ func (s *StatisticTaskSuite) TearDownSuite() {
 
 func (s *StatisticTaskSuite) TestStatisticTask_Timeout() {
 	ctx := context.Background()
-	task := s.getStatisticsTask(ctx)
+	task := s.GetStatisticsTask(ctx)
 
 	s.NoError(task.OnEnqueue())
 
@@ -173,8 +149,8 @@ func (s *StatisticTaskSuite) TestStatisticTask_Timeout() {
 	s.Greater(task.TimeoutTimestamp, typeutil.ZeroTimestamp)
 }
 
-func (s *StatisticTaskSuite) getStatisticsTask(ctx context.Context) *getStatisticsTask {
-	return &getStatisticsTask{
+func (s *StatisticTaskSuite) GetStatisticsTask(ctx context.Context) *GetStatisticsTask {
+	return &GetStatisticsTask{
 		baseTask:       baseTask{MetaCache: s.metaCache},
 		Condition:      NewTaskCondition(ctx),
 		ctx:            ctx,
@@ -197,7 +173,7 @@ func (s *StatisticTaskSuite) getStatisticsTask(ctx context.Context) *getStatisti
 
 func (s *StatisticTaskSuite) TestStatisticTask_NotShardLeader() {
 	ctx := context.Background()
-	task := s.getStatisticsTask(ctx)
+	task := s.GetStatisticsTask(ctx)
 
 	s.NoError(task.OnEnqueue())
 
@@ -215,7 +191,7 @@ func (s *StatisticTaskSuite) TestStatisticTask_NotShardLeader() {
 
 func (s *StatisticTaskSuite) TestStatisticTask_UnexpectedError() {
 	ctx := context.Background()
-	task := s.getStatisticsTask(ctx)
+	task := s.GetStatisticsTask(ctx)
 	s.NoError(task.OnEnqueue())
 
 	task.fromQueryNode = true
@@ -232,7 +208,7 @@ func (s *StatisticTaskSuite) TestStatisticTask_UnexpectedError() {
 
 func (s *StatisticTaskSuite) TestStatisticTask_Success() {
 	ctx := context.Background()
-	task := s.getStatisticsTask(ctx)
+	task := s.GetStatisticsTask(ctx)
 
 	s.NoError(task.OnEnqueue())
 	s.qn.EXPECT().GetStatistics(mock.Anything, mock.Anything).Return(nil, nil)

internal/proxy/dql/task_validator.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
@@ -14,7 +14,7 @@ type taskValidator[T any] interface {
 // searchTaskValidator validates search tasks
 type searchTaskValidator struct{}
 
-var searchTaskValidatorInstance taskValidator[*searchTask] = &searchTaskValidator{}
+var searchTaskValidatorInstance taskValidator[*SearchTask] = &searchTaskValidator{}
 
 func (v *searchTaskValidator) validateSubSearch(subReq *internalpb.SubSearchRequest) error {
 	maxResultEntries := paramtable.Get().ProxyCfg.MaxResultEntries.GetAsInt64()
@@ -33,7 +33,7 @@ func (v *searchTaskValidator) validateSubSearch(subReq *internalpb.SubSearchRequ
 	return nil
 }
 
-func (v *searchTaskValidator) validateSearch(search *searchTask) error {
+func (v *searchTaskValidator) validateSearch(search *SearchTask) error {
 	maxResultEntries := paramtable.Get().ProxyCfg.MaxResultEntries.GetAsInt64()
 	if maxResultEntries <= 0 {
 		return nil
@@ -50,7 +50,7 @@ func (v *searchTaskValidator) validateSearch(search *searchTask) error {
 	return nil
 }
 
-func (v *searchTaskValidator) validate(search *searchTask) error {
+func (v *searchTaskValidator) validate(search *SearchTask) error {
 	// if it is a hybrid search, check all sub-searches
 	if search.GetIsAdvanced() {
 		for _, subReq := range search.GetSubReqs() {
@@ -68,7 +68,7 @@ func (v *searchTaskValidator) validate(search *searchTask) error {
 
 func ValidateTask(task any) error {
 	switch t := task.(type) {
-	case *searchTask:
+	case *SearchTask:
 		return searchTaskValidatorInstance.validate(t)
 	default:
 		return nil

internal/proxy/dql/task_validator_test.go

@@ -1,4 +1,4 @@
-package proxy
+package dql
 
 import (
 	"testing"
@@ -23,7 +23,7 @@ func TestValidateTask(t *testing.T) {
 	}{
 		{
 			name: "valid search task",
-			task: &searchTask{
+			task: &SearchTask{
 				SearchRequest: &internalpb.SearchRequest{
 					Nq:         10,
 					Topk:       100,
@@ -34,7 +34,7 @@ func TestValidateTask(t *testing.T) {
 		},
 		{
 			name: "invalid search task",
-			task: &searchTask{
+			task: &SearchTask{
 				SearchRequest: &internalpb.SearchRequest{
 					Nq:         1000,
 					Topk:       2000,
@@ -46,7 +46,7 @@ func TestValidateTask(t *testing.T) {
 		},
 		{
 			name: "invalid search task with group size",
-			task: &searchTask{
+			task: &SearchTask{
 				SearchRequest: &internalpb.SearchRequest{
 					Nq:         100,
 					Topk:       200,
@@ -59,7 +59,7 @@ func TestValidateTask(t *testing.T) {
 		},
 		{
 			name: "invalid search task with sub-request",
-			task: &searchTask{
+			task: &SearchTask{
 				SearchRequest: &internalpb.SearchRequest{
 					IsAdvanced: true,
 					SubReqs: []*internalpb.SubSearchRequest{

internal/proxy/dql/testutil_test.go

@@ -0,0 +1,349 @@
+// 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 dql
+
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"math/rand"
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/bytedance/mockey"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
+	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
+	"github.com/milvus-io/milvus/internal/proxy/metacache"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
+	"github.com/milvus-io/milvus/pkg/v3/common"
+	"github.com/milvus-io/milvus/pkg/v3/util/merr"
+	"github.com/milvus-io/milvus/pkg/v3/util/testutils"
+	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
+)
+
+const testVecDim = 128
+
+const (
+	testInt64Field     = "int64"
+	testVarCharField   = "varChar"
+	testFloatVecField  = "fvec"
+	int64Field         = "int64"
+	floatVecField      = "fVec"
+	testBoolField      = "bool"
+	testInt32Field     = "int32"
+	testFloatField     = "float"
+	testDoubleField    = "double"
+	testBinaryVecField = "bvec"
+)
+
+// newTextSchemaForStorageV3Test builds a schema with a TEXT field for
+// StorageV3-gating tests.
+func newTextSchemaForStorageV3Test(collectionName string) *schemapb.CollectionSchema {
+	return &schemapb.CollectionSchema{
+		Name: collectionName,
+		Fields: []*schemapb.FieldSchema{
+			{FieldID: 100, Name: testInt64Field, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+			{FieldID: 101, Name: "text", DataType: schemapb.DataType_Text},
+			{
+				FieldID:  102,
+				Name:     testFloatVecField,
+				DataType: schemapb.DataType_FloatVector,
+				TypeParams: []*commonpb.KeyValuePair{
+					{Key: common.DimKey, Value: strconv.Itoa(testVecDim)},
+				},
+			},
+		},
+	}
+}
+
+// newScalarFieldData builds scalar field data for tests.
+func newScalarFieldData(fieldSchema *schemapb.FieldSchema, fieldName string, numRows int) *schemapb.FieldData {
+	return testutils.GenerateScalarFieldData(fieldSchema.GetDataType(), fieldName, numRows)
+}
+
+// newFloatVectorFieldData builds float-vector field data for tests.
+func newFloatVectorFieldData(fieldName string, numRows, dim int) *schemapb.FieldData {
+	return testutils.NewFloatVectorFieldData(fieldName, numRows, dim)
+}
+
+// generateFieldData builds field data for a data type.
+func generateFieldData(dataType schemapb.DataType, fieldName string, numRows int) *schemapb.FieldData {
+	if dataType < 100 {
+		return testutils.GenerateScalarFieldData(dataType, fieldName, numRows)
+	}
+	return testutils.GenerateVectorFieldData(dataType, fieldName, numRows, testVecDim)
+}
+
+// newTestCache returns an empty metacache whose methods are expected to be
+// patched by mockey in each test.
+func newTestCache() *metacache.MetaCache {
+	cache, err := metacache.NewMetaCache(nil)
+	if err != nil {
+		panic(err)
+	}
+	return cache
+}
+
+// mockTsoAllocator is a simple taskmodel.TsoAllocator for tests.
+type mockTsoAllocator struct {
+	mu        sync.Mutex
+	logicPart uint32
+}
+
+func (tso *mockTsoAllocator) AllocOne(ctx context.Context) (Timestamp, error) {
+	tso.mu.Lock()
+	defer tso.mu.Unlock()
+	tso.logicPart++
+	physical := uint64(time.Now().UnixMilli())
+	return (physical << 18) + uint64(tso.logicPart), nil
+}
+
+var _ taskmodel.TsoAllocator = (*mockTsoAllocator)(nil)
+
+const testMaxVarCharLength = 512
+
+// constructCollectionSchemaByDataType builds a collection schema for tests.
+func constructCollectionSchemaByDataType(collectionName string, fieldName2DataType map[string]schemapb.DataType, primaryFieldName string, autoID bool) *schemapb.CollectionSchema {
+	fieldsSchema := make([]*schemapb.FieldSchema, 0)
+
+	idx := int64(100)
+	for fieldName, dataType := range fieldName2DataType {
+		fieldSchema := &schemapb.FieldSchema{
+			FieldID:  idx,
+			Name:     fieldName,
+			DataType: dataType,
+		}
+		idx++
+		if typeutil.IsVectorType(dataType) {
+			fieldSchema.TypeParams = []*commonpb.KeyValuePair{
+				{
+					Key:   common.DimKey,
+					Value: strconv.Itoa(testVecDim),
+				},
+			}
+		}
+		if dataType == schemapb.DataType_VarChar {
+			fieldSchema.TypeParams = []*commonpb.KeyValuePair{
+				{
+					Key:   common.MaxLengthKey,
+					Value: strconv.Itoa(testMaxVarCharLength),
+				},
+			}
+		}
+		if fieldName == primaryFieldName {
+			fieldSchema.IsPrimaryKey = true
+			fieldSchema.AutoID = autoID
+		}
+
+		fieldsSchema = append(fieldsSchema, fieldSchema)
+	}
+
+	return &schemapb.CollectionSchema{
+		Name:   collectionName,
+		Fields: fieldsSchema,
+	}
+}
+
+// createTestSchema builds a test schemaInfo for query/search tests.
+func createTestSchema() *schemaInfo {
+	schema := &schemapb.CollectionSchema{
+		Name: "test_collection",
+		Fields: []*schemapb.FieldSchema{
+			{
+				FieldID:      100,
+				Name:         "id",
+				IsPrimaryKey: true,
+				DataType:     schemapb.DataType_Int64,
+			},
+			{
+				FieldID:  101,
+				Name:     "vector",
+				DataType: schemapb.DataType_FloatVector,
+				TypeParams: []*commonpb.KeyValuePair{
+					{Key: "dim", Value: "128"},
+				},
+			},
+			{
+				FieldID:  102,
+				Name:     "name",
+				DataType: schemapb.DataType_VarChar,
+			},
+		},
+	}
+	return mustNewSchemaInfo(schema)
+}
+
+// newTestSchema builds a schema covering every scalar/vector data type.
+func newTestSchema() *schemapb.CollectionSchema {
+	fields := []*schemapb.FieldSchema{
+		{FieldID: 0, Name: "FieldID", IsPrimaryKey: false, Description: "field no.1", DataType: schemapb.DataType_Int64},
+	}
+
+	for name, value := range schemapb.DataType_value {
+		dataType := schemapb.DataType(value)
+		if !typeutil.IsIntegerType(dataType) && !typeutil.IsFloatingType(dataType) && !typeutil.IsVectorType(dataType) && !typeutil.IsStringType(dataType) {
+			continue
+		}
+		newField := &schemapb.FieldSchema{
+			FieldID: int64(100 + value), Name: name + "Field", IsPrimaryKey: false, Description: "", DataType: dataType,
+		}
+		fields = append(fields, newField)
+	}
+
+	return &schemapb.CollectionSchema{
+		Name:               "test",
+		Description:        "schema for test used",
+		AutoID:             true,
+		Fields:             fields,
+		EnableDynamicField: true,
+	}
+}
+
+// constructCollectionSchema builds a minimal int64-pk + float-vector schema.
+func constructCollectionSchema(int64Field, floatVecField string, dim int, collectionName string) *schemapb.CollectionSchema {
+	pk := &schemapb.FieldSchema{
+		FieldID:      100,
+		Name:         int64Field,
+		IsPrimaryKey: true,
+		Description:  "",
+		DataType:     schemapb.DataType_Int64,
+		TypeParams:   nil,
+		IndexParams:  nil,
+		AutoID:       true,
+	}
+	fVec := &schemapb.FieldSchema{
+		FieldID:      101,
+		Name:         floatVecField,
+		IsPrimaryKey: false,
+		Description:  "",
+		DataType:     schemapb.DataType_FloatVector,
+		TypeParams: []*commonpb.KeyValuePair{
+			{
+				Key:   common.DimKey,
+				Value: strconv.Itoa(dim),
+			},
+		},
+		IndexParams: nil,
+		AutoID:      false,
+	}
+	return &schemapb.CollectionSchema{
+		Name:        collectionName,
+		Description: "",
+		AutoID:      false,
+		Fields: []*schemapb.FieldSchema{
+			pk,
+			fVec,
+		},
+		Properties: []*commonpb.KeyValuePair{
+			{
+				Key:   common.CollectionTTLConfigKey,
+				Value: "15",
+			},
+		},
+	}
+}
+
+// constructPlaceholderGroup builds a float-vector placeholder group for tests.
+func constructPlaceholderGroup(nq, dim int) *commonpb.PlaceholderGroup {
+	values := make([][]byte, 0, nq)
+	for i := 0; i < nq; i++ {
+		bs := make([]byte, 0, dim*4)
+		for j := 0; j < dim; j++ {
+			var buffer bytes.Buffer
+			f := rand.Float32()
+			err := binary.Write(&buffer, common.Endian, f)
+			if err != nil {
+				panic(err)
+			}
+			bs = append(bs, buffer.Bytes()...)
+		}
+		values = append(values, bs)
+	}
+
+	return &commonpb.PlaceholderGroup{
+		Placeholders: []*commonpb.PlaceholderValue{
+			{
+				Tag:    "$0",
+				Type:   commonpb.PlaceholderType_FloatVector,
+				Values: values,
+			},
+		},
+	}
+}
+
+// ConstructCollectionSchemaWithPartitionKey builds a schema with a partition key.
+func ConstructCollectionSchemaWithPartitionKey(collectionName string, fieldName2DataType map[string]schemapb.DataType, primaryFieldName string, partitionKeyFieldName string, autoID bool) *schemapb.CollectionSchema {
+	schema := constructCollectionSchemaByDataType(collectionName, fieldName2DataType, primaryFieldName, autoID)
+	for _, field := range schema.Fields {
+		if field.Name == partitionKeyFieldName {
+			field.IsPartitionKey = true
+		}
+	}
+
+	return schema
+}
+
+// fillFieldPropertiesOnly fills FieldId/Type from the schema for each column.
+func fillFieldPropertiesOnly(columns []*schemapb.FieldData, schema *schemaInfo) error {
+	for _, fieldData := range columns {
+		fieldSchema, err := schema.SchemaHelper.GetFieldFromNameDefaultJSON(fieldData.FieldName)
+		if err != nil {
+			return merr.WrapErrParameterInvalidMsg("fieldName %v not exist in collection schema", fieldData.FieldName)
+		}
+
+		fieldData.FieldId = fieldSchema.FieldID
+		fieldData.Type = fieldSchema.DataType
+
+		switch fieldData.Type {
+		case schemapb.DataType_Array:
+			fd, ok := fieldData.Field.(*schemapb.FieldData_Scalars)
+			if !ok || fd.Scalars.GetArrayData() == nil {
+				return merr.WrapErrParameterInvalidMsg("field convert FieldData_Scalars fail in fieldData, fieldName: %s, collectionName: %s",
+					fieldData.FieldName, schema.Name)
+			}
+			fd.Scalars.GetArrayData().ElementType = fieldSchema.ElementType
+		case schemapb.DataType_ArrayOfVector:
+			fd, ok := fieldData.Field.(*schemapb.FieldData_Vectors)
+			if !ok || fd.Vectors.GetVectorArray() == nil {
+				return merr.WrapErrParameterInvalidMsg("field convert FieldData_Vectors fail in fieldData, fieldName: %s, collectionName: %s",
+					fieldData.FieldName, schema.Name)
+			}
+			fd.Vectors.GetVectorArray().ElementType = fieldSchema.ElementType
+		}
+	}
+
+	return nil
+}
+
+// mockTest registers a mockey patch that is automatically unpatched when the
+// test finishes, so global mockey patches never leak between tests.
+func mockTest(t *testing.T, target any, rets ...any) *mockey.Mocker {
+	m := mockey.Mock(target).Return(rets...).Build()
+	t.Cleanup(func() { m.UnPatch() })
+	return m
+}
+
+// mockTestTo registers a mockey patch with a custom implementation that is
+// automatically unpatched when the test finishes.
+func mockTestTo(t *testing.T, target any, fn any) *mockey.Mocker {
+	m := mockey.Mock(target).To(fn).Build()
+	t.Cleanup(func() { m.UnPatch() })
+	return m
+}

internal/proxy/dql/translate_output_fields_test.go

@@ -0,0 +1,990 @@
+// 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 dql
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
+	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
+	"github.com/milvus-io/milvus/pkg/v3/common"
+	"github.com/milvus-io/milvus/pkg/v3/util/merr"
+	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
+)
+
+func TestTranslateOutputFields(t *testing.T) {
+	const (
+		idFieldName                = "id"
+		tsFieldName                = "timestamp"
+		floatVectorFieldName       = "float_vector"
+		binaryVectorFieldName      = "binary_vector"
+		float16VectorFieldName     = "float16_vector"
+		bfloat16VectorFieldName    = "bfloat16_vector"
+		sparseFloatVectorFieldName = "sparse_float_vector"
+	)
+	var outputFields []string
+	var userOutputFields []string
+	var userDynamicFields []string
+	var requestedPK bool
+	var err error
+
+	collSchema := &schemapb.CollectionSchema{
+		Name:        "TestTranslateOutputFields",
+		Description: "TestTranslateOutputFields",
+		AutoID:      false,
+		Fields: []*schemapb.FieldSchema{
+			{Name: idFieldName, FieldID: 0, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+			{Name: tsFieldName, FieldID: 1, DataType: schemapb.DataType_Int64},
+			{Name: floatVectorFieldName, FieldID: 100, DataType: schemapb.DataType_FloatVector},
+			{Name: binaryVectorFieldName, FieldID: 101, DataType: schemapb.DataType_BinaryVector},
+			{Name: float16VectorFieldName, FieldID: 102, DataType: schemapb.DataType_Float16Vector},
+			{Name: bfloat16VectorFieldName, FieldID: 103, DataType: schemapb.DataType_BFloat16Vector},
+			{Name: sparseFloatVectorFieldName, FieldID: 104, DataType: schemapb.DataType_SparseFloatVector, IsFunctionOutput: true},
+		},
+		Functions: []*schemapb.FunctionSchema{
+			{
+				Name:             "bm25",
+				Type:             schemapb.FunctionType_BM25,
+				OutputFieldNames: []string{sparseFloatVectorFieldName},
+				// omit other fields for brevity
+			},
+		},
+	}
+	schema := mustNewSchemaInfo(collSchema)
+
+	// Test empty output fields
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{}, schema, false)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{}, outputFields)
+	assert.ElementsMatch(t, []string{}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.False(t, requestedPK)
+
+	// Test single field
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName}, schema, false)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{idFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{idFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	// Test multiple fields
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName}, schema, false)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{idFieldName, tsFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{idFieldName, tsFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	// Test with vector field
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName, floatVectorFieldName}, schema, false)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	// Test without id field
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{tsFieldName, floatVectorFieldName}, schema, false)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.False(t, requestedPK)
+
+	// Test wildcard - should not include function output fields
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*"}, schema, false)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	// Test wildcard with spaces
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{" * "}, schema, false)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	_, _, _, _, _, err = translateOutputFields([]string{"*", sparseFloatVectorFieldName}, schema, false)
+	assert.Error(t, err)
+	_, _, _, _, _, err = translateOutputFields([]string{sparseFloatVectorFieldName}, schema, false)
+	assert.Error(t, err)
+	_, _, _, _, _, err = translateOutputFields([]string{sparseFloatVectorFieldName}, schema, true)
+	assert.Error(t, err)
+
+	// Test with removePkField=true
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{}, schema, true)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{}, outputFields)
+	assert.ElementsMatch(t, []string{}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.False(t, requestedPK)
+
+	// if removePkField is true, pk field should be removed from output fields
+
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*"}, schema, true)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName}, schema, true)
+	assert.Equal(t, nil, err)
+	assert.ElementsMatch(t, []string{tsFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{tsFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName, floatVectorFieldName}, schema, true)
+	assert.Equal(t, nil, err)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*"}, schema, true)
+	assert.Equal(t, nil, err)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*", tsFieldName}, schema, true)
+	assert.Equal(t, nil, err)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*", floatVectorFieldName}, schema, true)
+	assert.Equal(t, nil, err)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
+	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.True(t, requestedPK)
+
+	// Test non-existent field, dynamic field not enabled
+	_, _, _, _, _, err = translateOutputFields([]string{"A"}, schema, true)
+	assert.Error(t, err)
+
+	t.Run("enable dynamic schema", func(t *testing.T) {
+		collSchema := &schemapb.CollectionSchema{
+			Name:               "TestTranslateOutputFields",
+			Description:        "TestTranslateOutputFields",
+			AutoID:             false,
+			EnableDynamicField: true,
+			Fields: []*schemapb.FieldSchema{
+				{Name: idFieldName, FieldID: 1, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+				{Name: tsFieldName, FieldID: 2, DataType: schemapb.DataType_Int64},
+				{Name: floatVectorFieldName, FieldID: 100, DataType: schemapb.DataType_FloatVector},
+				{Name: binaryVectorFieldName, FieldID: 101, DataType: schemapb.DataType_BinaryVector},
+				{Name: common.MetaFieldName, FieldID: 102, DataType: schemapb.DataType_JSON, IsDynamic: true},
+			},
+		}
+		schema := mustNewSchemaInfo(collSchema)
+
+		outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"A", idFieldName}, schema, true)
+		assert.NoError(t, err)
+		assert.ElementsMatch(t, []string{common.MetaFieldName}, outputFields)
+		assert.ElementsMatch(t, []string{"A"}, userOutputFields)
+		assert.ElementsMatch(t, []string{"A"}, userDynamicFields)
+		assert.True(t, requestedPK)
+
+		outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"$meta[\"A\"]", idFieldName}, schema, true)
+		assert.NoError(t, err)
+		assert.ElementsMatch(t, []string{common.MetaFieldName}, outputFields)
+		assert.ElementsMatch(t, []string{"$meta[\"A\"]"}, userOutputFields)
+		assert.ElementsMatch(t, []string{"A"}, userDynamicFields)
+		assert.True(t, requestedPK)
+
+		// Test invalid dynamic field expressions
+		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, `$meta["A"]["B"]`}, schema, true)
+		assert.Error(t, err)
+
+		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "$meta[\"\"]"}, schema, true)
+		assert.Error(t, err)
+
+		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "$meta[]"}, schema, true)
+		assert.Error(t, err)
+
+		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "$meta["}, schema, true)
+		assert.Error(t, err)
+
+		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "[]"}, schema, true)
+		assert.Error(t, err)
+
+		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "A > 1"}, schema, true)
+		assert.Error(t, err)
+
+		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, ""}, schema, true)
+		assert.Error(t, err)
+	})
+}
+
+func TestTranslateOutputFields_StructArrayField(t *testing.T) {
+	const (
+		idFieldName                = "id"
+		tsFieldName                = "timestamp"
+		floatVectorFieldName       = "float_vector"
+		binaryVectorFieldName      = "binary_vector"
+		float16VectorFieldName     = "float16_vector"
+		bfloat16VectorFieldName    = "bfloat16_vector"
+		sparseFloatVectorFieldName = "sparse_float_vector"
+	)
+	var outputFields []string
+	var userOutputFields []string
+	var userDynamicFields []string
+	var requestedPK bool
+	var err error
+
+	collSchema := &schemapb.CollectionSchema{
+		Name:        "TestTranslateOutputFields",
+		Description: "TestTranslateOutputFields",
+		AutoID:      false,
+		Fields: []*schemapb.FieldSchema{
+			{Name: idFieldName, FieldID: 0, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+			{Name: tsFieldName, FieldID: 1, DataType: schemapb.DataType_Int64},
+			{Name: floatVectorFieldName, FieldID: 100, DataType: schemapb.DataType_FloatVector},
+		},
+		StructArrayFields: []*schemapb.StructArrayFieldSchema{
+			{
+				FieldID: 101,
+				Name:    "struct_array_field",
+				Fields: []*schemapb.FieldSchema{
+					{Name: "sub_field", FieldID: 102, DataType: schemapb.DataType_Array, ElementType: schemapb.DataType_Int64},
+					{Name: "sub_vector_field", FieldID: 103, DataType: schemapb.DataType_ArrayOfVector, ElementType: schemapb.DataType_FloatVector},
+				},
+			},
+		},
+	}
+	schema := mustNewSchemaInfo(collSchema)
+
+	// Test struct array field
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"sub_vector_field"}, schema, false)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{"sub_vector_field"}, outputFields)
+	assert.ElementsMatch(t, []string{"sub_vector_field"}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.False(t, requestedPK)
+
+	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"struct_array_field", "sub_field"}, schema, false)
+	assert.NoError(t, err)
+	assert.ElementsMatch(t, []string{"sub_vector_field", "sub_field"}, outputFields)
+	assert.ElementsMatch(t, []string{"sub_vector_field", "sub_field"}, userOutputFields)
+	assert.ElementsMatch(t, []string{}, userDynamicFields)
+	assert.False(t, requestedPK)
+}
+
+func Test_MaxQueryResultWindow(t *testing.T) {
+	paramtable.Init()
+	assert.Nil(t, validateMaxQueryResultWindow(0, 16384, false))
+	assert.Nil(t, validateMaxQueryResultWindow(0, 1, false))
+	assert.Error(t, validateMaxQueryResultWindow(0, 16385, false))
+	assert.Error(t, validateMaxQueryResultWindow(0, 0, false))
+	assert.Error(t, validateMaxQueryResultWindow(1, 0, false))
+
+	paramtable.Get().Save(paramtable.Get().QuotaConfig.LargeMaxQueryResultWindow.Key, "1000000")
+	defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.LargeMaxQueryResultWindow.Key)
+	assert.Nil(t, validateMaxQueryResultWindow(0, 16385, true))
+	assert.Nil(t, validateMaxQueryResultWindow(0, 1000000, true))
+	assert.Error(t, validateMaxQueryResultWindow(0, 1000001, true))
+}
+
+func Test_reconstructStructFieldData(t *testing.T) {
+	t.Run("count(*) query - should return early", func(t *testing.T) {
+		fieldsData := []*schemapb.FieldData{
+			{
+				FieldName: "count(*)",
+				FieldId:   0,
+				Type:      schemapb.DataType_Int64,
+			},
+		}
+		outputFields := []string{"count(*)"}
+
+		schema := &schemapb.CollectionSchema{
+			StructArrayFields: []*schemapb.StructArrayFieldSchema{
+				{
+					FieldID: 102,
+					Name:    "test_struct",
+					Fields: []*schemapb.FieldSchema{
+						{
+							FieldID:     1021,
+							Name:        "test_struct[sub_field]",
+							DataType:    schemapb.DataType_Array,
+							ElementType: schemapb.DataType_Int32,
+						},
+					},
+				},
+			},
+		}
+
+		originalFieldsData := make([]*schemapb.FieldData, len(fieldsData))
+		copy(originalFieldsData, fieldsData)
+		originalOutputFields := make([]string, len(outputFields))
+		copy(originalOutputFields, outputFields)
+
+		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
+
+		// Should not modify anything for count(*) query
+		assert.Equal(t, originalFieldsData, resultFieldsData)
+		assert.Equal(t, originalOutputFields, resultOutputFields)
+	})
+
+	t.Run("group by with count(*) should preserve aggregate field", func(t *testing.T) {
+		fieldsData := []*schemapb.FieldData{
+			{
+				FieldName: "id",
+				FieldId:   100,
+				Type:      schemapb.DataType_Int64,
+				Field: &schemapb.FieldData_Scalars{
+					Scalars: &schemapb.ScalarField{
+						Data: &schemapb.ScalarField_LongData{
+							LongData: &schemapb.LongArray{Data: []int64{1, 2}},
+						},
+					},
+				},
+			},
+			{
+				FieldName: "count(*)",
+				FieldId:   0,
+				Type:      schemapb.DataType_Int64,
+				Field: &schemapb.FieldData_Scalars{
+					Scalars: &schemapb.ScalarField{
+						Data: &schemapb.ScalarField_LongData{
+							LongData: &schemapb.LongArray{Data: []int64{3, 4}},
+						},
+					},
+				},
+			},
+		}
+		outputFields := []string{"id", "count(*)"}
+
+		schema := &schemapb.CollectionSchema{
+			Fields: []*schemapb.FieldSchema{
+				{
+					FieldID:      100,
+					Name:         "id",
+					IsPrimaryKey: true,
+					DataType:     schemapb.DataType_Int64,
+				},
+			},
+			StructArrayFields: []*schemapb.StructArrayFieldSchema{
+				{
+					FieldID: 102,
+					Name:    "test_struct",
+					Fields: []*schemapb.FieldSchema{
+						{
+							FieldID:     1021,
+							Name:        "test_struct[sub_field]",
+							DataType:    schemapb.DataType_Array,
+							ElementType: schemapb.DataType_Int32,
+						},
+					},
+				},
+			},
+		}
+
+		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
+
+		assert.Len(t, resultFieldsData, 2)
+		assert.Equal(t, "id", resultFieldsData[0].FieldName)
+		assert.Equal(t, int64(100), resultFieldsData[0].FieldId)
+		assert.Equal(t, "count(*)", resultFieldsData[1].FieldName)
+		assert.Equal(t, int64(0), resultFieldsData[1].FieldId)
+		assert.Equal(t, []string{"id", "count(*)"}, resultOutputFields)
+		assert.Equal(t, []int64{3, 4}, resultFieldsData[1].GetScalars().GetLongData().GetData())
+	})
+
+	t.Run("struct field query - should reconstruct struct field", func(t *testing.T) {
+		fieldsData := []*schemapb.FieldData{
+			{
+				FieldName: "test_struct[sub_field]",
+				FieldId:   1021, // Use the correct field ID that matches the schema
+				Type:      schemapb.DataType_Array,
+				Field: &schemapb.FieldData_Scalars{
+					Scalars: &schemapb.ScalarField{
+						Data: &schemapb.ScalarField_ArrayData{
+							ArrayData: &schemapb.ArrayArray{
+								ElementType: schemapb.DataType_Int32,
+								Data: []*schemapb.ScalarField{
+									{
+										Data: &schemapb.ScalarField_IntData{
+											IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		}
+		outputFields := []string{"test_struct[sub_field]"}
+
+		schema := &schemapb.CollectionSchema{
+			StructArrayFields: []*schemapb.StructArrayFieldSchema{
+				{
+					FieldID: 102,
+					Name:    "test_struct",
+					Fields: []*schemapb.FieldSchema{
+						{
+							FieldID:     1021,
+							Name:        "test_struct[sub_field]",
+							DataType:    schemapb.DataType_Array,
+							ElementType: schemapb.DataType_Int32,
+						},
+					},
+				},
+			},
+		}
+
+		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
+
+		// Should reconstruct the struct field with the restored field name
+		assert.Len(t, resultFieldsData, 1)
+		assert.Equal(t, "test_struct", resultFieldsData[0].FieldName)
+		assert.Equal(t, int64(102), resultFieldsData[0].FieldId)
+		assert.Equal(t, schemapb.DataType_ArrayOfStruct, resultFieldsData[0].Type)
+
+		// Check that the sub-field name has been restored
+		structArrayField := resultFieldsData[0].GetStructArrays()
+		assert.NotNil(t, structArrayField)
+		assert.Len(t, structArrayField.Fields, 1)
+		assert.Equal(t, "sub_field", structArrayField.Fields[0].FieldName) // Name should be restored
+
+		assert.Equal(t, []string{"test_struct"}, resultOutputFields)
+	})
+
+	t.Run("no struct array fields - should return early", func(t *testing.T) {
+		fieldsData := []*schemapb.FieldData{
+			{
+				FieldName: "field1",
+				FieldId:   100,
+				Type:      schemapb.DataType_Int64,
+			},
+			{
+				FieldName: "field2",
+				FieldId:   101,
+				Type:      schemapb.DataType_VarChar,
+			},
+		}
+		outputFields := []string{"field1", "field2"}
+
+		schema := &schemapb.CollectionSchema{
+			StructArrayFields: []*schemapb.StructArrayFieldSchema{},
+		}
+
+		originalFieldsData := make([]*schemapb.FieldData, len(fieldsData))
+		copy(originalFieldsData, fieldsData)
+		originalOutputFields := make([]string, len(outputFields))
+		copy(originalOutputFields, outputFields)
+
+		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
+
+		// Should not modify anything when no struct array fields
+		assert.Equal(t, originalFieldsData, resultFieldsData)
+		assert.Equal(t, originalOutputFields, resultOutputFields)
+	})
+
+	t.Run("reconstruct single struct field", func(t *testing.T) {
+		// Create mock data with transformed field names (as they would be internally)
+		subField1Data := &schemapb.FieldData{
+			FieldName: "test_struct[sub_int_array]",
+			FieldId:   1021,
+			Type:      schemapb.DataType_Array,
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_ArrayData{
+						ArrayData: &schemapb.ArrayArray{
+							ElementType: schemapb.DataType_Int32,
+							Data: []*schemapb.ScalarField{
+								{
+									Data: &schemapb.ScalarField_IntData{
+										IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		}
+
+		subField2Data := &schemapb.FieldData{
+			FieldName: "test_struct[sub_text_array]",
+			FieldId:   1022,
+			Type:      schemapb.DataType_Array,
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_ArrayData{
+						ArrayData: &schemapb.ArrayArray{
+							ElementType: schemapb.DataType_VarChar,
+							Data: []*schemapb.ScalarField{
+								{
+									Data: &schemapb.ScalarField_StringData{
+										StringData: &schemapb.StringArray{Data: []string{"hello", "world"}},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		}
+
+		fieldsData := []*schemapb.FieldData{subField1Data, subField2Data}
+		outputFields := []string{"test_struct[sub_int_array]", "test_struct[sub_text_array]"}
+
+		schema := &schemapb.CollectionSchema{
+			Fields: []*schemapb.FieldSchema{
+				{
+					FieldID:      100,
+					Name:         "pk",
+					IsPrimaryKey: true,
+					DataType:     schemapb.DataType_Int64,
+				},
+			},
+			StructArrayFields: []*schemapb.StructArrayFieldSchema{
+				{
+					FieldID: 102,
+					Name:    "test_struct",
+					Fields: []*schemapb.FieldSchema{
+						{
+							FieldID:     1021,
+							Name:        "test_struct[sub_int_array]",
+							DataType:    schemapb.DataType_Array,
+							ElementType: schemapb.DataType_Int32,
+						},
+						{
+							FieldID:     1022,
+							Name:        "test_struct[sub_text_array]",
+							DataType:    schemapb.DataType_Array,
+							ElementType: schemapb.DataType_VarChar,
+						},
+					},
+				},
+			},
+		}
+
+		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
+
+		// Check result
+		assert.Len(t, resultFieldsData, 1, "Should only have one reconstructed struct field")
+		assert.Len(t, resultOutputFields, 1, "Output fields should only have one")
+
+		structField := resultFieldsData[0]
+		assert.Equal(t, "test_struct", structField.FieldName)
+		assert.Equal(t, int64(102), structField.FieldId)
+		assert.Equal(t, schemapb.DataType_ArrayOfStruct, structField.Type)
+		assert.Equal(t, "test_struct", resultOutputFields[0])
+
+		// Check fields inside struct
+		structArrays := structField.GetStructArrays()
+		assert.NotNil(t, structArrays)
+		assert.Len(t, structArrays.Fields, 2, "Struct should contain 2 sub fields")
+
+		// Check sub fields
+		var foundIntField, foundTextField bool
+		for _, field := range structArrays.Fields {
+			switch field.FieldId {
+			case 1021:
+				assert.Equal(t, "sub_int_array", field.FieldName)
+				assert.Equal(t, schemapb.DataType_Array, field.Type)
+				foundIntField = true
+			case 1022:
+				assert.Equal(t, "sub_text_array", field.FieldName)
+				assert.Equal(t, schemapb.DataType_Array, field.Type)
+				foundTextField = true
+			}
+		}
+		assert.True(t, foundIntField, "Should find int array field")
+		assert.True(t, foundTextField, "Should find text array field")
+	})
+
+	t.Run("mixed regular and struct fields", func(t *testing.T) {
+		// Create regular field data
+		regularField := &schemapb.FieldData{
+			FieldName: "regular_field",
+			FieldId:   100,
+			Type:      schemapb.DataType_Int64,
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_LongData{
+						LongData: &schemapb.LongArray{Data: []int64{1, 2, 3}},
+					},
+				},
+			},
+		}
+
+		// Create struct sub field data with transformed name
+		subFieldData := &schemapb.FieldData{
+			FieldName: "test_struct[sub_field]",
+			FieldId:   1021,
+			Type:      schemapb.DataType_Array,
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_ArrayData{
+						ArrayData: &schemapb.ArrayArray{
+							ElementType: schemapb.DataType_Int32,
+							Data: []*schemapb.ScalarField{
+								{
+									Data: &schemapb.ScalarField_IntData{
+										IntData: &schemapb.IntArray{Data: []int32{10, 20}},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		}
+
+		fieldsData := []*schemapb.FieldData{regularField, subFieldData}
+		outputFields := []string{"regular_field", "test_struct[sub_field]"}
+
+		schema := &schemapb.CollectionSchema{
+			Fields: []*schemapb.FieldSchema{
+				{
+					FieldID:  100,
+					Name:     "regular_field",
+					DataType: schemapb.DataType_Int64,
+				},
+			},
+			StructArrayFields: []*schemapb.StructArrayFieldSchema{
+				{
+					FieldID: 102,
+					Name:    "test_struct",
+					Fields: []*schemapb.FieldSchema{
+						{
+							FieldID:     1021,
+							Name:        "test_struct[sub_field]",
+							DataType:    schemapb.DataType_Array,
+							ElementType: schemapb.DataType_Int32,
+						},
+					},
+				},
+			},
+		}
+
+		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
+
+		// Check result: should have 2 fields (1 regular + 1 reconstructed struct)
+		assert.Len(t, resultFieldsData, 2)
+		assert.Len(t, resultOutputFields, 2)
+
+		// Check regular and struct fields both exist
+		var foundRegularField, foundStructField bool
+		for i, field := range resultFieldsData {
+			switch field.FieldId {
+			case 100:
+				assert.Equal(t, "regular_field", field.FieldName)
+				assert.Equal(t, schemapb.DataType_Int64, field.Type)
+				assert.Equal(t, "regular_field", resultOutputFields[i])
+				foundRegularField = true
+			case 102:
+				assert.Equal(t, "test_struct", field.FieldName)
+				assert.Equal(t, schemapb.DataType_ArrayOfStruct, field.Type)
+				assert.Equal(t, "test_struct", resultOutputFields[i])
+				foundStructField = true
+			}
+		}
+		assert.True(t, foundRegularField, "Should find regular field")
+		assert.True(t, foundStructField, "Should find reconstructed struct field")
+	})
+
+	t.Run("multiple struct fields", func(t *testing.T) {
+		// Create sub field for first struct
+		struct1SubField := &schemapb.FieldData{
+			FieldName: "struct1[struct1_sub]",
+			FieldId:   1021,
+			Type:      schemapb.DataType_Array,
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_ArrayData{
+						ArrayData: &schemapb.ArrayArray{
+							ElementType: schemapb.DataType_Int32,
+							Data:        []*schemapb.ScalarField{},
+						},
+					},
+				},
+			},
+		}
+
+		// Create sub fields for second struct
+		struct2SubField1 := &schemapb.FieldData{
+			FieldName: "struct2[struct2_sub1]",
+			FieldId:   1031,
+			Type:      schemapb.DataType_Array,
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_ArrayData{
+						ArrayData: &schemapb.ArrayArray{
+							ElementType: schemapb.DataType_Int32,
+							Data:        []*schemapb.ScalarField{},
+						},
+					},
+				},
+			},
+		}
+
+		struct2SubField2 := &schemapb.FieldData{
+			FieldName: "struct2[struct2_sub2]",
+			FieldId:   1032,
+			Type:      schemapb.DataType_VarChar,
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_StringData{
+						StringData: &schemapb.StringArray{Data: []string{"test"}},
+					},
+				},
+			},
+		}
+
+		fieldsData := []*schemapb.FieldData{struct1SubField, struct2SubField1, struct2SubField2}
+		outputFields := []string{"struct1[struct1_sub]", "struct2[struct2_sub1]", "struct2[struct2_sub2]"}
+
+		schema := &schemapb.CollectionSchema{
+			Fields: []*schemapb.FieldSchema{
+				{
+					FieldID:      100,
+					Name:         "pk",
+					IsPrimaryKey: true,
+					DataType:     schemapb.DataType_Int64,
+				},
+			},
+			StructArrayFields: []*schemapb.StructArrayFieldSchema{
+				{
+					FieldID: 102,
+					Name:    "struct1",
+					Fields: []*schemapb.FieldSchema{
+						{
+							FieldID:     1021,
+							Name:        "struct1[struct1_sub]",
+							DataType:    schemapb.DataType_Array,
+							ElementType: schemapb.DataType_Int32,
+						},
+					},
+				},
+				{
+					FieldID: 103,
+					Name:    "struct2",
+					Fields: []*schemapb.FieldSchema{
+						{
+							FieldID:     1031,
+							Name:        "struct2[struct2_sub1]",
+							DataType:    schemapb.DataType_Array,
+							ElementType: schemapb.DataType_Int32,
+						},
+						{
+							FieldID:  1032,
+							Name:     "struct2[struct2_sub2]",
+							DataType: schemapb.DataType_VarChar,
+						},
+					},
+				},
+			},
+		}
+
+		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
+
+		// Check result: should have 2 struct fields
+		assert.Len(t, resultFieldsData, 2)
+		assert.Len(t, resultOutputFields, 2)
+
+		// Check both struct fields
+		var foundStruct1, foundStruct2 bool
+		for _, field := range resultFieldsData {
+			switch field.FieldId {
+			case 102:
+				assert.Equal(t, "struct1", field.FieldName)
+				assert.Equal(t, schemapb.DataType_ArrayOfStruct, field.Type)
+				foundStruct1 = true
+				structArrays := field.GetStructArrays()
+				assert.NotNil(t, structArrays)
+				assert.Len(t, structArrays.Fields, 1)
+			case 103:
+				assert.Equal(t, "struct2", field.FieldName)
+				assert.Equal(t, schemapb.DataType_ArrayOfStruct, field.Type)
+				foundStruct2 = true
+				structArrays := field.GetStructArrays()
+				assert.NotNil(t, structArrays)
+				assert.Len(t, structArrays.Fields, 2)
+			}
+		}
+		assert.True(t, foundStruct1, "Should find struct1")
+		assert.True(t, foundStruct2, "Should find struct2")
+	})
+
+	t.Run("partial struct fields query - only return queried fields", func(t *testing.T) {
+		// Create a struct with 3 fields, but only query 2 of them
+		// This tests that we only return what the user requested
+
+		// Create mock data for only 2 out of 3 struct fields
+		clipStrData := &schemapb.FieldData{
+			FieldName: "clip[str]",
+			FieldId:   2001,
+			Type:      schemapb.DataType_VarChar,
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_StringData{
+						StringData: &schemapb.StringArray{Data: []string{"text1", "text2"}},
+					},
+				},
+			},
+		}
+
+		clipIntData := &schemapb.FieldData{
+			FieldName: "clip[int]",
+			FieldId:   2002,
+			Type:      schemapb.DataType_Int32,
+			Field: &schemapb.FieldData_Scalars{
+				Scalars: &schemapb.ScalarField{
+					Data: &schemapb.ScalarField_IntData{
+						IntData: &schemapb.IntArray{Data: []int32{100, 200}},
+					},
+				},
+			},
+		}
+
+		// Note: clip[embedding] is NOT included in query results
+		fieldsData := []*schemapb.FieldData{clipStrData, clipIntData}
+		outputFields := []string{"clip[str]", "clip[int]"}
+
+		schema := &schemapb.CollectionSchema{
+			Fields: []*schemapb.FieldSchema{
+				{
+					FieldID:      100,
+					Name:         "pk",
+					IsPrimaryKey: true,
+					DataType:     schemapb.DataType_Int64,
+				},
+			},
+			StructArrayFields: []*schemapb.StructArrayFieldSchema{
+				{
+					FieldID: 200,
+					Name:    "clip",
+					Fields: []*schemapb.FieldSchema{
+						{
+							FieldID:  2001,
+							Name:     "clip[str]",
+							DataType: schemapb.DataType_VarChar,
+						},
+						{
+							FieldID:  2002,
+							Name:     "clip[int]",
+							DataType: schemapb.DataType_Int32,
+						},
+						{
+							FieldID:  2003,
+							Name:     "clip[embedding]",
+							DataType: schemapb.DataType_FloatVector,
+							TypeParams: []*commonpb.KeyValuePair{
+								{Key: "dim", Value: "128"},
+							},
+						},
+					},
+				},
+			},
+		}
+
+		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
+
+		// Check result
+		assert.Len(t, resultFieldsData, 1, "Should have one reconstructed struct field")
+		assert.Len(t, resultOutputFields, 1, "Output fields should have one")
+
+		structField := resultFieldsData[0]
+		assert.Equal(t, "clip", structField.FieldName)
+		assert.Equal(t, int64(200), structField.FieldId)
+		assert.Equal(t, schemapb.DataType_ArrayOfStruct, structField.Type)
+		assert.Equal(t, "clip", resultOutputFields[0])
+
+		// Check that struct only contains the 2 queried fields, NOT the embedding field
+		structArrays := structField.GetStructArrays()
+		assert.NotNil(t, structArrays)
+		assert.Len(t, structArrays.Fields, 2, "Struct should only contain 2 queried fields, not 3")
+
+		// Verify the field names have been restored to original names
+		var foundStr, foundInt bool
+		for _, field := range structArrays.Fields {
+			switch field.FieldId {
+			case 2001:
+				assert.Equal(t, "str", field.FieldName, "Field name should be restored to original")
+				assert.Equal(t, schemapb.DataType_VarChar, field.Type)
+				foundStr = true
+			case 2002:
+				assert.Equal(t, "int", field.FieldName, "Field name should be restored to original")
+				assert.Equal(t, schemapb.DataType_Int32, field.Type)
+				foundInt = true
+			case 2003:
+				assert.Fail(t, "Should not include embedding field as it was not queried")
+			}
+		}
+		assert.True(t, foundStr, "Should find str field")
+		assert.True(t, foundInt, "Should find int field")
+	})
+}
+
+func TestResolveTimezone(t *testing.T) {
+	ctx := context.Background()
+	colInfoWithTz := &collectionInfo{
+		Properties: []*commonpb.KeyValuePair{
+			{Key: common.TimezoneKey, Value: "America/New_York"},
+		},
+	}
+	colInfoWithoutTz := &collectionInfo{}
+
+	t.Run("request timezone wins over collection timezone", func(t *testing.T) {
+		params := []*commonpb.KeyValuePair{{Key: common.TimezoneKey, Value: "Asia/Shanghai"}}
+		tz, err := resolveTimezone(ctx, params, colInfoWithTz)
+		assert.NoError(t, err)
+		assert.Equal(t, "Asia/Shanghai", tz)
+	})
+
+	t.Run("invalid request timezone is rejected as ParameterInvalid", func(t *testing.T) {
+		params := []*commonpb.KeyValuePair{{Key: common.TimezoneKey, Value: "Not/AZone"}}
+		_, err := resolveTimezone(ctx, params, colInfoWithTz)
+		assert.Error(t, err)
+		assert.ErrorIs(t, err, merr.ErrParameterInvalid)
+	})
+
+	t.Run("absent request timezone falls back to collection timezone", func(t *testing.T) {
+		tz, err := resolveTimezone(ctx, nil, colInfoWithTz)
+		assert.NoError(t, err)
+		assert.Equal(t, "America/New_York", tz)
+	})
+
+	t.Run("empty request timezone is treated as unspecified", func(t *testing.T) {
+		params := []*commonpb.KeyValuePair{{Key: common.TimezoneKey, Value: ""}}
+		tz, err := resolveTimezone(ctx, params, colInfoWithTz)
+		assert.NoError(t, err)
+		assert.Equal(t, "America/New_York", tz)
+	})
+
+	t.Run("no request or collection timezone defaults to UTC", func(t *testing.T) {
+		tz, err := resolveTimezone(ctx, nil, colInfoWithoutTz)
+		assert.NoError(t, err)
+		assert.Equal(t, common.DefaultTimezone, tz)
+	})
+}

internal/proxy/dql/util_dql.go

@@ -0,0 +1,783 @@
+// 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 dql
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/samber/lo"
+
+	"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"
+	typeutil2 "github.com/milvus-io/milvus/internal/util/typeutil"
+	"github.com/milvus-io/milvus/pkg/v3/common"
+	"github.com/milvus-io/milvus/pkg/v3/mlog"
+	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
+	"github.com/milvus-io/milvus/pkg/v3/util/contextutil"
+	"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/timestamptz"
+	"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
+	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
+)
+
+// strongTS / boundedTS are magic guarantee-timestamp sentinels for the
+// Bounded and Strong consistency mappings.
+const (
+	strongTS  = 0
+	boundedTS = 2
+)
+
+// isAlpha returns true if the byte is an ASCII letter.
+func isAlpha(c uint8) bool {
+	return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
+}
+
+// isNumber returns true if the byte is an ASCII digit.
+func isNumber(c uint8) bool {
+	return c >= '0' && c <= '9'
+}
+
+// validateLimit checks that a search/query limit is within the configured range.
+func validateLimit(limit int64, largeTopKEnabled bool) error {
+	topKLimit := paramtable.Get().QuotaConfig.TopKLimit.GetAsInt64()
+	if largeTopKEnabled {
+		topKLimit = paramtable.Get().QuotaConfig.LargeTopKLimit.GetAsInt64()
+	}
+	if limit <= 0 || limit > topKLimit {
+		return merr.WrapErrParameterInvalidMsg("it should be in range [1, %d], but got %d", topKLimit, limit)
+	}
+	return nil
+}
+
+// validatePartitionTag validates a partition tag string. It is duplicated here
+// because the DDL and DML task groups also need it and each sub-package keeps
+// its own copy of small shared helpers (see DEPENDENCIES.md).
+func validatePartitionTag(partitionTag string, strictCheck bool) error {
+	partitionTag = strings.TrimSpace(partitionTag)
+
+	invalidMsg := "Invalid partition name: " + partitionTag + ". "
+	if partitionTag == "" {
+		msg := invalidMsg + "Partition name should not be empty."
+		return merr.WrapErrParameterInvalidMsg("%s", msg)
+	}
+	if len(partitionTag) > paramtable.Get().ProxyCfg.MaxNameLength.GetAsInt() {
+		msg := invalidMsg + "The length of a partition name must be less than " + paramtable.Get().ProxyCfg.MaxNameLength.GetValue() + " characters."
+		return merr.WrapErrParameterInvalidMsg("%s", msg)
+	}
+
+	if strictCheck {
+		firstChar := partitionTag[0]
+		if firstChar != '_' && !isAlpha(firstChar) && !isNumber(firstChar) {
+			msg := invalidMsg + "The first character of a partition name must be an underscore or letter."
+			return merr.WrapErrParameterInvalidMsg("%s", msg)
+		}
+
+		tagSize := len(partitionTag)
+		for i := 1; i < tagSize; i++ {
+			c := partitionTag[i]
+			if c != '_' && !isAlpha(c) && !isNumber(c) && c != '-' {
+				msg := invalidMsg + "Partition name can only contain numbers, letters and underscores."
+				return merr.WrapErrParameterInvalidMsg("%s", msg)
+			}
+		}
+	}
+
+	return nil
+}
+
+// namespaceForPlan returns nil when partition-mode namespacing is enabled so
+// plan building skips the namespace field, else passes it through.
+func namespaceForPlan(schema *schemapb.CollectionSchema, namespace *string) *string {
+	if namespacePartitionModeEnabled(schema) {
+		return nil
+	}
+	return namespace
+}
+
+// namespacePartitionModeEnabled reports whether the collection runs in
+// partition-mode namespacing, which renders the plan namespace redundant.
+func namespacePartitionModeEnabled(schema *schemapb.CollectionSchema) bool {
+	return schema != nil && schema.GetEnableNamespace() && common.IsNamespaceModePartition(schema.GetProperties()...)
+}
+
+func validateMaxQueryResultWindow(offset int64, limit int64, largeTopKEnabled bool) error {
+	if offset < 0 {
+		return merr.WrapErrParameterInvalidMsg("%s [%d] is invalid, should be gte than 0", OffsetKey, offset)
+	}
+	if limit <= 0 {
+		return merr.WrapErrParameterInvalidMsg("%s [%d] is invalid, should be greater than 0", LimitKey, limit)
+	}
+
+	depth := offset + limit
+	maxQueryResultWindow := paramtable.Get().QuotaConfig.MaxQueryResultWindow.GetAsInt64()
+	if largeTopKEnabled {
+		maxQueryResultWindow = paramtable.Get().QuotaConfig.LargeMaxQueryResultWindow.GetAsInt64()
+	}
+	if depth <= 0 || depth > maxQueryResultWindow {
+		return merr.WrapErrParameterInvalidMsg("(offset+limit) should be in range [1, %d], but got %d", maxQueryResultWindow, depth)
+	}
+	return nil
+}
+
+func validateCollectionNameOrAlias(entity, entityType string) error {
+	if entity == "" {
+		return merr.WrapErrParameterInvalidMsg("collection %s should not be empty", entityType)
+	}
+
+	invalidMsg := fmt.Sprintf("Invalid collection %s: %s. ", entityType, entity)
+	if len(entity) > paramtable.Get().ProxyCfg.MaxNameLength.GetAsInt() {
+		return merr.WrapErrParameterInvalidMsg("%s the length of a collection %s must be less than %s characters", invalidMsg, entityType,
+			paramtable.Get().ProxyCfg.MaxNameLength.GetValue())
+	}
+
+	firstChar := entity[0]
+	if firstChar != '_' && !isAlpha(firstChar) {
+		return merr.WrapErrParameterInvalidMsg("%s the first character of a collection %s must be an underscore or letter", invalidMsg, entityType)
+	}
+
+	for i := 1; i < len(entity); i++ {
+		c := entity[i]
+		if c != '_' && !isAlpha(c) && !isNumber(c) {
+			return merr.WrapErrParameterInvalidMsg("%s collection %s can only contain numbers, letters and underscores", invalidMsg, entityType)
+		}
+	}
+	return nil
+}
+
+func validateCollectionName(collName string) error {
+	return validateCollectionNameOrAlias(collName, "name")
+}
+
+// return value.
+func translateOutputFields(outputFields []string, schema *schemaInfo, removePkField bool) ([]string, []string, []string, []agg.AggregateBase, bool, error) {
+	var primaryFieldName string
+	allFieldNameMap := make(map[string]*schemapb.FieldSchema)
+	resultFieldNameMap := make(map[string]bool)
+	resultFieldNames := make([]string, 0)
+	userOutputFieldsMap := make(map[string]bool)
+	userOutputFields := make([]string, 0)
+	userDynamicFieldsMap := make(map[string]bool)
+	userDynamicFields := make([]string, 0)
+	useAllDyncamicFields := false
+	aggregates := make([]agg.AggregateBase, 0)
+	for _, field := range schema.Fields {
+		if field.IsPrimaryKey {
+			primaryFieldName = field.Name
+		}
+		allFieldNameMap[field.Name] = field
+	}
+
+	// User may specify a struct array field or some specific fields in the struct array field
+	for _, subStruct := range schema.StructArrayFields {
+		for _, field := range subStruct.Fields {
+			allFieldNameMap[field.Name] = field
+		}
+	}
+
+	structArrayNameToFields := make(map[string][]*schemapb.FieldSchema)
+	for _, subStruct := range schema.StructArrayFields {
+		structArrayNameToFields[subStruct.Name] = subStruct.Fields
+	}
+
+	userRequestedPkFieldExplicitly := false
+
+	for _, outputFieldName := range outputFields {
+		outputFieldName = strings.TrimSpace(outputFieldName)
+		if outputFieldName == primaryFieldName {
+			userRequestedPkFieldExplicitly = true
+		}
+		if outputFieldName == "*" {
+			userRequestedPkFieldExplicitly = true
+			for fieldName, field := range allFieldNameMap {
+				if schema.CanRetrieveRawFieldData(field) {
+					resultFieldNameMap[fieldName] = true
+					userOutputFieldsMap[fieldName] = true
+				}
+			}
+			useAllDyncamicFields = true
+		} else {
+			if isAgg, aggregateName, aggFieldName := agg.MatchAggregationExpression(outputFieldName); isAgg {
+				if aggField, ok := allFieldNameMap[aggFieldName]; ok {
+					aggFuncs, aggErr := agg.NewAggregate(aggregateName, aggField.GetFieldID(), outputFieldName, aggField.GetDataType())
+					if aggErr != nil {
+						return nil, nil, nil, nil, false, aggErr
+					}
+					aggregates = append(aggregates, aggFuncs...)
+				} else if aggFieldName == "*" {
+					// only count(*) is allowed
+					if aggregateName != "count" {
+						return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("%s(*) is not supported, only count(*) is allowed", aggregateName)
+					}
+					if err := agg.ValidateAggFieldType(aggregateName, schemapb.DataType_None); err != nil {
+						return nil, nil, nil, nil, false, err
+					}
+					aggFuncs, aggErr := agg.NewAggregate(aggregateName, 0, outputFieldName, schemapb.DataType_None)
+					if aggErr != nil {
+						return nil, nil, nil, nil, false, aggErr
+					}
+					aggregates = append(aggregates, aggFuncs...)
+				} else {
+					return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("target field %s for aggregation:%s does not exist", aggFieldName, aggregateName)
+				}
+				userOutputFieldsMap[outputFieldName] = true
+				continue
+			}
+
+			if structArrayField, ok := structArrayNameToFields[outputFieldName]; ok {
+				for _, field := range structArrayField {
+					if schema.CanRetrieveRawFieldData(field) {
+						resultFieldNameMap[field.Name] = true
+						userOutputFieldsMap[field.Name] = true
+					}
+				}
+				continue
+			}
+			if field, ok := allFieldNameMap[outputFieldName]; ok {
+				if !schema.CanRetrieveRawFieldData(field) {
+					return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("not allowed to retrieve raw data of field %s", outputFieldName)
+				}
+				resultFieldNameMap[outputFieldName] = true
+				userOutputFieldsMap[outputFieldName] = true
+			} else {
+				if schema.EnableDynamicField {
+					dynamicNestedPath := outputFieldName
+					err := planparserv2.ParseIdentifier(schema.SchemaHelper, outputFieldName, func(expr *planpb.Expr) error {
+						columnInfo := expr.GetColumnExpr().GetInfo()
+						// there must be no error here
+						dynamicField, _ := schema.SchemaHelper.GetDynamicField()
+						// only $meta["xxx"] is allowed for now
+						if dynamicField.GetFieldID() != columnInfo.GetFieldId() {
+							return merr.WrapErrParameterInvalidMsg("not support getting subkeys of json field yet")
+						}
+						nestedPaths := columnInfo.GetNestedPath()
+						// $meta["A"]["B"] not allowed for now
+						if len(nestedPaths) != 1 {
+							return merr.WrapErrParameterInvalidMsg("not support getting multiple level of dynamic field for now")
+						}
+						// $meta["dyn_field"], output field name could be:
+						// 1. "dyn_field", outputFieldName == nestedPath
+						// 2. `$meta["dyn_field"]` explicit form
+						if nestedPaths[0] != outputFieldName {
+							// use "dyn_field" as userDynamicFieldsMap when outputField = `$meta["dyn_field"]`
+							dynamicNestedPath = nestedPaths[0]
+						}
+						return nil
+					})
+					if err != nil {
+						mlog.Info(context.TODO(), "parse output field name failed", mlog.String("field name", outputFieldName), mlog.Err(err))
+						return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("parse output field name failed: %s", outputFieldName)
+					}
+					resultFieldNameMap[common.MetaFieldName] = true
+					userOutputFieldsMap[outputFieldName] = true
+					userDynamicFieldsMap[dynamicNestedPath] = true
+				} else {
+					return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("field %s not exist", outputFieldName)
+				}
+			}
+		}
+	}
+
+	if removePkField {
+		delete(resultFieldNameMap, primaryFieldName)
+		delete(userOutputFieldsMap, primaryFieldName)
+	}
+
+	for fieldName := range resultFieldNameMap {
+		resultFieldNames = append(resultFieldNames, fieldName)
+	}
+	for fieldName := range userOutputFieldsMap {
+		userOutputFields = append(userOutputFields, fieldName)
+	}
+	if !useAllDyncamicFields {
+		for fieldName := range userDynamicFieldsMap {
+			userDynamicFields = append(userDynamicFields, fieldName)
+		}
+	}
+
+	return resultFieldNames, userOutputFields, userDynamicFields, aggregates, userRequestedPkFieldExplicitly, nil
+}
+
+func isPartitionKeyMode(ctx context.Context, metaCache Cache, dbName string, colName string) (bool, error) {
+	colSchema, err := metaCache.GetCollectionSchema(ctx, dbName, colName)
+	if err != nil {
+		return false, err
+	}
+
+	for _, fieldSchema := range colSchema.GetFields() {
+		if fieldSchema.IsPartitionKey {
+			return true, nil
+		}
+	}
+
+	return false, nil
+}
+
+func assignNamespacePartitionKey(ctx context.Context, metaCache Cache, dbName string, collName string, namespace *string) ([]string, error) {
+	if namespace == nil {
+		return nil, nil
+	}
+
+	return assignPartitionKeys(ctx, metaCache, dbName, collName, []*planpb.GenericValue{
+		{Val: &planpb.GenericValue_StringVal{StringVal: *namespace}},
+	})
+}
+
+func namespacePartitionKeyMode(schema *schemapb.CollectionSchema) bool {
+	return schema != nil && schema.GetEnableNamespace() && common.IsNamespaceModePartitionKey(schema.GetProperties()...)
+}
+
+func namespacePartitionKeyModeEnabled(schema *schemapb.CollectionSchema) bool {
+	return namespaceShardingEnabled(schema) && namespacePartitionKeyMode(schema)
+}
+
+func resolveNamespacePartitionNames(schema *schemapb.CollectionSchema, namespace *string, partitionNames []string) ([]string, bool, error) {
+	if err := common.CheckNamespace(schema, namespace); err != nil {
+		return nil, false, err
+	}
+	if !namespacePartitionModeEnabled(schema) {
+		return partitionNames, false, nil
+	}
+
+	namespacePartitionName := *namespace
+	if err := validatePartitionTag(namespacePartitionName, true); err != nil {
+		return nil, true, err
+	}
+	if len(partitionNames) == 0 {
+		return []string{namespacePartitionName}, true, nil
+	}
+	if len(partitionNames) == 1 && partitionNames[0] == namespacePartitionName {
+		return partitionNames, true, nil
+	}
+	return nil, true, merr.WrapErrParameterInvalidMsg("partition names %v mismatch namespace %q", partitionNames, namespacePartitionName)
+}
+
+// filtering and result formatting use the same timezone.
+func resolveTimezone(ctx context.Context, params []*commonpb.KeyValuePair, colInfo *collectionInfo) (string, error) {
+	timezone, _ := funcutil.TryGetAttrByKeyFromRepeatedKV(common.TimezoneKey, params)
+	if timezone != "" {
+		if !timestamptz.IsTimezoneValid(timezone) {
+			mlog.Info(ctx, "get invalid timezone from request", mlog.String("timezone", timezone))
+			return "", merr.WrapErrParameterInvalidMsg("unknown or invalid IANA Time Zone ID: %s", timezone)
+		}
+		mlog.Debug(ctx, "determine timezone from request", mlog.String("user defined timezone", timezone))
+		return timezone, nil
+	}
+	timezone = getColTimezone(colInfo)
+	mlog.Debug(ctx, "determine timezone from collection", mlog.String("collection timezone", timezone))
+	return timezone, nil
+}
+
+func validateTextStorageV3Enabled(schema *schemapb.CollectionSchema) error {
+	if err := typeutil.ValidateTextRequiresStorageV3(schema, paramtable.Get().CommonCfg.UseLoonFFI.GetAsBool()); err != nil {
+		return merr.WrapErrParameterInvalidMsg("%s", err.Error())
+	}
+	return nil
+}
+
+func assignPartitionKeys(ctx context.Context, metaCache Cache, dbName string, collName string, keys []*planpb.GenericValue) ([]string, error) {
+	partitionNames, err := metaCache.GetPartitionsIndex(ctx, dbName, collName)
+	if err != nil {
+		return nil, err
+	}
+
+	schema, err := metaCache.GetCollectionSchema(ctx, dbName, collName)
+	if err != nil {
+		return nil, err
+	}
+
+	partitionKeyFieldSchema, err := typeutil.GetPartitionKeyFieldSchema(schema.CollectionSchema)
+	if err != nil {
+		return nil, err
+	}
+
+	hashedPartitionNames, err := typeutil2.HashKey2Partitions(partitionKeyFieldSchema, keys, partitionNames)
+	return hashedPartitionNames, err
+}
+
+func GetCurUserFromContext(ctx context.Context) (string, error) {
+	return contextutil.GetCurUserFromContext(ctx)
+}
+
+func parseGuaranteeTsFromConsistency(ts, tMax typeutil.Timestamp, consistency commonpb.ConsistencyLevel) typeutil.Timestamp {
+	switch consistency {
+	case commonpb.ConsistencyLevel_Strong:
+		ts = tMax
+	case commonpb.ConsistencyLevel_Bounded:
+		ratio := paramtable.Get().CommonCfg.GracefulTime.GetAsDuration(time.Millisecond)
+		ts = tsoutil.AddPhysicalDurationOnTs(tMax, -ratio)
+	case commonpb.ConsistencyLevel_Eventually:
+		ts = 1
+	}
+	return ts
+}
+
+func parseGuaranteeTs(ts, tMax typeutil.Timestamp) typeutil.Timestamp {
+	switch ts {
+	case strongTS:
+		ts = tMax
+	case boundedTS:
+		ratio := paramtable.Get().CommonCfg.GracefulTime.GetAsDuration(time.Millisecond)
+		ts = tsoutil.AddPhysicalDurationOnTs(tMax, -ratio)
+	}
+	return ts
+}
+
+func namespaceShardingChannel(schema *schemapb.CollectionSchema, namespace *string, channelNames []string) (string, bool, error) {
+	channelID, ok, err := namespaceShardingChannelID(schema, namespace, channelNames)
+	if !ok || err != nil {
+		return "", ok, err
+	}
+	return channelNames[channelID], true, nil
+}
+
+func preferredNodeForChannel(preferredNodes map[string]int64, channel string) int64 {
+	if preferredNodes == nil {
+		return 0
+	}
+	preferredNodeID, ok := preferredNodes[channel]
+	if !ok {
+		return 0
+	}
+	return preferredNodeID
+}
+
+func reconstructStructFieldDataForQuery(results *milvuspb.QueryResults, schema *schemapb.CollectionSchema) {
+	fieldsData, outputFields := reconstructStructFieldData(
+		results.FieldsData,
+		results.OutputFields,
+		schema,
+	)
+	results.FieldsData = fieldsData
+	results.OutputFields = outputFields
+}
+
+func getMaxMvccTsFromChannels(channelsTs map[string]uint64, beginTs typeutil.Timestamp) typeutil.Timestamp {
+	maxTs := typeutil.Timestamp(0)
+	for _, ts := range channelsTs {
+		if ts > maxTs {
+			maxTs = ts
+		}
+	}
+
+	if maxTs == 0 {
+		mlog.Warn(context.TODO(), "no channel ts found, use beginTs instead")
+		return beginTs
+	}
+
+	return maxTs
+}
+
+func extractFieldsFromResults(results []*schemapb.FieldData, timezone string, fieldList []string) error {
+	targetLocation, err := time.LoadLocation(timezone)
+	if err != nil {
+		mlog.Error(context.TODO(), "invalid timezone", mlog.String("timezone", timezone), mlog.Err(err))
+		return merr.WrapErrParameterInvalidMsg("got invalid timezone: %s", timezone)
+	}
+
+	for _, fieldData := range results {
+		if fieldData.GetType() != schemapb.DataType_Timestamptz {
+			continue
+		}
+
+		scalarField := fieldData.GetScalars()
+		if scalarField == nil || scalarField.GetTimestamptzData() == nil {
+			if longData := scalarField.GetLongData(); longData != nil && len(longData.GetData()) > 0 {
+				mlog.Warn(context.TODO(), "field data is not Timestamptz data, but found LongData instead", mlog.String("fieldName", fieldData.GetFieldName()))
+				return merr.WrapErrParameterInvalidMsg("field data for '%s' is not Timestamptz data", fieldData.GetFieldName())
+			}
+			continue
+		}
+
+		utcTimestamps := scalarField.GetTimestamptzData().GetData()
+		extractedResults := make([]*schemapb.ScalarField, 0, len(fieldList))
+
+		for _, ts := range utcTimestamps {
+			t := time.UnixMicro(ts).UTC()
+			localTime := t.In(targetLocation)
+
+			values, err := extractFields(localTime, fieldList)
+			if err != nil {
+				return err
+			}
+			valuesScalarField := &schemapb.ScalarField_LongData{
+				LongData: &schemapb.LongArray{
+					Data: values,
+				},
+			}
+			extractedResults = append(extractedResults, &schemapb.ScalarField{
+				Data: valuesScalarField,
+			})
+		}
+
+		fieldData.GetScalars().Data = &schemapb.ScalarField_ArrayData{
+			ArrayData: &schemapb.ArrayArray{
+				Data:        extractedResults,
+				ElementType: schemapb.DataType_Int64,
+			},
+		}
+		fieldData.Type = schemapb.DataType_Array
+	}
+	return nil
+}
+
+func timestamptzUTC2IsoStr(results []*schemapb.FieldData, colTimezone string) error {
+	location, err := time.LoadLocation(colTimezone)
+	if err != nil {
+		mlog.Error(context.TODO(), "invalid timezone", mlog.String("timezone", colTimezone), mlog.Err(err))
+		return merr.WrapErrParameterInvalidMsg("got invalid default timezone: %s", colTimezone)
+	}
+
+	for _, fieldData := range results {
+		if fieldData.GetType() != schemapb.DataType_Timestamptz {
+			continue
+		}
+
+		scalarField := fieldData.GetScalars()
+
+		// Guard against nil scalars or missing timestamp data
+		if scalarField == nil || scalarField.GetTimestamptzData() == nil {
+			if longData := scalarField.GetLongData(); longData != nil && len(longData.GetData()) > 0 {
+				mlog.Warn(context.TODO(), "field data is not Timestamptz data", mlog.String("fieldName", fieldData.GetFieldName()))
+				return merr.WrapErrParameterInvalidMsg("field data for '%s' is not Timestamptz data", fieldData.GetFieldName())
+			}
+			// Handle the case of an empty field (e.g., all nulls), skip if no data to process.
+			continue
+		}
+
+		utcTimestamps := scalarField.GetTimestamptzData().GetData()
+		isoStrings := make([]string, len(utcTimestamps))
+
+		// CORE CHANGE: Use the optimized formatting function
+		for i, ts := range utcTimestamps {
+			// 1. Convert Unix Microsecond (UTC) to a time.Time object (still in UTC).
+			t := time.UnixMicro(ts).UTC()
+
+			// 2. Adjust the time object to the target location.
+			localTime := t.In(location)
+
+			// 3. Format using the optimized logic (max 6 digits, no trailing zeros)
+			isoStrings[i] = timestamptz.FormatTimeMicroWithoutTrailingZeros(localTime)
+		}
+
+		// Replace the TimestamptzData with the new StringData in place.
+		fieldData.GetScalars().Data = &schemapb.ScalarField_StringData{
+			StringData: &schemapb.StringArray{
+				Data: isoStrings,
+			},
+		}
+	}
+	return nil
+}
+
+func validateNQLimit(limit int64) error {
+	nqLimit := paramtable.Get().QuotaConfig.NQLimit.GetAsInt64()
+	if limit <= 0 || limit > nqLimit {
+		return merr.WrapErrParameterInvalidMsg("nq (number of search vector per search request) should be in range [1, %d], but got %d", nqLimit, limit)
+	}
+	return nil
+}
+
+func getBM25FunctionOfAnnsField(fieldID int64, functions []*schemapb.FunctionSchema) (*schemapb.FunctionSchema, bool) {
+	return lo.Find(functions, func(function *schemapb.FunctionSchema) bool {
+		return function.GetType() == schemapb.FunctionType_BM25 && function.OutputFieldIds[0] == fieldID
+	})
+}
+
+func extractFields(t time.Time, fieldList []string) ([]int64, error) {
+	extractedValues := make([]int64, 0, len(fieldList))
+	for _, field := range fieldList {
+		var val int64
+		switch strings.ToLower(field) {
+		case common.TszYear:
+			val = int64(t.Year())
+		case common.TszMonth:
+			val = int64(t.Month())
+		case common.TszDay:
+			val = int64(t.Day())
+		case common.TszHour:
+			val = int64(t.Hour())
+		case common.TszMinute:
+			val = int64(t.Minute())
+		case common.TszSecond:
+			val = int64(t.Second())
+		case common.TszMicrosecond:
+			val = int64(t.Nanosecond() / 1000)
+		default:
+			return nil, merr.WrapErrParameterInvalidMsg("unsupported field for extraction: %s, fields should be seprated by ',' or ' '", field)
+		}
+		extractedValues = append(extractedValues, val)
+	}
+	return extractedValues, nil
+}
+
+func reconstructStructFieldData(
+	fieldsData []*schemapb.FieldData,
+	outputFields []string,
+	schema *schemapb.CollectionSchema,
+) ([]*schemapb.FieldData, []string) {
+	if len(outputFields) == 1 && outputFields[0] == "count(*)" {
+		return fieldsData, outputFields
+	}
+
+	if len(schema.StructArrayFields) == 0 {
+		return fieldsData, outputFields
+	}
+
+	regularFieldIDs := make(map[int64]interface{})
+	subFieldToStructMap := make(map[int64]int64)
+	groupedStructFields := make(map[int64][]*schemapb.FieldData)
+	structFieldNames := make(map[int64]string)
+	reconstructedOutputFields := make([]string, 0, len(fieldsData))
+
+	// record all regular field IDs
+	for _, field := range schema.Fields {
+		regularFieldIDs[field.GetFieldID()] = nil
+	}
+
+	// build the mapping from sub-field ID to struct field ID
+	for _, structField := range schema.StructArrayFields {
+		for _, subField := range structField.GetFields() {
+			subFieldToStructMap[subField.GetFieldID()] = structField.GetFieldID()
+		}
+		structFieldNames[structField.GetFieldID()] = structField.GetName()
+	}
+
+	newFieldsData := make([]*schemapb.FieldData, 0, len(fieldsData))
+	for _, field := range fieldsData {
+		fieldID := field.GetFieldId()
+		if _, ok := regularFieldIDs[fieldID]; ok {
+			newFieldsData = append(newFieldsData, field)
+			reconstructedOutputFields = append(reconstructedOutputFields, field.GetFieldName())
+		} else if structFieldID, ok := subFieldToStructMap[fieldID]; ok {
+			groupedStructFields[structFieldID] = append(groupedStructFields[structFieldID], field)
+		} else {
+			newFieldsData = append(newFieldsData, field)
+			reconstructedOutputFields = append(reconstructedOutputFields, field.GetFieldName())
+		}
+	}
+
+	for structFieldID, fields := range groupedStructFields {
+		// Restore original field names (from "structName[fieldName]" to "fieldName")
+		// for the user-facing response.
+		for _, field := range fields {
+			originalName, err := extractOriginalFieldName(field.FieldName)
+			if err != nil {
+				mlog.Error(context.TODO(), "failed to extract original field name from struct field",
+					mlog.String("fieldName", field.FieldName),
+					mlog.Err(err))
+			} else {
+				field.FieldName = originalName
+			}
+		}
+
+		newFieldsData = append(newFieldsData, &schemapb.FieldData{
+			FieldName: structFieldNames[structFieldID],
+			FieldId:   structFieldID,
+			Type:      schemapb.DataType_ArrayOfStruct,
+			Field:     &schemapb.FieldData_StructArrays{StructArrays: &schemapb.StructArrayField{Fields: fields}},
+		})
+		reconstructedOutputFields = append(reconstructedOutputFields, structFieldNames[structFieldID])
+	}
+
+	return newFieldsData, reconstructedOutputFields
+}
+
+// enableMultipleVectorFields indicates whether to enable multiple vector fields.
+const enableMultipleVectorFields = true
+
+// defaultMaxSearchRequest is the maximum number of ann search requests in a hybrid search.
+const defaultMaxSearchRequest = 1024
+
+func namespaceShardingEnabled(schema *schemapb.CollectionSchema) bool {
+	if schema == nil || !schema.GetEnableNamespace() {
+		return false
+	}
+	enabled, err := common.IsNamespaceShardingEnabled(schema.GetProperties()...)
+	return err == nil && enabled
+}
+
+func getColTimezone(colInfo *collectionInfo) string {
+	timezone, _ := funcutil.TryGetAttrByKeyFromRepeatedKV(common.TimezoneKey, colInfo.Properties)
+	if timezone == "" {
+		timezone = common.DefaultTimezone
+	}
+	return timezone
+}
+
+func namespaceShardingChannelID(schema *schemapb.CollectionSchema, namespace *string, channelNames []string) (uint32, bool, error) {
+	if namespace == nil || !namespacePartitionKeyModeEnabled(schema) {
+		return 0, false, nil
+	}
+	if len(channelNames) == 0 {
+		return 0, false, merr.WrapErrServiceInternalMsg("no virtual channels available for namespace sharding")
+	}
+	return typeutil.HashNamespace2Channels(*namespace, channelNames), true, nil
+}
+
+func extractOriginalFieldName(transformedName string) (string, error) {
+	idx := strings.Index(transformedName, "[")
+	if idx == -1 {
+		return "", merr.WrapErrParameterInvalidMsg("not a transformed struct field name: %s", transformedName)
+	}
+
+	if !strings.HasSuffix(transformedName, "]") {
+		return "", merr.WrapErrParameterInvalidMsg("invalid struct field format: %s, missing closing bracket", transformedName)
+	}
+
+	if idx == 0 {
+		return "", merr.WrapErrParameterInvalidMsg("invalid struct field format: %s, missing struct name", transformedName)
+	}
+
+	fieldName := transformedName[idx+1 : len(transformedName)-1]
+	if fieldName == "" {
+		return "", merr.WrapErrParameterInvalidMsg("invalid struct field format: %s, empty field name", transformedName)
+	}
+
+	return fieldName, nil
+}
+
+func preferredNodeFromConcurrentMap(preferredNodes *typeutil.ConcurrentMap[string, int64], channel string) int64 {
+	if preferredNodes == nil {
+		return 0
+	}
+	preferredNodeID, ok := preferredNodes.Get(channel)
+	if !ok {
+		return 0
+	}
+	return preferredNodeID
+}
+
+func reconstructStructFieldDataForSearch(results *milvuspb.SearchResults, schema *schemapb.CollectionSchema) {
+	if results.Results == nil {
+		return
+	}
+	fieldsData, outputFields := reconstructStructFieldData(
+		results.Results.FieldsData,
+		results.Results.OutputFields,
+		schema,
+	)
+	results.Results.FieldsData = fieldsData
+	results.Results.OutputFields = outputFields
+}

internal/proxy/dql/vector_type_convert.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"encoding/binary"
@@ -147,7 +147,7 @@ func convertPlaceholder(
 	return proto.Marshal(phg)
 }
 
-func normalizeFP32ToFP16BF16VectorFieldData(columns []*schemapb.FieldData, schema *schemaInfo) error {
+func NormalizeFP32ToFP16BF16VectorFieldData(columns []*schemapb.FieldData, schema *schemaInfo) error {
 	for _, fieldData := range columns {
 		fieldSchema, err := schema.SchemaHelper.GetFieldFromNameDefaultJSON(fieldData.GetFieldName())
 		if err != nil {

internal/proxy/dql/vector_type_convert_test.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package proxy
+package dql
 
 import (
 	"fmt"
@@ -371,7 +371,7 @@ func TestNormalizeFp32ToFp16Bf16VectorFieldData(t *testing.T) {
 			err := fillFieldPropertiesOnly([]*schemapb.FieldData{tt.fieldData}, schema)
 			assert.NoError(t, err)
 
-			err = normalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{tt.fieldData}, schema)
+			err = NormalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{tt.fieldData}, schema)
 			assert.NoError(t, err)
 			assert.Equal(t, tt.dataType, tt.fieldData.GetType())
 			assert.Equal(t, int64(2), tt.fieldData.GetVectors().GetDim())
@@ -388,7 +388,7 @@ func TestNormalizeFp32ToFp16Bf16VectorFieldData(t *testing.T) {
 		schema := newTestSchema(schemapb.DataType_FloatVector)
 		err := fillFieldPropertiesOnly([]*schemapb.FieldData{fieldData}, schema)
 		assert.NoError(t, err)
-		err = normalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{fieldData}, schema)
+		err = NormalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{fieldData}, schema)
 		assert.NoError(t, err)
 		assert.NotNil(t, fieldData.GetVectors().GetFloatVector())
 	})
@@ -396,13 +396,13 @@ func TestNormalizeFp32ToFp16Bf16VectorFieldData(t *testing.T) {
 	t.Run("skip non float vector payload", func(t *testing.T) {
 		fieldData := &schemapb.FieldData{FieldName: "vector"}
 		schema := newTestSchema(schemapb.DataType_Float16Vector)
-		err := normalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{fieldData}, schema)
+		err := NormalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{fieldData}, schema)
 		assert.NoError(t, err)
 	})
 
 	t.Run("unknown field", func(t *testing.T) {
 		fieldData := testFloatVectorFieldData("missing", []float32{0.1, 0.2}, 2)
-		err := normalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{fieldData}, newTestSchema(schemapb.DataType_Float16Vector))
+		err := NormalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{fieldData}, newTestSchema(schemapb.DataType_Float16Vector))
 		assert.Error(t, err)
 	})
 
@@ -411,7 +411,7 @@ func TestNormalizeFp32ToFp16Bf16VectorFieldData(t *testing.T) {
 		schema := newTestSchema(schemapb.DataType_Float16Vector)
 		err := fillFieldPropertiesOnly([]*schemapb.FieldData{fieldData}, schema)
 		assert.NoError(t, err)
-		err = normalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{fieldData}, schema)
+		err = NormalizeFP32ToFP16BF16VectorFieldData([]*schemapb.FieldData{fieldData}, schema)
 		assert.Error(t, err)
 	})
 }

internal/proxy/impl.go

@@ -52,6 +52,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proxy/privilege"
 	"github.com/milvus-io/milvus/internal/proxy/replicate"
 	"github.com/milvus-io/milvus/internal/proxy/rls"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/fileresource"
 	"github.com/milvus-io/milvus/internal/util/hookutil"
@@ -1160,15 +1161,7 @@ func (node *Proxy) GetStatistics(ctx context.Context, request *milvuspb.GetStati
 	defer sp.End()
 	method := "GetStatistics"
 	tr := timerecord.NewTimeRecorder(method)
-	g := &getStatisticsTask{
-		baseTask:  baseTask{MetaCache: node.GetMetaCache()},
-		request:   request,
-		Condition: NewTaskCondition(ctx),
-		ctx:       ctx,
-		tr:        tr,
-		mixc:      node.mixCoord,
-		lb:        node.lbPolicy,
-	}
+	g := NewGetStatisticsTask(ctx, node, request, tr)
 
 	mlog.Debug(ctx,
 		rpcReceived(method),
@@ -1210,7 +1203,7 @@ func (node *Proxy) GetStatistics(ctx context.Context, request *milvuspb.GetStati
 		mlog.Uint64("EndTS", g.EndTs()))
 
 	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
-	return g.result, nil
+	return g.Result(), nil
 }
 
 // GetCollectionStatistics get the collection statistics, such as `num_rows`.
@@ -1225,13 +1218,7 @@ func (node *Proxy) GetCollectionStatistics(ctx context.Context, request *milvusp
 	defer sp.End()
 	method := "GetCollectionStatistics"
 	tr := timerecord.NewTimeRecorder(method)
-	g := &getCollectionStatisticsTask{
-		baseTask:                       baseTask{MetaCache: node.GetMetaCache()},
-		ctx:                            ctx,
-		Condition:                      NewTaskCondition(ctx),
-		GetCollectionStatisticsRequest: request,
-		mixCoord:                       node.mixCoord,
-	}
+	g := NewGetCollectionStatisticsTask(ctx, node, request)
 
 	mlog.Debug(ctx, rpcReceived(method))
 
@@ -1268,7 +1255,7 @@ func (node *Proxy) GetCollectionStatistics(ctx context.Context, request *milvusp
 		mlog.Uint64("EndTS", g.EndTs()))
 
 	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
-	return g.result, nil
+	return g.Result(), nil
 }
 
 // ShowCollections list all collections in Milvus.
@@ -1813,13 +1800,7 @@ func (node *Proxy) GetPartitionStatistics(ctx context.Context, request *milvuspb
 	method := "GetPartitionStatistics"
 	tr := timerecord.NewTimeRecorder(method)
 
-	g := &getPartitionStatisticsTask{
-		baseTask:                      baseTask{MetaCache: node.GetMetaCache()},
-		ctx:                           ctx,
-		Condition:                     NewTaskCondition(ctx),
-		GetPartitionStatisticsRequest: request,
-		mixCoord:                      node.mixCoord,
-	}
+	g := NewGetPartitionStatisticsTask(ctx, node, request)
 
 	mlog.Debug(ctx, rpcReceived(method))
 
@@ -1856,7 +1837,7 @@ func (node *Proxy) GetPartitionStatistics(ctx context.Context, request *milvuspb
 		mlog.Uint64("EndTS", g.EndTs()))
 
 	metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
-	return g.result, nil
+	return g.Result(), nil
 }
 
 // ShowPartitions list all partitions in the specific collection.
@@ -3041,31 +3022,7 @@ func (node *Proxy) search(ctx context.Context, request *milvuspb.SearchRequest,
 		}, false, false, false, nil
 	}
 
-	qt := &searchTask{
-		baseTask: baseTask{
-			MetaCache: node.GetMetaCache(),
-		},
-		ctx:       ctx,
-		Condition: NewTaskCondition(ctx),
-		SearchRequest: &internalpb.SearchRequest{
-			Base: commonpbutil.NewMsgBase(
-				commonpbutil.WithMsgType(commonpb.MsgType_Search),
-				commonpbutil.WithSourceID(paramtable.GetNodeID()),
-			),
-			ReqID:              paramtable.GetNodeID(),
-			IsTopkReduce:       optimizedSearch,
-			IsRecallEvaluation: isRecallEvaluation,
-		},
-		request:                request,
-		tr:                     timerecord.NewTimeRecorder("search"),
-		mixCoord:               node.mixCoord,
-		node:                   node,
-		lb:                     node.lbPolicy,
-		shardClientMgr:         node.shardMgr,
-		enableMaterializedView: node.enableMaterializedView,
-		mustUsePartitionKey:    Params.ProxyCfg.MustUsePartitionKey.GetAsBool(),
-		chMgr:                  node.chMgr,
-	}
+	qt := NewSearchTask(ctx, node, node.sched, request, optimizedSearch, isRecallEvaluation, tr)
 
 	succeeded := false
 	defer func() {
@@ -3141,7 +3098,7 @@ func (node *Proxy) search(ctx context.Context, request *milvuspb.SearchRequest,
 
 	metrics.ProxySearchVectors.
 		WithLabelValues(nodeID, dbName, collectionName).
-		Add(float64(qt.result.GetResults().GetNumQueries()))
+		Add(float64(qt.Result().GetResults().GetNumQueries()))
 
 	searchDur := tr.ElapseSpan().Milliseconds()
 	metrics.ProxySQLatency.WithLabelValues(
@@ -3166,40 +3123,40 @@ func (node *Proxy) search(ctx context.Context, request *milvuspb.SearchRequest,
 			metrics.SearchLabel,
 			dbName,
 			collectionName,
-		).Add(float64(qt.storageCost.ScannedRemoteBytes) / 1024 / 1024)
+		).Add(float64(qt.StorageCost().ScannedRemoteBytes) / 1024 / 1024)
 
 		metrics.ProxyScannedTotalMB.WithLabelValues(
 			nodeID,
 			metrics.SearchLabel,
 			dbName,
 			collectionName,
-		).Add(float64(qt.storageCost.ScannedTotalBytes) / 1024 / 1024)
+		).Add(float64(qt.StorageCost().ScannedTotalBytes) / 1024 / 1024)
 	}
 
-	if qt.result != nil {
+	if qt.Result() != nil {
 		username := GetCurUserFromContextOrDefault(ctx)
-		sentSize := proto.Size(qt.result)
+		sentSize := proto.Size(qt.Result())
 		v := hookutil.GetExtension().Report(map[string]any{
 			hookutil.OpTypeKey:          hookutil.OpTypeSearch,
 			hookutil.DatabaseKey:        dbName,
 			hookutil.UsernameKey:        username,
 			hookutil.ResultDataSizeKey:  sentSize,
-			hookutil.RelatedDataSizeKey: qt.relatedDataSize,
-			hookutil.RelatedCntKey:      qt.result.GetResults().GetAllSearchCount(),
+			hookutil.RelatedDataSizeKey: qt.RelatedDataSize(),
+			hookutil.RelatedCntKey:      qt.Result().GetResults().GetAllSearchCount(),
 		})
-		SetReportValue(qt.result.GetStatus(), v)
-		SetStorageCost(qt.result.GetStatus(), qt.storageCost)
-		if merr.Ok(qt.result.GetStatus()) {
+		SetReportValue(qt.Result().GetStatus(), v)
+		SetStorageCost(qt.Result().GetStatus(), qt.StorageCost())
+		if merr.Ok(qt.Result().GetStatus()) {
 			metrics.ProxyReportValue.WithLabelValues(nodeID, hookutil.OpTypeSearch, dbName, username).Add(float64(v))
 		}
 	}
 
-	if qt.result != nil && qt.result.Results != nil && len(validData) > 0 {
-		adjustSearchResultsForNullVectors(qt.result, validData)
+	if qt.Result() != nil && qt.Result().Results != nil && len(validData) > 0 {
+		adjustSearchResultsForNullVectors(qt.Result(), validData)
 	}
 
 	succeeded = true
-	return qt.result, qt.resultSizeInsufficient, qt.isTopkReduce, qt.isRecallEvaluation, nil
+	return qt.Result(), qt.ResultSizeInsufficient(), qt.IsTopkReduce(), qt.IsRecallEvaluation(), nil
 }
 
 func (node *Proxy) HybridSearch(ctx context.Context, request *milvuspb.HybridSearchRequest) (*milvuspb.SearchResults, error) {
@@ -3277,30 +3234,8 @@ func (node *Proxy) hybridSearch(ctx context.Context, request *milvuspb.HybridSea
 
 	ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-HybridSearch")
 	defer sp.End()
-	newSearchReq := convertHybridSearchToSearch(request)
-	qt := &searchTask{
-		baseTask: baseTask{
-			MetaCache: node.GetMetaCache(),
-		},
-		ctx:       ctx,
-		Condition: NewTaskCondition(ctx),
-		SearchRequest: &internalpb.SearchRequest{
-			Base: commonpbutil.NewMsgBase(
-				commonpbutil.WithMsgType(commonpb.MsgType_Search),
-				commonpbutil.WithSourceID(paramtable.GetNodeID()),
-			),
-			ReqID:        paramtable.GetNodeID(),
-			IsTopkReduce: optimizedSearch,
-		},
-		request:             newSearchReq,
-		tr:                  timerecord.NewTimeRecorder(method),
-		mixCoord:            node.mixCoord,
-		node:                node,
-		lb:                  node.lbPolicy,
-		shardClientMgr:      node.shardMgr,
-		mustUsePartitionKey: Params.ProxyCfg.MustUsePartitionKey.GetAsBool(),
-		chMgr:               node.chMgr,
-	}
+	newSearchReq := ConvertHybridSearchToSearch(request)
+	qt := NewSearchTask(ctx, node, node.sched, newSearchReq, optimizedSearch, false, tr)
 
 	succeeded := false
 	defer func() {
@@ -3404,35 +3339,35 @@ func (node *Proxy) hybridSearch(ctx context.Context, request *milvuspb.HybridSea
 			metrics.HybridSearchLabel,
 			dbName,
 			collectionName,
-		).Add(float64(qt.storageCost.ScannedRemoteBytes) / 1024 / 1024)
+		).Add(float64(qt.StorageCost().ScannedRemoteBytes) / 1024 / 1024)
 
 		metrics.ProxyScannedTotalMB.WithLabelValues(
 			nodeID,
 			metrics.HybridSearchLabel,
 			dbName,
 			collectionName,
-		).Add(float64(qt.storageCost.ScannedTotalBytes) / 1024 / 1024)
+		).Add(float64(qt.StorageCost().ScannedTotalBytes) / 1024 / 1024)
 	}
 
-	if qt.result != nil {
-		sentSize := proto.Size(qt.result)
+	if qt.Result() != nil {
+		sentSize := proto.Size(qt.Result())
 		username := GetCurUserFromContextOrDefault(ctx)
 		v := hookutil.GetExtension().Report(map[string]any{
 			hookutil.OpTypeKey:          hookutil.OpTypeHybridSearch,
 			hookutil.DatabaseKey:        dbName,
 			hookutil.UsernameKey:        username,
 			hookutil.ResultDataSizeKey:  sentSize,
-			hookutil.RelatedDataSizeKey: qt.relatedDataSize,
-			hookutil.RelatedCntKey:      qt.result.GetResults().GetAllSearchCount(),
+			hookutil.RelatedDataSizeKey: qt.RelatedDataSize(),
+			hookutil.RelatedCntKey:      qt.Result().GetResults().GetAllSearchCount(),
 		})
-		SetReportValue(qt.result.GetStatus(), v)
-		SetStorageCost(qt.result.GetStatus(), qt.storageCost)
-		if merr.Ok(qt.result.GetStatus()) {
+		SetReportValue(qt.Result().GetStatus(), v)
+		SetStorageCost(qt.Result().GetStatus(), qt.StorageCost())
+		if merr.Ok(qt.Result().GetStatus()) {
 			metrics.ProxyReportValue.WithLabelValues(nodeID, hookutil.OpTypeHybridSearch, dbName, username).Add(float64(v))
 		}
 	}
 	succeeded = true
-	return qt.result, qt.resultSizeInsufficient, qt.isTopkReduce, nil
+	return qt.Result(), qt.ResultSizeInsufficient(), qt.IsTopkReduce(), nil
 }
 
 func adjustSearchResultsForNullVectors(results *milvuspb.SearchResults, validData []bool) {
@@ -3582,31 +3517,15 @@ func (node *Proxy) handleIfSearchByPK(ctx context.Context, request *milvuspb.Sea
 	}
 
 	// Create queryTask to execute the retrieval
-	qt := &queryTask{
-		baseTask: baseTask{
-			MetaCache: node.GetMetaCache(),
-		},
-		ctx:       ctx,
-		Condition: NewTaskCondition(ctx),
-		RetrieveRequest: &internalpb.RetrieveRequest{
-			Base: commonpbutil.NewMsgBase(
-				commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
-				commonpbutil.WithSourceID(paramtable.GetNodeID()),
-			),
-			ReqID:            paramtable.GetNodeID(),
-			ConsistencyLevel: request.ConsistencyLevel,
-			QueryLabel:       metrics.QueryLabel,
-		},
-		request:             queryReq,
-		plan:                plan,
-		mixCoord:            node.mixCoord,
-		lb:                  node.lbPolicy,
-		shardclientMgr:      node.shardMgr,
-		mustUsePartitionKey: Params.ProxyCfg.MustUsePartitionKey.GetAsBool(),
-		// reQuery defaults to false - we need full query processing:
-		// partition conversion, struct field reconstruction, timestamp handling etc
-		chMgr: node.chMgr,
-	}
+	qt := NewQueryTask(ctx, node, queryReq, plan, &internalpb.RetrieveRequest{
+		Base: commonpbutil.NewMsgBase(
+			commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
+			commonpbutil.WithSourceID(paramtable.GetNodeID()),
+		),
+		ReqID:            paramtable.GetNodeID(),
+		ConsistencyLevel: request.ConsistencyLevel,
+		QueryLabel:       metrics.QueryLabel,
+	}, node.GetMetaCache(), paramtable.Get().ProxyCfg.MustUsePartitionKey.GetAsBool())
 
 	// Execute query
 	queryResult, _, err := node.query(ctx, qt, nil)
@@ -3673,7 +3592,7 @@ func (node *Proxy) handleIfSearchByPK(ctx context.Context, request *milvuspb.Sea
 	for i := 0; i < returnedPKCount; i++ {
 		pkOffset[pkItr(i)] = i
 	}
-	orderedFieldsData, err := pickFieldData(ids, pkOffset, queryResult.GetFieldsData(),
+	orderedFieldsData, err := PickFieldData(ids, pkOffset, queryResult.GetFieldsData(),
 		collectionInfo.Schema.CollectionSchema, collectionInfo.CollID)
 	if err != nil {
 		return nil, err
@@ -3807,11 +3726,23 @@ func (node *Proxy) Flush(ctx context.Context, request *milvuspb.FlushRequest) (*
 	return ft.result, nil
 }
 
+// ExecuteQuery implements taskmodel.QueryRunner. It is the single seam through
+// which requery operators in the extracted task packages execute a query task
+// on the composition root; the concrete task type is asserted here because
+// root imports the concrete packages, not the reverse.
+func (node *Proxy) ExecuteQuery(ctx context.Context, qt taskmodel.Task, sp trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
+	queryTask, ok := qt.(*queryTask)
+	if !ok {
+		return nil, segcore.StorageCost{}, merr.WrapErrServiceInternalMsg("query task of unexpected type: %T", qt)
+	}
+	return node.query(ctx, queryTask, sp)
+}
+
 // Query get the records by primary keys.
 func (node *Proxy) query(ctx context.Context, qt *queryTask, sp trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
-	request := qt.request
+	request := qt.Request()
 	method := "Query"
-	queryLabel := qt.getQueryLabel()
+	queryLabel := qt.GetMetricLabel()
 	if sp != nil {
 		sp.SetAttributes(attribute.String("queryLabel", queryLabel))
 	}
@@ -3887,7 +3818,7 @@ func (node *Proxy) query(ctx context.Context, qt *queryTask, sp trace.Span) (*mi
 		}, segcore.StorageCost{}, nil
 	}
 
-	if !qt.reQuery {
+	if !qt.ReQuery() {
 		span := tr.CtxRecord(ctx, "wait query result")
 		metrics.ProxyWaitForSearchResultLatency.WithLabelValues(
 			strconv.FormatInt(paramtable.GetNodeID(), 10),
@@ -3913,33 +3844,20 @@ func (node *Proxy) query(ctx context.Context, qt *queryTask, sp trace.Span) (*mi
 	}
 
 	succeeded = true
-	return qt.result, qt.storageCost, nil
+	return qt.Result(), qt.StorageCost(), nil
 }
 
 // Query get the records by primary keys.
 func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
-	qt := &queryTask{
-		baseTask: baseTask{
-			MetaCache: node.GetMetaCache(),
-		},
-		ctx:       ctx,
-		Condition: NewTaskCondition(ctx),
-		RetrieveRequest: &internalpb.RetrieveRequest{
-			Base: commonpbutil.NewMsgBase(
-				commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
-				commonpbutil.WithSourceID(paramtable.GetNodeID()),
-			),
-			ReqID:            paramtable.GetNodeID(),
-			ConsistencyLevel: request.ConsistencyLevel,
-			QueryLabel:       metrics.QueryLabel,
-		},
-		request:             request,
-		mixCoord:            node.mixCoord,
-		lb:                  node.lbPolicy,
-		shardclientMgr:      node.shardMgr,
-		mustUsePartitionKey: Params.ProxyCfg.MustUsePartitionKey.GetAsBool(),
-		chMgr:               node.chMgr,
-	}
+	qt := NewQueryTask(ctx, node, request, nil, &internalpb.RetrieveRequest{
+		Base: commonpbutil.NewMsgBase(
+			commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
+			commonpbutil.WithSourceID(paramtable.GetNodeID()),
+		),
+		ReqID:            paramtable.GetNodeID(),
+		ConsistencyLevel: request.ConsistencyLevel,
+		QueryLabel:       metrics.QueryLabel,
+	}, node.GetMetaCache(), paramtable.Get().ProxyCfg.MustUsePartitionKey.GetAsBool())
 
 	subLabel := GetCollectionRateSubLabel(request)
 	metrics.GetStats(ctx).
@@ -3977,14 +3895,14 @@ func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (*
 			metrics.QueryLabel,
 			request.DbName,
 			request.CollectionName,
-		).Add(float64(qt.storageCost.ScannedRemoteBytes) / 1024 / 1024)
+		).Add(float64(qt.StorageCost().ScannedRemoteBytes) / 1024 / 1024)
 
 		metrics.ProxyScannedTotalMB.WithLabelValues(
 			strconv.FormatInt(paramtable.GetNodeID(), 10),
 			metrics.QueryLabel,
 			request.DbName,
 			request.CollectionName,
-		).Add(float64(qt.storageCost.ScannedTotalBytes) / 1024 / 1024)
+		).Add(float64(qt.StorageCost().ScannedTotalBytes) / 1024 / 1024)
 	}
 
 	if err != nil || !merr.Ok(res.Status) {
@@ -4000,14 +3918,14 @@ func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (*
 		hookutil.DatabaseKey:        request.DbName,
 		hookutil.UsernameKey:        username,
 		hookutil.ResultDataSizeKey:  proto.Size(res),
-		hookutil.RelatedDataSizeKey: qt.totalRelatedDataSize,
-		hookutil.RelatedCntKey:      qt.allQueryCnt,
+		hookutil.RelatedDataSizeKey: qt.TotalRelatedDataSize(),
+		hookutil.RelatedCntKey:      qt.AllQueryCnt(),
 	})
 	SetReportValue(res.Status, v)
 	SetStorageCost(res.Status, storageCost)
 	metrics.ProxyReportValue.WithLabelValues(nodeID, hookutil.OpTypeQuery, request.DbName, username).Add(float64(v))
 
-	if mlog.LevelEnabled(mlog.DebugLevel) && matchCountRule(request.OutputFields) {
+	if mlog.LevelEnabled(mlog.DebugLevel) && MatchCountRule(request.OutputFields) {
 		r, _ := protojson.Marshal(res)
 		mlog.Debug(ctx, "Count result", mlog.String("result", string(r)))
 	}

internal/proxy/impl_test.go

@@ -2597,7 +2597,7 @@ func TestHandleIfSearchByPK_PreservesNamespaceInInternalQuery(t *testing.T) {
 
 		var capturedNamespace *string
 		mockey.Mock((*Proxy).query).To(func(_ *Proxy, _ context.Context, qt *queryTask, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
-			capturedNamespace = qt.request.Namespace
+			capturedNamespace = qt.Request().Namespace
 			return &milvuspb.QueryResults{
 				Status: merr.Success(),
 				FieldsData: []*schemapb.FieldData{

internal/proxy/proxy.go

@@ -35,6 +35,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proxy/rls"
 	"github.com/milvus-io/milvus/internal/proxy/scheduler"
 	"github.com/milvus-io/milvus/internal/proxy/shardclient"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/adminauth"
 	"github.com/milvus-io/milvus/internal/util/dependency"
@@ -131,6 +132,13 @@ type Proxy struct {
 	slowQueries *expirable.LRU[Timestamp, *metricsinfo.SlowQuery]
 }
 
+// Compile-time assertions that *Proxy satisfies the task-model contracts the
+// extracted task packages consume through the composition root.
+var (
+	_ taskmodel.TaskNode    = (*Proxy)(nil)
+	_ taskmodel.QueryRunner = (*Proxy)(nil)
+)
+
 // NewProxy returns a Proxy struct.
 func NewProxy(ctx context.Context, factory dependency.Factory) (*Proxy, error) {
 	rand.Seed(time.Now().UnixNano())
@@ -183,6 +191,36 @@ func (node *Proxy) GetMetaCache() Cache {
 	return node.getMetaCache()
 }
 
+// MixCoord returns the MixCoord client consumed by concrete tasks through the
+// taskmodel.TaskNode contract.
+func (node *Proxy) MixCoord() types.MixCoordClient {
+	return node.mixCoord
+}
+
+// LBPolicy returns the replica load-balance policy consumed by concrete tasks
+// through the taskmodel.TaskNode contract.
+func (node *Proxy) LBPolicy() shardclient.LBPolicy {
+	return node.lbPolicy
+}
+
+// ShardMgr returns the shard client manager consumed by concrete tasks through
+// the taskmodel.TaskNode contract.
+func (node *Proxy) ShardMgr() shardclient.ShardClientMgr {
+	return node.shardMgr
+}
+
+// ChMgr returns the channel manager consumed by concrete tasks through the
+// taskmodel.TaskNode contract.
+func (node *Proxy) ChMgr() channelmgr.ChannelsMgr {
+	return node.chMgr
+}
+
+// TsoAllocator returns the timestamp allocator consumed by concrete tasks
+// through the taskmodel.TaskNode contract.
+func (node *Proxy) TsoAllocator() taskmodel.TsoAllocator {
+	return node.tsoAllocator
+}
+
 // IsDQLQueueFull reports whether the next DQL enqueue would be rejected with
 // TooManyRequests. The REST layer probes it (via interface assertion, like
 // GetMetaCache) to reject search/query before paying for body decoding.

internal/proxy/proxy_test.go

@@ -4970,3 +4970,9 @@ func TestProxy_IsDQLQueueFull(t *testing.T) {
 	}))
 	assert.True(t, node.IsDQLQueueFull())
 }
+
+func getMixCoordClient() *mocks.MockMixCoordClient {
+	mixc := &mocks.MockMixCoordClient{}
+	mixc.EXPECT().Close().Return(nil)
+	return mixc
+}

internal/proxy/task.go

@@ -19,7 +19,6 @@ package proxy
 import (
 	"context"
 	"fmt"
-	"math"
 	"strconv"
 
 	"github.com/cockroachdb/errors"
@@ -42,52 +41,11 @@ import (
 	"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/timestamptz"
 	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
 )
 
 const (
-	SumScorer string = "sum"
-	MaxScorer string = "max"
-	AvgScorer string = "avg"
-)
-
-const (
-	IgnoreGrowingKey     = "ignore_growing"
-	ReduceStopForBestKey = "reduce_stop_for_best"
-	IteratorField        = "iterator"
-	CollectionID         = "collection_id"
-	GroupByFieldKey      = "group_by_field"
-	GroupSizeKey         = "group_size"
-	StrictGroupSize      = "strict_group_size"
-	JSONPath             = "json_path"
-	JSONType             = "json_type"
-	StrictCastKey        = "strict_cast"
-	RankGroupScorer      = "rank_group_scorer"
-	AnnsFieldKey         = "anns_field"
-	AnalyzerKey          = "analyzer_name"
-	TopKKey              = "topk"
-	NQKey                = "nq"
-	MetricTypeKey        = common.MetricTypeKey
-	ParamsKey            = common.ParamsKey
-	ExprParamsKey        = "expr_params"
-	RoundDecimalKey      = "round_decimal"
-	OffsetKey            = "offset"
-	LimitKey             = "limit"
-	// key for timestamptz translation
-	TimefieldsKey = "time_fields"
-
-	SearchIterV2Key        = "search_iter_v2"
-	SearchIterBatchSizeKey = "search_iter_batch_size"
-	SearchIterLastBoundKey = "search_iter_last_bound"
-	SearchIterIdKey        = "search_iter_id"
-	QueryIterLastPKKey     = "query_iter_last_pk"
-	QueryIterLastOffsetKey = "query_iter_last_element_offset"
-	GroupByFieldsKey       = "group_by_fields"
-	OrderByFieldsKey       = "order_by_fields"
-	PipelineTraceKey       = "pipeline_trace"
-
 	InsertTaskName                = "InsertTask"
 	CreateCollectionTaskName      = "CreateCollectionTask"
 	DropCollectionTaskName        = "DropCollectionTask"
@@ -134,14 +92,6 @@ const (
 	AddFieldTaskName              = "AddFieldTaskName"
 	AddStructFieldTaskName        = "AddStructFieldTaskName"
 	AlterCollectionSchemaTaskName = "AlterCollectionSchemaTaskName"
-
-	// minFloat32 minimum float.
-	minFloat32 = -1 * float32(math.MaxFloat32)
-
-	RankTypeKey      = "strategy"
-	RRFParamsKey     = "k"
-	WeightsParamsKey = "weights"
-	NormScoreKey     = "norm_score"
 )
 
 func validateTextStorageV3Enabled(schema *schemapb.CollectionSchema) error {
@@ -4224,107 +4174,3 @@ func (t *RunAnalyzerTask) Execute(ctx context.Context) error {
 func (t *RunAnalyzerTask) PostExecute(ctx context.Context) error {
 	return nil
 }
-
-// git highlight after search
-type HighlightTask struct {
-	baseTask
-	Condition
-	*querypb.GetHighlightRequest
-	ctx            context.Context
-	collectionName string
-	collectionID   typeutil.UniqueID
-	dbName         string
-	lb             shardclient.LBPolicy
-
-	result *querypb.GetHighlightResponse
-}
-
-func (t *HighlightTask) TraceCtx() context.Context {
-	return t.ctx
-}
-
-func (t *HighlightTask) ID() UniqueID {
-	return t.Base.MsgID
-}
-
-func (t *HighlightTask) SetID(uid UniqueID) {
-	t.Base.MsgID = uid
-}
-
-func (t *HighlightTask) Name() string {
-	return HighlightTaskName
-}
-
-func (t *HighlightTask) Type() commonpb.MsgType {
-	return t.Base.MsgType
-}
-
-func (t *HighlightTask) BeginTs() Timestamp {
-	return t.Base.Timestamp
-}
-
-func (t *HighlightTask) EndTs() Timestamp {
-	return t.Base.Timestamp
-}
-
-func (t *HighlightTask) SetTs(ts Timestamp) {
-	t.Base.Timestamp = ts
-}
-
-func (t *HighlightTask) OnEnqueue() error {
-	if t.Base == nil {
-		t.Base = commonpbutil.NewMsgBase()
-	}
-	t.Base.MsgType = commonpb.MsgType_Undefined
-	t.Base.SourceID = paramtable.GetNodeID()
-	return nil
-}
-
-func (t *HighlightTask) PreExecute(ctx context.Context) error {
-	return nil
-}
-
-func (t *HighlightTask) getHighlightOnShardleader(ctx context.Context, nodeID int64, qn types.QueryNodeClient, channel string) error {
-	ctx = retry.WithMaxAttemptsContext(ctx, 1)
-	t.Channel = channel
-	resp, err := qn.GetHighlight(ctx, t.GetHighlightRequest)
-	if err != nil {
-		return err
-	}
-
-	if err := merr.Error(resp.GetStatus()); err != nil {
-		return err
-	}
-	t.result = resp
-	return nil
-}
-
-func (t *HighlightTask) Execute(ctx context.Context) error {
-	err := t.lb.ExecuteOneChannel(ctx, shardclient.CollectionWorkLoad{
-		Db:             t.dbName,
-		CollectionName: t.collectionName,
-		CollectionID:   t.collectionID,
-		Nq:             int64(len(t.GetTopks()) * len(t.GetTasks())),
-		Exec:           t.getHighlightOnShardleader,
-	})
-
-	return err
-}
-
-func (t *HighlightTask) PostExecute(ctx context.Context) error {
-	return nil
-}
-
-// isIgnoreGrowing is used to check if the request should ignore growing
-func isIgnoreGrowing(params []*commonpb.KeyValuePair) (bool, error) {
-	for _, kv := range params {
-		if kv.GetKey() == IgnoreGrowingKey {
-			ignoreGrowing, err := strconv.ParseBool(kv.GetValue())
-			if err != nil {
-				return false, merr.WrapErrParameterInvalidMsg("parse ignore growing field failed")
-			}
-			return ignoreGrowing, nil
-		}
-	}
-	return false, nil
-}

internal/proxy/task_delete.go

@@ -374,7 +374,7 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
 	dr.plan, err = planparserv2.CreateRetrievePlanArgs(dr.schema.SchemaHelper, dr.req.GetExpr(), dr.req.GetExprTemplateValues(), visitorArgs)
 	if err != nil {
 		metrics.ProxyParseExpressionLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), "delete", metrics.FailLabel).Observe(float64(time.Since(start).Microseconds()) / 1000.0)
-		return merr.WrapErrAsInputError(wrapPlanCreationError(err, "failed to create delete plan"))
+		return merr.WrapErrAsInputError(WrapPlanCreationError(err, "failed to create delete plan"))
 	}
 	metrics.ProxyParseExpressionLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), "delete", metrics.SuccessLabel).Observe(float64(time.Since(start).Microseconds()) / 1000.0)
 
@@ -403,7 +403,7 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
 		if err != nil {
 			return err
 		}
-		dr.partitionIDs, err = getPartitionIDs(ctx, dr.GetMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames)
+		dr.partitionIDs, err = GetPartitionIDs(ctx, dr.GetMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames)
 		if err != nil {
 			return err
 		}
@@ -420,7 +420,7 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
 		if err != nil {
 			return err
 		}
-		dr.partitionIDs, err = getPartitionIDs(ctx, dr.GetMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames)
+		dr.partitionIDs, err = GetPartitionIDs(ctx, dr.GetMetaCache(), dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames)
 		if err != nil {
 			return err
 		}
@@ -637,7 +637,7 @@ func (dr *deleteRunner) complexDelete(ctx context.Context, plan *planpb.PlanNode
 	// Budget and marshal once before LB fan-out. A roaring plan can be tens of
 	// MiB; marshaling inside each concurrent shard callback would multiply that
 	// transient buffer by shard count and repeat it on retries.
-	serializedPlan, _, err := marshalPlanWithMembershipFilterSizeLimit(plan, 0)
+	serializedPlan, _, err := MarshalPlanWithMembershipFilterSizeLimit(plan, 0)
 	if err != nil {
 		return err
 	}

internal/proxy/task_delete_test.go

@@ -1247,3 +1247,21 @@ func TestDeleteRunner_Run(t *testing.T) {
 		assert.Equal(t, int64(3), dr.result.DeleteCnt)
 	})
 }
+
+func namespaceEnabledSchema(fields ...*schemapb.FieldSchema) *schemapb.CollectionSchema {
+	return &schemapb.CollectionSchema{
+		Name: "test_collection",
+		Fields: append(fields, &schemapb.FieldSchema{
+			FieldID:        999,
+			Name:           common.NamespaceFieldName,
+			IsPartitionKey: true,
+			DataType:       schemapb.DataType_VarChar,
+		}),
+		EnableNamespace: true,
+	}
+}
+
+func expectedNamespacePartitionID(namespace string, partitionNames []string, partitionIDs map[string]int64) int64 {
+	idx := typeutil.HashString2Uint32(namespace) % uint32(len(partitionNames))
+	return partitionIDs[partitionNames[idx]]
+}

internal/proxy/task_insert.go

@@ -280,7 +280,7 @@ func (it *insertTask) PreExecute(ctx context.Context) error {
 		log.Info(ctx, "fill field properties failed", mlog.Err(err))
 		return err
 	}
-	err = normalizeFP32ToFP16BF16VectorFieldData(it.insertMsg.GetFieldsData(), schema)
+	err = NormalizeFP32ToFP16BF16VectorFieldData(it.insertMsg.GetFieldsData(), schema)
 	if err != nil {
 		log.Info(ctx, "normalize fp32 to fp16/bf16 vector field data failed", mlog.Err(err))
 		return err

internal/proxy/task_test.go

@@ -43,7 +43,6 @@ import (
 	"github.com/milvus-io/milvus/internal/json"
 	"github.com/milvus-io/milvus/internal/mocks"
 	"github.com/milvus-io/milvus/internal/proxy/channelmgr"
-	"github.com/milvus-io/milvus/internal/proxy/shardclient"
 	"github.com/milvus-io/milvus/internal/querycoordv2/meta"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/function/embedding"
@@ -60,7 +59,6 @@ import (
 	"github.com/milvus-io/milvus/pkg/v3/util/metric"
 	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/v3/util/testutils"
-	"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
 	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
 	"github.com/milvus-io/milvus/pkg/v3/util/uniquegenerator"
 )
@@ -691,273 +689,6 @@ func constructSearchRequest(
 	}
 }
 
-func TestTranslateOutputFields(t *testing.T) {
-	const (
-		idFieldName                = "id"
-		tsFieldName                = "timestamp"
-		floatVectorFieldName       = "float_vector"
-		binaryVectorFieldName      = "binary_vector"
-		float16VectorFieldName     = "float16_vector"
-		bfloat16VectorFieldName    = "bfloat16_vector"
-		sparseFloatVectorFieldName = "sparse_float_vector"
-	)
-	var outputFields []string
-	var userOutputFields []string
-	var userDynamicFields []string
-	var requestedPK bool
-	var err error
-
-	collSchema := &schemapb.CollectionSchema{
-		Name:        "TestTranslateOutputFields",
-		Description: "TestTranslateOutputFields",
-		AutoID:      false,
-		Fields: []*schemapb.FieldSchema{
-			{Name: idFieldName, FieldID: 0, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
-			{Name: tsFieldName, FieldID: 1, DataType: schemapb.DataType_Int64},
-			{Name: floatVectorFieldName, FieldID: 100, DataType: schemapb.DataType_FloatVector},
-			{Name: binaryVectorFieldName, FieldID: 101, DataType: schemapb.DataType_BinaryVector},
-			{Name: float16VectorFieldName, FieldID: 102, DataType: schemapb.DataType_Float16Vector},
-			{Name: bfloat16VectorFieldName, FieldID: 103, DataType: schemapb.DataType_BFloat16Vector},
-			{Name: sparseFloatVectorFieldName, FieldID: 104, DataType: schemapb.DataType_SparseFloatVector, IsFunctionOutput: true},
-		},
-		Functions: []*schemapb.FunctionSchema{
-			{
-				Name:             "bm25",
-				Type:             schemapb.FunctionType_BM25,
-				OutputFieldNames: []string{sparseFloatVectorFieldName},
-				// omit other fields for brevity
-			},
-		},
-	}
-	schema := mustNewSchemaInfo(collSchema)
-
-	// Test empty output fields
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{}, schema, false)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{}, outputFields)
-	assert.ElementsMatch(t, []string{}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.False(t, requestedPK)
-
-	// Test single field
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName}, schema, false)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{idFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{idFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	// Test multiple fields
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName}, schema, false)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{idFieldName, tsFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{idFieldName, tsFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	// Test with vector field
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName, floatVectorFieldName}, schema, false)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	// Test without id field
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{tsFieldName, floatVectorFieldName}, schema, false)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.False(t, requestedPK)
-
-	// Test wildcard - should not include function output fields
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*"}, schema, false)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	// Test wildcard with spaces
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{" * "}, schema, false)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	_, _, _, _, _, err = translateOutputFields([]string{"*", sparseFloatVectorFieldName}, schema, false)
-	assert.Error(t, err)
-	_, _, _, _, _, err = translateOutputFields([]string{sparseFloatVectorFieldName}, schema, false)
-	assert.Error(t, err)
-	_, _, _, _, _, err = translateOutputFields([]string{sparseFloatVectorFieldName}, schema, true)
-	assert.Error(t, err)
-
-	// Test with removePkField=true
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{}, schema, true)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{}, outputFields)
-	assert.ElementsMatch(t, []string{}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.False(t, requestedPK)
-
-	// if removePkField is true, pk field should be removed from output fields
-
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*"}, schema, true)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName}, schema, true)
-	assert.Equal(t, nil, err)
-	assert.ElementsMatch(t, []string{tsFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{tsFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName, floatVectorFieldName}, schema, true)
-	assert.Equal(t, nil, err)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*"}, schema, true)
-	assert.Equal(t, nil, err)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*", tsFieldName}, schema, true)
-	assert.Equal(t, nil, err)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*", floatVectorFieldName}, schema, true)
-	assert.Equal(t, nil, err)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
-	assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.True(t, requestedPK)
-
-	// Test non-existent field, dynamic field not enabled
-	_, _, _, _, _, err = translateOutputFields([]string{"A"}, schema, true)
-	assert.Error(t, err)
-
-	t.Run("enable dynamic schema", func(t *testing.T) {
-		collSchema := &schemapb.CollectionSchema{
-			Name:               "TestTranslateOutputFields",
-			Description:        "TestTranslateOutputFields",
-			AutoID:             false,
-			EnableDynamicField: true,
-			Fields: []*schemapb.FieldSchema{
-				{Name: idFieldName, FieldID: 1, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
-				{Name: tsFieldName, FieldID: 2, DataType: schemapb.DataType_Int64},
-				{Name: floatVectorFieldName, FieldID: 100, DataType: schemapb.DataType_FloatVector},
-				{Name: binaryVectorFieldName, FieldID: 101, DataType: schemapb.DataType_BinaryVector},
-				{Name: common.MetaFieldName, FieldID: 102, DataType: schemapb.DataType_JSON, IsDynamic: true},
-			},
-		}
-		schema := mustNewSchemaInfo(collSchema)
-
-		outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"A", idFieldName}, schema, true)
-		assert.NoError(t, err)
-		assert.ElementsMatch(t, []string{common.MetaFieldName}, outputFields)
-		assert.ElementsMatch(t, []string{"A"}, userOutputFields)
-		assert.ElementsMatch(t, []string{"A"}, userDynamicFields)
-		assert.True(t, requestedPK)
-
-		outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"$meta[\"A\"]", idFieldName}, schema, true)
-		assert.NoError(t, err)
-		assert.ElementsMatch(t, []string{common.MetaFieldName}, outputFields)
-		assert.ElementsMatch(t, []string{"$meta[\"A\"]"}, userOutputFields)
-		assert.ElementsMatch(t, []string{"A"}, userDynamicFields)
-		assert.True(t, requestedPK)
-
-		// Test invalid dynamic field expressions
-		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, `$meta["A"]["B"]`}, schema, true)
-		assert.Error(t, err)
-
-		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "$meta[\"\"]"}, schema, true)
-		assert.Error(t, err)
-
-		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "$meta[]"}, schema, true)
-		assert.Error(t, err)
-
-		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "$meta["}, schema, true)
-		assert.Error(t, err)
-
-		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "[]"}, schema, true)
-		assert.Error(t, err)
-
-		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "A > 1"}, schema, true)
-		assert.Error(t, err)
-
-		_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, ""}, schema, true)
-		assert.Error(t, err)
-	})
-}
-
-func TestTranslateOutputFields_StructArrayField(t *testing.T) {
-	const (
-		idFieldName                = "id"
-		tsFieldName                = "timestamp"
-		floatVectorFieldName       = "float_vector"
-		binaryVectorFieldName      = "binary_vector"
-		float16VectorFieldName     = "float16_vector"
-		bfloat16VectorFieldName    = "bfloat16_vector"
-		sparseFloatVectorFieldName = "sparse_float_vector"
-	)
-	var outputFields []string
-	var userOutputFields []string
-	var userDynamicFields []string
-	var requestedPK bool
-	var err error
-
-	collSchema := &schemapb.CollectionSchema{
-		Name:        "TestTranslateOutputFields",
-		Description: "TestTranslateOutputFields",
-		AutoID:      false,
-		Fields: []*schemapb.FieldSchema{
-			{Name: idFieldName, FieldID: 0, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
-			{Name: tsFieldName, FieldID: 1, DataType: schemapb.DataType_Int64},
-			{Name: floatVectorFieldName, FieldID: 100, DataType: schemapb.DataType_FloatVector},
-		},
-		StructArrayFields: []*schemapb.StructArrayFieldSchema{
-			{
-				FieldID: 101,
-				Name:    "struct_array_field",
-				Fields: []*schemapb.FieldSchema{
-					{Name: "sub_field", FieldID: 102, DataType: schemapb.DataType_Array, ElementType: schemapb.DataType_Int64},
-					{Name: "sub_vector_field", FieldID: 103, DataType: schemapb.DataType_ArrayOfVector, ElementType: schemapb.DataType_FloatVector},
-				},
-			},
-		},
-	}
-	schema := mustNewSchemaInfo(collSchema)
-
-	// Test struct array field
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"sub_vector_field"}, schema, false)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{"sub_vector_field"}, outputFields)
-	assert.ElementsMatch(t, []string{"sub_vector_field"}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.False(t, requestedPK)
-
-	outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"struct_array_field", "sub_field"}, schema, false)
-	assert.NoError(t, err)
-	assert.ElementsMatch(t, []string{"sub_vector_field", "sub_field"}, outputFields)
-	assert.ElementsMatch(t, []string{"sub_vector_field", "sub_field"}, userOutputFields)
-	assert.ElementsMatch(t, []string{}, userDynamicFields)
-	assert.False(t, requestedPK)
-}
-
 func TestAddFieldTask(t *testing.T) {
 	rc := NewMixCoordMock()
 	ctx := context.Background()
@@ -5033,47 +4764,6 @@ func TestPartitionKey(t *testing.T) {
 		assert.NoError(t, dt.Execute(ctx))
 		assert.NoError(t, dt.PostExecute(ctx))
 	})
-
-	t.Run("search", func(t *testing.T) {
-		searchTask := &searchTask{
-			baseTask: baseTask{MetaCache: cache},
-			ctx:      ctx,
-			SearchRequest: &internalpb.SearchRequest{
-				Base: &commonpb.MsgBase{},
-			},
-			request: &milvuspb.SearchRequest{
-				CollectionName: collectionName,
-				Nq:             1,
-			},
-			mixCoord: qc,
-			tr:       timerecord.NewTimeRecorder("test-search"),
-		}
-
-		// don't support specify partition name if use partition key
-		searchTask.request.PartitionNames = partitionNames
-		err = searchTask.PreExecute(ctx)
-		assert.Error(t, err)
-	})
-
-	t.Run("query", func(t *testing.T) {
-		queryTask := &queryTask{
-			baseTask: baseTask{MetaCache: cache},
-			ctx:      ctx,
-			RetrieveRequest: &internalpb.RetrieveRequest{
-				QueryLabel: "query",
-				Base:       &commonpb.MsgBase{},
-			},
-			request: &milvuspb.QueryRequest{
-				CollectionName: collectionName,
-			},
-			mixCoord: qc,
-		}
-
-		// don't support specify partition name if use partition key
-		queryTask.request.PartitionNames = partitionNames
-		err = queryTask.PreExecute(ctx)
-		assert.Error(t, err)
-	})
 }
 
 func TestDefaultPartition(t *testing.T) {
@@ -7185,240 +6875,6 @@ func TestAlterCollection_AllowInsertAutoID_AutoIDFalse(t *testing.T) {
 	assert.Equal(t, merr.Code(merr.ErrParameterInvalid), merr.Code(err))
 }
 
-func TestHighlightTask(t *testing.T) {
-	paramtable.Init()
-	ctx := context.Background()
-
-	t.Run("basic methods", func(t *testing.T) {
-		task := &HighlightTask{
-			ctx: ctx,
-			GetHighlightRequest: &querypb.GetHighlightRequest{
-				Base: commonpbutil.NewMsgBase(),
-			},
-			Condition: NewTaskCondition(ctx),
-		}
-
-		t.Run("traceCtx", func(t *testing.T) {
-			traceCtx := task.TraceCtx()
-			assert.NotNil(t, traceCtx)
-			assert.Equal(t, ctx, traceCtx)
-		})
-
-		t.Run("id", func(t *testing.T) {
-			id := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
-			task.SetID(id)
-			assert.Equal(t, id, task.ID())
-		})
-
-		t.Run("name", func(t *testing.T) {
-			assert.Equal(t, HighlightTaskName, task.Name())
-		})
-
-		t.Run("type", func(t *testing.T) {
-			assert.Equal(t, commonpb.MsgType_Undefined, task.Type())
-		})
-
-		t.Run("ts", func(t *testing.T) {
-			ts := Timestamp(time.Now().UnixNano())
-			task.SetTs(ts)
-			assert.Equal(t, ts, task.BeginTs())
-			assert.Equal(t, ts, task.EndTs())
-		})
-	})
-
-	t.Run("OnEnqueue", func(t *testing.T) {
-		task := &HighlightTask{
-			ctx:                 ctx,
-			GetHighlightRequest: &querypb.GetHighlightRequest{},
-			Condition:           NewTaskCondition(ctx),
-		}
-
-		err := task.OnEnqueue()
-		assert.NoError(t, err)
-		assert.NotNil(t, task.Base)
-		assert.Equal(t, commonpb.MsgType_Undefined, task.Base.MsgType)
-		assert.Equal(t, paramtable.GetNodeID(), task.Base.SourceID)
-	})
-
-	t.Run("PreExecute", func(t *testing.T) {
-		task := &HighlightTask{
-			ctx: ctx,
-			GetHighlightRequest: &querypb.GetHighlightRequest{
-				Base: commonpbutil.NewMsgBase(),
-			},
-			Condition: NewTaskCondition(ctx),
-		}
-
-		err := task.PreExecute(ctx)
-		assert.NoError(t, err)
-	})
-
-	t.Run("getHighlightOnShardleader success", func(t *testing.T) {
-		task := &HighlightTask{
-			ctx: ctx,
-			GetHighlightRequest: &querypb.GetHighlightRequest{
-				Base:  commonpbutil.NewMsgBase(),
-				Topks: []int64{10},
-				Tasks: []*querypb.HighlightTask{
-					{},
-				},
-			},
-			Condition: NewTaskCondition(ctx),
-		}
-
-		mockQN := mocks.NewMockQueryNodeClient(t)
-		expectedResp := &querypb.GetHighlightResponse{
-			Status: merr.Success(),
-			Results: []*querypb.HighlightResult{
-				{
-					Fragments: []*querypb.HighlightFragment{
-						{
-							StartOffset: 0,
-							EndOffset:   10,
-							Offsets:     []int64{0, 5, 5, 10},
-						},
-					},
-				},
-			},
-		}
-
-		mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(expectedResp, nil)
-
-		err := task.getHighlightOnShardleader(ctx, 1, mockQN, "test_channel")
-		assert.NoError(t, err)
-		assert.NotNil(t, task.result)
-		assert.Equal(t, expectedResp, task.result)
-		assert.Equal(t, "test_channel", task.Channel)
-	})
-
-	t.Run("getHighlightOnShardleader rpc error", func(t *testing.T) {
-		task := &HighlightTask{
-			ctx: ctx,
-			GetHighlightRequest: &querypb.GetHighlightRequest{
-				Base: commonpbutil.NewMsgBase(),
-			},
-			Condition: NewTaskCondition(ctx),
-		}
-
-		mockQN := mocks.NewMockQueryNodeClient(t)
-		mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(nil, errors.New("rpc error"))
-
-		err := task.getHighlightOnShardleader(ctx, 1, mockQN, "test_channel")
-		assert.Error(t, err)
-		assert.Contains(t, err.Error(), "rpc error")
-	})
-
-	t.Run("getHighlightOnShardleader status error", func(t *testing.T) {
-		task := &HighlightTask{
-			ctx: ctx,
-			GetHighlightRequest: &querypb.GetHighlightRequest{
-				Base: commonpbutil.NewMsgBase(),
-			},
-			Condition: NewTaskCondition(ctx),
-		}
-
-		mockQN := mocks.NewMockQueryNodeClient(t)
-		expectedResp := &querypb.GetHighlightResponse{
-			Status: merr.Status(errors.New("status error")),
-		}
-
-		mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(expectedResp, nil)
-
-		err := task.getHighlightOnShardleader(ctx, 1, mockQN, "test_channel")
-		assert.Error(t, err)
-	})
-
-	t.Run("Execute success", func(t *testing.T) {
-		task := &HighlightTask{
-			ctx:            ctx,
-			collectionName: "test_collection",
-			collectionID:   100,
-			dbName:         "default",
-			GetHighlightRequest: &querypb.GetHighlightRequest{
-				Base:  commonpbutil.NewMsgBase(),
-				Topks: []int64{10, 20},
-				Tasks: []*querypb.HighlightTask{
-					{},
-					{},
-				},
-			},
-			Condition: NewTaskCondition(ctx),
-		}
-
-		mockLB := shardclient.NewMockLBPolicy(t)
-		mockQN := mocks.NewMockQueryNodeClient(t)
-
-		expectedResp := &querypb.GetHighlightResponse{
-			Status: merr.Success(),
-			Results: []*querypb.HighlightResult{
-				{
-					Fragments: []*querypb.HighlightFragment{
-						{
-							StartOffset: 0,
-							EndOffset:   10,
-							Offsets:     []int64{0, 5, 5, 10},
-						},
-					},
-				},
-			},
-		}
-
-		mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(expectedResp, nil)
-
-		mockLB.EXPECT().ExecuteOneChannel(mock.Anything, mock.Anything).Run(
-			func(ctx context.Context, workload shardclient.CollectionWorkLoad) {
-				assert.Equal(t, "default", workload.Db)
-				assert.Equal(t, "test_collection", workload.CollectionName)
-				assert.Equal(t, int64(100), workload.CollectionID)
-				assert.Equal(t, int64(4), workload.Nq) // len(topks) * len(tasks) = 2 * 2 = 4
-				err := workload.Exec(ctx, 1, mockQN, "test_channel")
-				assert.NoError(t, err)
-			},
-		).Return(nil)
-
-		task.lb = mockLB
-
-		err := task.Execute(ctx)
-		assert.NoError(t, err)
-		assert.NotNil(t, task.result)
-	})
-
-	t.Run("Execute lb error", func(t *testing.T) {
-		task := &HighlightTask{
-			ctx:            ctx,
-			collectionName: "test_collection",
-			collectionID:   100,
-			dbName:         "default",
-			GetHighlightRequest: &querypb.GetHighlightRequest{
-				Base: commonpbutil.NewMsgBase(),
-			},
-			Condition: NewTaskCondition(ctx),
-		}
-
-		mockLB := shardclient.NewMockLBPolicy(t)
-		mockLB.EXPECT().ExecuteOneChannel(mock.Anything, mock.Anything).Return(errors.New("lb error"))
-
-		task.lb = mockLB
-
-		err := task.Execute(ctx)
-		assert.Error(t, err)
-		assert.Contains(t, err.Error(), "lb error")
-	})
-
-	t.Run("PostExecute", func(t *testing.T) {
-		task := &HighlightTask{
-			ctx: ctx,
-			GetHighlightRequest: &querypb.GetHighlightRequest{
-				Base: commonpbutil.NewMsgBase(),
-			},
-			Condition: NewTaskCondition(ctx),
-		}
-
-		err := task.PostExecute(ctx)
-		assert.NoError(t, err)
-	})
-}
-
 func TestValidateCollectionTTL(t *testing.T) {
 	cases := []struct {
 		name      string

internal/proxy/task_upsert.go

@@ -32,7 +32,7 @@ import (
 	"github.com/milvus-io/milvus/internal/parser/planparserv2"
 	"github.com/milvus-io/milvus/internal/proxy/channelmgr"
 	"github.com/milvus-io/milvus/internal/proxy/fieldvalidator"
-	"github.com/milvus-io/milvus/internal/types"
+	"github.com/milvus-io/milvus/internal/proxy/taskmodel"
 	"github.com/milvus-io/milvus/internal/util/segcore"
 	"github.com/milvus-io/milvus/pkg/v3/common"
 	"github.com/milvus-io/milvus/pkg/v3/metrics"
@@ -75,7 +75,7 @@ type upsertTask struct {
 	schemaVersion   int32
 
 	// write after read, generate write part by queryPreExecute
-	node types.ProxyComponent
+	node taskmodel.TaskNode
 
 	deletePKs       *schemapb.IDs
 	insertFieldData []*schemapb.FieldData
@@ -229,36 +229,27 @@ func retrieveByPKs(ctx context.Context, t *upsertTask, ids *schemapb.IDs, output
 
 	plan := planparserv2.CreateRequeryPlan(pkField, ids)
 	plan.Namespace = namespaceForPlan(t.schema.CollectionSchema, t.req.Namespace)
-	qt := &queryTask{
-		baseTask: baseTask{
-			MetaCache: t.GetMetaCache(),
-		},
-		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,
-			MvccTimestamp:    0,
-			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,
-		actualChannelsMvcc: channelReadTs,
-	}
+	qt := NewQueryTask(t.ctx, t.node, queryReq, plan, &internalpb.RetrieveRequest{
+		Base: commonpbutil.NewMsgBase(
+			commonpbutil.WithMsgType(commonpb.MsgType_Retrieve),
+			commonpbutil.WithSourceID(paramtable.GetNodeID()),
+		),
+		ReqID:            paramtable.GetNodeID(),
+		PartitionIDs:     partitionIDs,
+		ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
+		MvccTimestamp:    0,
+		QueryLabel:       metrics.UpsertQueryLabel,
+	}, t.GetMetaCache(), false)
+	qt.SetActualChannelsMvcc(channelReadTs)
 	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)
+	queryRunner, ok := t.node.(taskmodel.QueryRunner)
+	if !ok {
+		return nil, segcore.StorageCost{}, merr.WrapErrServiceInternalMsg("host node does not implement QueryRunner")
+	}
+	queryResult, storageCost, err := queryRunner.ExecuteQuery(ctx, qt, sp)
 	if err := merr.CheckRPCCall(queryResult, err); err != nil {
 		return nil, storageCost, err
 	}
@@ -1744,7 +1735,7 @@ func (it *upsertTask) insertPreExecute(ctx context.Context) error {
 		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)
+	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)

internal/proxy/task_upsert_test.go

@@ -1869,13 +1869,13 @@ func TestPartialUpdateRetriesAfterCASConflict(t *testing.T) {
 	// The retry must not allocate a separate fixed timestamp before querying.
 	task.node.(*Proxy).tsoAllocator = nil
 	queryPatch := mockey.Mock((*Proxy).query).To(func(_ *Proxy, _ context.Context, qt *queryTask, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
-		require.Equal(t, commonpb.ConsistencyLevel_Strong, qt.request.GetConsistencyLevel())
-		require.Zero(t, qt.request.GetGuaranteeTimestamp())
+		require.Equal(t, commonpb.ConsistencyLevel_Strong, qt.Request().GetConsistencyLevel())
+		require.Zero(t, qt.Request().GetGuaranteeTimestamp())
 		require.Zero(t, qt.GetMvccTimestamp())
 		require.False(t, qt.CanSkipAllocTimestamp())
 		for channel, meta := range task.partialUpdateCASGroups {
 			require.EqualValues(t, 10, meta.GetObservedPchannelTerm())
-			qt.actualChannelsMvcc.Insert(channel, 2000)
+			qt.RecordChannelMvcc(channel, 2000)
 		}
 		return &milvuspb.QueryResults{Status: merr.Success()}, segcore.StorageCost{}, nil
 	}).Build()
@@ -2466,20 +2466,20 @@ func TestRetrieveByPKs_Success(t *testing.T) {
 			}).Build()
 			defer partition.UnPatch()
 			query := mockey.Mock((*Proxy).query).To(func(_ *Proxy, _ context.Context, qt *queryTask, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
-				require.Equal(t, commonpb.ConsistencyLevel_Strong, qt.request.GetConsistencyLevel())
+				require.Equal(t, commonpb.ConsistencyLevel_Strong, qt.Request().GetConsistencyLevel())
 				require.Equal(t, commonpb.ConsistencyLevel_Strong, qt.GetConsistencyLevel())
-				require.Zero(t, qt.request.GetBase().GetTimestamp())
-				require.Zero(t, qt.request.GetGuaranteeTimestamp())
+				require.Zero(t, qt.Request().GetBase().GetTimestamp())
+				require.Zero(t, qt.Request().GetGuaranteeTimestamp())
 				require.Zero(t, qt.GetMvccTimestamp())
 				if partitionKeyMode {
 					require.Equal(t, []int64{common.AllPartitionsID}, qt.GetPartitionIDs())
-					require.Empty(t, qt.request.GetPartitionNames())
+					require.Empty(t, qt.Request().GetPartitionNames())
 				} else {
 					require.Equal(t, []int64{1002}, qt.GetPartitionIDs())
-					require.Equal(t, []string{"_default"}, qt.request.GetPartitionNames())
+					require.Equal(t, []string{"_default"}, qt.Request().GetPartitionNames())
 				}
 				for channel := range task.partialUpdateCASGroups {
-					qt.actualChannelsMvcc.Insert(channel, 70)
+					qt.RecordChannelMvcc(channel, 70)
 				}
 				return &milvuspb.QueryResults{
 					Status: merr.Success(), FieldsData: []*schemapb.FieldData{partialUpdateCASPKFieldData([]int64{10, 20})},
@@ -2529,16 +2529,16 @@ func TestRetrieveByPKsUsesStrongQueryForFullAutoID(t *testing.T) {
 
 	require.NoError(t, err)
 	require.NotNil(t, captured)
-	require.Equal(t, commonpb.ConsistencyLevel_Strong, captured.request.GetConsistencyLevel())
+	require.Equal(t, commonpb.ConsistencyLevel_Strong, captured.Request().GetConsistencyLevel())
 	require.Equal(t, commonpb.ConsistencyLevel_Strong, captured.GetConsistencyLevel())
-	require.False(t, captured.request.GetUseDefaultConsistency())
-	require.Zero(t, captured.request.GetGuaranteeTimestamp())
+	require.False(t, captured.Request().GetUseDefaultConsistency())
+	require.Zero(t, captured.Request().GetGuaranteeTimestamp())
 	require.Zero(t, captured.GetGuaranteeTimestamp())
-	require.Zero(t, captured.request.GetBase().GetTimestamp())
+	require.Zero(t, captured.Request().GetBase().GetTimestamp())
 	require.False(t, captured.CanSkipAllocTimestamp(), "the Query must obtain its own Strong-read timestamp")
-	require.Equal(t, []string{"id"}, captured.request.GetOutputFields())
+	require.Equal(t, []string{"id"}, captured.Request().GetOutputFields())
 	require.Zero(t, captured.GetMvccTimestamp())
-	require.Nil(t, captured.actualChannelsMvcc)
+	require.False(t, captured.HasRecordedChannelMvcc())
 	require.Empty(t, task.partialUpdateCASGroups)
 	require.Equal(t, beginTS, task.BeginTs())
 }
@@ -3281,12 +3281,12 @@ func TestRetrieveByPKsStrongReadBindsActualSnapshots(t *testing.T) {
 	require.NoError(t, task.preparePartialUpdateCASGroups(context.Background()))
 
 	patch := mockey.Mock((*Proxy).query).To(func(_ *Proxy, ctx context.Context, qt *queryTask, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
-		require.Equal(t, commonpb.ConsistencyLevel_Strong, qt.request.GetConsistencyLevel())
-		require.Zero(t, qt.request.GetGuaranteeTimestamp())
+		require.Equal(t, commonpb.ConsistencyLevel_Strong, qt.Request().GetConsistencyLevel())
+		require.Zero(t, qt.Request().GetGuaranteeTimestamp())
 		require.Zero(t, qt.GetMvccTimestamp())
 		require.False(t, qt.CanSkipAllocTimestamp())
 		for i, ch := range partialUpdateCASTestVChannels {
-			qt.actualChannelsMvcc.Insert(ch, uint64(70+i*20))
+			qt.RecordChannelMvcc(ch, uint64(70+i*20))
 		}
 		return &milvuspb.QueryResults{Status: merr.Success()}, segcore.StorageCost{}, nil
 	}).Build()
@@ -3320,17 +3320,17 @@ func TestRetrieveByPKsStrongReadRejectsMissingProofBeforeMerge(t *testing.T) {
 			patch := mockey.Mock((*Proxy).query).To(func(_ *Proxy, ctx context.Context, qt *queryTask, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
 				calls++
 				cost := segcore.StorageCost{ScannedRemoteBytes: 3, ScannedTotalBytes: 5}
-				require.NotNil(t, qt.actualChannelsMvcc)
+				require.True(t, qt.HasRecordedChannelMvcc())
 				for ch := range task.partialUpdateCASGroups {
 					if scenario == "zero" {
-						qt.actualChannelsMvcc.Insert(ch, 0)
+						qt.RecordChannelMvcc(ch, 0)
 					}
 					if scenario == "query_error" {
-						qt.actualChannelsMvcc.Insert(ch, 70)
+						qt.RecordChannelMvcc(ch, 70)
 					}
 				}
 				if scenario == "missing" {
-					qt.actualChannelsMvcc.Insert("unrelated-channel", 70)
+					qt.RecordChannelMvcc("unrelated-channel", 70)
 				}
 				if scenario == "query_error" {
 					return nil, cost, expectedErr
@@ -4235,7 +4235,7 @@ func TestPartialUpdateAutoIDDestinationSnapshotFromQueryRPC(t *testing.T) {
 			patch((*MetaCache).GetCollectionSchema, task.schema, nil)
 			patch((*MetaCache).GetPartitionID, int64(100), nil)
 			patch((*MetaCache).GetPartitionsIndex, []string{"_default_0"}, nil)
-			patch(getPartitionIDs, []int64{100}, nil)
+			patch(GetPartitionIDs, []int64{100}, nil)
 			patch(getDefaultPartitionsInPartitionKeyMode, []string{"_default_0"}, nil)
 			allocated := 0
 			alloc := mockey.Mock((*allocator.IDAllocator).Alloc).To(func(_ *allocator.IDAllocator, count uint32) (int64, int64, error) {
@@ -4300,7 +4300,7 @@ func TestPartialUpdateAutoIDDestinationSnapshotFromQueryRPC(t *testing.T) {
 						return nil, segcore.StorageCost{}, err
 					}
 				}
-				return qt.result, qt.storageCost, nil
+				return qt.Result(), qt.StorageCost(), nil
 			}).Build()
 			t.Cleanup(func() { query.UnPatch() })
 
@@ -7180,9 +7180,9 @@ func TestPartialUpdateRetryStrongReadRejectsMissingSnapshot(t *testing.T) {
 	fakeWAL.term = 10
 	task.node.(*Proxy).tsoAllocator = nil
 	query := mockey.Mock((*Proxy).query).To(func(_ *Proxy, _ context.Context, qt *queryTask, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
-		require.Equal(t, commonpb.ConsistencyLevel_Strong, qt.request.GetConsistencyLevel())
+		require.Equal(t, commonpb.ConsistencyLevel_Strong, qt.Request().GetConsistencyLevel())
 		require.Zero(t, qt.GetMvccTimestamp())
-		require.NotNil(t, qt.actualChannelsMvcc)
+		require.True(t, qt.HasRecordedChannelMvcc())
 		for _, meta := range task.partialUpdateCASGroups {
 			require.Zero(t, meta.GetReadTs())
 			require.EqualValues(t, 10, meta.GetObservedPchannelTerm())

internal/proxy/tasks_alias.go

@@ -0,0 +1,85 @@
+// 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 (
+	"github.com/milvus-io/milvus/internal/proxy/dql"
+)
+
+// Aliases to the extracted dql package, kept so the composition root can keep
+// using the original (unexported) type names. This is the only place in the
+// root package allowed to reference the dql package.
+
+type (
+	searchTask                  = dql.SearchTask
+	queryTask                   = dql.QueryTask
+	getStatisticsTask           = dql.GetStatisticsTask
+	getCollectionStatisticsTask = dql.GetCollectionStatisticsTask
+	getPartitionStatisticsTask  = dql.GetPartitionStatisticsTask
+)
+
+// Constants re-exported from the dql package so external consumers of the root
+// proxy package (e.g. the REST httpserver) keep compiling. They are true consts
+// because dql declares them as consts; a var alias would break callers that use
+// them in compile-time constant contexts.
+const (
+	AnnsFieldKey           = dql.AnnsFieldKey
+	AnalyzerKey            = dql.AnalyzerKey
+	ExprParamsKey          = dql.ExprParamsKey
+	GroupByFieldKey        = dql.GroupByFieldKey
+	GroupByFieldsKey       = dql.GroupByFieldsKey
+	GroupSizeKey           = dql.GroupSizeKey
+	IgnoreGrowingKey       = dql.IgnoreGrowingKey
+	LimitKey               = dql.LimitKey
+	MetricTypeKey          = dql.MetricTypeKey
+	NQKey                  = dql.NQKey
+	NormScoreKey           = dql.NormScoreKey
+	OffsetKey              = dql.OffsetKey
+	OrderByFieldsKey       = dql.OrderByFieldsKey
+	ParamsKey              = dql.ParamsKey
+	PipelineTraceKey       = dql.PipelineTraceKey
+	QueryIterLastOffsetKey = dql.QueryIterLastOffsetKey
+	QueryIterLastPKKey     = dql.QueryIterLastPKKey
+	RRFParamsKey           = dql.RRFParamsKey
+	RankTypeKey            = dql.RankTypeKey
+	ReduceStopForBestKey   = dql.ReduceStopForBestKey
+	RoundDecimalKey        = dql.RoundDecimalKey
+	SearchIterBatchSizeKey = dql.SearchIterBatchSizeKey
+	SearchIterIdKey        = dql.SearchIterIdKey
+	SearchIterLastBoundKey = dql.SearchIterLastBoundKey
+	StrictCastKey          = dql.StrictCastKey
+	StrictGroupSize        = dql.StrictGroupSize
+	TimefieldsKey          = dql.TimefieldsKey
+	TopKKey                = dql.TopKKey
+	WeightsParamsKey       = dql.WeightsParamsKey
+	CollectionID           = dql.CollectionID
+)
+
+var (
+	NewSearchTask                            = dql.NewSearchTask
+	NewQueryTask                             = dql.NewQueryTask
+	NewGetStatisticsTask                     = dql.NewGetStatisticsTask
+	NewGetCollectionStatisticsTask           = dql.NewGetCollectionStatisticsTask
+	NewGetPartitionStatisticsTask            = dql.NewGetPartitionStatisticsTask
+	ConvertHybridSearchToSearch              = dql.ConvertHybridSearchToSearch
+	PickFieldData                            = dql.PickFieldData
+	MatchCountRule                           = dql.MatchCountRule
+	WrapPlanCreationError                    = dql.WrapPlanCreationError
+	GetPartitionIDs                          = dql.GetPartitionIDs
+	MarshalPlanWithMembershipFilterSizeLimit = dql.MarshalPlanWithMembershipFilterSizeLimit
+	NormalizeFP32ToFP16BF16VectorFieldData   = dql.NormalizeFP32ToFP16BF16VectorFieldData
+)

internal/proxy/util.go

@@ -35,9 +35,7 @@ import (
 	"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/json"
-	"github.com/milvus-io/milvus/internal/parser/planparserv2"
 	"github.com/milvus-io/milvus/internal/proxy/fieldvalidator"
 	"github.com/milvus-io/milvus/internal/proxy/privilege"
 	"github.com/milvus-io/milvus/internal/types"
@@ -64,7 +62,6 @@ import (
 	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/v3/util/rbacutil"
 	"github.com/milvus-io/milvus/pkg/v3/util/requestutil"
-	"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"
 )
@@ -119,30 +116,6 @@ func restoreStructFieldNames(schema *schemapb.CollectionSchema) error {
 	return nil
 }
 
-// extractOriginalFieldName extracts the original field name from structName[fieldName] format
-// This function should only be called on transformed struct field names
-func extractOriginalFieldName(transformedName string) (string, error) {
-	idx := strings.Index(transformedName, "[")
-	if idx == -1 {
-		return "", merr.WrapErrParameterInvalidMsg("not a transformed struct field name: %s", transformedName)
-	}
-
-	if !strings.HasSuffix(transformedName, "]") {
-		return "", merr.WrapErrParameterInvalidMsg("invalid struct field format: %s, missing closing bracket", transformedName)
-	}
-
-	if idx == 0 {
-		return "", merr.WrapErrParameterInvalidMsg("invalid struct field format: %s, missing struct name", transformedName)
-	}
-
-	fieldName := transformedName[idx+1 : len(transformedName)-1]
-	if fieldName == "" {
-		return "", merr.WrapErrParameterInvalidMsg("invalid struct field format: %s, empty field name", transformedName)
-	}
-
-	return fieldName, nil
-}
-
 // isAlpha check if c is alpha.
 func isAlpha(c uint8) bool {
 	if (c < 'A' || c > 'Z') && (c < 'a' || c > 'z') {
@@ -179,25 +152,6 @@ func validateRunAnalyzer(req *milvuspb.RunAnalyzerRequest) error {
 	return nil
 }
 
-func validateMaxQueryResultWindow(offset int64, limit int64, largeTopKEnabled bool) error {
-	if offset < 0 {
-		return merr.WrapErrParameterInvalidMsg("%s [%d] is invalid, should be gte than 0", OffsetKey, offset)
-	}
-	if limit <= 0 {
-		return merr.WrapErrParameterInvalidMsg("%s [%d] is invalid, should be greater than 0", LimitKey, limit)
-	}
-
-	depth := offset + limit
-	maxQueryResultWindow := Params.QuotaConfig.MaxQueryResultWindow.GetAsInt64()
-	if largeTopKEnabled {
-		maxQueryResultWindow = Params.QuotaConfig.LargeMaxQueryResultWindow.GetAsInt64()
-	}
-	if depth <= 0 || depth > maxQueryResultWindow {
-		return merr.WrapErrParameterInvalidMsg("(offset+limit) should be in range [1, %d], but got %d", maxQueryResultWindow, depth)
-	}
-	return nil
-}
-
 func validateLimit(limit int64, largeTopKEnabled bool) error {
 	topKLimit := Params.QuotaConfig.TopKLimit.GetAsInt64()
 	if largeTopKEnabled {
@@ -1773,167 +1727,6 @@ func computeRecall(results *schemapb.SearchResultData, gts *schemapb.SearchResul
 	}
 }
 
-// Support wildcard in output fields:
-//
-//	"*" - all fields
-//
-// For example, A and B are scalar fields, C and D are vector fields, duplicated fields will automatically be removed.
-//
-//	output_fields=["*"] 	 ==> [A,B,C,D]
-//	output_fields=["*",A] 	 ==> [A,B,C,D]
-//	output_fields=["*",C]    ==> [A,B,C,D]
-//
-// 4th return value is true if user requested pk field explicitly or using wildcard.
-// if removePkField is true, pk field will not be include in the first(resultFieldNames)/second(userOutputFields)
-// return value.
-func translateOutputFields(outputFields []string, schema *schemaInfo, removePkField bool) ([]string, []string, []string, []agg.AggregateBase, bool, error) {
-	var primaryFieldName string
-	allFieldNameMap := make(map[string]*schemapb.FieldSchema)
-	resultFieldNameMap := make(map[string]bool)
-	resultFieldNames := make([]string, 0)
-	userOutputFieldsMap := make(map[string]bool)
-	userOutputFields := make([]string, 0)
-	userDynamicFieldsMap := make(map[string]bool)
-	userDynamicFields := make([]string, 0)
-	useAllDyncamicFields := false
-	aggregates := make([]agg.AggregateBase, 0)
-	for _, field := range schema.Fields {
-		if field.IsPrimaryKey {
-			primaryFieldName = field.Name
-		}
-		allFieldNameMap[field.Name] = field
-	}
-
-	// User may specify a struct array field or some specific fields in the struct array field
-	for _, subStruct := range schema.StructArrayFields {
-		for _, field := range subStruct.Fields {
-			allFieldNameMap[field.Name] = field
-		}
-	}
-
-	structArrayNameToFields := make(map[string][]*schemapb.FieldSchema)
-	for _, subStruct := range schema.StructArrayFields {
-		structArrayNameToFields[subStruct.Name] = subStruct.Fields
-	}
-
-	userRequestedPkFieldExplicitly := false
-
-	for _, outputFieldName := range outputFields {
-		outputFieldName = strings.TrimSpace(outputFieldName)
-		if outputFieldName == primaryFieldName {
-			userRequestedPkFieldExplicitly = true
-		}
-		if outputFieldName == "*" {
-			userRequestedPkFieldExplicitly = true
-			for fieldName, field := range allFieldNameMap {
-				if schema.CanRetrieveRawFieldData(field) {
-					resultFieldNameMap[fieldName] = true
-					userOutputFieldsMap[fieldName] = true
-				}
-			}
-			useAllDyncamicFields = true
-		} else {
-			if isAgg, aggregateName, aggFieldName := agg.MatchAggregationExpression(outputFieldName); isAgg {
-				if aggField, ok := allFieldNameMap[aggFieldName]; ok {
-					aggFuncs, aggErr := agg.NewAggregate(aggregateName, aggField.GetFieldID(), outputFieldName, aggField.GetDataType())
-					if aggErr != nil {
-						return nil, nil, nil, nil, false, aggErr
-					}
-					aggregates = append(aggregates, aggFuncs...)
-				} else if aggFieldName == "*" {
-					// only count(*) is allowed
-					if aggregateName != "count" {
-						return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("%s(*) is not supported, only count(*) is allowed", aggregateName)
-					}
-					if err := agg.ValidateAggFieldType(aggregateName, schemapb.DataType_None); err != nil {
-						return nil, nil, nil, nil, false, err
-					}
-					aggFuncs, aggErr := agg.NewAggregate(aggregateName, 0, outputFieldName, schemapb.DataType_None)
-					if aggErr != nil {
-						return nil, nil, nil, nil, false, aggErr
-					}
-					aggregates = append(aggregates, aggFuncs...)
-				} else {
-					return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("target field %s for aggregation:%s does not exist", aggFieldName, aggregateName)
-				}
-				userOutputFieldsMap[outputFieldName] = true
-				continue
-			}
-
-			if structArrayField, ok := structArrayNameToFields[outputFieldName]; ok {
-				for _, field := range structArrayField {
-					if schema.CanRetrieveRawFieldData(field) {
-						resultFieldNameMap[field.Name] = true
-						userOutputFieldsMap[field.Name] = true
-					}
-				}
-				continue
-			}
-			if field, ok := allFieldNameMap[outputFieldName]; ok {
-				if !schema.CanRetrieveRawFieldData(field) {
-					return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("not allowed to retrieve raw data of field %s", outputFieldName)
-				}
-				resultFieldNameMap[outputFieldName] = true
-				userOutputFieldsMap[outputFieldName] = true
-			} else {
-				if schema.EnableDynamicField {
-					dynamicNestedPath := outputFieldName
-					err := planparserv2.ParseIdentifier(schema.SchemaHelper, outputFieldName, func(expr *planpb.Expr) error {
-						columnInfo := expr.GetColumnExpr().GetInfo()
-						// there must be no error here
-						dynamicField, _ := schema.SchemaHelper.GetDynamicField()
-						// only $meta["xxx"] is allowed for now
-						if dynamicField.GetFieldID() != columnInfo.GetFieldId() {
-							return merr.WrapErrParameterInvalidMsg("not support getting subkeys of json field yet")
-						}
-						nestedPaths := columnInfo.GetNestedPath()
-						// $meta["A"]["B"] not allowed for now
-						if len(nestedPaths) != 1 {
-							return merr.WrapErrParameterInvalidMsg("not support getting multiple level of dynamic field for now")
-						}
-						// $meta["dyn_field"], output field name could be:
-						// 1. "dyn_field", outputFieldName == nestedPath
-						// 2. `$meta["dyn_field"]` explicit form
-						if nestedPaths[0] != outputFieldName {
-							// use "dyn_field" as userDynamicFieldsMap when outputField = `$meta["dyn_field"]`
-							dynamicNestedPath = nestedPaths[0]
-						}
-						return nil
-					})
-					if err != nil {
-						mlog.Info(context.TODO(), "parse output field name failed", mlog.String("field name", outputFieldName), mlog.Err(err))
-						return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("parse output field name failed: %s", outputFieldName)
-					}
-					resultFieldNameMap[common.MetaFieldName] = true
-					userOutputFieldsMap[outputFieldName] = true
-					userDynamicFieldsMap[dynamicNestedPath] = true
-				} else {
-					return nil, nil, nil, nil, false, merr.WrapErrParameterInvalidMsg("field %s not exist", outputFieldName)
-				}
-			}
-		}
-	}
-
-	if removePkField {
-		delete(resultFieldNameMap, primaryFieldName)
-		delete(userOutputFieldsMap, primaryFieldName)
-	}
-
-	for fieldName := range resultFieldNameMap {
-		resultFieldNames = append(resultFieldNames, fieldName)
-	}
-	for fieldName := range userOutputFieldsMap {
-		userOutputFields = append(userOutputFields, fieldName)
-	}
-	if !useAllDyncamicFields {
-		for fieldName := range userDynamicFieldsMap {
-			userDynamicFields = append(userDynamicFields, fieldName)
-		}
-	}
-
-	return resultFieldNames, userOutputFields, userDynamicFields, aggregates, userRequestedPkFieldExplicitly, nil
-}
-
 func validCharInIndexName(c byte) bool {
 	return c == '_' || c == '[' || c == ']' || isAlpha(c) || isNumber(c)
 }
@@ -3452,104 +3245,6 @@ func getCollectionTTL(pairs []*commonpb.KeyValuePair) uint64 {
 	return uint64(ttl)
 }
 
-// reconstructStructFieldData regroups flattened sub-fields (named "structName[fieldName]")
-// back into StructArrayField entries, restoring original field names for the user-facing response.
-// It modifies sub-field FieldName in place; callers must not reuse the input slice afterwards.
-func reconstructStructFieldData(
-	fieldsData []*schemapb.FieldData,
-	outputFields []string,
-	schema *schemapb.CollectionSchema,
-) ([]*schemapb.FieldData, []string) {
-	if len(outputFields) == 1 && outputFields[0] == "count(*)" {
-		return fieldsData, outputFields
-	}
-
-	if len(schema.StructArrayFields) == 0 {
-		return fieldsData, outputFields
-	}
-
-	regularFieldIDs := make(map[int64]interface{})
-	subFieldToStructMap := make(map[int64]int64)
-	groupedStructFields := make(map[int64][]*schemapb.FieldData)
-	structFieldNames := make(map[int64]string)
-	reconstructedOutputFields := make([]string, 0, len(fieldsData))
-
-	// record all regular field IDs
-	for _, field := range schema.Fields {
-		regularFieldIDs[field.GetFieldID()] = nil
-	}
-
-	// build the mapping from sub-field ID to struct field ID
-	for _, structField := range schema.StructArrayFields {
-		for _, subField := range structField.GetFields() {
-			subFieldToStructMap[subField.GetFieldID()] = structField.GetFieldID()
-		}
-		structFieldNames[structField.GetFieldID()] = structField.GetName()
-	}
-
-	newFieldsData := make([]*schemapb.FieldData, 0, len(fieldsData))
-	for _, field := range fieldsData {
-		fieldID := field.GetFieldId()
-		if _, ok := regularFieldIDs[fieldID]; ok {
-			newFieldsData = append(newFieldsData, field)
-			reconstructedOutputFields = append(reconstructedOutputFields, field.GetFieldName())
-		} else if structFieldID, ok := subFieldToStructMap[fieldID]; ok {
-			groupedStructFields[structFieldID] = append(groupedStructFields[structFieldID], field)
-		} else {
-			newFieldsData = append(newFieldsData, field)
-			reconstructedOutputFields = append(reconstructedOutputFields, field.GetFieldName())
-		}
-	}
-
-	for structFieldID, fields := range groupedStructFields {
-		// Restore original field names (from "structName[fieldName]" to "fieldName")
-		// for the user-facing response.
-		for _, field := range fields {
-			originalName, err := extractOriginalFieldName(field.FieldName)
-			if err != nil {
-				mlog.Error(context.TODO(), "failed to extract original field name from struct field",
-					mlog.String("fieldName", field.FieldName),
-					mlog.Err(err))
-			} else {
-				field.FieldName = originalName
-			}
-		}
-
-		newFieldsData = append(newFieldsData, &schemapb.FieldData{
-			FieldName: structFieldNames[structFieldID],
-			FieldId:   structFieldID,
-			Type:      schemapb.DataType_ArrayOfStruct,
-			Field:     &schemapb.FieldData_StructArrays{StructArrays: &schemapb.StructArrayField{Fields: fields}},
-		})
-		reconstructedOutputFields = append(reconstructedOutputFields, structFieldNames[structFieldID])
-	}
-
-	return newFieldsData, reconstructedOutputFields
-}
-
-func reconstructStructFieldDataForQuery(results *milvuspb.QueryResults, schema *schemapb.CollectionSchema) {
-	fieldsData, outputFields := reconstructStructFieldData(
-		results.FieldsData,
-		results.OutputFields,
-		schema,
-	)
-	results.FieldsData = fieldsData
-	results.OutputFields = outputFields
-}
-
-func reconstructStructFieldDataForSearch(results *milvuspb.SearchResults, schema *schemapb.CollectionSchema) {
-	if results.Results == nil {
-		return
-	}
-	fieldsData, outputFields := reconstructStructFieldData(
-		results.Results.FieldsData,
-		results.Results.OutputFields,
-		schema,
-	)
-	results.Results.FieldsData = fieldsData
-	results.Results.OutputFields = outputFields
-}
-
 func getColTimezone(colInfo *collectionInfo) string {
 	timezone, _ := funcutil.TryGetAttrByKeyFromRepeatedKV(common.TimezoneKey, colInfo.Properties)
 	if timezone == "" {
@@ -3558,160 +3253,6 @@ func getColTimezone(colInfo *collectionInfo) string {
 	return timezone
 }
 
-// resolveTimezone returns the effective timezone for a request: the validated
-// request-level timezone when one is specified in params, otherwise the
-// collection timezone (UTC by default). Callers must resolve the timezone
-// before generating any plan that parses naive TIMESTAMPTZ literals, so that
-// filtering and result formatting use the same timezone.
-func resolveTimezone(ctx context.Context, params []*commonpb.KeyValuePair, colInfo *collectionInfo) (string, error) {
-	timezone, _ := funcutil.TryGetAttrByKeyFromRepeatedKV(common.TimezoneKey, params)
-	if timezone != "" {
-		if !timestamptz.IsTimezoneValid(timezone) {
-			mlog.Info(ctx, "get invalid timezone from request", mlog.String("timezone", timezone))
-			return "", merr.WrapErrParameterInvalidMsg("unknown or invalid IANA Time Zone ID: %s", timezone)
-		}
-		mlog.Debug(ctx, "determine timezone from request", mlog.String("user defined timezone", timezone))
-		return timezone, nil
-	}
-	timezone = getColTimezone(colInfo)
-	mlog.Debug(ctx, "determine timezone from collection", mlog.String("collection timezone", timezone))
-	return timezone, nil
-}
-
-// timestamptzUTC2IsoStr converts Timestamptz (Unix Microsecond) data
-// within FieldData results into ISO-8601 strings, applying the correct
-// timezone offset and using the optimized format (microsecond precision, no trailing zeros).
-func timestamptzUTC2IsoStr(results []*schemapb.FieldData, colTimezone string) error {
-	location, err := time.LoadLocation(colTimezone)
-	if err != nil {
-		mlog.Error(context.TODO(), "invalid timezone", mlog.String("timezone", colTimezone), mlog.Err(err))
-		return merr.WrapErrParameterInvalidMsg("got invalid default timezone: %s", colTimezone)
-	}
-
-	for _, fieldData := range results {
-		if fieldData.GetType() != schemapb.DataType_Timestamptz {
-			continue
-		}
-
-		scalarField := fieldData.GetScalars()
-
-		// Guard against nil scalars or missing timestamp data
-		if scalarField == nil || scalarField.GetTimestamptzData() == nil {
-			if longData := scalarField.GetLongData(); longData != nil && len(longData.GetData()) > 0 {
-				mlog.Warn(context.TODO(), "field data is not Timestamptz data", mlog.String("fieldName", fieldData.GetFieldName()))
-				return merr.WrapErrParameterInvalidMsg("field data for '%s' is not Timestamptz data", fieldData.GetFieldName())
-			}
-			// Handle the case of an empty field (e.g., all nulls), skip if no data to process.
-			continue
-		}
-
-		utcTimestamps := scalarField.GetTimestamptzData().GetData()
-		isoStrings := make([]string, len(utcTimestamps))
-
-		// CORE CHANGE: Use the optimized formatting function
-		for i, ts := range utcTimestamps {
-			// 1. Convert Unix Microsecond (UTC) to a time.Time object (still in UTC).
-			t := time.UnixMicro(ts).UTC()
-
-			// 2. Adjust the time object to the target location.
-			localTime := t.In(location)
-
-			// 3. Format using the optimized logic (max 6 digits, no trailing zeros)
-			isoStrings[i] = timestamptz.FormatTimeMicroWithoutTrailingZeros(localTime)
-		}
-
-		// Replace the TimestamptzData with the new StringData in place.
-		fieldData.GetScalars().Data = &schemapb.ScalarField_StringData{
-			StringData: &schemapb.StringArray{
-				Data: isoStrings,
-			},
-		}
-	}
-	return nil
-}
-
-// extractFields is a helper function to extract specific integer fields from a time.Time object.
-// Supported fields are: "year", "month", "day", "hour", "minute", "second", "microsecond", "nanosecond".
-func extractFields(t time.Time, fieldList []string) ([]int64, error) {
-	extractedValues := make([]int64, 0, len(fieldList))
-	for _, field := range fieldList {
-		var val int64
-		switch strings.ToLower(field) {
-		case common.TszYear:
-			val = int64(t.Year())
-		case common.TszMonth:
-			val = int64(t.Month())
-		case common.TszDay:
-			val = int64(t.Day())
-		case common.TszHour:
-			val = int64(t.Hour())
-		case common.TszMinute:
-			val = int64(t.Minute())
-		case common.TszSecond:
-			val = int64(t.Second())
-		case common.TszMicrosecond:
-			val = int64(t.Nanosecond() / 1000)
-		default:
-			return nil, merr.WrapErrParameterInvalidMsg("unsupported field for extraction: %s, fields should be seprated by ',' or ' '", field)
-		}
-		extractedValues = append(extractedValues, val)
-	}
-	return extractedValues, nil
-}
-
-func extractFieldsFromResults(results []*schemapb.FieldData, timezone string, fieldList []string) error {
-	targetLocation, err := time.LoadLocation(timezone)
-	if err != nil {
-		mlog.Error(context.TODO(), "invalid timezone", mlog.String("timezone", timezone), mlog.Err(err))
-		return merr.WrapErrParameterInvalidMsg("got invalid timezone: %s", timezone)
-	}
-
-	for _, fieldData := range results {
-		if fieldData.GetType() != schemapb.DataType_Timestamptz {
-			continue
-		}
-
-		scalarField := fieldData.GetScalars()
-		if scalarField == nil || scalarField.GetTimestamptzData() == nil {
-			if longData := scalarField.GetLongData(); longData != nil && len(longData.GetData()) > 0 {
-				mlog.Warn(context.TODO(), "field data is not Timestamptz data, but found LongData instead", mlog.String("fieldName", fieldData.GetFieldName()))
-				return merr.WrapErrParameterInvalidMsg("field data for '%s' is not Timestamptz data", fieldData.GetFieldName())
-			}
-			continue
-		}
-
-		utcTimestamps := scalarField.GetTimestamptzData().GetData()
-		extractedResults := make([]*schemapb.ScalarField, 0, len(fieldList))
-
-		for _, ts := range utcTimestamps {
-			t := time.UnixMicro(ts).UTC()
-			localTime := t.In(targetLocation)
-
-			values, err := extractFields(localTime, fieldList)
-			if err != nil {
-				return err
-			}
-			valuesScalarField := &schemapb.ScalarField_LongData{
-				LongData: &schemapb.LongArray{
-					Data: values,
-				},
-			}
-			extractedResults = append(extractedResults, &schemapb.ScalarField{
-				Data: valuesScalarField,
-			})
-		}
-
-		fieldData.GetScalars().Data = &schemapb.ScalarField_ArrayData{
-			ArrayData: &schemapb.ArrayArray{
-				Data:        extractedResults,
-				ElementType: schemapb.DataType_Int64,
-			},
-		}
-		fieldData.Type = schemapb.DataType_Array
-	}
-	return nil
-}
-
 func genFunctionFields(ctx context.Context, insertMsg *msgstream.InsertMsg, schema *schemaInfo, partialUpdate bool) error {
 	functions := schema.GetFunctions()
 	allowNonBM25Outputs := common.GetCollectionAllowInsertNonBM25FunctionOutputs(schema.Properties)

internal/proxy/util_test.go

@@ -46,7 +46,6 @@ import (
 	"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/proto/planpb"
 	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
 	"github.com/milvus-io/milvus/pkg/v3/proto/rootcoordpb"
 	"github.com/milvus-io/milvus/pkg/v3/util"
@@ -76,39 +75,6 @@ func captureProxyLogs(t *testing.T) *mlog.TestSink {
 	})
 }
 
-func TestSearchInfoDetermineSearchTypeWithPluralGroupByFieldIDs(t *testing.T) {
-	info := &SearchInfo{
-		planInfo: &planpb.QueryInfo{
-			GroupByFieldIds: []int64{101},
-		},
-	}
-
-	assert.Equal(t, internalpb.SearchType_DEFAULT, info.DetermineSearchType(false))
-}
-
-func TestParseGroupByInfoLegacyFieldPrecedence(t *testing.T) {
-	schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
-		{FieldID: 101, Name: "brand", DataType: schemapb.DataType_VarChar},
-		{FieldID: 102, Name: "category", DataType: schemapb.DataType_VarChar},
-	}}
-
-	info, err := parseGroupByInfo([]*commonpb.KeyValuePair{
-		{Key: GroupByFieldKey, Value: "brand"},
-		{Key: GroupByFieldsKey, Value: "category"},
-	}, schema)
-	assert.NoError(t, err)
-	assert.Equal(t, []int64{101}, info.groupByFieldIds)
-	assert.Equal(t, []string{"brand"}, info.groupByFieldNames)
-
-	info, err = parseGroupByInfo([]*commonpb.KeyValuePair{
-		{Key: GroupByFieldKey, Value: " "},
-		{Key: GroupByFieldsKey, Value: "brand, category"},
-	}, schema)
-	assert.NoError(t, err)
-	assert.Equal(t, []int64{101, 102}, info.groupByFieldIds)
-	assert.Equal(t, []string{"brand", "category"}, info.groupByFieldNames)
-}
-
 func TestValidateCollectionName(t *testing.T) {
 	assert.Nil(t, validateCollectionName("abc"))
 	assert.Nil(t, validateCollectionName("_123abc"))
@@ -309,7 +275,7 @@ func TestConvertHybridSearchToSearchCopiesNamespace(t *testing.T) {
 		},
 	}
 
-	searchReq := convertHybridSearchToSearch(req)
+	searchReq := ConvertHybridSearchToSearch(req)
 	require.NotNil(t, searchReq.Namespace)
 	assert.Equal(t, namespace, searchReq.GetNamespace())
 }
@@ -2557,21 +2523,6 @@ func Test_LargeTopKLimit(t *testing.T) {
 	assert.Error(t, validateLimit(0, true))
 }
 
-func Test_MaxQueryResultWindow(t *testing.T) {
-	paramtable.Init()
-	assert.Nil(t, validateMaxQueryResultWindow(0, 16384, false))
-	assert.Nil(t, validateMaxQueryResultWindow(0, 1, false))
-	assert.Error(t, validateMaxQueryResultWindow(0, 16385, false))
-	assert.Error(t, validateMaxQueryResultWindow(0, 0, false))
-	assert.Error(t, validateMaxQueryResultWindow(1, 0, false))
-
-	Params.Save(Params.QuotaConfig.LargeMaxQueryResultWindow.Key, "1000000")
-	defer Params.Reset(Params.QuotaConfig.LargeMaxQueryResultWindow.Key)
-	assert.Nil(t, validateMaxQueryResultWindow(0, 16385, true))
-	assert.Nil(t, validateMaxQueryResultWindow(0, 1000000, true))
-	assert.Error(t, validateMaxQueryResultWindow(0, 1000001, true))
-}
-
 func Test_GetPartitionProgressFailed(t *testing.T) {
 	qc := mocks.NewMockQueryCoordClient(t)
 	qc.EXPECT().ShowLoadPartitions(mock.Anything, mock.Anything).Return(&querypb.ShowPartitionsResponse{
@@ -4890,641 +4841,6 @@ func TestValidateStructArrayField_MaxCapacity(t *testing.T) {
 	})
 }
 
-func Test_reconstructStructFieldData(t *testing.T) {
-	t.Run("count(*) query - should return early", func(t *testing.T) {
-		fieldsData := []*schemapb.FieldData{
-			{
-				FieldName: "count(*)",
-				FieldId:   0,
-				Type:      schemapb.DataType_Int64,
-			},
-		}
-		outputFields := []string{"count(*)"}
-
-		schema := &schemapb.CollectionSchema{
-			StructArrayFields: []*schemapb.StructArrayFieldSchema{
-				{
-					FieldID: 102,
-					Name:    "test_struct",
-					Fields: []*schemapb.FieldSchema{
-						{
-							FieldID:     1021,
-							Name:        "test_struct[sub_field]",
-							DataType:    schemapb.DataType_Array,
-							ElementType: schemapb.DataType_Int32,
-						},
-					},
-				},
-			},
-		}
-
-		originalFieldsData := make([]*schemapb.FieldData, len(fieldsData))
-		copy(originalFieldsData, fieldsData)
-		originalOutputFields := make([]string, len(outputFields))
-		copy(originalOutputFields, outputFields)
-
-		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
-
-		// Should not modify anything for count(*) query
-		assert.Equal(t, originalFieldsData, resultFieldsData)
-		assert.Equal(t, originalOutputFields, resultOutputFields)
-	})
-
-	t.Run("group by with count(*) should preserve aggregate field", func(t *testing.T) {
-		fieldsData := []*schemapb.FieldData{
-			{
-				FieldName: "id",
-				FieldId:   100,
-				Type:      schemapb.DataType_Int64,
-				Field: &schemapb.FieldData_Scalars{
-					Scalars: &schemapb.ScalarField{
-						Data: &schemapb.ScalarField_LongData{
-							LongData: &schemapb.LongArray{Data: []int64{1, 2}},
-						},
-					},
-				},
-			},
-			{
-				FieldName: "count(*)",
-				FieldId:   0,
-				Type:      schemapb.DataType_Int64,
-				Field: &schemapb.FieldData_Scalars{
-					Scalars: &schemapb.ScalarField{
-						Data: &schemapb.ScalarField_LongData{
-							LongData: &schemapb.LongArray{Data: []int64{3, 4}},
-						},
-					},
-				},
-			},
-		}
-		outputFields := []string{"id", "count(*)"}
-
-		schema := &schemapb.CollectionSchema{
-			Fields: []*schemapb.FieldSchema{
-				{
-					FieldID:      100,
-					Name:         "id",
-					IsPrimaryKey: true,
-					DataType:     schemapb.DataType_Int64,
-				},
-			},
-			StructArrayFields: []*schemapb.StructArrayFieldSchema{
-				{
-					FieldID: 102,
-					Name:    "test_struct",
-					Fields: []*schemapb.FieldSchema{
-						{
-							FieldID:     1021,
-							Name:        "test_struct[sub_field]",
-							DataType:    schemapb.DataType_Array,
-							ElementType: schemapb.DataType_Int32,
-						},
-					},
-				},
-			},
-		}
-
-		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
-
-		assert.Len(t, resultFieldsData, 2)
-		assert.Equal(t, "id", resultFieldsData[0].FieldName)
-		assert.Equal(t, int64(100), resultFieldsData[0].FieldId)
-		assert.Equal(t, "count(*)", resultFieldsData[1].FieldName)
-		assert.Equal(t, int64(0), resultFieldsData[1].FieldId)
-		assert.Equal(t, []string{"id", "count(*)"}, resultOutputFields)
-		assert.Equal(t, []int64{3, 4}, resultFieldsData[1].GetScalars().GetLongData().GetData())
-	})
-
-	t.Run("struct field query - should reconstruct struct field", func(t *testing.T) {
-		fieldsData := []*schemapb.FieldData{
-			{
-				FieldName: "test_struct[sub_field]",
-				FieldId:   1021, // Use the correct field ID that matches the schema
-				Type:      schemapb.DataType_Array,
-				Field: &schemapb.FieldData_Scalars{
-					Scalars: &schemapb.ScalarField{
-						Data: &schemapb.ScalarField_ArrayData{
-							ArrayData: &schemapb.ArrayArray{
-								ElementType: schemapb.DataType_Int32,
-								Data: []*schemapb.ScalarField{
-									{
-										Data: &schemapb.ScalarField_IntData{
-											IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
-										},
-									},
-								},
-							},
-						},
-					},
-				},
-			},
-		}
-		outputFields := []string{"test_struct[sub_field]"}
-
-		schema := &schemapb.CollectionSchema{
-			StructArrayFields: []*schemapb.StructArrayFieldSchema{
-				{
-					FieldID: 102,
-					Name:    "test_struct",
-					Fields: []*schemapb.FieldSchema{
-						{
-							FieldID:     1021,
-							Name:        "test_struct[sub_field]",
-							DataType:    schemapb.DataType_Array,
-							ElementType: schemapb.DataType_Int32,
-						},
-					},
-				},
-			},
-		}
-
-		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
-
-		// Should reconstruct the struct field with the restored field name
-		assert.Len(t, resultFieldsData, 1)
-		assert.Equal(t, "test_struct", resultFieldsData[0].FieldName)
-		assert.Equal(t, int64(102), resultFieldsData[0].FieldId)
-		assert.Equal(t, schemapb.DataType_ArrayOfStruct, resultFieldsData[0].Type)
-
-		// Check that the sub-field name has been restored
-		structArrayField := resultFieldsData[0].GetStructArrays()
-		assert.NotNil(t, structArrayField)
-		assert.Len(t, structArrayField.Fields, 1)
-		assert.Equal(t, "sub_field", structArrayField.Fields[0].FieldName) // Name should be restored
-
-		assert.Equal(t, []string{"test_struct"}, resultOutputFields)
-	})
-
-	t.Run("no struct array fields - should return early", func(t *testing.T) {
-		fieldsData := []*schemapb.FieldData{
-			{
-				FieldName: "field1",
-				FieldId:   100,
-				Type:      schemapb.DataType_Int64,
-			},
-			{
-				FieldName: "field2",
-				FieldId:   101,
-				Type:      schemapb.DataType_VarChar,
-			},
-		}
-		outputFields := []string{"field1", "field2"}
-
-		schema := &schemapb.CollectionSchema{
-			StructArrayFields: []*schemapb.StructArrayFieldSchema{},
-		}
-
-		originalFieldsData := make([]*schemapb.FieldData, len(fieldsData))
-		copy(originalFieldsData, fieldsData)
-		originalOutputFields := make([]string, len(outputFields))
-		copy(originalOutputFields, outputFields)
-
-		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
-
-		// Should not modify anything when no struct array fields
-		assert.Equal(t, originalFieldsData, resultFieldsData)
-		assert.Equal(t, originalOutputFields, resultOutputFields)
-	})
-
-	t.Run("reconstruct single struct field", func(t *testing.T) {
-		// Create mock data with transformed field names (as they would be internally)
-		subField1Data := &schemapb.FieldData{
-			FieldName: "test_struct[sub_int_array]",
-			FieldId:   1021,
-			Type:      schemapb.DataType_Array,
-			Field: &schemapb.FieldData_Scalars{
-				Scalars: &schemapb.ScalarField{
-					Data: &schemapb.ScalarField_ArrayData{
-						ArrayData: &schemapb.ArrayArray{
-							ElementType: schemapb.DataType_Int32,
-							Data: []*schemapb.ScalarField{
-								{
-									Data: &schemapb.ScalarField_IntData{
-										IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
-									},
-								},
-							},
-						},
-					},
-				},
-			},
-		}
-
-		subField2Data := &schemapb.FieldData{
-			FieldName: "test_struct[sub_text_array]",
-			FieldId:   1022,
-			Type:      schemapb.DataType_Array,
-			Field: &schemapb.FieldData_Scalars{
-				Scalars: &schemapb.ScalarField{
-					Data: &schemapb.ScalarField_ArrayData{
-						ArrayData: &schemapb.ArrayArray{
-							ElementType: schemapb.DataType_VarChar,
-							Data: []*schemapb.ScalarField{
-								{
-									Data: &schemapb.ScalarField_StringData{
-										StringData: &schemapb.StringArray{Data: []string{"hello", "world"}},
-									},
-								},
-							},
-						},
-					},
-				},
-			},
-		}
-
-		fieldsData := []*schemapb.FieldData{subField1Data, subField2Data}
-		outputFields := []string{"test_struct[sub_int_array]", "test_struct[sub_text_array]"}
-
-		schema := &schemapb.CollectionSchema{
-			Fields: []*schemapb.FieldSchema{
-				{
-					FieldID:      100,
-					Name:         "pk",
-					IsPrimaryKey: true,
-					DataType:     schemapb.DataType_Int64,
-				},
-			},
-			StructArrayFields: []*schemapb.StructArrayFieldSchema{
-				{
-					FieldID: 102,
-					Name:    "test_struct",
-					Fields: []*schemapb.FieldSchema{
-						{
-							FieldID:     1021,
-							Name:        "test_struct[sub_int_array]",
-							DataType:    schemapb.DataType_Array,
-							ElementType: schemapb.DataType_Int32,
-						},
-						{
-							FieldID:     1022,
-							Name:        "test_struct[sub_text_array]",
-							DataType:    schemapb.DataType_Array,
-							ElementType: schemapb.DataType_VarChar,
-						},
-					},
-				},
-			},
-		}
-
-		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
-
-		// Check result
-		assert.Len(t, resultFieldsData, 1, "Should only have one reconstructed struct field")
-		assert.Len(t, resultOutputFields, 1, "Output fields should only have one")
-
-		structField := resultFieldsData[0]
-		assert.Equal(t, "test_struct", structField.FieldName)
-		assert.Equal(t, int64(102), structField.FieldId)
-		assert.Equal(t, schemapb.DataType_ArrayOfStruct, structField.Type)
-		assert.Equal(t, "test_struct", resultOutputFields[0])
-
-		// Check fields inside struct
-		structArrays := structField.GetStructArrays()
-		assert.NotNil(t, structArrays)
-		assert.Len(t, structArrays.Fields, 2, "Struct should contain 2 sub fields")
-
-		// Check sub fields
-		var foundIntField, foundTextField bool
-		for _, field := range structArrays.Fields {
-			switch field.FieldId {
-			case 1021:
-				assert.Equal(t, "sub_int_array", field.FieldName)
-				assert.Equal(t, schemapb.DataType_Array, field.Type)
-				foundIntField = true
-			case 1022:
-				assert.Equal(t, "sub_text_array", field.FieldName)
-				assert.Equal(t, schemapb.DataType_Array, field.Type)
-				foundTextField = true
-			}
-		}
-		assert.True(t, foundIntField, "Should find int array field")
-		assert.True(t, foundTextField, "Should find text array field")
-	})
-
-	t.Run("mixed regular and struct fields", func(t *testing.T) {
-		// Create regular field data
-		regularField := &schemapb.FieldData{
-			FieldName: "regular_field",
-			FieldId:   100,
-			Type:      schemapb.DataType_Int64,
-			Field: &schemapb.FieldData_Scalars{
-				Scalars: &schemapb.ScalarField{
-					Data: &schemapb.ScalarField_LongData{
-						LongData: &schemapb.LongArray{Data: []int64{1, 2, 3}},
-					},
-				},
-			},
-		}
-
-		// Create struct sub field data with transformed name
-		subFieldData := &schemapb.FieldData{
-			FieldName: "test_struct[sub_field]",
-			FieldId:   1021,
-			Type:      schemapb.DataType_Array,
-			Field: &schemapb.FieldData_Scalars{
-				Scalars: &schemapb.ScalarField{
-					Data: &schemapb.ScalarField_ArrayData{
-						ArrayData: &schemapb.ArrayArray{
-							ElementType: schemapb.DataType_Int32,
-							Data: []*schemapb.ScalarField{
-								{
-									Data: &schemapb.ScalarField_IntData{
-										IntData: &schemapb.IntArray{Data: []int32{10, 20}},
-									},
-								},
-							},
-						},
-					},
-				},
-			},
-		}
-
-		fieldsData := []*schemapb.FieldData{regularField, subFieldData}
-		outputFields := []string{"regular_field", "test_struct[sub_field]"}
-
-		schema := &schemapb.CollectionSchema{
-			Fields: []*schemapb.FieldSchema{
-				{
-					FieldID:  100,
-					Name:     "regular_field",
-					DataType: schemapb.DataType_Int64,
-				},
-			},
-			StructArrayFields: []*schemapb.StructArrayFieldSchema{
-				{
-					FieldID: 102,
-					Name:    "test_struct",
-					Fields: []*schemapb.FieldSchema{
-						{
-							FieldID:     1021,
-							Name:        "test_struct[sub_field]",
-							DataType:    schemapb.DataType_Array,
-							ElementType: schemapb.DataType_Int32,
-						},
-					},
-				},
-			},
-		}
-
-		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
-
-		// Check result: should have 2 fields (1 regular + 1 reconstructed struct)
-		assert.Len(t, resultFieldsData, 2)
-		assert.Len(t, resultOutputFields, 2)
-
-		// Check regular and struct fields both exist
-		var foundRegularField, foundStructField bool
-		for i, field := range resultFieldsData {
-			switch field.FieldId {
-			case 100:
-				assert.Equal(t, "regular_field", field.FieldName)
-				assert.Equal(t, schemapb.DataType_Int64, field.Type)
-				assert.Equal(t, "regular_field", resultOutputFields[i])
-				foundRegularField = true
-			case 102:
-				assert.Equal(t, "test_struct", field.FieldName)
-				assert.Equal(t, schemapb.DataType_ArrayOfStruct, field.Type)
-				assert.Equal(t, "test_struct", resultOutputFields[i])
-				foundStructField = true
-			}
-		}
-		assert.True(t, foundRegularField, "Should find regular field")
-		assert.True(t, foundStructField, "Should find reconstructed struct field")
-	})
-
-	t.Run("multiple struct fields", func(t *testing.T) {
-		// Create sub field for first struct
-		struct1SubField := &schemapb.FieldData{
-			FieldName: "struct1[struct1_sub]",
-			FieldId:   1021,
-			Type:      schemapb.DataType_Array,
-			Field: &schemapb.FieldData_Scalars{
-				Scalars: &schemapb.ScalarField{
-					Data: &schemapb.ScalarField_ArrayData{
-						ArrayData: &schemapb.ArrayArray{
-							ElementType: schemapb.DataType_Int32,
-							Data:        []*schemapb.ScalarField{},
-						},
-					},
-				},
-			},
-		}
-
-		// Create sub fields for second struct
-		struct2SubField1 := &schemapb.FieldData{
-			FieldName: "struct2[struct2_sub1]",
-			FieldId:   1031,
-			Type:      schemapb.DataType_Array,
-			Field: &schemapb.FieldData_Scalars{
-				Scalars: &schemapb.ScalarField{
-					Data: &schemapb.ScalarField_ArrayData{
-						ArrayData: &schemapb.ArrayArray{
-							ElementType: schemapb.DataType_Int32,
-							Data:        []*schemapb.ScalarField{},
-						},
-					},
-				},
-			},
-		}
-
-		struct2SubField2 := &schemapb.FieldData{
-			FieldName: "struct2[struct2_sub2]",
-			FieldId:   1032,
-			Type:      schemapb.DataType_VarChar,
-			Field: &schemapb.FieldData_Scalars{
-				Scalars: &schemapb.ScalarField{
-					Data: &schemapb.ScalarField_StringData{
-						StringData: &schemapb.StringArray{Data: []string{"test"}},
-					},
-				},
-			},
-		}
-
-		fieldsData := []*schemapb.FieldData{struct1SubField, struct2SubField1, struct2SubField2}
-		outputFields := []string{"struct1[struct1_sub]", "struct2[struct2_sub1]", "struct2[struct2_sub2]"}
-
-		schema := &schemapb.CollectionSchema{
-			Fields: []*schemapb.FieldSchema{
-				{
-					FieldID:      100,
-					Name:         "pk",
-					IsPrimaryKey: true,
-					DataType:     schemapb.DataType_Int64,
-				},
-			},
-			StructArrayFields: []*schemapb.StructArrayFieldSchema{
-				{
-					FieldID: 102,
-					Name:    "struct1",
-					Fields: []*schemapb.FieldSchema{
-						{
-							FieldID:     1021,
-							Name:        "struct1[struct1_sub]",
-							DataType:    schemapb.DataType_Array,
-							ElementType: schemapb.DataType_Int32,
-						},
-					},
-				},
-				{
-					FieldID: 103,
-					Name:    "struct2",
-					Fields: []*schemapb.FieldSchema{
-						{
-							FieldID:     1031,
-							Name:        "struct2[struct2_sub1]",
-							DataType:    schemapb.DataType_Array,
-							ElementType: schemapb.DataType_Int32,
-						},
-						{
-							FieldID:  1032,
-							Name:     "struct2[struct2_sub2]",
-							DataType: schemapb.DataType_VarChar,
-						},
-					},
-				},
-			},
-		}
-
-		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
-
-		// Check result: should have 2 struct fields
-		assert.Len(t, resultFieldsData, 2)
-		assert.Len(t, resultOutputFields, 2)
-
-		// Check both struct fields
-		var foundStruct1, foundStruct2 bool
-		for _, field := range resultFieldsData {
-			switch field.FieldId {
-			case 102:
-				assert.Equal(t, "struct1", field.FieldName)
-				assert.Equal(t, schemapb.DataType_ArrayOfStruct, field.Type)
-				foundStruct1 = true
-				structArrays := field.GetStructArrays()
-				assert.NotNil(t, structArrays)
-				assert.Len(t, structArrays.Fields, 1)
-			case 103:
-				assert.Equal(t, "struct2", field.FieldName)
-				assert.Equal(t, schemapb.DataType_ArrayOfStruct, field.Type)
-				foundStruct2 = true
-				structArrays := field.GetStructArrays()
-				assert.NotNil(t, structArrays)
-				assert.Len(t, structArrays.Fields, 2)
-			}
-		}
-		assert.True(t, foundStruct1, "Should find struct1")
-		assert.True(t, foundStruct2, "Should find struct2")
-	})
-
-	t.Run("partial struct fields query - only return queried fields", func(t *testing.T) {
-		// Create a struct with 3 fields, but only query 2 of them
-		// This tests that we only return what the user requested
-
-		// Create mock data for only 2 out of 3 struct fields
-		clipStrData := &schemapb.FieldData{
-			FieldName: "clip[str]",
-			FieldId:   2001,
-			Type:      schemapb.DataType_VarChar,
-			Field: &schemapb.FieldData_Scalars{
-				Scalars: &schemapb.ScalarField{
-					Data: &schemapb.ScalarField_StringData{
-						StringData: &schemapb.StringArray{Data: []string{"text1", "text2"}},
-					},
-				},
-			},
-		}
-
-		clipIntData := &schemapb.FieldData{
-			FieldName: "clip[int]",
-			FieldId:   2002,
-			Type:      schemapb.DataType_Int32,
-			Field: &schemapb.FieldData_Scalars{
-				Scalars: &schemapb.ScalarField{
-					Data: &schemapb.ScalarField_IntData{
-						IntData: &schemapb.IntArray{Data: []int32{100, 200}},
-					},
-				},
-			},
-		}
-
-		// Note: clip[embedding] is NOT included in query results
-		fieldsData := []*schemapb.FieldData{clipStrData, clipIntData}
-		outputFields := []string{"clip[str]", "clip[int]"}
-
-		schema := &schemapb.CollectionSchema{
-			Fields: []*schemapb.FieldSchema{
-				{
-					FieldID:      100,
-					Name:         "pk",
-					IsPrimaryKey: true,
-					DataType:     schemapb.DataType_Int64,
-				},
-			},
-			StructArrayFields: []*schemapb.StructArrayFieldSchema{
-				{
-					FieldID: 200,
-					Name:    "clip",
-					Fields: []*schemapb.FieldSchema{
-						{
-							FieldID:  2001,
-							Name:     "clip[str]",
-							DataType: schemapb.DataType_VarChar,
-						},
-						{
-							FieldID:  2002,
-							Name:     "clip[int]",
-							DataType: schemapb.DataType_Int32,
-						},
-						{
-							FieldID:  2003,
-							Name:     "clip[embedding]",
-							DataType: schemapb.DataType_FloatVector,
-							TypeParams: []*commonpb.KeyValuePair{
-								{Key: "dim", Value: "128"},
-							},
-						},
-					},
-				},
-			},
-		}
-
-		resultFieldsData, resultOutputFields := reconstructStructFieldData(fieldsData, outputFields, schema)
-
-		// Check result
-		assert.Len(t, resultFieldsData, 1, "Should have one reconstructed struct field")
-		assert.Len(t, resultOutputFields, 1, "Output fields should have one")
-
-		structField := resultFieldsData[0]
-		assert.Equal(t, "clip", structField.FieldName)
-		assert.Equal(t, int64(200), structField.FieldId)
-		assert.Equal(t, schemapb.DataType_ArrayOfStruct, structField.Type)
-		assert.Equal(t, "clip", resultOutputFields[0])
-
-		// Check that struct only contains the 2 queried fields, NOT the embedding field
-		structArrays := structField.GetStructArrays()
-		assert.NotNil(t, structArrays)
-		assert.Len(t, structArrays.Fields, 2, "Struct should only contain 2 queried fields, not 3")
-
-		// Verify the field names have been restored to original names
-		var foundStr, foundInt bool
-		for _, field := range structArrays.Fields {
-			switch field.FieldId {
-			case 2001:
-				assert.Equal(t, "str", field.FieldName, "Field name should be restored to original")
-				assert.Equal(t, schemapb.DataType_VarChar, field.Type)
-				foundStr = true
-			case 2002:
-				assert.Equal(t, "int", field.FieldName, "Field name should be restored to original")
-				assert.Equal(t, schemapb.DataType_Int32, field.Type)
-				foundInt = true
-			case 2003:
-				assert.Fail(t, "Should not include embedding field as it was not queried")
-			}
-		}
-		assert.True(t, foundStr, "Should find str field")
-		assert.True(t, foundInt, "Should find int field")
-	})
-}
-
 func TestLackOfFieldsDataBySchema(t *testing.T) {
 	schema := &schemapb.CollectionSchema{
 		Fields: []*schemapb.FieldSchema{
@@ -7084,46 +6400,3 @@ func TestFailMetricLabel(t *testing.T) {
 	assertLabels(metrics.CauseCancel, context.Canceled)
 	assertLabels(metrics.CauseCancel, errors.Wrap(context.Canceled, "rpc aborted"))
 }
-
-func TestResolveTimezone(t *testing.T) {
-	ctx := context.Background()
-	colInfoWithTz := &collectionInfo{
-		Properties: []*commonpb.KeyValuePair{
-			{Key: common.TimezoneKey, Value: "America/New_York"},
-		},
-	}
-	colInfoWithoutTz := &collectionInfo{}
-
-	t.Run("request timezone wins over collection timezone", func(t *testing.T) {
-		params := []*commonpb.KeyValuePair{{Key: common.TimezoneKey, Value: "Asia/Shanghai"}}
-		tz, err := resolveTimezone(ctx, params, colInfoWithTz)
-		assert.NoError(t, err)
-		assert.Equal(t, "Asia/Shanghai", tz)
-	})
-
-	t.Run("invalid request timezone is rejected as ParameterInvalid", func(t *testing.T) {
-		params := []*commonpb.KeyValuePair{{Key: common.TimezoneKey, Value: "Not/AZone"}}
-		_, err := resolveTimezone(ctx, params, colInfoWithTz)
-		assert.Error(t, err)
-		assert.ErrorIs(t, err, merr.ErrParameterInvalid)
-	})
-
-	t.Run("absent request timezone falls back to collection timezone", func(t *testing.T) {
-		tz, err := resolveTimezone(ctx, nil, colInfoWithTz)
-		assert.NoError(t, err)
-		assert.Equal(t, "America/New_York", tz)
-	})
-
-	t.Run("empty request timezone is treated as unspecified", func(t *testing.T) {
-		params := []*commonpb.KeyValuePair{{Key: common.TimezoneKey, Value: ""}}
-		tz, err := resolveTimezone(ctx, params, colInfoWithTz)
-		assert.NoError(t, err)
-		assert.Equal(t, "America/New_York", tz)
-	})
-
-	t.Run("no request or collection timezone defaults to UTC", func(t *testing.T) {
-		tz, err := resolveTimezone(ctx, nil, colInfoWithoutTz)
-		assert.NoError(t, err)
-		assert.Equal(t, common.DefaultTimezone, tz)
-	})
-}