feat: derive SortMergeJoin sort options from input ordering and allow SMJ for a single partition

datafusion/core/src/physical_planner.rs

@@ -1781,8 +1781,8 @@ impl DefaultPhysicalPlanner {
                             None,
                         )?)
                     }
-                } else if session_state.config().target_partitions() > 1
-                    && session_state.config().repartition_joins()
+                } else if (session_state.config().target_partitions() == 1
+                    || session_state.config().repartition_joins())
                     && !prefer_hash_join
                     && !*null_aware
                 // Null-aware joins (e.g. `NOT IN` with a nullable subquery) must
@@ -1791,14 +1791,18 @@ impl DefaultPhysicalPlanner {
                 // results when the right side contains a null join key.
                 {
                     // Use SortMergeJoin if hash join is not preferred
-                    let join_on_len = join_on.len();
+                    let sort_options = sort_merge_join_sort_options(
+                        &physical_left,
+                        &physical_right,
+                        &join_on,
+                    )?;
                     Arc::new(SortMergeJoinExec::try_new(
                         physical_left,
                         physical_right,
                         join_on,
                         join_filter,
                         *join_type,
-                        vec![SortOptions::default(); join_on_len],
+                        sort_options,
                         *null_equality,
                     )?)
                 } else {
@@ -3405,6 +3409,80 @@ enum PlanAsyncExpr {
     Async(AsyncMapper, PlannedExprResult),
 }
 
+/// Chooses the per-key [`SortOptions`] of a [`SortMergeJoinExec`].
+///
+/// When both inputs already declare an ordering on the join keys, the join adopts it, so no
+/// `SortExec` is needed just because of a different direction or null ordering. The same options
+/// apply to the left and the right input, so a key is only adopted if both inputs agree on it.
+/// Otherwise this falls back to [`SortOptions::default`] for all keys.
+fn sort_merge_join_sort_options(
+    left: &Arc<dyn ExecutionPlan>,
+    right: &Arc<dyn ExecutionPlan>,
+    on: &join_utils::JoinOn,
+) -> Result<Vec<SortOptions>> {
+    let mut chosen = Vec::with_capacity(on.len());
+    if find_sort_merge_join_sort_options(left, right, on, &mut chosen)? {
+        Ok(chosen)
+    } else {
+        Ok(vec![SortOptions::default(); on.len()])
+    }
+}
+
+/// Backtracking search that appends options for the keys `chosen.len()..on.len()`.
+/// The default options are tried first, so plans that already match are unchanged.
+fn find_sort_merge_join_sort_options(
+    left: &Arc<dyn ExecutionPlan>,
+    right: &Arc<dyn ExecutionPlan>,
+    on: &join_utils::JoinOn,
+    chosen: &mut Vec<SortOptions>,
+) -> Result<bool> {
+    if chosen.len() == on.len() {
+        return Ok(true);
+    }
+
+    let default = SortOptions::default();
+    let candidates = [
+        default,
+        SortOptions {
+            nulls_first: !default.nulls_first,
+            ..default
+        },
+        SortOptions {
+            descending: !default.descending,
+            ..default
+        },
+        SortOptions {
+            descending: !default.descending,
+            nulls_first: !default.nulls_first,
+        },
+    ];
+
+    for options in candidates {
+        chosen.push(options);
+        if keys_ordered(left, on, chosen, true)?
+            && keys_ordered(right, on, chosen, false)?
+            && find_sort_merge_join_sort_options(left, right, on, chosen)?
+        {
+            return Ok(true);
+        }
+        chosen.pop();
+    }
+    Ok(false)
+}
+
+/// Whether `plan` is ordered on the first `options.len()` join keys with the given options.
+fn keys_ordered(
+    plan: &Arc<dyn ExecutionPlan>,
+    on: &join_utils::JoinOn,
+    options: &[SortOptions],
+    is_left: bool,
+) -> Result<bool> {
+    let ordering = on.iter().zip(options).map(|((l, r), options)| {
+        PhysicalSortExpr::new(Arc::clone(if is_left { l } else { r }), *options)
+    });
+    plan.equivalence_properties().ordering_satisfy(ordering)
+}
+
 fn tuple_err<T, R>(value: (Result<T>, Result<R>)) -> Result<(T, R)> {
     match value {
         (Ok(e), Ok(e1)) => Ok((e, e1)),

datafusion/sqllogictest/test_files/sort_merge_join.slt

@@ -1029,6 +1029,34 @@ DROP TABLE c;
 statement ok
 set datafusion.execution.target_partitions = 4;
 
+# prefer_hash_join = false also selects a sort merge join with a single target partition
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+query TT
+EXPLAIN SELECT t1.a, t1.b, t2.b FROM t1 JOIN t2 ON t1.a = t2.a
+----
+logical_plan
+01)Projection: t1.a, t1.b, t2.b
+02)--Inner Join: t1.a = t2.a
+03)----TableScan: t1 projection=[a, b]
+04)----TableScan: t2 projection=[a, b]
+physical_plan
+01)SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)], projection=[a@0, b@1, b@3]
+02)--SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+04)--SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
+05)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+query III rowsort
+SELECT t1.a, t1.b, t2.b FROM t1 JOIN t2 ON t1.a = t2.a
+----
+2 200 250
+3 300 201
+
+statement ok
+set datafusion.execution.target_partitions = 4;
+
 # Config reset
 statement ok
 RESET datafusion.optimizer.prefer_hash_join;

