enhance: move the search utilities from querynode into new module (#3…

…7531) issue: #33285 - The search utilities will be shared between query node and streaming node. Signed-off-by: chyezh <[email protected]>
milvus-io · Nov 11, 2024 · fca946d · fca946d
1 parent 5e90f34
commit fca946d
Show file tree

Hide file tree

Showing 19 changed files with 26 additions and 22 deletions.
diff --git a/internal/querynodev2/delegator/delegator.go b/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"

diff --git a/internal/querynodev2/server.go b/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,
 		)
 

diff --git a/internal/querynodev2/server_test.go b/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"
 )

diff --git a/internal/querynodev2/services.go b/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 {

diff --git a/internal/querynodev2/tasks/query_stream_task.go b/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,

diff --git a/internal/querynodev2/tasks/query_task.go b/internal/querynodev2/tasks/query_task.go
@@ -15,14 +15,15 @@ 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"
 	"github.com/milvus-io/milvus/pkg/util/timerecord"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
-var _ Task = &QueryTask{}
+var _ scheduler.Task = &QueryTask{}
 
 func NewQueryTask(ctx context.Context,
 	collection *segments.Collection,

diff --git a/internal/querynodev2/tasks/search_task.go b/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
 }
 

diff --git a/...querynodev2/optimizers/mock_query_hook.go → .../searchutil/optimizers/mock_query_hook.go b/...querynodev2/optimizers/mock_query_hook.go → .../searchutil/optimizers/mock_query_hook.go
diff --git a/...rnal/querynodev2/optimizers/query_hook.go → .../util/searchutil/optimizers/query_hook.go b/...rnal/querynodev2/optimizers/query_hook.go → .../util/searchutil/optimizers/query_hook.go
diff --git a/...querynodev2/optimizers/query_hook_test.go → .../searchutil/optimizers/query_hook_test.go b/...querynodev2/optimizers/query_hook_test.go → .../searchutil/optimizers/query_hook_test.go
diff --git a/...nodev2/tasks/concurrent_safe_scheduler.go → ...il/scheduler/concurrent_safe_scheduler.go b/...nodev2/tasks/concurrent_safe_scheduler.go → ...il/scheduler/concurrent_safe_scheduler.go
@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"fmt"

diff --git a/...2/tasks/concurrent_safe_scheduler_test.go → ...heduler/concurrent_safe_scheduler_test.go b/...2/tasks/concurrent_safe_scheduler_test.go → ...heduler/concurrent_safe_scheduler_test.go
@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"context"

diff --git a/internal/querynodev2/tasks/fifo_policy.go → .../util/searchutil/scheduler/fifo_policy.go b/internal/querynodev2/tasks/fifo_policy.go → .../util/searchutil/scheduler/fifo_policy.go
@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"github.com/milvus-io/milvus/pkg/util/paramtable"

diff --git a/internal/querynodev2/tasks/mock_task_test.go → ...il/searchutil/scheduler/mock_task_test.go b/internal/querynodev2/tasks/mock_task_test.go → ...il/searchutil/scheduler/mock_task_test.go
@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"context"

diff --git a/internal/querynodev2/tasks/policy_test.go → .../util/searchutil/scheduler/policy_test.go b/internal/querynodev2/tasks/policy_test.go → .../util/searchutil/scheduler/policy_test.go
@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"fmt"

diff --git a/internal/querynodev2/tasks/queues.go → internal/util/searchutil/scheduler/queues.go b/internal/querynodev2/tasks/queues.go → internal/util/searchutil/scheduler/queues.go
@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"container/ring"

diff --git a/internal/querynodev2/tasks/queues_test.go → .../util/searchutil/scheduler/queues_test.go b/internal/querynodev2/tasks/queues_test.go → .../util/searchutil/scheduler/queues_test.go
@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"fmt"

diff --git a/internal/querynodev2/tasks/tasks.go → internal/util/searchutil/scheduler/tasks.go b/internal/querynodev2/tasks/tasks.go → internal/util/searchutil/scheduler/tasks.go
@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import "github.com/milvus-io/milvus/internal/proto/internalpb"
 

diff --git a/...ynodev2/tasks/user_task_polling_policy.go → ...til/scheduler/user_task_polling_policy.go b/...ynodev2/tasks/user_task_polling_policy.go → ...til/scheduler/user_task_polling_policy.go
@@ -1,4 +1,4 @@
-package tasks
+package scheduler
 
 import (
 	"fmt"