Do not add redundant subquery ordering into plan

datafusion/sql/src/relation/mod.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
+use datafusion_common::tree_node::{Transformed, TreeNode};
 use datafusion_common::{not_impl_err, plan_err, DFSchema, Result, TableReference};
 use datafusion_expr::{expr::Unnest, Expr, LogicalPlan, LogicalPlanBuilder};
 use sqlparser::ast::{FunctionArg, FunctionArgExpr, TableFactor};
@@ -142,10 +143,36 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 );
             }
         };
+
+        let optimized_plan = optimize_subquery_sort(plan)?.data;
         if let Some(alias) = alias {
-            self.apply_table_alias(plan, alias)
+            self.apply_table_alias(optimized_plan, alias)
         } else {
-            Ok(plan)
+            Ok(optimized_plan)
         }
     }
 }
+
+fn optimize_subquery_sort(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+    // When a subquery is initialized, we look for the sort options since they might be redundant.
+    // It's only important if the subquery result is affected with the order by statement.
+    // Which are the cases of:
+    // 1. DISTINCT ON / ARRAY_AGG ... => It's handled as an Aggregate plan and keep the sorting
+    // 2. RANK / ROW_NUMBER ... => It's handled as a WindowAggr plan sorting is preserved
+    // 3. LIMIT => It's handled as a Sort plan type so that we need to search for it
+    let mut has_limit = false;
+    let new_plan = plan.clone().transform_down(|c| {
+        if let LogicalPlan::Limit(_) = c {
+            has_limit = true;
+            return Ok(Transformed::no(c));
+        }
+        if let LogicalPlan::Sort(_) = c {
+            if !has_limit {
+                has_limit = false;
+                return Ok(Transformed::yes(c.inputs()[0].clone()));
+            }
+        }
+        Ok(Transformed::no(c))
+    });
+    new_plan
+}

datafusion/sqllogictest/test_files/group_by.slt

@@ -2527,20 +2527,24 @@ EXPLAIN SELECT s.country, ARRAY_AGG(s.amount ORDER BY s.amount DESC) AS amounts,
         SUM(s.amount) AS sum1
           FROM (SELECT *
             FROM sales_global
-            ORDER BY country) AS s
+            ORDER BY country
+            LIMIT 10) AS s
           GROUP BY s.country
 ----
 logical_plan
 01)Projection: s.country, array_agg(s.amount) ORDER BY [s.amount DESC NULLS FIRST] AS amounts, sum(s.amount) AS sum1
 02)--Aggregate: groupBy=[[s.country]], aggr=[[array_agg(s.amount) ORDER BY [s.amount DESC NULLS FIRST], sum(CAST(s.amount AS Float64))]]
 03)----SubqueryAlias: s
-04)------Sort: sales_global.country ASC NULLS LAST
-05)--------TableScan: sales_global projection=[country, amount]
+04)------Limit: skip=0, fetch=10
+05)--------Sort: sales_global.country ASC NULLS LAST, fetch=10
+06)----------TableScan: sales_global projection=[country, amount]
 physical_plan
 01)ProjectionExec: expr=[country@0 as country, array_agg(s.amount) ORDER BY [s.amount DESC NULLS FIRST]@1 as amounts, sum(s.amount)@2 as sum1]
 02)--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[array_agg(s.amount) ORDER BY [s.amount DESC NULLS FIRST], sum(s.amount)], ordering_mode=Sorted
 03)----SortExec: expr=[country@0 ASC NULLS LAST,amount@1 DESC], preserve_partitioning=[false]
-04)------MemoryExec: partitions=1, partition_sizes=[1]
+04)------SortExec: TopK(fetch=10), expr=[country@0 ASC NULLS LAST], preserve_partitioning=[false]
+05)--------MemoryExec: partitions=1, partition_sizes=[1]
+
 
 query T?R rowsort
 SELECT s.country, ARRAY_AGG(s.amount ORDER BY s.amount DESC) AS amounts,
@@ -2563,20 +2567,23 @@ EXPLAIN SELECT s.country, s.zip_code, ARRAY_AGG(s.amount ORDER BY s.amount DESC)
         SUM(s.amount) AS sum1
           FROM (SELECT *
             FROM sales_global
-            ORDER BY country) AS s
+            ORDER BY country
+            LIMIT 10) AS s
           GROUP BY s.country, s.zip_code
 ----
 logical_plan
 01)Projection: s.country, s.zip_code, array_agg(s.amount) ORDER BY [s.amount DESC NULLS FIRST] AS amounts, sum(s.amount) AS sum1
 02)--Aggregate: groupBy=[[s.country, s.zip_code]], aggr=[[array_agg(s.amount) ORDER BY [s.amount DESC NULLS FIRST], sum(CAST(s.amount AS Float64))]]
 03)----SubqueryAlias: s