datafusion/sqllogictest/test_files/sort_merge_join_sort_options.slt

@@ -0,0 +1,102 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+##########
+## The sort options of a SortMergeJoin follow the ordering both inputs declare on the join keys,
+## so no SortExec is needed on top of inputs that are already sorted that way.
+##########
+
+statement ok
+set datafusion.optimizer.prefer_hash_join = false;
+
+statement ok
+set datafusion.execution.target_partitions = 2;
+
+statement ok
+COPY (SELECT * FROM (VALUES (3, 'c', 30), (2, 'b', 20), (1, 'a', 10)) AS t(k1, k2, v) ORDER BY k1 DESC NULLS LAST, k2 ASC NULLS FIRST)
+TO 'test_files/scratch/sort_merge_join_sort_options/left.parquet';
+
+statement ok
+COPY (SELECT * FROM (VALUES (3, 'c', 300), (2, 'b', 200), (4, 'd', 400)) AS t(k1, k2, w) ORDER BY k1 DESC NULLS LAST, k2 ASC NULLS FIRST)
+TO 'test_files/scratch/sort_merge_join_sort_options/right.parquet';
+
+# both inputs sorted `k1 DESC NULLS LAST, k2 ASC NULLS FIRST`: the join adopts these per key options
+statement ok
+CREATE EXTERNAL TABLE l (k1 INT, k2 VARCHAR, v INT) STORED AS PARQUET
+WITH ORDER (k1 DESC NULLS LAST, k2 ASC NULLS FIRST)
+LOCATION 'test_files/scratch/sort_merge_join_sort_options/left.parquet';
+
+statement ok
+CREATE EXTERNAL TABLE r (k1 INT, k2 VARCHAR, w INT) STORED AS PARQUET
+WITH ORDER (k1 DESC NULLS LAST, k2 ASC NULLS FIRST)
+LOCATION 'test_files/scratch/sort_merge_join_sort_options/right.parquet';
+
+query TT
+EXPLAIN SELECT l.k1, l.k2, v, w FROM l JOIN r ON l.k1 = r.k1 AND l.k2 = r.k2
+----
+logical_plan
+01)Projection: l.k1, l.k2, l.v, r.w
+02)--Inner Join: l.k1 = r.k1, l.k2 = r.k2
+03)----TableScan: l projection=[k1, k2, v]
+04)----TableScan: r projection=[k1, k2, w]
+physical_plan
+01)SortMergeJoinExec: join_type=Inner, on=[(k1@0, k1@0), (k2@1, k2@1)], projection=[k1@0, k2@1, v@2, w@5]
+02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_merge_join_sort_options/left.parquet]]}, projection=[k1, k2, v], output_ordering=[k1@0 DESC NULLS LAST, k2@1 ASC], file_type=parquet
+03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_merge_join_sort_options/right.parquet]]}, projection=[k1, k2, w], output_ordering=[k1@0 DESC NULLS LAST, k2@1 ASC], file_type=parquet
+
+query ITII rowsort
+SELECT l.k1, l.k2, v, w FROM l JOIN r ON l.k1 = r.k1 AND l.k2 = r.k2
+----
+2 b 20 200
+3 c 30 300
+
+# the right input is sorted differently: the join keeps the default options and sorts the inputs
+statement ok
+CREATE EXTERNAL TABLE r_asc (k1 INT, k2 VARCHAR, w INT) STORED AS PARQUET
+WITH ORDER (k1 ASC NULLS LAST, k2 ASC NULLS FIRST)
+LOCATION 'test_files/scratch/sort_merge_join_sort_options/right.parquet';
+
+query TT
+EXPLAIN SELECT l.k1, l.k2, v, w FROM l JOIN r_asc r ON l.k1 = r.k1 AND l.k2 = r.k2
+----
+logical_plan
+01)Projection: l.k1, l.k2, l.v, r.w
+02)--Inner Join: l.k1 = r.k1, l.k2 = r.k2
+03)----TableScan: l projection=[k1, k2, v]
+04)----SubqueryAlias: r
+05)------TableScan: r_asc projection=[k1, k2, w]
+physical_plan
+01)SortMergeJoinExec: join_type=Inner, on=[(k1@0, k1@0), (k2@1, k2@1)], projection=[k1@0, k2@1, v@2, w@5]
+02)--SortExec: expr=[k1@0 ASC, k2@1 ASC], preserve_partitioning=[false]
+03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_merge_join_sort_options/left.parquet]]}, projection=[k1, k2, v], file_type=parquet, sort_order_for_reorder=[k1@0 ASC, k2@1 ASC]
+04)--SortExec: expr=[k1@0 ASC, k2@1 ASC], preserve_partitioning=[false]
+05)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_merge_join_sort_options/right.parquet]]}, projection=[k1, k2, w], file_type=parquet, sort_order_for_reorder=[k1@0 ASC, k2@1 ASC]
+
+statement ok
+DROP TABLE l;
+
+statement ok
+DROP TABLE r;
+
+statement ok
+DROP TABLE r_asc;
+
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+reset datafusion.optimizer.prefer_hash_join;