enhance: move the search utilities from querynode into new module (#37531)

issue: #33285

- The search utilities will be shared between query node and streaming
node.

Signed-off-by: chyezh <chyezh@outlook.com>

internal/querynodev2/delegator/delegator.go

@@ -39,13 +39,13 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/querypb"
 	"github.com/milvus-io/milvus/internal/querynodev2/cluster"
 	"github.com/milvus-io/milvus/internal/querynodev2/delegator/deletebuffer"
-	"github.com/milvus-io/milvus/internal/querynodev2/optimizers"
 	"github.com/milvus-io/milvus/internal/querynodev2/pkoracle"
 	"github.com/milvus-io/milvus/internal/querynodev2/segments"
 	"github.com/milvus-io/milvus/internal/querynodev2/tsafe"
 	"github.com/milvus-io/milvus/internal/storage"
 	"github.com/milvus-io/milvus/internal/util/function"
 	"github.com/milvus-io/milvus/internal/util/reduce"
+	"github.com/milvus-io/milvus/internal/util/searchutil/optimizers"
 	"github.com/milvus-io/milvus/internal/util/streamrpc"
 	"github.com/milvus-io/milvus/pkg/common"
 	"github.com/milvus-io/milvus/pkg/log"

internal/querynodev2/server.go

@@ -50,16 +50,16 @@ import (
 	grpcquerynodeclient "github.com/milvus-io/milvus/internal/distributed/querynode/client"
 	"github.com/milvus-io/milvus/internal/querynodev2/cluster"
 	"github.com/milvus-io/milvus/internal/querynodev2/delegator"
-	"github.com/milvus-io/milvus/internal/querynodev2/optimizers"
 	"github.com/milvus-io/milvus/internal/querynodev2/pipeline"
 	"github.com/milvus-io/milvus/internal/querynodev2/segments"
-	"github.com/milvus-io/milvus/internal/querynodev2/tasks"
 	"github.com/milvus-io/milvus/internal/querynodev2/tsafe"
 	"github.com/milvus-io/milvus/internal/registry"
 	"github.com/milvus-io/milvus/internal/storage"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/dependency"
 	"github.com/milvus-io/milvus/internal/util/initcore"
+	"github.com/milvus-io/milvus/internal/util/searchutil/optimizers"
+	"github.com/milvus-io/milvus/internal/util/searchutil/scheduler"
 	"github.com/milvus-io/milvus/internal/util/sessionutil"
 	"github.com/milvus-io/milvus/pkg/config"
 	"github.com/milvus-io/milvus/pkg/log"
@@ -113,7 +113,7 @@ type QueryNode struct {
 	loader segments.Loader
 
 	// Search/Query
-	scheduler tasks.Scheduler
+	scheduler scheduler.Scheduler
 
 	// etcd client
 	etcdCli *clientv3.Client
@@ -339,7 +339,7 @@ func (node *QueryNode) Init() error {
 		}
 
 		schedulePolicy := paramtable.Get().QueryNodeCfg.SchedulePolicyName.GetValue()
-		node.scheduler = tasks.NewScheduler(
+		node.scheduler = scheduler.NewScheduler(
 			schedulePolicy,
 		)
 

internal/querynodev2/server_test.go

@@ -34,10 +34,10 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proto/querypb"
-	"github.com/milvus-io/milvus/internal/querynodev2/optimizers"
 	"github.com/milvus-io/milvus/internal/querynodev2/segments"
 	"github.com/milvus-io/milvus/internal/storage"
 	"github.com/milvus-io/milvus/internal/util/dependency"
+	"github.com/milvus-io/milvus/internal/util/searchutil/optimizers"
 	"github.com/milvus-io/milvus/pkg/util/etcd"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 )

internal/querynodev2/services.go

@@ -42,6 +42,7 @@ import (
 	"github.com/milvus-io/milvus/internal/querynodev2/segments"
 	"github.com/milvus-io/milvus/internal/querynodev2/tasks"
 	"github.com/milvus-io/milvus/internal/storage"
+	"github.com/milvus-io/milvus/internal/util/searchutil/scheduler"
 	"github.com/milvus-io/milvus/internal/util/streamrpc"
 	"github.com/milvus-io/milvus/pkg/common"
 	"github.com/milvus-io/milvus/pkg/log"
@@ -689,7 +690,7 @@ func (node *QueryNode) SearchSegments(ctx context.Context, req *querypb.SearchRe
 		return resp, nil
 	}
 
-	var task tasks.Task
+	var task scheduler.Task
 	if paramtable.Get().QueryNodeCfg.UseStreamComputing.GetAsBool() {
 		task = tasks.NewStreamingSearchTask(searchCtx, collection, node.manager, req, node.serverID)
 	} else {

internal/querynodev2/tasks/query_stream_task.go

@@ -6,10 +6,11 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/proto/querypb"
 	"github.com/milvus-io/milvus/internal/querynodev2/segments"
+	"github.com/milvus-io/milvus/internal/util/searchutil/scheduler"
 	"github.com/milvus-io/milvus/internal/util/streamrpc"
 )
 
-var _ Task = &QueryStreamTask{}
+var _ scheduler.Task = &QueryStreamTask{}
 
 func NewQueryStreamTask(ctx context.Context,
 	collection *segments.Collection,

internal/querynodev2/tasks/query_task.go

@@ -15,6 +15,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/querypb"
 	"github.com/milvus-io/milvus/internal/proto/segcorepb"
 	"github.com/milvus-io/milvus/internal/querynodev2/segments"
+	"github.com/milvus-io/milvus/internal/util/searchutil/scheduler"
 	"github.com/milvus-io/milvus/pkg/metrics"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
@@ -22,7 +23,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
-var _ Task = &QueryTask{}
+var _ scheduler.Task = &QueryTask{}
 
 func NewQueryTask(ctx context.Context,
 	collection *segments.Collection,

internal/querynodev2/tasks/search_task.go

@@ -20,6 +20,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/proto/querypb"
 	"github.com/milvus-io/milvus/internal/querynodev2/segments"
+	"github.com/milvus-io/milvus/internal/util/searchutil/scheduler"
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/metrics"
 	"github.com/milvus-io/milvus/pkg/util/funcutil"
@@ -30,8 +31,8 @@ import (
 )
 
 var (
-	_ Task      = &SearchTask{}
-	_ MergeTask = &SearchTask{}
+	_ scheduler.Task      = &SearchTask{}
+	_ scheduler.MergeTask = &SearchTask{}
 )
 
 type SearchTask struct {
@@ -346,7 +347,7 @@ func (t *SearchTask) NQ() int64 {
 	return t.nq
 }
 
-func (t *SearchTask) MergeWith(other Task) bool {
+func (t *SearchTask) MergeWith(other scheduler.Task) bool {
 	switch other := other.(type) {
 	case *SearchTask:
 		return t.Merge(other)
@@ -416,7 +417,7 @@ func NewStreamingSearchTask(ctx context.Context,
 	}
 }
 
-func (t *StreamingSearchTask) MergeWith(other Task) bool {
+func (t *StreamingSearchTask) MergeWith(other scheduler.Task) bool {
 	return false
 }
 

internal/util/searchutil/scheduler/concurrent_safe_scheduler.go

@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"fmt"

internal/util/searchutil/scheduler/concurrent_safe_scheduler_test.go

@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"context"

internal/util/searchutil/scheduler/fifo_policy.go

@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"github.com/milvus-io/milvus/pkg/util/paramtable"

internal/util/searchutil/scheduler/mock_task_test.go

@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"context"

internal/util/searchutil/scheduler/policy_test.go

@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"fmt"

internal/util/searchutil/scheduler/queues.go

@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"container/ring"

internal/util/searchutil/scheduler/queues_test.go

@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"fmt"

internal/util/searchutil/scheduler/tasks.go

@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import "github.com/milvus-io/milvus/internal/proto/internalpb"
 

internal/util/searchutil/scheduler/user_task_polling_policy.go

@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"fmt"