-04)------Sort: sales_global.country ASC NULLS LAST
-05)--------TableScan: sales_global projection=[zip_code, country, amount]
+04)------Limit: skip=0, fetch=10
+05)--------Sort: sales_global.country ASC NULLS LAST, fetch=10
+06)----------TableScan: sales_global projection=[zip_code, country, amount]
 physical_plan
 01)ProjectionExec: expr=[country@0 as country, zip_code@1 as zip_code, array_agg(s.amount) ORDER BY [s.amount DESC NULLS FIRST]@2 as amounts, sum(s.amount)@3 as sum1]
 02)--AggregateExec: mode=Single, gby=[country@1 as country, zip_code@0 as zip_code], aggr=[array_agg(s.amount) ORDER BY [s.amount DESC NULLS FIRST], sum(s.amount)], ordering_mode=PartiallySorted([0])
 03)----SortExec: expr=[country@1 ASC NULLS LAST,amount@2 DESC], preserve_partitioning=[false]
-04)------MemoryExec: partitions=1, partition_sizes=[1]
+04)------SortExec: TopK(fetch=10), expr=[country@1 ASC NULLS LAST], preserve_partitioning=[false]
+05)--------MemoryExec: partitions=1, partition_sizes=[1]
 
 query TI?R rowsort
 SELECT s.country, s.zip_code, ARRAY_AGG(s.amount ORDER BY s.amount DESC) AS amounts,
@@ -2599,19 +2606,21 @@ EXPLAIN SELECT s.country, ARRAY_AGG(s.amount ORDER BY s.country DESC) AS amounts
         SUM(s.amount) AS sum1
           FROM (SELECT *
             FROM sales_global
-            ORDER BY country) AS s
+            ORDER BY country
+            LIMIT 10) AS s
           GROUP BY s.country
 ----
 logical_plan
 01)Projection: s.country, array_agg(s.amount) ORDER BY [s.country DESC NULLS FIRST] AS amounts, sum(s.amount) AS sum1
 02)--Aggregate: groupBy=[[s.country]], aggr=[[array_agg(s.amount) ORDER BY [s.country DESC NULLS FIRST], sum(CAST(s.amount AS Float64))]]
 03)----SubqueryAlias: s
-04)------Sort: sales_global.country ASC NULLS LAST
-05)--------TableScan: sales_global projection=[country, amount]
+04)------Limit: skip=0, fetch=10
+05)--------Sort: sales_global.country ASC NULLS LAST, fetch=10
+06)----------TableScan: sales_global projection=[country, amount]
 physical_plan
 01)ProjectionExec: expr=[country@0 as country, array_agg(s.amount) ORDER BY [s.country DESC NULLS FIRST]@1 as amounts, sum(s.amount)@2 as sum1]
 02)--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[array_agg(s.amount) ORDER BY [s.country DESC NULLS FIRST], sum(s.amount)], ordering_mode=Sorted
-03)----SortExec: expr=[country@0 ASC NULLS LAST], preserve_partitioning=[false]
+03)----SortExec: TopK(fetch=10), expr=[country@0 ASC NULLS LAST], preserve_partitioning=[false]
 04)------MemoryExec: partitions=1, partition_sizes=[1]
 
 query T?R rowsort
@@ -2634,20 +2643,24 @@ EXPLAIN SELECT s.country, ARRAY_AGG(s.amount ORDER BY s.country DESC, s.amount D
         SUM(s.amount) AS sum1
           FROM (SELECT *
             FROM sales_global
-            ORDER BY country) AS s
+            ORDER BY country
+            LIMIT 10) AS s
           GROUP BY s.country
 ----
 logical_plan
 01)Projection: s.country, array_agg(s.amount) ORDER BY [s.country DESC NULLS FIRST, s.amount DESC NULLS FIRST] AS amounts, sum(s.amount) AS sum1
 02)--Aggregate: groupBy=[[s.country]], aggr=[[array_agg(s.amount) ORDER BY [s.country DESC NULLS FIRST, s.amount DESC NULLS FIRST], sum(CAST(s.amount AS Float64))]]
 03)----SubqueryAlias: s
