fix: partitionIDs should be union of partitionIDs of each subquery

Signed-off-by: zhenshan.cao <[email protected]>
czs007 · Apr 11, 2024 · 6470349 · 6470349
1 parent 35f3959
commit 6470349
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go
@@ -2995,7 +2995,7 @@ func (node *Proxy) hybridSearch(ctx context.Context, request *milvuspb.HybridSea
 
 	metrics.ProxySearchVectors.
 		WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), request.GetDbName(), request.GetCollectionName()).
-		Add(float64(len(request.GetRequests())))
+		Add(float64(len(request.GetRequests()) * int(qt.SearchRequest.GetNq())))
 
 	searchDur := tr.ElapseSpan().Milliseconds()
 	metrics.ProxySQLatency.WithLabelValues(

diff --git a/internal/proxy/task_search.go b/internal/proxy/task_search.go
@@ -348,6 +348,10 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error {
 			zap.Int64s("plan.OutputFieldIds", plan.GetOutputFieldIds()),
 			zap.Stringer("plan", plan)) // may be very large if large term passed.
 	}
+	// used for requery
+	if t.partitionKeyMode {
+		t.SearchRequest.PartitionIDs = t.partitionIDsSet.Collect()
+	}
 	var err error
 	t.reScorers, err = NewReScorers(len(t.request.GetSubReqs()), t.request.GetSearchParams())
 	if err != nil {