fix: choose the forced shuffled hash join at planning time so Spark keeps its sorts

.github/labeler.yml

@@ -87,7 +87,7 @@ area:joins:
   - changed-files:
       - any-glob-to-any-file:
           - 'native/core/src/execution/operators/dynamic_filter/**'
-          - 'spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala'
+          - 'spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala'
           - 'spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala'
           - 'spark/src/main/**/*Join*'
 

docs/source/user-guide/latest/compatibility/operators.md

@@ -133,19 +133,26 @@ express falls back to Spark and records a fallback reason (see
 an operator above the join relies on the order.
 
 An outer join falls back to Spark when Comet cannot check the order Spark reports and that order
-reaches the join through an input from the JVM or through a nested loop join. Spark keeps the rows
-in the order they arrive. One such input is an RDD checkpoint. It reaches a native join only with
+reaches the join through an input from the JVM. Spark keeps the rows in the order they arrive. One
+such input is an RDD checkpoint. It reaches a native join only with
 `spark.comet.convert.rdd.enabled=true`, which defaults to `false`. Without that setting the
 checkpoint stays a Spark input and the join stays in Spark anyway. Another is a cache written by an
 outer hash join whose own probe input was not a scan
-([#6922](https://github.com/apache/datafusion-comet/issues/6922)). The output of a nested loop
-join can already be out of order when it reaches the join
+([#6922](https://github.com/apache/datafusion-comet/issues/6922)).
+
+The native nested loop join falls back to Spark when Spark reports an output order for it, since
+it emits its rows build row by build row
 ([#6898](https://github.com/apache/datafusion-comet/issues/6898)).
 
 When the order Comet cannot check comes only from another native hash join in the same plan, the
 join does not fall back. It runs natively and declares no order, since DataFusion keeps the order
 it can see there.
 
+Hash joins of the other types that report an order, such as inner joins, keep their streamed rows
+in place either way. They declare the order when Comet can check and express it, so a native
+operator above them, a window for instance, sees it. Otherwise they declare nothing and still run
+natively.
+
 ## Window Functions
 
 Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of
@@ -203,6 +210,10 @@ supported range of 0001-01-01 to 9999-12-31, such as one produced by `date_add`
 native frame can treat the bound as the edge of the partition where Spark compares against it
 ([#5022](https://github.com/apache/datafusion-comet/issues/5022)).
 
+Spark drops the sort below a window when it believes the window's input is already ordered. For a window with
+`PARTITION BY` keys, when Comet cannot verify that order, as with rows read from a checkpoint, the native window sorts
+its input by the `PARTITION BY` keys and then the `ORDER BY` keys. The results match Spark, at the cost of a sort.
+
 `WindowGroupLimitExec` (window-based limit pushdown for `ROW_NUMBER`, `RANK`, and `DENSE_RANK`)
 runs natively; it is controlled by `spark.comet.exec.windowGroupLimit.enabled` (default: true).
 

docs/source/user-guide/latest/operators.md

@@ -96,12 +96,12 @@ omitted from the tables below and may be reconsidered based on demand:
 
 ## Joins
 
-| Operator                      | Status | Notes                                                                                                                                                                         |
-| ----------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `BroadcastHashJoinExec`       | ✅     | Falls back to Spark for a `NOT IN` (null-aware anti join) whose streamed side reports an ordering.                                                                            |
-| `ShuffledHashJoinExec`        | ✅     |                                                                                                                                                                               |
-| `SortMergeJoinExec`           | ✅     | Supports `BINARY` join keys. Nested-type (struct, array, map) and collated-string join keys fall back to Spark.                                                               |
-| `BroadcastNestedLoopJoinExec` | ✅     | Falls back to Spark when the preserved side is broadcast (for example LEFT OUTER with BROADCAST on the left) ([#4429](https://github.com/apache/datafusion-comet/pull/4429)). |
+| Operator                      | Status | Notes                                                                                                                                                                                                                                                                                                                                        |
+| ----------------------------- | ------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `BroadcastHashJoinExec`       | ✅     | Falls back to Spark for a `NOT IN` (null-aware anti join) whose streamed side reports an ordering.                                                                                                                                                                                                                                           |
+| `ShuffledHashJoinExec`        | ✅     |                                                                                                                                                                                                                                                                                                                                              |
+| `SortMergeJoinExec`           | ✅     | Supports `BINARY` join keys. Nested-type (struct, array, map) and collated-string join keys fall back to Spark.                                                                                                                                                                                                                              |
+| `BroadcastNestedLoopJoinExec` | ✅     | Falls back to Spark when the preserved side is broadcast (for example LEFT OUTER with BROADCAST on the left) ([#4429](https://github.com/apache/datafusion-comet/pull/4429)), and when Spark reports an output ordering for the join, which the native join does not keep ([#6898](https://github.com/apache/datafusion-comet/issues/6898)). |
 
 `ExistenceJoin` (produced when `EXISTS` / `IN` is combined with another predicate via `OR`) runs
 natively on `BroadcastHashJoinExec` and `ShuffledHashJoinExec`. Existence sort-merge joins, residual

docs/source/user-guide/latest/tuning/operators.md

@@ -30,9 +30,12 @@ to configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`. Comet does
 to test with both for your specific workloads.
 
 To configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`, set `spark.comet.exec.forceShuffledHashJoin=true`.
-The conversion only happens when the build side is under a size limit, and a join whose build side has no statistics
-is left as a `SortMergeJoin`. The size is Spark's planning estimate of the build side, or under AQE the materialized
-shuffle size of the build side. By default the limit is Spark's own rule for choosing a `ShuffledHashJoin`:
+Comet makes the choice while Spark plans the query, so a join Spark would broadcast, or one with a join strategy hint
+such as `MERGE`, is planned as Spark would plan it. The choice also applies to joins of static tables inside a
+streaming query, which then run as Spark's own `ShuffledHashJoin`. The conversion only happens when the build side is
+under a size limit, and a join whose build side has no statistics is left as a `SortMergeJoin`. The size is Spark's
+planning estimate of the build side, or under AQE the materialized shuffle size of the build side. By default the
+limit is Spark's own rule for choosing a `ShuffledHashJoin`:
 `spark.sql.autoBroadcastJoinThreshold` times the initial shuffle partition count
 (`spark.sql.adaptive.coalescePartitions.initialPartitionNum` when AQE and partition coalescing are both on and it is
 set, else `spark.sql.shuffle.partitions`). When broadcasts are disabled with a non-positive threshold, Spark's default

docs/source/user-guide/latest/understanding-comet-plans.md

@@ -224,6 +224,8 @@ Keep the following in mind when reading the reports:
   native plan still counts as accelerated, so treat the percentage as an upper bound.
 - Comet's split Iceberg V2 write (`spark.comet.write.iceberg.enabled`) is declined in
   plan-only mode, so such writes run on, and are reported as, Spark.
+- `spark.comet.exec.forceShuffledHashJoin` is not applied in plan-only mode, so the report shows
+  the sort-merge joins Spark plans rather than the hash joins a real run would use.
 - Under AQE the report describes the plan before any adaptive re-planning, so coverage of the
   plan that finally executes can differ. In particular, AQE plans subqueries into the outer
   query only after the report is produced, so the outer report counts their operators as

native/core/src/execution/operators/dynamic_filter/early.rs

@@ -28,11 +28,12 @@ use datafusion::common::{internal_err, Result};
 use datafusion::physical_expr::expressions::{Column, DynamicFilterPhysicalExpr};
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
-use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions};
 use datafusion_comet_operators::CometFilterExec;
 
 use super::parquet_reader::is_direct_column_null_checks;
 use super::{DynamicFilterExec, DynamicFilterJoinExec};
+use crate::execution::operators::SortedInputExec;
 use datafusion_comet_operators::CometProjectionExec;
 
 pub(super) fn place_early_filter(
@@ -126,6 +127,19 @@ fn place(
             {
                 return Ok(Some(filter.with_execution_input(child)));
             }
+        } else if let Some(sorted) = input.downcast_ref::<SortedInputExec>() {
+            // It only declares its input's order, which dropping rows below it keeps.
+            if let Some(child) = place(
+                sorted.input(),
+                Arc::clone(&predicate),
+                metrics,
+                crossed_join,
+            )? {
+                return Ok(Some(Arc::clone(input).replace_children(
+                    vec![child],
+                    ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+                )?));
+            }
         }
     }
     // A terminal consumer is useful only after crossing a join. Above other

native/core/src/execution/operators/dynamic_filter/join/tests.rs

@@ -2010,3 +2010,98 @@ fn early_filter_stops_at_intermediate_join_fetch() {
         &place_early_filter(&wrapped, predicate, &ExecutionPlanMetricsSet::new()).unwrap()
     ));
 }
+
+/// An order that `key@index` declares, as the planner declares a sorted probe's.
+fn declared_order(input: Arc<dyn ExecutionPlan>, index: usize) -> Arc<dyn ExecutionPlan> {
+    use crate::execution::operators::SortedInputExec;
+    use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
+    let key = Arc::new(Column::new(input.schema().field(index).name(), index));
+    let ordering = LexOrdering::new([PhysicalSortExpr::new_default(key)]).unwrap();
+    Arc::new(SortedInputExec::new(input, ordering))
+}
+
+#[test]
+fn reader_filter_crosses_a_declared_probe_order() {
+    // The order wrapper only passes rows through, so the reader below it can drop rows too.
+    let session = Arc::new(SessionContext::new_with_config(
+        SessionConfig::new().with_target_partitions(1),
+    ));
+    let (_file, scan) = parquet_probe((0..400).collect(), &session, 100);
+    let probe = declared_order(filtered_probe(&scan), 0);
+    let ordering = probe.properties().output_ordering().unwrap().to_string();
+    let join = single_key_join_plans(two_batch_build(), probe, PartitionMode::Partitioned);
+    let wrapper =
+        DynamicFilterJoinExec::new(&join, session.copied_config().options().as_ref().clone())
+            .unwrap();
+    let runtime = wrapper.build_runtime_join().unwrap();
+    assert!(runtime.reader_filter_attached);
+    let consumer = runtime
+        .join
+        .right()
+        .downcast_ref::<DynamicFilterExec>()
+        .unwrap();
+    let sorted = consumer.children()[0];
+    assert_eq!("SortedInputExec", sorted.name());
+    assert_eq!(
+        ordering,
+        sorted.properties().output_ordering().unwrap().to_string()
+    );
+    let filter = sorted.children()[0]
+        .downcast_ref::<CometFilterExec>()
+        .unwrap();
+    let reader = filter.input().downcast_ref::<DataSourceExec>().unwrap();
+    let (_, source) = reader.downcast_to_file_source::<ParquetSource>().unwrap();
+    let reader_filter = source.filter().unwrap();
+    let reader_filter = find_dynamic_filter(&reader_filter).unwrap();
+    assert_eq!(
+        reader_filter.expression_id().unwrap(),
+        produced_join_filter(&runtime.join).expression_id().unwrap()
+    );
+    // The rebuilt join still keeps the declared order in its own output.
+    assert!(runtime.join.properties().output_ordering().is_some());
+}
+
+#[tokio::test]
+async fn ancestor_filter_crosses_a_declared_probe_order() {
+    // As in the test above without the wrapper: the ancestor's filter still reaches the
+    // intermediate join's probe.
+    let ctx = SessionContext::new();
+    let fact = input(
+        vec![Some(0), Some(1), Some(2), Some(2), Some(3), None],
+        &DataType::Int32,
+        0,
+    );
+    let dimension = input(
+        vec![Some(0), Some(1), Some(2), Some(2), Some(3)],
+        &DataType::Int32,
+        0,
+    );
+    let ancestor_build = input(vec![Some(2), Some(2)], &DataType::Int32, 0);
+    let inner = chain_join(dimension, fact, 0, 0);
+    let inner = Arc::new(
+        DynamicFilterJoinExec::try_new(&inner, ctx.copied_config().options())
+            .unwrap()
+            .unwrap(),
+    ) as Arc<dyn ExecutionPlan>;
+    let probe = declared_order(inner, 2);
+    let outer = chain_join(ancestor_build, probe, 0, 2);
+    let expected = collect(
+        Arc::new(outer.builder().reset_state().build().unwrap()),
+        ctx.task_ctx(),
+    )
+    .await
+    .unwrap();
+    assert_eq!(row_count(&expected), 8);
+    let outer = Arc::new(
+        DynamicFilterJoinExec::try_new(&outer, ctx.copied_config().options())
+            .unwrap()
+            .unwrap(),
+    ) as Arc<dyn ExecutionPlan>;
+    let actual = collect(Arc::clone(&outer), ctx.task_ctx()).await.unwrap();
+    assert_eq!(
+        batches_to_sort_string(&actual),
+        batches_to_sort_string(&expected)
+    );
+    assert_eq!(metric(&outer, "dynamic_filter_early_rows_evaluated"), 6);
+    assert_eq!(metric(&outer, "dynamic_filter_early_rows_pruned"), 4);
+}

native/core/src/execution/operators/dynamic_filter/parquet_reader.rs

@@ -29,11 +29,13 @@ use datafusion::physical_expr::expressions::{
 };
 use datafusion::physical_expr::utils::collect_columns;
 use datafusion::physical_expr::PhysicalExpr;
-use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions};
 
 use crate::parquet::file_error_context::ParquetErrorContext;
 use datafusion_comet_operators::CometFilterExec;
 
+use crate::execution::operators::SortedInputExec;
+
 mod schema_adapter;
 
 use schema_adapter::RuntimeFilterSchemaAdapterFactory;
@@ -98,6 +100,17 @@ pub(super) fn try_attach_parquet_reader_filter(
                 .map(|reader| filter.with_execution_input(reader)),
         );
     }
+    // It only declares its input's order, which dropping rows below it keeps.
+    if let Some(sorted) = input.downcast_ref::<SortedInputExec>() {
+        return try_attach_parquet_reader_filter(sorted.input(), predicate, config)?
+            .map(|reader| {
+                Arc::clone(input).replace_children(
+                    vec![reader],
+                    ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+                )
+            })
+            .transpose();
+    }
     // Spark inserts IS NOT NULL residuals above equijoin inputs, including AND
     // chains of inferred null checks. A reader predicate can cross those direct
     // checks because both operations only discard rows. Broader deterministic

native/core/src/execution/operators/sorted_input.rs

@@ -37,13 +37,14 @@ use datafusion::physical_plan::{
 /// sorted before it was cached for instance, but the `ScanExec` it arrives through declares none,
 /// and DataFusion's hash join keeps the unmatched probe rows of a right outer join in place only
 /// when the probe input declares an order. Otherwise it moves them after the matched rows of each
-/// batch, which loses an order Spark has already dropped a sort for. The planner puts this only
-/// on that join's probe input, with the order Spark reports for it, and only when the rows reach
-/// that input from the JVM through operators that keep their order. The join keeps the order in
-/// its own output, so the operators above it in the native plan see it too: an aggregate grouped
-/// by the key runs in sorted, streaming mode, a sort aggregate takes the order as already
-/// satisfied, and a top-k sort whose keys start with it stops reading early. A wrong order would
-/// therefore give wrong groups or wrong rows, not only rows out of order.
+/// batch, which loses an order Spark has already dropped a sort for. The planner puts this on a
+/// hash join's probe input, with the order Spark reports for it, when the rows reach that input
+/// from the JVM or a file scan through operators that keep their order, and on the native range,
+/// whose order is real, so that the join and the operators above it in the native plan see the
+/// order: an aggregate grouped by the key runs in sorted, streaming mode, a window finds the order
+/// it requires, a sort aggregate takes the order as already satisfied, and a top-k sort whose keys
+/// start with it stops reading early. A wrong order would therefore give wrong groups or wrong
+/// rows, not only rows out of order.
 #[derive(Debug)]
 pub struct SortedInputExec {
     input: Arc<dyn ExecutionPlan>,
@@ -70,6 +71,10 @@ impl SortedInputExec {
             cache,
         }
     }
+
+    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
+        &self.input
+    }
 }
 
 impl DisplayAs for SortedInputExec {