-04)------Sort: sales_global.country ASC NULLS LAST
-05)--------TableScan: sales_global projection=[country, amount]
+04)------Limit: skip=0, fetch=10
+05)--------Sort: sales_global.country ASC NULLS LAST, fetch=10
+06)----------TableScan: sales_global projection=[country, amount]
 physical_plan
 01)ProjectionExec: expr=[country@0 as country, array_agg(s.amount) ORDER BY [s.country DESC NULLS FIRST, s.amount DESC NULLS FIRST]@1 as amounts, sum(s.amount)@2 as sum1]
 02)--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[array_agg(s.amount) ORDER BY [s.country DESC NULLS FIRST, s.amount DESC NULLS FIRST], sum(s.amount)], ordering_mode=Sorted
 03)----SortExec: expr=[country@0 ASC NULLS LAST,amount@1 DESC], preserve_partitioning=[false]
-04)------MemoryExec: partitions=1, partition_sizes=[1]
+04)------SortExec: TopK(fetch=10), expr=[country@0 ASC NULLS LAST], preserve_partitioning=[false]
+05)--------MemoryExec: partitions=1, partition_sizes=[1]
+
 
 query T?R rowsort
 SELECT s.country, ARRAY_AGG(s.amount ORDER BY s.country DESC, s.amount DESC) AS amounts,
@@ -2801,8 +2814,7 @@ EXPLAIN SELECT country, FIRST_VALUE(amount ORDER BY ts DESC) as fv1,
 logical_plan
 01)Projection: sales_global.country, first_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST] AS fv1, last_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST] AS lv1, sum(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST] AS sum1
 02)--Aggregate: groupBy=[[sales_global.country]], aggr=[[first_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST], last_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST], sum(CAST(sales_global.amount AS Float64)) ORDER BY [sales_global.ts DESC NULLS FIRST]]]
-03)----Sort: sales_global.ts ASC NULLS LAST
-04)------TableScan: sales_global projection=[country, ts, amount]
+03)----TableScan: sales_global projection=[country, ts, amount]
 physical_plan
 01)ProjectionExec: expr=[country@0 as country, first_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@1 as fv1, last_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@2 as lv1, sum(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@3 as sum1]
 02)--AggregateExec: mode=Single, gby=[country@0 as country], aggr=[first_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST], last_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST], sum(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]]

datafusion/sqllogictest/test_files/joins.slt

@@ -3898,6 +3898,7 @@ SELECT * FROM (
 ) as lhs RIGHT JOIN (
     SELECT * from right_table_no_nulls
     ORDER BY b
+    LIMIT 10
 ) AS rhs ON lhs.b=rhs.b
 ----
 11 1 21 1
@@ -3911,23 +3912,26 @@ EXPLAIN SELECT * FROM (
 ) as lhs RIGHT JOIN (
     SELECT * from right_table_no_nulls
     ORDER BY b
+    LIMIT 10
 ) AS rhs ON lhs.b=rhs.b
 ----
 logical_plan
 01)Right Join: lhs.b = rhs.b
 02)--SubqueryAlias: lhs
 03)----TableScan: left_table_no_nulls projection=[a, b]
 04)--SubqueryAlias: rhs
-05)----Sort: right_table_no_nulls.b ASC NULLS LAST
-06)------TableScan: right_table_no_nulls projection=[a, b]
+05)----Limit: skip=0, fetch=10
+06)------Sort: right_table_no_nulls.b ASC NULLS LAST, fetch=10
+07)--------TableScan: right_table_no_nulls projection=[a, b]
 physical_plan
 01)CoalesceBatchesExec: target_batch_size=3
 02)--HashJoinExec: mode=CollectLeft, join_type=Right, on=[(b@1, b@1)]
 03)----MemoryExec: partitions=1, partition_sizes=[1]
-04)----SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false]
+04)----SortExec: TopK(fetch=10), expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false]
 05)------MemoryExec: partitions=1, partition_sizes=[1]
 
 
+
 # Missing probe index in the middle of the batch:
 statement ok
 CREATE TABLE left_table_missing_probe(a INT UNSIGNED, b INT UNSIGNED)
@@ -3953,6 +3957,7 @@ SELECT * FROM (
 ) as lhs RIGHT JOIN (
     SELECT * from right_table_missing_probe
     ORDER BY b
+    LIMIT 10
 ) AS rhs ON lhs.b=rhs.b
 ----
 11 1 21 1
@@ -3974,14 +3979,12 @@ logical_plan
 02)--SubqueryAlias: lhs
 03)----TableScan: left_table_no_nulls projection=[a, b]
 04)--SubqueryAlias: rhs
-05)----Sort: right_table_no_nulls.b ASC NULLS LAST
-06)------TableScan: right_table_no_nulls projection=[a, b]
+05)----TableScan: right_table_no_nulls projection=[a, b]
 physical_plan
 01)CoalesceBatchesExec: target_batch_size=3
 02)--HashJoinExec: mode=CollectLeft, join_type=Right, on=[(b@1, b@1)]
 03)----MemoryExec: partitions=1, partition_sizes=[1]
