SERVER-60254 Enable group pushdown $unionWith explain tests

2 files changed, 79 insertions, 58 deletions

diff --git a/jstests/aggregation/extras/utils.js b/jstests/aggregation/extras/utils.js
index 626a7a6939b..f25f77d5a7f 100644
--- a/jstests/aggregation/extras/utils.js
+++ b/jstests/aggregation/extras/utils.js
@@ -100,15 +100,15 @@ function documentEq(dl, dr, verbose = false, valueComparator, fieldsToSkip = [])
         if (!dl.hasOwnProperty(propertyName))
             continue;
 
+        if (fieldsToSkip.includes(propertyName))
+            continue;
+
         // The documents aren't equal if they don't both have the property.
         if (!dr.hasOwnProperty(propertyName)) {
             debug('documentEq: dr doesn\'t have property ' + propertyName);
             return false;
         }
 
-        if (fieldsToSkip.includes(propertyName))
-            continue;
-
         if (!anyEq(dl[propertyName], dr[propertyName], verbose, valueComparator, fieldsToSkip)) {
             return false;
         }
@@ -119,6 +119,9 @@ function documentEq(dl, dr, verbose = false, valueComparator, fieldsToSkip = [])
         if (!dr.hasOwnProperty(propertyName))
             continue;
 
+        if (fieldsToSkip.includes(propertyName))
+            continue;
+
         // If dl doesn't have this they are not equal; if it does, we compared it above and know it
         // to be equal.
         if (!dl.hasOwnProperty(propertyName)) {
diff --git a/jstests/aggregation/sources/unionWith/unionWith_explain.js b/jstests/aggregation/sources/unionWith/unionWith_explain.js
index 8044881e3a8..664844c6159 100644
--- a/jstests/aggregation/sources/unionWith/unionWith_explain.js
+++ b/jstests/aggregation/sources/unionWith/unionWith_explain.js
@@ -28,6 +28,27 @@ for (let i = 0; i < docsPerColl; i++) {
     assert.commandWorked(collB.insert({b: i, val: i * 2, groupKey: i}));
     assert.commandWorked(collC.insert({c: i, val: 10 - i, groupKey: i}));
 }
+
+const executionStatsIngoredFields = [
+    "executionTimeMillis",
+    "executionTimeMillisEstimate",
+    "saveState",
+    "restoreState",
+];
+
+const stagesIgnoredFields = [
+    "slots",
+];
+
+const mongosIgnoredFields = [
+    "works",
+    "needTime",
+].concat(executionStatsIngoredFields, stagesIgnoredFields);
+
+const queryPlannerIgnoredFields = [
+    "optimizedPipeline",
+].concat(stagesIgnoredFields);
+
 function getUnionWithStage(explain) {
     if (explain.splitPipeline != null) {
         // If there is only one shard, the whole pipeline will run on that shard.
@@ -49,71 +70,67 @@ 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",
-        "slots",
-    ]);
+function anyEqWithIgnoredFields(union, regular, ignoredFields) {
+    return anyEq(union, regular, false /* verbose */, null /* valueComparator */, ignoredFields);
 }
 
-function arrayEqWithIgnoredFields(union, regular) {
-    return arrayEq(union, regular, false /* verbose */, null /* valueComparator */, [
-        "slots",
-    ]);
+function documentEqWithIgnoredFields(union, regular, ignoredFields) {
+    return documentEq(
+        union, regular, false /* verbose */, null /* valueComparator */, ignoredFields);
 }
 
-function assertExplainEq(unionExplain, regularExplain) {
-    const unionStage = getUnionWithStage(unionExplain);
-    assert(unionStage);
-    const unionSubExplain = unionStage.$unionWith.pipeline;
+function arrayEqWithIgnoredFields(union, regular, ignoredFields) {
+    return arrayEq(union, regular, false /* verbose */, null /* valueComparator */, ignoredFields);
+}
+
+function assertExplainEq(union, regular) {
     if (FixtureHelpers.isMongos(testDB)) {
-        const splitPipe = unionExplain.splitPipeline;
-        if (splitPipe === null) {
-            assert.eq(unionSubExplain.splitPipeline,
-                      regularExplain.splitPipeline,
-                      buildErrorString(unionSubExplain, regularExplain, "splitPipeline"));
-        } else {
-            assert(
-                docEqWithIgnoredFields(unionSubExplain.splitPipeline, regularExplain.splitPipeline),
-                buildErrorString(unionSubExplain, regularExplain, "splitPipeline"));
-        }
-        assert.eq(unionSubExplain.mergeType,
-                  regularExplain.mergeType,
-                  buildErrorString(unionSubExplain, regularExplain, "mergeType"));
-        assert(docEqWithIgnoredFields(unionSubExplain.shards, regularExplain.shards),
-               buildErrorString(unionSubExplain, regularExplain, "shards"));
+        assert(
+            anyEqWithIgnoredFields(union.splitPipeline, regular.splitPipeline, mongosIgnoredFields),
+            buildErrorString(union, regular, "splitPipeline"));
+
+        assert.eq(
+            union.mergeType, regular.mergeType, buildErrorString(union, regular, "mergeType"));
+
+        assert(documentEqWithIgnoredFields(union.shards, regular.shards, mongosIgnoredFields),
+               buildErrorString(union, regular, "shards"));
+    } else if ("executionStats" in regular) {
+        const unionStats = union[0].$cursor.executionStats;
+        const regularStats = regular.executionStats;
+
+        assert(documentEqWithIgnoredFields(unionStats.executionStages,
+                                           regularStats.executionStages,
+                                           executionStatsIngoredFields),
+               buildErrorString(unionStats, regularStats, "executionStages"));
+    } else if ("stages" in regular) {
+        assert(arrayEqWithIgnoredFields(union, regular.stages, stagesIgnoredFields),
+               buildErrorString(union, regular, "stages"));
+    } else if ("queryPlanner" in regular) {
+        assert.eq(union.length, 1, "Expected single union stage");
+        const unionCursor = union[0].$cursor;
+        assert(documentEqWithIgnoredFields(
+                   regular.queryPlanner, unionCursor.queryPlanner, queryPlannerIgnoredFields),
+               buildErrorString(unionCursor, regular, "queryPlanner"));
     } else {
-        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, "stages"));
-        } else {
-            const realExplain = regularExplain.stages;
-            assert(arrayEqWithIgnoredFields(unionSubExplain, realExplain),
-                   buildErrorString(unionSubExplain, realExplain, "stages"));
-        }
+        assert(false,
+               "Don't know how to compare following explains.\n" +
+                   "regular: " + tojson(regularExplain) + "\n" +
+                   "union: " + tojson(unionSubExplain) + "\n");
     }
 }
 
