SERVER-50246 Fix $unionWith explain with mode 'executionStats' to account for pushed down stages

3 files changed, 44 insertions, 9 deletions

diff --git a/src/mongo/db/pipeline/document_source_union_with.cpp b/src/mongo/db/pipeline/document_source_union_with.cpp
index 63d7e124a64..294b16accc3 100644
--- a/src/mongo/db/pipeline/document_source_union_with.cpp
+++ b/src/mongo/db/pipeline/document_source_union_with.cpp
@@ -66,7 +66,7 @@ std::unique_ptr<Pipeline, PipelineDeleter> buildPipelineFromViewDefinition(
     auto unionExpCtx = expCtx->copyForSubPipeline(resolvedNs.ns);
 
     if (resolvedNs.pipeline.empty()) {
-        return Pipeline::parse(std::move(currentPipeline), unionExpCtx, validatorCallback);
+        return Pipeline::parse(currentPipeline, unionExpCtx, validatorCallback);
     }
     auto resolvedPipeline = std::move(resolvedNs.pipeline);
     resolvedPipeline.reserve(currentPipeline.size() + resolvedPipeline.size());
@@ -251,18 +251,40 @@ void DocumentSourceUnionWith::doDispose() {
 
 Value DocumentSourceUnionWith::serialize(boost::optional<ExplainOptions::Verbosity> explain) const {
     if (explain) {
+        // There are several different possible states depending on the explain verbosity as well as
+        // the other stages in the pipeline:
+        //  * If verbosity is queryPlanner, then the sub-pipeline should be untouched and we can
+        //  explain it directly.
+        //  * If verbosity is execStats or allPlansExecution, then whether or not to explain the
+        //  sub-pipeline depends on if we've started reading from it. For instance, there could be a
+        //  $limit stage after the $unionWith which results in only reading from the base collection
+        //  branch and not the sub-pipeline.
+        Pipeline* pipeCopy = nullptr;
+        if (*explain == ExplainOptions::Verbosity::kQueryPlanner) {
+            pipeCopy = Pipeline::create(_pipeline->getSources(), _pipeline->getContext()).release();
+        } else if (*explain >= ExplainOptions::Verbosity::kExecStats &&
+                   _executionState > ExecutionProgress::kIteratingSource) {
+            // We've either exhausted the sub-pipeline or at least started iterating it. Use the
+            // cached pipeline to get the explain output since the '_pipeline' may have been
+            // modified for any optimizations or pushdowns into the initial $cursor stage.
+            pipeCopy = _cachedPipeline;
+        } else {
+            // The plan does not require reading from the sub-pipeline, so just include the
+            // serialization in the explain output.
+            BSONArrayBuilder bab;
+            for (auto&& stage : _pipeline->serialize())
+                bab << stage;
+            return Value(DOC(getSourceName() << DOC("coll" << _pipeline->getContext()->ns.coll()
+                                                           << "pipeline" << bab.arr())));
+        }
 
-        auto pipeCopy = Pipeline::create(_pipeline->getSources(), _pipeline->getContext());
-
-        // If we have already started getting documents from the sub-pipeline, this is an explain
-        // that has done some execution. We don't want to serialize the mergeCursors stage, so if
-        // we have a cursor stage we tell the process interface to remove it in the case it is a
-        // mergeCursors stage.
+        invariant(pipeCopy);
         BSONObj explainLocal =
-            pExpCtx->mongoProcessInterface->preparePipelineAndExplain(pipeCopy.release(), *explain);
+            pExpCtx->mongoProcessInterface->preparePipelineAndExplain(pipeCopy, *explain);
         LOGV2_DEBUG(4553501, 3, "$unionWith attached cursor to pipeline for explain");
         // We expect this to be an explanation of a pipeline -- there should only be one field.
         invariant(explainLocal.nFields() == 1);
+
         return Value(
             DOC(getSourceName() << DOC("coll" << _pipeline->getContext()->ns.coll() << "pipeline"
                                               << explainLocal.firstElement())));
diff --git a/src/mongo/db/pipeline/document_source_union_with.h b/src/mongo/db/pipeline/document_source_union_with.h
index 05c3d5d8505..5686f569374 100644
--- a/src/mongo/db/pipeline/document_source_union_with.h
+++ b/src/mongo/db/pipeline/document_source_union_with.h
@@ -62,7 +62,14 @@ public:
 
     DocumentSourceUnionWith(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                             std::unique_ptr<Pipeline, PipelineDeleter> pipeline)
-        : DocumentSource(kStageName, expCtx), _pipeline(std::move(pipeline)) {}
+        : DocumentSource(kStageName, expCtx), _pipeline(std::move(pipeline)) {
+        // If this pipeline is being run as part of explain, then cache a copy to use later during
+        // serialization.
+        if (expCtx->explain >= ExplainOptions::Verbosity::kExecStats) {
+            _cachedPipeline =
+                Pipeline::create(_pipeline->getSources(), _pipeline->getContext()).release();
+        }
+    }
 
     ~DocumentSourceUnionWith();
 
@@ -155,6 +162,7 @@ private:
     void addViewDefinition(NamespaceString nss, std::vector<BSONObj> viewPipeline);
 
     std::unique_ptr<Pipeline, PipelineDeleter> _pipeline;
+    Pipeline* _cachedPipeline = nullptr;
     bool _usedDisk = false;
     ExecutionProgress _executionState = ExecutionProgress::kIteratingSource;
 };
diff --git a/src/mongo/db/pipeline/process_interface/non_shardsvr_process_interface.cpp b/src/mongo/db/pipeline/process_interface/non_shardsvr_process_interface.cpp
index 3cae0550907..537d89749d1 100644
--- a/src/mongo/db/pipeline/process_interface/non_shardsvr_process_interface.cpp
+++ b/src/mongo/db/pipeline/process_interface/non_shardsvr_process_interface.cpp
@@ -187,6 +187,11 @@ BSONObj NonShardServerProcessInterface::preparePipelineAndExplain(
         ownedPipeline = nullptr;
     } else {
         auto pipelineWithCursor = attachCursorSourceToPipelineForLocalRead(ownedPipeline);
+        // If we need execution stats, this runs the plan in order to gather the stats.
+        if (verbosity >= ExplainOptions::Verbosity::kExecStats) {
+            while (pipelineWithCursor->getNext()) {
+            }
+        }
         pipelineVec = pipelineWithCursor->writeExplainOps(verbosity);
     }
     BSONArrayBuilder bab;