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,6 +57,10 @@ function assertResultsMatch({
             result = getWinningPlan(explain.stages[0].$cursor.queryPlanner);
         }
 
+        let optimizer = getOptimizer(explain);
+
+        switch (optimizer) {
+            case "classic": {
                 // Check that $project uses the query system.
                 assert.eq(expectProjectToCoalesce,
                           planHasStage(db, result, "PROJECTION_DEFAULT") ||
@@ -70,6 +75,14 @@ function assertResultsMatch({
                             assert.neq(removedProjectStage, stage["$project"], explain);
                     });
                 }
+                break;
+            }
+            case "CQF": {
+                // TODO SERVER-77719: Implement the assertion for projection optimization rules for
+                // CQF.
+                break;
+            }
+        }
     }
 
     // Again without an index.

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,
  * ]
  */

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(!planHasStage(db, explainOutput, "FETCH"), explainOutput);
+    // TODO SERVER-77719: Ensure that all stages are defined for all optimizers.
-        assert(planHasStage(db, explainOutput, "IXSCAN"), explainOutput);
+    let stage = {"classic": "FETCH"};
-    } else {
+    if (stage[optimizer]) {
-        assert(isAggregationPlan(explainOutput), explainOutput);
+        assert(!planHasStage(db, explainOutput, stage[optimizer]), explainOutput);
-        assert(!aggPlanHasStage(explainOutput, "FETCH"), explainOutput);
-        assert(aggPlanHasStage(explainOutput, "IXSCAN"), explainOutput);
     }
+    // TODO SERVER-77719: Ensure that all stages are defined for all optimizers.
+    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;
+    }
     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(planHasStage(db, explainOutput, "FETCH") || aggPlanHasStage("COLLSCAN"),
+    // TODO SERVER-77719: Ensure that all stages are defined for all optimizers.
-               explainOutput);
+    let stage1 = {"classic": "FETCH"};
-    } else {
+    let stage2 = {"classic": "COLLSCAN"};
-        assert(isAggregationPlan(explainOutput), explainOutput);
+    if (stage1[optimizer] && stage2[optimizer]) {
-        assert(aggPlanHasStage(explainOutput, "FETCH") || aggPlanHasStage("COLLSCAN"),
+        assert(planHasStage(db, explainOutput, stage1[optimizer]) ||
+                   aggPlanHasStage(stage2[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;
+    }
     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.