diff --git a/internal/querynode/task_search.go b/internal/querynode/task_search.go index f86f5af4df..65dced852e 100644 --- a/internal/querynode/task_search.go +++ b/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 diff --git a/internal/querynode/task_search_test.go b/internal/querynode/task_search_test.go new file mode 100644 index 0000000000..19f0c5c23b --- /dev/null +++ b/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)) +}