Fix slice collectors to leaves association with profile enabled

Signed-off-by: Ticheng Lin <[email protected]>

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

@@ -65,6 +65,25 @@ public AbstractProfileBreakdown<QueryTimingType> context(Object context) {
  @Override
  public Map<String, Long> toBreakdownMap() {
  final Map<String, Long> topLevelBreakdownMapWithWeightTime = super.toBreakdownMap();
+ int totalLeafCount = 0;
+ for (List<LeafReaderContext> leaves : sliceCollectorsToLeaves.values()) {
+ totalLeafCount += leaves.size();
+ }
+ if (contexts.size() != totalLeafCount) {
+ System.out.println("$$$$$");
+ // If the leaf counts in contexts map don't match those in sliceCollectorsToLeaves, this means that the current
+ // node is outside the concurrent search path, which occurred before the contextIndexSearcher::searchLeaf.
+ // For example, in the post filter query, we may create a profile node during filtered collector context creation.
+ // In this case, as there is no concurrent implementation, we can simply return the accumulated breakdown result
+ // instead of the statistical result.
+ final Map<String, Long> map = new HashMap<>(topLevelBreakdownMapWithWeightTime);
+ for (final AbstractProfileBreakdown<QueryTimingType> context : contexts.values()) {
+ for (final Map.Entry<String, Long> entry : context.toBreakdownMap().entrySet()) {
+ map.merge(entry.getKey(), entry.getValue(), Long::sum);
+ }
+ }
+ return map;
+ }
  final long createWeightStartTime = topLevelBreakdownMapWithWeightTime.get(
  QueryTimingType.CREATE_WEIGHT + TIMING_TYPE_START_TIME_SUFFIX
  );
@@ -83,7 +102,7 @@ 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.
  AbstractProfileBreakdown<QueryTimingType> breakdown = contexts.values().iterator().next();
  queryNodeTime = breakdown.toNodeTime() + createWeightTime;
  maxSliceNodeTime = 0L;