From 822e25344a7158bb38ec9a7126e4b5e90ddb359f Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Thu, 9 Nov 2023 20:33:33 -0500
Subject: [PATCH] Fix slice collectors to leaves association with post filter
 (#11134) (#11150)

* Fix slice collectors to leaves association with profile enabled (#11134)



* Update server/src/main/java/org/opensearch/search/profile/query/ConcurrentQueryProfileBreakdown.java





---------




(cherry picked from commit 3eda422ae1979d3e1e7b158abab7d179f919f868)

Signed-off-by: Ticheng Lin <ticheng@amazon.com>
Signed-off-by: Ticheng Lin <51488860+ticheng-aws@users.noreply.github.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Andriy Redko <drreta@gmail.com>
---
 .../opensearch/common/lucene/search/FilteredCollector.java   | 4 ++++
 .../profile/query/ConcurrentQueryProfileBreakdown.java       | 5 ++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/server/src/main/java/org/opensearch/common/lucene/search/FilteredCollector.java b/server/src/main/java/org/opensearch/common/lucene/search/FilteredCollector.java
index 0d7a8866f7788..b5c0e84a10308 100644
--- a/server/src/main/java/org/opensearch/common/lucene/search/FilteredCollector.java
+++ b/server/src/main/java/org/opensearch/common/lucene/search/FilteredCollector.java
@@ -40,6 +40,7 @@
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
 import org.opensearch.common.lucene.Lucene;
+import org.opensearch.search.profile.query.ProfileWeight;
 
 import java.io.IOException;
 
@@ -64,6 +65,9 @@ public Collector getCollector() {
 
     @Override
     public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+        if (filter instanceof ProfileWeight) {
+            ((ProfileWeight) filter).associateCollectorToLeaves(context, collector);
+        }
         final ScorerSupplier filterScorerSupplier = filter.scorerSupplier(context);
         final LeafCollector in = collector.getLeafCollector(context);
         final Bits bits = Lucene.asSequentialAccessBits(context.reader().maxDoc(), filterScorerSupplier);
diff --git a/server/src/main/java/org/opensearch/search/profile/query/ConcurrentQueryProfileBreakdown.java b/server/src/main/java/org/opensearch/search/profile/query/ConcurrentQueryProfileBreakdown.java
index 59ef01f9f947a..e1d41227a22f7 100644
--- a/server/src/main/java/org/opensearch/search/profile/query/ConcurrentQueryProfileBreakdown.java
+++ b/server/src/main/java/org/opensearch/search/profile/query/ConcurrentQueryProfileBreakdown.java
@@ -83,7 +83,10 @@ public Map<String, Long> toBreakdownMap() {
             // creates a new weight and breakdown map for each rewritten query. This new breakdown map captures the timing information for
             // the new rewritten query. The sliceCollectorsToLeaves is empty because this breakdown for rewritten query gets created later
             // in search leaf path which doesn't have collector. Also, this is not needed since this breakdown is per leaf and there is no
-            // concurrency involved. An empty sliceCollectorsToLeaves could also happen in the case of early termination.
+            // concurrency involved.
+            assert contexts.size() == 1 : "Unexpected size: "
+                + contexts.size()
+                + " of leaves breakdown in ConcurrentQueryProfileBreakdown of rewritten query for a leaf.";
             AbstractProfileBreakdown<QueryTimingType> breakdown = contexts.values().iterator().next();
             queryNodeTime = breakdown.toNodeTime() + createWeightTime;
             maxSliceNodeTime = 0L;