SERVER-58888 Account for pushed down stages in union explain with mode 'executionStats'

(cherry picked from commit 432d231645aa6496a291996a1385409b08bcd7dc)

2 files changed, 21 insertions, 2 deletions

diff --git a/jstests/aggregation/sources/unionWith/unionWith_explain.js b/jstests/aggregation/sources/unionWith/unionWith_explain.js
index f7841a5a14a..507a899e602 100644
--- a/jstests/aggregation/sources/unionWith/unionWith_explain.js
+++ b/jstests/aggregation/sources/unionWith/unionWith_explain.js
@@ -81,11 +81,11 @@ function assertExplainEq(unionExplain, regularExplain) {
                 unionStage.$unionWith.pipeline[0].$cursor.executionStats.executionStages;
             const realStats = regularExplain.executionStats.executionStages;
             assert(docEqWithIgnoredFields(unionSubStats, realStats),
-                   buildErrorString(unionSubStats, realStats));
+                   buildErrorString(unionSubStats, realStats, "stages"));
         } else {
             const realExplain = regularExplain.stages;
             assert(arrayEq(unionSubExplain, realExplain),
-                   buildErrorString(unionSubExplain, realExplain));
+                   buildErrorString(unionSubExplain, realExplain, "stages"));
         }
     }
 }
@@ -212,4 +212,15 @@ assertExplainEq(result, expectedResult);
 
 // Test a nested $unionWith which itself should perform an index scan.
 testPipeline([{$unionWith: {coll: indexedColl.getName(), pipeline: [{$match: {val: {$gt: 0}}}]}}]);
+
+// Similar test as above, except the $match is pushed down to the inner pipeline as part of a
+// rewrite optimization.
+const res = db.adminCommand({getParameter: 1, "failpoint.disablePipelineOptimization": 1});
+assert.commandWorked(res);
+if (!res["failpoint.disablePipelineOptimization"].mode) {
+    result = collA.explain("executionStats")
+                 .aggregate([{$unionWith: indexedColl.getName()}, {$match: {val: {$gt: 2}}}]);
+    expectedResult = indexedColl.explain("executionStats").aggregate([{$match: {val: {$gt: 2}}}]);
+    assertExplainEq(result, expectedResult);
+}
 })();
diff --git a/src/mongo/db/pipeline/document_source_union_with.cpp b/src/mongo/db/pipeline/document_source_union_with.cpp
index d54e95b8425..cdea4c1985a 100644
--- a/src/mongo/db/pipeline/document_source_union_with.cpp
+++ b/src/mongo/db/pipeline/document_source_union_with.cpp
@@ -267,6 +267,14 @@ Pipeline::SourceContainer::iterator DocumentSourceUnionWith::doOptimizeAt(
     Pipeline::SourceContainer::iterator itr, Pipeline::SourceContainer* container) {
     auto duplicateAcrossUnion = [&](auto&& nextStage) {
         _pipeline->addFinalSource(nextStage->clone());
+        // Apply the same rewrite to the cached pipeline if available.
+        if (pExpCtx->explain >= ExplainOptions::Verbosity::kExecStats) {
+            auto cloneForExplain = nextStage->clone();
+            if (!_cachedPipeline.empty()) {
+                cloneForExplain->setSource(_cachedPipeline.back().get());
+            }
+            _cachedPipeline.push_back(std::move(cloneForExplain));
+        }
         auto newStageItr = container->insert(itr, std::move(nextStage));
         container->erase(std::next(itr));
         return newStageItr == container->begin() ? newStageItr : std::prev(newStageItr);