-04)----SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false]
-05)------MemoryExec: partitions=1, partition_sizes=[1]
+04)----MemoryExec: partitions=1, partition_sizes=[1]
 
 
 # Null build indices:
@@ -4009,6 +4012,7 @@ SELECT * FROM (
 ) as lhs RIGHT JOIN (
     SELECT * from right_table_append_null_build
     ORDER BY b
+    LIMIT 10
 ) AS rhs ON lhs.b=rhs.b
 ----
 NULL NULL 21 4
@@ -4025,18 +4029,20 @@ EXPLAIN SELECT * FROM (
 ) as lhs RIGHT JOIN (
     SELECT * from right_table_no_nulls
     ORDER BY b
+    LIMIT 10
 ) AS rhs ON lhs.b=rhs.b
 ----
 logical_plan
 01)Right Join: lhs.b = rhs.b
 02)--SubqueryAlias: lhs
 03)----TableScan: left_table_no_nulls projection=[a, b]
 04)--SubqueryAlias: rhs
-05)----Sort: right_table_no_nulls.b ASC NULLS LAST
-06)------TableScan: right_table_no_nulls projection=[a, b]
+05)----Limit: skip=0, fetch=10
+06)------Sort: right_table_no_nulls.b ASC NULLS LAST, fetch=10
+07)--------TableScan: right_table_no_nulls projection=[a, b]
 physical_plan
 01)CoalesceBatchesExec: target_batch_size=3
 02)--HashJoinExec: mode=CollectLeft, join_type=Right, on=[(b@1, b@1)]
 03)----MemoryExec: partitions=1, partition_sizes=[1]
-04)----SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false]
+04)----SortExec: TopK(fetch=10), expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false]
 05)------MemoryExec: partitions=1, partition_sizes=[1]

datafusion/sqllogictest/test_files/select.slt

@@ -990,13 +990,13 @@ FROM (
     ) AS a
 ) AS b
 ----
-a 5 -101
-a 5 -54
 a 5 -38
+a 5 -54
+a 6 36
+a 6 -31
 a 5 65
+a 5 -101
 a 6 -101
-a 6 -31
-a 6 36
 
 # nested select without aliases
 query TII
@@ -1011,13 +1011,13 @@ FROM (
     )
 )
 ----
-a 5 -101
-a 5 -54
 a 5 -38
+a 5 -54
+a 6 36
+a 6 -31
 a 5 65
+a 5 -101
 a 6 -101
-a 6 -31
-a 6 36
 
 # select with join unaliased subqueries
 query TIITII
@@ -1461,22 +1461,26 @@ query TT
 EXPLAIN SELECT c2, COUNT(*)
 FROM (SELECT c2
 FROM aggregate_test_100
-ORDER BY c1, c2)
+ORDER BY c1, c2
+LIMIT 4)
 GROUP BY c2;
 ----
 logical_plan
 01)Aggregate: groupBy=[[aggregate_test_100.c2]], aggr=[[count(Int64(1)) AS count(*)]]
 02)--Projection: aggregate_test_100.c2
-03)----Sort: aggregate_test_100.c1 ASC NULLS LAST, aggregate_test_100.c2 ASC NULLS LAST
-04)------Projection: aggregate_test_100.c2, aggregate_test_100.c1
-05)--------TableScan: aggregate_test_100 projection=[c1, c2]
+03)----Limit: skip=0, fetch=4
+04)------Sort: aggregate_test_100.c1 ASC NULLS LAST, aggregate_test_100.c2 ASC NULLS LAST, fetch=4
+05)--------Projection: aggregate_test_100.c2, aggregate_test_100.c1
+06)----------TableScan: aggregate_test_100 projection=[c1, c2]
 physical_plan
 01)AggregateExec: mode=FinalPartitioned, gby=[c2@0 as c2], aggr=[count(*)]
 02)--CoalesceBatchesExec: target_batch_size=8192
 03)----RepartitionExec: partitioning=Hash([c2@0], 2), input_partitions=2
 04)------AggregateExec: mode=Partial, gby=[c2@0 as c2], aggr=[count(*)]
 05)--------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
-06)----------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c2], has_header=true
+06)----------ProjectionExec: expr=[c2@0 as c2]
+07)------------SortExec: TopK(fetch=4), expr=[c1@1 ASC NULLS LAST,c2@0 ASC NULLS LAST], preserve_partitioning=[false]
+08)--------------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c2, c1], has_header=true
 
 # FilterExec can track equality of non-column expressions.
 # plan below shouldn't have a SortExec because given column 'a' is ordered.