SERVER-83102: Make existing jstest work with the new CQF explain (core/aggregation) (3)

jstests/aggregation/optimize_away_pipeline.js

@@ -5,11 +5,13 @@
 //
 // Relies on the pipeline stages to be collapsed into a single $cursor stage, so pipelines cannot be
 // wrapped into a facet stage to not prevent this optimization. Also, this test is not prepared to
-// handle explain output for sharded collections.
+// handle explain output for sharded collections. The assertions made in this test are irrelevant
-// This test makes assumptions about how the explain output will be formatted, so cannot be run when
+// for CQF, since it has no concept of a "query layer" or "aggregation layer". This test makes
-// pipeline optimization is disabled.
+// assumptions about how the explain output will be formatted, so cannot be run when pipeline
+// optimization is disabled.
 // @tags: [
 //   assumes_unsharded_collection,
+//   cqf_incompatible,
 //   do_not_wrap_aggregations_in_facets,
 //   requires_pipeline_optimization,
 //   requires_profiling,

jstests/aggregation/sources/project/remove_redundant_projects.js

@@ -6,6 +6,7 @@
 // ]
 import {orderedArrayEq} from "jstests/aggregation/extras/utils.js";
 import {
+    getOptimizer,
     getWinningPlan,
     isAggregationPlan,
     isQueryPlan,
@@ -56,19 +57,31 @@ function assertResultsMatch({
             result = getWinningPlan(explain.stages[0].$cursor.queryPlanner);
         }
 
-        // Check that $project uses the query system.
+        let optimizer = getOptimizer(explain);
-        assert.eq(expectProjectToCoalesce,
-                  planHasStage(db, result, "PROJECTION_DEFAULT") ||
-                      planHasStage(db, result, "PROJECTION_COVERED") ||
-                      planHasStage(db, result, "PROJECTION_SIMPLE"),
-                  explain);
 
-        if (!pipelineOptimizedAway) {
+        switch (optimizer) {
-            // Check that $project was removed from pipeline and pushed to the query system.
+            case "classic": {
-            explain.stages.forEach(function(stage) {
+                // Check that $project uses the query system.
-                if (stage.hasOwnProperty("$project"))
+                assert.eq(expectProjectToCoalesce,
-                    assert.neq(removedProjectStage, stage["$project"], explain);
+                          planHasStage(db, result, "PROJECTION_DEFAULT") ||
-            });
+                              planHasStage(db, result, "PROJECTION_COVERED") ||
+                              planHasStage(db, result, "PROJECTION_SIMPLE"),
+                          explain);
+
+                if (!pipelineOptimizedAway) {
+                    // Check that $project was removed from pipeline and pushed to the query system.
+                    explain.stages.forEach(function(stage) {
+                        if (stage.hasOwnProperty("$project"))
+                            assert.neq(removedProjectStage, stage["$project"], explain);
+                    });
+                }
+                break;
+            }
+            case "CQF": {
+                // TODO SERVER-77719: Implement the assertion for projection optimization rules for
+                // CQF.
+                break;
+            }
         }
     }
 
@@ -176,4 +189,4 @@ assertResultsMatch({
     index: indexSpec,
     pipelineOptimizedAway: true,
     removedProjectStage: {'_id.a': 1},
-});
+});

jstests/aggregation/sources/setWindowFields/optimize.js

@@ -221,4 +221,4 @@ const explain13 = coll.explain().aggregate([
     },
     {$sort: {a: {$meta: "textScore"}}},
 ]);
-assert.eq(2, numberOfStages(explain13, '$sort'), explain13);
+assert.eq(2, numberOfStages(explain13, '$sort'), explain13);

jstests/aggregation/sources/unionWith/unionWith_explain.js

@@ -1,7 +1,10 @@
 /**
  * Test that $unionWith's pipeline argument returns the same explain as an equivalent normal
- * pipeline.
+ * pipeline. The assertions in this test assume that the optimizer for $unionWith queries is
+ * the same as the optimizer for the "normal" pipeline. This assumption is not strictly true when
+ * CQF is enabled.
  * @tags: [
+ *   cqf_incompatible,
  *   do_not_wrap_aggregations_in_facets,
  * ]
  */
@@ -256,4 +259,4 @@ if (!res["failpoint.disablePipelineOptimization"].mode) {
                  .aggregate([{$unionWith: indexedColl.getName()}, {$match: {val: {$gt: 2}}}]);
     expectedResult = indexedColl.explain("executionStats").aggregate([{$match: {val: {$gt: 2}}}]);
     assertExplainMatch(result, expectedResult);
-}
+}

jstests/aggregation/use_query_projection.js

