SERVER-49469 Fix union inner pipeline explain

(cherry picked from commit 274192ab7d429710e0f69f0b1d922005fca9d516)

2 files changed, 64 insertions, 14 deletions

diff --git a/jstests/aggregation/extras/utils.js b/jstests/aggregation/extras/utils.js
index 2dd25601f60..3ea7d4a9724 100644
--- a/jstests/aggregation/extras/utils.js
+++ b/jstests/aggregation/extras/utils.js
@@ -28,7 +28,7 @@ function testExpressionWithCollation(coll, expression, result, collationSpec) {
  * with the exception of '_id'. If 'al' and 'ar' are neither object nor arrays, they must compare
  * equal using 'valueComparator', or == if not provided.
  */
-function anyEq(al, ar, verbose = false, valueComparator) {
+function anyEq(al, ar, verbose = false, valueComparator, fieldsToSkip = []) {
     const debug = msg => verbose ? print(msg) : null;  // Helper to log 'msg' iff 'verbose' is true.
 
     if (al instanceof Array) {
@@ -49,7 +49,7 @@ function anyEq(al, ar, verbose = false, valueComparator) {
             return false;
         }
 
-        if (!documentEq(al, ar, verbose, valueComparator)) {
+        if (!documentEq(al, ar, verbose, valueComparator, fieldsToSkip)) {
             debug(`anyEq: documentEq(al, ar): false; al=${tojson(al)}, ar=${tojson(ar)}`);
             return false;
         }
@@ -74,9 +74,10 @@ function customDocumentEq({left, right, verbose, valueComparator}) {
 
 /**
  * Compare two documents for equality. Returns true or false. Only equal if they have the exact same
- * set of properties, and all the properties' values match.
+ * set of properties, and all the properties' values match except the values with names in the
+ * fieldsToSkip array. The fields in fieldsToSkip will be skipped at all levels of the document.
  */
-function documentEq(dl, dr, verbose = false, valueComparator) {
+function documentEq(dl, dr, verbose = false, valueComparator, fieldsToSkip = []) {
     const debug = msg => verbose ? print(msg) : null;  // Helper to log 'msg' iff 'verbose' is true.
 
     // Make sure these are both objects.
@@ -102,10 +103,10 @@ function documentEq(dl, dr, verbose = false, valueComparator) {
         }
 
         // If the property is the _id, they don't have to be equal.
-        if (propertyName == '_id')
+        if (propertyName == '_id' || fieldsToSkip.includes(propertyName))
             continue;
 
-        if (!anyEq(dl[propertyName], dr[propertyName], verbose, valueComparator)) {
+        if (!anyEq(dl[propertyName], dr[propertyName], verbose, valueComparator, fieldsToSkip)) {
             return false;
         }
     }
@@ -134,7 +135,7 @@ function documentEq(dl, dr, verbose = false, valueComparator) {
  * This is a predicate, not an assertion; use assert.sameMembers() in assert.js for the
  * equivalent assertion.
  */
-function arrayEq(al, ar, verbose = false, valueComparator) {
+function arrayEq(al, ar, verbose = false, valueComparator, fieldsToSkip = []) {
     const debug = msg => verbose ? print(msg) : null;  // Helper to log 'msg' iff 'verbose' is true.
 
     // Check that these are both arrays.
@@ -159,7 +160,7 @@ function arrayEq(al, ar, verbose = false, valueComparator) {
         let foundMatch = false;
         for (let i = 0; i < ar.length; ++i) {
             if (!matchedElementsInRight.has(i) &&
-                anyEq(leftElem, ar[i], verbose, valueComparator)) {
+                anyEq(leftElem, ar[i], verbose, valueComparator, fieldsToSkip)) {
                 matchedElementsInRight.add(i);  // Don't use the same value each time.
                 foundMatch = true;
                 break;
diff --git a/jstests/aggregation/sources/unionWith/unionWith_explain.js b/jstests/aggregation/sources/unionWith/unionWith_explain.js
index 39085fa4657..eace8e715c8 100644
--- a/jstests/aggregation/sources/unionWith/unionWith_explain.js
+++ b/jstests/aggregation/sources/unionWith/unionWith_explain.js
@@ -35,6 +35,17 @@ function buildErrorString(unionExplain, realExplain, field) {
         "\nRegular:\n" + tojson(realExplain);
 }
 
+function docEqWithIgnoredFields(union, regular) {
+    return documentEq(union, regular, false /* verbose */, null /* valueComparator */, [
+        "executionTimeMillis",
+        "executionTimeMillisEstimate",
+        "saveState",
+        "restoreState",
+        "works",
+        "needTime",
+    ]);
+}
+
 function assertExplainEq(unionExplain, regularExplain) {
     if (FixtureHelpers.isMongos(testDB)) {
         const splitPipe = unionExplain.splitPipeline;
@@ -48,22 +59,32 @@ function assertExplainEq(unionExplain, regularExplain) {
                       regularExplain.splitPipeline,
                       buildErrorString(unionSubExplain, regularExplain, "splitPipeline"));
         } else {
-            assert(documentEq(unionSubExplain.splitPipeline, regularExplain.splitPipeline),
-                   buildErrorString(unionSubExplain, regularExplain, "splitPipeline"));
+            assert(
+                docEqWithIgnoredFields(unionSubExplain.splitPipeline, regularExplain.splitPipeline),
+                buildErrorString(unionSubExplain, regularExplain, "splitPipeline"));
         }
         assert.eq(unionSubExplain.mergeType,
                   regularExplain.mergeType,
                   buildErrorString(unionSubExplain, regularExplain, "mergeType"));
-        assert(documentEq(unionSubExplain.shards, regularExplain.shards),
+        assert(docEqWithIgnoredFields(unionSubExplain.shards, regularExplain.shards),
                buildErrorString(unionSubExplain, regularExplain, "shards"));
     } else {
         const unionStage = getUnionWithStage(unionExplain.stages);
         const unionSubExplain = unionStage.$unionWith.pipeline;
-        const realExplain = regularExplain.stages;
-        assert(arrayEq(unionSubExplain, realExplain),
-               buildErrorString(unionSubExplain, realExplain));
+        if ("executionStats" in unionSubExplain[0].$cursor) {
+            const unionSubStats =
+                unionStage.$unionWith.pipeline[0].$cursor.executionStats.executionStages;
+            const realStats = regularExplain.executionStats.executionStages;
+            assert(docEqWithIgnoredFields(unionSubStats, realStats),
+                   buildErrorString(unionSubStats, realStats));
+        } else {
+            const realExplain = regularExplain.stages;
+            assert(arrayEq(unionSubExplain, realExplain),
+                   buildErrorString(unionSubExplain, realExplain));
+        }
     }
 }
+
 function testPipeline(pipeline) {
     let unionResult = collA.aggregate([{$unionWith: {coll: collB.getName(), pipeline: pipeline}}],
                                       {explain: true});
@@ -126,4 +147,32 @@ var result = assert.commandWorked(testDB.runCommand({
         cursor: {},
     }
 }));
+
+// Test that execution stats inner cursor is populated.
+result = collA.explain("executionStats").aggregate([{"$unionWith": collB.getName()}]);
+var expectedResult = collB.explain("executionStats").aggregate([]);
+assert(result.ok, result);
+assert(expectedResult.ok, result);
+// If we attached a fresh cursor stage, the number returned would still be zero.
+if (FixtureHelpers.isMongos(testDB)) {
+    if (result.splitPipeline != null) {
+        const pipeline = result.splitPipeline.mergerPart;
+        const unionStage = getUnionWithStage(pipeline);
+        assert(docEqWithIgnoredFields(expectedResult.shards, unionStage.$unionWith.pipeline.shards),
+               buildErrorString(unionStage, expectedResult));
+    }
+} else {
+    assert(result.stages[1].$unionWith.pipeline[0].$cursor.executionStats.nreturned != 0, result);
+}
+
+// Test an index scan.
+const indexedColl = testDB.indexed;
+assert.commandWorked(indexedColl.createIndex({val: 1}));
+indexedColl.insert([{val: 0}, {val: 1}, {val: 2}, {val: 3}]);
+result = collA.explain("executionStats").aggregate([
+    {$unionWith: {coll: indexedColl.getName(), pipeline: [{$match: {val: {$gt: 2}}}]}}
+]);
+expectedResult = indexedColl.explain("executionStats").aggregate([{$match: {val: {$gt: 2}}}]);
+
+assertExplainEq(result, expectedResult);
 })();