SERVER-113145 Enable $sort+$group optimization for compatible sort orders for featureFlagShardFilteringDistinctScan (#44881)

GitOrigin-RevId: ef80dbacc5bf03bc59d68c49387920d5eb94df60

etc/backports_required_for_multiversion_tests.yml

@@ -87,6 +87,8 @@ last-continuous:
       ticket: SERVER-113468
     - test_file: jstests/sharding/disable_resumable_range_deleter.js
       ticket: SERVER-112357
+    - test_file: jstests/noPassthrough/query/agg/sort_group_to_distinct_scan.js
+      ticket: SERVER-113145
     - test_file: jstests/aggregation/accumulators/accumulator_js_size_limits.js
       ticket: SERVER-112267
     - test_file: jstests/sharding/resharding_error_during_critical_section.js
@@ -690,6 +692,8 @@ last-lts:
       ticket: SERVER-113468
     - test_file: jstests/sharding/disable_resumable_range_deleter.js
       ticket: SERVER-112357
+    - test_file: jstests/noPassthrough/query/agg/sort_group_to_distinct_scan.js
+      ticket: SERVER-113145
     - test_file: jstests/aggregation/accumulators/accumulator_js_size_limits.js
       ticket: SERVER-112267
     - test_file: jstests/concurrency/fsm_workloads/timeseries/timeseries_create_indexes.js

jstests/noPassthrough/query/agg/sort_group_to_distinct_scan.js

@@ -71,8 +71,7 @@ const indexes = [
     {season: 1, year: 1},
 ];
 
-// TODO SERVER-113145: DISTINCT_SCAN optimization cases enabled when featureFlagShardFilteringDistinctScan is enabled.
-const alwaysEnabledDistictScanTestCases = [
+const distinctScanTestCases = [
     {
         pipeline: [
             {$sort: {season: 1, year: 1}},
@@ -87,10 +86,6 @@ const alwaysEnabledDistictScanTestCases = [
         ],
         index: {season: 1, year: 1},
     },
-];
-
-// TODO SERVER-113145: DISTINCT_SCAN optimization cases disabled when featureFlagShardFilteringDistinctScan is enabled.
-const distictScanTestCases = [
     {
         pipeline: [
             {$sort: {season: 1, year: 1}},
@@ -260,6 +255,20 @@ const distictScanTestCases = [
 ];
 
 const indexScanTestCases = [
+    {
+        pipeline: [
+            {$sort: {season: 1, year: 1}},
+            {$group: {_id: "$season", year: {$bottom: {output: "$year", sortBy: {season: -1, year: -1}}}}},
+        ],
+        index: {season: 1, year: 1},
+    },
+    {
+        pipeline: [
+            {$sort: {season: -1, year: -1}},
+            {$group: {_id: "$season", year: {$bottom: {output: "$year", sortBy: {season: 1, year: 1}}}}},
+        ],
+        index: {season: 1, year: 1},
+    },
     {
         pipeline: [
             {$sort: {season: 1, year: 1}},
@@ -358,7 +367,7 @@ const indexScanTestCases = [
     },
 ];
 
-function runTests(db, isDistinctScanOptimizationEnabled) {
+function runTests(db) {
     const coll = db[jsTest.name()];
     coll.drop();
     assert.commandWorked(coll.insertMany(docs));
@@ -388,19 +397,10 @@ function runTests(db, isDistinctScanOptimizationEnabled) {
         assertPipeline(pipeline, validateExplain);
     }
 
-    for (const testCase of alwaysEnabledDistictScanTestCases) {
+    for (const testCase of distinctScanTestCases) {
         assertDistinctScan(testCase.pipeline, testCase.index);
     }
 
-    for (const testCase of distictScanTestCases) {
-        // TODO SERVER-113145: Some distinct scan optimizations are suppressed when shard filtering feature is enabled.
-        if (isDistinctScanOptimizationEnabled) {
-            assertDistinctScan(testCase.pipeline, testCase.index);
-        } else {
-            assertIndexScan(testCase.pipeline, testCase.index);
-        }
-    }
-
     for (const testCase of indexScanTestCases) {
         assertIndexScan(testCase.pipeline, testCase.index);
     }
@@ -413,8 +413,7 @@ function runMongodAndTest(featureFlagShardFilteringDistinctScan) {
     const conn = MongoRunner.runMongod(mongodOptions);
     try {
         const db = conn.getDB(`${jsTest.name()}_db`);
-        const isDistinctScanOptimizationEnabled = !featureFlagShardFilteringDistinctScan;
-        runTests(db, isDistinctScanOptimizationEnabled);
+        runTests(db);
     } finally {
         MongoRunner.stopMongod(conn);
     }

src/mongo/db/pipeline/document_source_group_base.cpp

@@ -520,8 +520,8 @@ std::pair<size_t, SortPatternDirectionComparison> findMatchedSortingInfix(
 }
 
 /**
- *  Compare the the directions of the sorts specified by the preceding $sort stage and the
- * accumulators of the current $group stage. Returns whether they have the same or opptite
+ * Compare the the directions of the sorts specified by the preceding $sort stage and the
+ * accumulators of the current $group stage. Returns whether they have the same or opposite
  * directions, or the sort patterns are incompatible.
  * For example, sortStagePattern {a: 1, b: 1, c: 1}. groupPatterns
  * {b: 1, c: 1} or {b: 1} return sameDirection, groupPatterns {b: -1, c: -1} or {b: -1} return

src/mongo/db/query/distinct_access.cpp

@@ -597,18 +597,20 @@ bool finalizeDistinctScan(const CanonicalQuery& canonicalQuery,
             "Expected a strict distinct scan when the query has a sortRequirement",
             !hasSortRequirement || strictDistinctOnly);
 
-    // If there are other factors that affect the scan direction, don't attempt to reverse it. Note
-    // that 'flipDistinctScanDirection' is intentionally ignored here, since its purpose is to
-    // implement $last/$bottom with a distinct scan by reversing the sort direction of the query.
-    const bool reverseScanIfNeededToSatisfySortRequirement =
-        !canonicalQuery.getSortPattern() && !plannerParams.traversalPreference;
-
-    if (hasSortRequirement &&
+    // Only validate the sort requirement when there is no sort pattern on 'canonicalQuery'. Because
+    // when there is a sort pattern, 'analyzeSort()' has already validated that the required sort
+    // pattern is provided. Additionally, for the $sort + $group case, the sort patterns have
+    // already been checked for compatibility when attempting the $group rewrite.
+    if (!canonicalQuery.getSortPattern() && hasSortRequirement &&
         !QueryPlannerAnalysis::analyzeNonBlockingSort(
             plannerParams,
             canonicalQuery.getDistinct()->getSerializedSortRequirement(),
             canonicalQuery.getFindCommandRequest().getHint(),
-            reverseScanIfNeededToSatisfySortRequirement,
+            // If there are other factors that affect the scan direction, don't attempt to reverse
+            // it. Note that 'flipDistinctScanDirection' is intentionally ignored here, since its
+            // purpose is to implement $last/$bottom with a distinct scan by reversing the sort
+            // direction of the query.
+            !plannerParams.traversalPreference /*reverseScanIfNeeded*/,
             indexScanNode ? static_cast<QuerySolutionNode*>(indexScanNode)
                           : static_cast<QuerySolutionNode*>(distinctScanNode))) {
         return false;