Improve search task merge policy (#22006)

Signed-off-by: yah01 <yang.cen@zilliz.com>

internal/querynode/task_search.go

@@ -348,12 +348,16 @@ func (s *searchTask) CanMergeWith(t readTask) bool {
 		return false
 	}
 
-	pre := s.NQ * s.TopK * 1.0
-	newTopK := s.TopK
-	if newTopK < s2.TopK {
-		newTopK = s2.TopK
+	pre := s.NQ * s.TopK
+	if s2.NQ*s2.TopK < pre {
+		pre = s2.NQ * s2.TopK
 	}
-	after := (s.NQ + s2.NQ) * newTopK
+
+	maxTopk := s.TopK
+	if maxTopk < s2.TopK {
+		maxTopk = s2.TopK
+	}
+	after := (s.NQ + s2.NQ) * maxTopk
 
 	if pre == 0 {
 		return false

internal/querynode/task_search_test.go

@@ -0,0 +1,51 @@
+// 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 querynode
+
+import (
+	"testing"
+
+	"github.com/milvus-io/milvus/internal/proto/planpb"
+	"github.com/stretchr/testify/suite"
+)
+
+type SearchTaskSuite struct {
+	suite.Suite
+}
+
+func (s *SearchTaskSuite) TestMerge() {
+	plan := &planpb.PlanNode{
+		Node: &planpb.PlanNode_VectorAnns{},
+	}
+
+	s1 := &searchTask{
+		NQ:   1,
+		TopK: 1000,
+		plan: plan,
+	}
+	s2 := &searchTask{
+		NQ:   1,
+		TopK: 1,
+		plan: plan,
+	}
+
+	s.Equal(s1.CanMergeWith(s2), s2.CanMergeWith(s1))
+}
+
+func TestSearchTask(t *testing.T) {
+	suite.Run(t, new(SearchTaskSuite))
+}