@@ -9,6 +9,7 @@
 // ]
 import {
     aggPlanHasStage,
+    getOptimizer,
     hasRejectedPlans,
     isAggregationPlan,
     isQueryPlan,
@@ -24,36 +25,55 @@ for (let i = 0; i < 100; ++i) {
 }
 assert.commandWorked(bulk.execute());
 
-function assertQueryCoversProjection(
+function assertQueryCoversProjection({pipeline = [], options = {}} = {}) {
-    {pipeline = [], pipelineOptimizedAway = true, options = {}} = {}) {
     const explainOutput = coll.explain().aggregate(pipeline, options);
+    const optimizer = getOptimizer(explainOutput);
 
-    if (pipelineOptimizedAway) {
+    assert(isQueryPlan(explainOutput), explainOutput);
-        assert(isQueryPlan(explainOutput), explainOutput);
+    // TODO SERVER-77719: Ensure that all stages are defined for all optimizers.
-        assert(!planHasStage(db, explainOutput, "FETCH"), explainOutput);
+    let stage = {"classic": "FETCH"};
-        assert(planHasStage(db, explainOutput, "IXSCAN"), explainOutput);
+    if (stage[optimizer]) {
-    } else {
+        assert(!planHasStage(db, explainOutput, stage[optimizer]), explainOutput);
-        assert(isAggregationPlan(explainOutput), explainOutput);
+    }
-        assert(!aggPlanHasStage(explainOutput, "FETCH"), explainOutput);
+    // TODO SERVER-77719: Ensure that all stages are defined for all optimizers.
-        assert(aggPlanHasStage(explainOutput, "IXSCAN"), explainOutput);
+    stage = {"classic": "IXSCAN"};
+    if (stage[optimizer]) {
+        assert(planHasStage(db, explainOutput, stage[optimizer]), explainOutput);
+    }
+
+    switch (optimizer) {
+        case "classic":
+            assert(!hasRejectedPlans(explainOutput), explainOutput);
+            break;
+        case "CQF":
+            // TODO SERVER-77719: Address the existence of rejected plans in CQF.
+            break;
     }
-    assert(!hasRejectedPlans(explainOutput), explainOutput);
     return explainOutput;
 }
 
-function assertQueryDoesNotCoverProjection({pipeline = [], pipelineOptimizedAway = true} = {}) {
+function assertQueryDoesNotCoverProjection({pipeline = []} = {}) {
     const explainOutput = coll.explain().aggregate(pipeline);
+    const optimizer = getOptimizer(explainOutput);
 
-    if (pipelineOptimizedAway) {
+    assert(isQueryPlan(explainOutput), explainOutput);
-        assert(isQueryPlan(explainOutput), explainOutput);
+    // TODO SERVER-77719: Ensure that all stages are defined for all optimizers.
-        assert(planHasStage(db, explainOutput, "FETCH") || aggPlanHasStage("COLLSCAN"),
+    let stage1 = {"classic": "FETCH"};
-               explainOutput);
+    let stage2 = {"classic": "COLLSCAN"};
-    } else {
+    if (stage1[optimizer] && stage2[optimizer]) {
-        assert(isAggregationPlan(explainOutput), explainOutput);
+        assert(planHasStage(db, explainOutput, stage1[optimizer]) ||
-        assert(aggPlanHasStage(explainOutput, "FETCH") || aggPlanHasStage("COLLSCAN"),
+                   aggPlanHasStage(stage2[optimizer]),
                explainOutput);
     }
-    assert(!hasRejectedPlans(explainOutput), explainOutput);
+
+    switch (optimizer) {
+        case "classic":
+            assert(!hasRejectedPlans(explainOutput), explainOutput);
+            break;
+        case "CQF":
+            // TODO SERVER-77719: Address the existence of rejected plans in CQF.
+            break;
+    }
     return explainOutput;
 }
 

jstests/libs/analyze_plan.js

@@ -193,6 +193,7 @@ export function getAllPlanStages(root) {
  * Asserts that no more than one stage is a match.
  */
 export function getPlanStage(root, stage) {
+    assert(stage, "Stage was not defined in getPlanStage.")
     var planStageList = getPlanStages(root, stage);
 
     if (planStageList.length === 0) {
@@ -341,6 +342,7 @@ export function getShardQueryPlans(root) {
  * structure matches expected format.
  */
 export function getAggPlanStages(root, stage, useQueryPlannerSection = false) {
+    assert(stage, "Stage was not defined in getAggPlanStages.");
     let results = [];
 
     function getDocumentSources(docSourceArray) {
@@ -434,6 +436,7 @@ export function getAggPlanStages(root, stage, useQueryPlannerSection = false) {
  * will be used to lookup the given 'stage', even if 'executionStats' section is available.
  */
 export function getAggPlanStage(root, stage, useQueryPlannerSection = false) {
+    assert(stage, "Stage was not defined in getAggPlanStage.")
     let planStageList = getAggPlanStages(root, stage, useQueryPlannerSection);
 
     if (planStageList.length === 0) {
@@ -465,6 +468,7 @@ export function aggPlanHasStage(root, stage) {
  * on one node's query plan, an error will be thrown.
  */
 export function planHasStage(db, root, stage) {
+    assert(stage, "Stage was not defined in planHasStage.")
     const matchingStages = getPlanStages(root, stage);
 
     // If we are executing against a mongos, we may get more than one occurrence of the stage.