SERVER-46233 Relocate second call to optimizePipeline()

The second call to optimizePipeline() can sometimes catch
optimizations that were missed in the first call. The first call
simplifies filter expressions, opening up some optimization
opportunities that were previously blocked.

We want to catch these opportunities before we push down any stages
into the PlanExecutor.

1 files changed, 67 insertions, 0 deletions

diff --git a/jstests/aggregation/split_match_and_swap_with_sort.js b/jstests/aggregation/split_match_and_swap_with_sort.js
new file mode 100644
index 00000000000..1905c7be1b7
--- /dev/null
+++ b/jstests/aggregation/split_match_and_swap_with_sort.js
@@ -0,0 +1,67 @@
+// When a $sort is followed by a $match, the pipeline optimizer should be able to move the match in
+// front of the $sort, so as to reduce the number of documents that need to be sorted.
+//
+// @tags: [
+//   # Wrapping the pipeline in a $facet will prevent its $match from getting absorbed by the query
+//   # system, which is what we are testing for.
+//   do_not_wrap_aggregations_in_facets,
+//   # Likewise, splitting up the pipeline for sharding may also prevent the $match from getting
+//   # absorbed.
+//   assumes_unsharded_collection,
+//   # Don't disable the thing we are specifically testing for!
+//   requires_pipeline_optimization,
+// ]
+load('jstests/libs/analyze_plan.js');
+
+(function() {
+"use strict";
+
+const coll = db.getSiblingDB("split_match_and_swap_with_sort")["test"];
+coll.drop();
+
+assert.commandWorked(
+    coll.insert([{_id: 1, a: 1, b: 3}, {_id: 2, a: 2, b: 2}, {_id: 3, a: 3, b: 1}]));
+
+{
+    const pipeline = [{$sort: {b: 1}}, {$match: {a: {$ne: 2}}}];
+
+    assert.eq([{_id: 3, a: 3, b: 1}, {_id: 1, a: 1, b: 3}], coll.aggregate(pipeline).toArray());
+
+    const pipelineExplained = coll.explain().aggregate(pipeline);
+    const collScanStage = getPlanStage(pipelineExplained, "COLLSCAN");
+
+    // After moving the $match to the front of the pipeline, we expect the pipeline optimizer to
+    // push it down into the PlanExecutor, so that it becomes a filter on the COLLSCAN.
+    assert.neq(null, collScanStage, pipelineExplained);
+    assert.eq({a: {"$not": {"$eq": 2}}}, collScanStage.filter, collScanStage);
+}
+
+{
+    // Note that a $expr cannot be moved. However, the optimizer can split the non-$expr part of the
+    // $match and move that to be in front of the $sort.
+    const pipeline =
+        [{$sort: {b: 1}}, {$match: {$and: [{a: {$ne: 2}}, {$expr: {$ne: ["$a", "$b"]}}]}}];
+
+    assert.eq([{_id: 3, a: 3, b: 1}, {_id: 1, a: 1, b: 3}], coll.aggregate(pipeline).toArray());
+
+    const pipelineExplained = coll.explain().aggregate(pipeline);
+    const collScanStage = getAggPlanStage(pipelineExplained, "COLLSCAN");
+    assert.neq(null, collScanStage, pipelineExplained);
+    assert.eq({a: {"$not": {"$eq": 2}}}, collScanStage.filter, collScanStage);
+}
+
+{
+    // SERVER-46233: Normally a $or at the root of a $match expression prevents it from getting
+    // split. However, When the $or has only one child, the optimizer should be able to eliminate
+    // the $or and then recognize that the simplified result _is_ eligible to be split.
+    const pipeline =
+        [{$sort: {b: 1}}, {$match: {$or: [{$and: [{a: {$ne: 2}}, {$expr: {$ne: ["$a", "$b"]}}]}]}}];
+
+    assert.eq([{_id: 3, a: 3, b: 1}, {_id: 1, a: 1, b: 3}], coll.aggregate(pipeline).toArray());
+
+    const pipelineExplained = coll.explain().aggregate(pipeline);
+    const collScanStage = getAggPlanStage(pipelineExplained, "COLLSCAN");
+    assert.neq(null, collScanStage, pipelineExplained);
+    assert.eq({a: {"$not": {"$eq": 2}}}, collScanStage.filter, collScanStage);
+}
+}());
\ No newline at end of file