-function testPipeline(pipeline) {
-    // When the SBE $group pushdown feature is enabled, a $group alone is pushed down but it is not
-    // when it's in $unionWith sub-pipeline. So, we don't need test such scenarios for now.
-    // Eventually such scenarios should be enabled.
-    if (groupPushdownEnabled && pipeline.some(stage => stage.hasOwnProperty("$group"))) {
-        return;
-    }
+function assertExplainMatch(unionExplain, regularExplain) {
+    const unionStage = getUnionWithStage(unionExplain);
+    assert(unionStage);
+    const unionSubExplain = unionStage.$unionWith.pipeline;
+    assertExplainEq(unionSubExplain, regularExplain);
+}
 
+function testPipeline(pipeline) {
     let unionResult = collA.aggregate([{$unionWith: {coll: collB.getName(), pipeline: pipeline}}],
                                       {explain: true});
     let queryResult = collB.aggregate(pipeline, {explain: true});
-    assertExplainEq(unionResult, queryResult);
+    assertExplainMatch(unionResult, queryResult);
 }
 
 testPipeline([{$addFields: {bump: true}}]);
@@ -180,7 +197,8 @@ assert.commandWorked(expectedResult);
 let unionStage = getUnionWithStage(result);
 assert(unionStage, result);
 if (FixtureHelpers.isMongos(testDB)) {
-    assert(docEqWithIgnoredFields(expectedResult.shards, unionStage.$unionWith.pipeline.shards),
+    assert(documentEqWithIgnoredFields(
+               expectedResult.shards, unionStage.$unionWith.pipeline.shards, mongosIgnoredFields),
            buildErrorString(unionStage, expectedResult));
     // TODO SERVER-50597 Fix unionWith nReturned stat in sharded cluster
     // assert.eq(unionStage.nReturned, docsPerColl, unionStage);
@@ -227,7 +245,7 @@ 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);
+assertExplainMatch(result, expectedResult);
 
 // Test a nested $unionWith which itself should perform an index scan.
 testPipeline([{$unionWith: {coll: indexedColl.getName(), pipeline: [{$match: {val: {$gt: 0}}}]}}]);
@@ -240,6 +258,6 @@ 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);
+    assertExplainMatch(result, expectedResult);
 }
 })();