diff --git a/.github/labeler.yml b/.github/labeler.yml index 83db8b4c6ac..1f7fec4b70a 100644 --- a/.github/labeler.yml +++ b/.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*' diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index b2beae1740b..82a3ca1ef82 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -372,11 +372,16 @@ So an operator that reports an ordering must keep it natively, and the native op by returning `true` from `maintains_input_order` for that child, as DataFusion's filter and projection do. The planner declares the order through an operator only when the operator says it keeps it. An operator that cannot keep the order natively must not report it, or has to be added to -`mayReportFalseOrder` in `operators.scala`. The hash join and window serdes then send no order over -it, an outer hash join above it that has to keep that order falls back to Spark, and a window that -runs in Spark above it gets its input sorted again. Comet's hash joins are the exception within one -native plan: an outer hash join over one of them stays native, even over one that declared no order, -since DataFusion keeps the order it can see there. +`mayReportFalseOrder` in `operators.scala`. The hash join and window serdes then declare no order +over it, an outer hash join above it that has to keep that order falls back to Spark, and a window +that runs in Spark above it gets its input sorted again. A native window with `PARTITION BY` keys +above it sorts its input by the whole order it reports, since Spark relies on that order above the +window too, and falls back to Spark when Comet cannot express all of it. The window's own keys +lead that order, in place of the reported keys they match: a reported key can match one of them +through an equal column only, such as an alias or the other key of an inner join. +Comet's hash joins are the exception within one native plan: an outer hash join over one of them +stays native, even over one that declared no order, since DataFusion keeps the order it can see +there. A leaf that reports an ordering is where the order starts, so the planner has to know that the native operator produces its rows in that order. It knows that of an input from the JVM and of a diff --git a/docs/source/contributor-guide/debugging.md b/docs/source/contributor-guide/debugging.md index 51d74a1ea94..754852cdb39 100644 --- a/docs/source/contributor-guide/debugging.md +++ b/docs/source/contributor-guide/debugging.md @@ -260,8 +260,9 @@ so `spark.comet.debug.enabled` and `COMET_LOG_LEVEL` do not affect it. #### Example: dumping window operator output In `WindowBuilder::build` in `native/core/src/execution/planner/window.rs`, -the window operator builds `window_agg` over `input`, the child plan with -Spark's input order declared, either a `BoundedWindowAggExec` or a +the window operator builds `window_agg` over `window_input`, the child plan +with Spark's input order declared and sorted when it is not in the order the +window relies on, either a `BoundedWindowAggExec` or a `WindowAggExec` depending on whether every window expression runs with bounded memory. It then optionally wraps `window_agg` in a casting `ProjectionExec` and builds the `SparkPlan` from either the projection or `window_agg`. To see the raw @@ -272,14 +273,14 @@ projection: let window_agg: Arc = if all_bounded { Arc::new(BoundedWindowAggExec::try_new( window_expr, - input, + window_input, InputOrderMode::Sorted, !partition_exprs.is_empty(), )?) } else { Arc::new(WindowAggExec::try_new( window_expr, - input, + window_input, !partition_exprs.is_empty(), )?) }; @@ -291,12 +292,13 @@ let window_agg: Arc = { }; // ... the casting `ProjectionExec` is built from `window_agg`, and -// `SparkPlan::new_with_additional(spark_plan.plan_id, projection, vec![child], -// vec![window_agg])` wraps it. Without a cast, `SparkPlan::new` wraps `window_agg`. +// `SparkPlan::new_with_additional` wraps it, with `window_agg` and any input +// sort as additional plans. Without a cast, `SparkPlan::new_with_additional` +// wraps `window_agg`, with any input sort as an additional plan. ``` Wrapping `window_agg` leaves the rest of the builder unchanged. To see the -output after the cast, wrap `projection` the same way just before the +output after the cast, wrap `projection` the same way just before the first `SparkPlan::new_with_additional` call. The `let` rebinding is only compiled in when `debug_assertions` is on, so release builds use the original binding. diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index e54594bac66..4f5ea76cf66 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -155,19 +155,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 @@ -231,6 +238,13 @@ 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 whole order the input reports, since Spark also relies on that order above the window. The results +match Spark, at the cost of a sort. When that order has a key Comet cannot serialize, or one with a non-default string +collation, the window falls back to Spark with the reason "the order of the window's input, which Spark relies on, +cannot be fully expressed natively". + `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). diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 2e1a49d23a1..627a2bf0ed6 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -96,13 +96,13 @@ 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)). | -| `CartesianProductExec` | ⚠️ | Opt-in (`spark.comet.exec.cartesianProduct.enabled`). Runs natively only without a join condition and when each input reports exactly one partition. Row order can differ from Spark's. | +| 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)). | +| `CartesianProductExec` | ⚠️ | Opt-in (`spark.comet.exec.cartesianProduct.enabled`). Runs natively only without a join condition and when each input reports exactly one partition. Row order can differ from Spark's. | `ExistenceJoin` (produced when `EXISTS` / `IN` is combined with another predicate via `OR`) runs natively on `BroadcastHashJoinExec` and `ShuffledHashJoinExec`. Existence sort-merge joins, residual diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 058a6a00cd7..4532a23355c 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/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 diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index d86f8076bde..9d1fd3593de 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -226,6 +226,8 @@ Keep the following in mind when reading the reports: Rust 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 diff --git a/native/core/src/execution/operators/dynamic_filter/early.rs b/native/core/src/execution/operators/dynamic_filter/early.rs index 9b101044bd9..2a51c652f87 100644 --- a/native/core/src/execution/operators/dynamic_filter/early.rs +++ b/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::() { + // 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 diff --git a/native/core/src/execution/operators/dynamic_filter/join/tests.rs b/native/core/src/execution/operators/dynamic_filter/join/tests.rs index 54ae6737edf..e55ab5b38af 100644 --- a/native/core/src/execution/operators/dynamic_filter/join/tests.rs +++ b/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, index: usize) -> Arc { + 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::() + .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::() + .unwrap(); + let reader = filter.input().downcast_ref::().unwrap(); + let (_, source) = reader.downcast_to_file_source::().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; + 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; + 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); +} diff --git a/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs index 6afd49cb226..13be5efc2b2 100644 --- a/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs +++ b/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::() { + 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 diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs index ad561cf49e3..eb5470384ed 100644 --- a/native/core/src/execution/operators/sorted_input.rs +++ b/native/core/src/execution/operators/sorted_input.rs @@ -37,14 +37,16 @@ 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. A window rejects an input -/// that does not declare its partition keys' order. The planner puts this only on that join's -/// probe input or that window's 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. Both keep the +/// batch, which loses an order Spark has already dropped a sort for. A window sorts an input +/// that declares no order on its partition keys, a sort Spark has already left out. The planner +/// puts this on a hash join's probe input or a window's 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. The join and the window keep the /// order in their own output, so the operators above them 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. +/// 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, @@ -71,6 +73,10 @@ impl SortedInputExec { cache, } } + + pub fn input(&self) -> &Arc { + &self.input + } } impl DisplayAs for SortedInputExec { diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 8ea033d11da..80f52bc0c18 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3130,6 +3130,7 @@ mod tests { mod empty_native_scan; mod multi_store_scan; mod semi_join_reduction; + mod window_input_order; use futures::{poll, StreamExt}; use std::{ @@ -4113,7 +4114,7 @@ mod tests { .expect("a hash join") } - fn sorted_outer_hash_join( + fn sorted_hash_join( join_type: spark_operator::JoinType, build_side: spark_operator::BuildSide, children: Vec, @@ -4147,7 +4148,7 @@ mod tests { ), ] { let operator = - sorted_outer_hash_join(join_type, build_side, vec![create_scan(), create_scan()]); + sorted_hash_join(join_type, build_side, vec![create_scan(), create_scan()]); let (_, _, planned) = PhysicalPlanner::default() .create_plan(&operator, &mut vec![], 1) .unwrap(); @@ -4170,6 +4171,39 @@ mod tests { } } + #[test] + fn inner_hash_join_declares_streamed_order_on_probe_input() { + // An inner join keeps its probe order whether or not the probe declares it, but the + // operators above it, a window for one, see that order only once the probe declares it. + for build_side in [ + spark_operator::BuildSide::BuildLeft, + spark_operator::BuildSide::BuildRight, + ] { + let operator = sorted_hash_join( + spark_operator::JoinType::Inner, + build_side, + vec![create_scan(), create_scan()], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + let join = planned_hash_join(&planned.native_plan); + assert_eq!(&datafusion::common::JoinType::Inner, join.join_type()); + assert_eq!("ScanExec", join.left().name(), "{build_side:?}"); + let probe = join.right(); + assert_eq!("SortedInputExec", probe.name(), "{build_side:?}"); + assert_eq!("ScanExec", probe.children()[0].name(), "{build_side:?}"); + assert_eq!( + "col_0@0 ASC", + probe.properties().output_ordering().unwrap().to_string() + ); + assert!( + planned.native_plan.properties().output_ordering().is_some(), + "{build_side:?}" + ); + } + } + #[test] fn hash_join_declares_order_through_operators_that_keep_it() { let projection = Operator { @@ -4179,7 +4213,7 @@ mod tests { })), ..Default::default() }; - let operator = sorted_outer_hash_join( + let operator = sorted_hash_join( spark_operator::JoinType::RightOuter, spark_operator::BuildSide::BuildLeft, vec![create_scan(), projection], @@ -4228,7 +4262,7 @@ mod tests { (LeftAnti, BuildRight, DF::RightAnti), (Existence, BuildRight, DF::RightMark), ] { - let operator = sorted_outer_hash_join( + let operator = sorted_hash_join( spark_operator::JoinType::RightOuter, BuildLeft, vec![ @@ -4279,7 +4313,7 @@ mod tests { ), ] { let operator = - sorted_outer_hash_join(RightOuter, BuildLeft, vec![create_scan(), probe_input]); + sorted_hash_join(RightOuter, BuildLeft, vec![create_scan(), probe_input]); let (_, _, planned) = PhysicalPlanner::default() .create_plan(&operator, &mut vec![], 1) .unwrap(); @@ -4300,7 +4334,7 @@ mod tests { })), ..Default::default() }; - let operator = sorted_outer_hash_join( + let operator = sorted_hash_join( spark_operator::JoinType::RightOuter, spark_operator::BuildSide::BuildLeft, vec![create_scan(), sort], @@ -4551,7 +4585,8 @@ mod tests { #[test] fn window_declares_input_order_on_a_scan() { // Both window operators need their partition keys' order, which a `ScanExec` declares - // only through the order Spark sent. `percent_rank` runs in `WindowAggExec`. + // only through the order Spark sent. Without it the planner sorts the input instead. + // `percent_rank` runs in `WindowAggExec`. let partition_by: Vec> = vec![Arc::new(Column::new("col_0", 0))]; for (func, operator) in [ ("row_number", "BoundedWindowAggExec"), @@ -4581,8 +4616,10 @@ mod tests { assert_eq!(1, ordered.len(), "{func}"); let input = window_input(vec![]); - assert_eq!("ScanExec", input.name()); - let ordered = get_ordered_partition_by_indices(&partition_by, &input).unwrap(); + assert_eq!("SortExec", input.name(), "{func}"); + assert_eq!("ScanExec", input.children()[0].name(), "{func}"); + let ordered = + get_ordered_partition_by_indices(&partition_by, input.children()[0]).unwrap(); assert!(ordered.is_empty(), "{func}"); } } diff --git a/native/core/src/execution/planner/join.rs b/native/core/src/execution/planner/join.rs index f135c6afa70..63923e79cf8 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -538,8 +538,9 @@ impl PhysicalPlanner { /// keeps it natively and says so, and a Comet leaf that reports one is a source this walk /// knows or declares its order itself, as the native range does. One that does not lands /// here with `plan` unchanged, and the join then moves its unmatched rows after Spark has - /// dropped the sort above it, or the window rejects its input. Nothing here can fall back at - /// that point, so the invariant is kept at the operator: see "Adding a New Operator" in the + /// dropped the sort above it, or the window sorts its input by its own keys only, which + /// loses any longer order Spark relies on above it. Nothing here can fall back at that + /// point, so the invariant is kept at the operator: see "Adding a New Operator" in the /// contributor guide. pub(super) fn with_declared_order( &self, diff --git a/native/core/src/execution/planner/tests/window_input_order.rs b/native/core/src/execution/planner/tests/window_input_order.rs new file mode 100644 index 00000000000..d4d10383dfa --- /dev/null +++ b/native/core/src/execution/planner/tests/window_input_order.rs @@ -0,0 +1,599 @@ +// 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. + +//! A native window sorts its input unless it is in the order the window relies on. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, Float64Array, Int32Array, UInt64Array}; +use arrow::compute::cast; +use arrow::compute::SortOptions; +use arrow::datatypes::DataType; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalSortExpr; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use datafusion_comet_proto::spark_expression::{ + self, data_type::DataTypeId, expr::ExprStruct, Expr, +}; +use datafusion_comet_proto::spark_operator::{ + self, lower_window_frame_bound::LowerFrameBoundStruct, operator::OpStruct, + upper_window_frame_bound::UpperFrameBoundStruct, CurrentRow, LowerWindowFrameBound, Operator, + UnboundedPreceding, UpperWindowFrameBound, WindowFrame, WindowFrameType, WindowSpecDefinition, +}; + +use super::collect_with_inputs; +use crate::execution::metrics::utils::to_native_metric_node; +use crate::execution::operators::InputBatch; +use crate::execution::planner::PhysicalPlanner; + +/// `row_number` runs in `BoundedWindowAggExec` and `percent_rank` in `WindowAggExec`. +const FUNCTIONS: [(&str, &str); 2] = [ + ("row_number", "BoundedWindowAggExec"), + ("percent_rank", "WindowAggExec"), +]; + +fn data_type(type_id: DataTypeId) -> spark_expression::DataType { + spark_expression::DataType { + type_id: type_id as i32, + type_info: None, + } +} + +fn int32() -> spark_expression::DataType { + data_type(DataTypeId::Int32) +} + +fn typed_column(index: i32, type_id: DataTypeId) -> Expr { + Expr { + expr_struct: Some(ExprStruct::Bound(spark_expression::BoundReference { + index, + datatype: Some(data_type(type_id)), + })), + ..Default::default() + } +} + +fn column(index: i32) -> Expr { + typed_column(index, DataTypeId::Int32) +} + +/// Spark's default null ordering: ASC NULLS FIRST, DESC NULLS LAST. +fn sort_order_of(child: Expr, descending: bool) -> Expr { + Expr { + expr_struct: Some(ExprStruct::SortOrder(Box::new( + spark_expression::SortOrder { + child: Some(Box::new(child)), + direction: i32::from(descending), + null_ordering: i32::from(descending), + }, + ))), + ..Default::default() + } +} + +fn sort_order(index: i32, descending: bool) -> Expr { + sort_order_of(column(index), descending) +} + +/// Columns from the JVM, which declares no order. +fn scan_of(fields: Vec) -> Operator { + Operator { + op_struct: Some(OpStruct::Scan(spark_operator::Scan { + fields, + source: String::new(), + })), + ..Default::default() + } +} + +/// Three Int32 columns, (a, b, v). +fn scan() -> Operator { + scan_of(vec![int32(); 3]) +} + +fn sort_by(child: Operator, sort_orders: Vec) -> Operator { + Operator { + children: vec![child], + op_struct: Some(OpStruct::Sort(spark_operator::Sort { + sort_orders, + ..Default::default() + })), + ..Default::default() + } +} + +fn sort(child: Operator, order: &[(i32, bool)]) -> Operator { + let orders = order.iter().map(|&(i, desc)| sort_order(i, desc)).collect(); + sort_by(child, orders) +} + +fn window(child: Operator, function: &str, partition_by: &[i32], order_by: &[i32]) -> Operator { + let partition_by = partition_by.iter().map(|&i| column(i)).collect(); + let order_by = order_by.iter().map(|&i| sort_order(i, false)).collect(); + window_by(child, function, partition_by, order_by, None) +} + +/// A window with one function. A `result_type` that differs from the native one makes the +/// planner cast the result in a projection above the window. +fn window_by( + child: Operator, + function: &str, + partition_by_list: Vec, + order_by_list: Vec, + result_type: Option, +) -> Operator { + let expr = spark_operator::WindowExpr { + built_in_window_function: Some(Expr { + expr_struct: Some(ExprStruct::ScalarFunc(spark_expression::ScalarFunc { + func: function.to_string(), + ..Default::default() + })), + ..Default::default() + }), + spec: Some(WindowSpecDefinition { + frame_specification: Some(WindowFrame { + frame_type: WindowFrameType::Rows as i32, + lower_bound: Some(LowerWindowFrameBound { + lower_frame_bound_struct: Some(LowerFrameBoundStruct::UnboundedPreceding( + UnboundedPreceding {}, + )), + }), + upper_bound: Some(UpperWindowFrameBound { + upper_frame_bound_struct: Some(UpperFrameBoundStruct::CurrentRow( + CurrentRow {}, + )), + }), + }), + ..Default::default() + }), + result_type, + ..Default::default() + }; + Operator { + children: vec![child], + op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window { + window_expr: vec![expr], + partition_by_list, + order_by_list, + ..Default::default() + }))), + ..Default::default() + } +} + +fn find_window(plan: &Arc) -> Arc { + if matches!(plan.name(), "BoundedWindowAggExec" | "WindowAggExec") { + return Arc::clone(plan); + } + find_window(plan.children().first().expect("no window in the plan")) +} + +fn count_sorts(plan: &Arc) -> usize { + usize::from(plan.downcast_ref::().is_some()) + + plan.children().into_iter().map(count_sorts).sum::() +} + +/// The keys of the sort directly below the window as (column, descending, nulls first). +fn input_sort_keys(window: &Arc) -> Option> { + let sort = window.children()[0].downcast_ref::()?; + Some( + sort.expr() + .iter() + .map(|e| { + let column = e.expr.downcast_ref::().expect("a column sort key"); + (column.index(), e.options.descending, e.options.nulls_first) + }) + .collect(), + ) +} + +struct Case { + name: &'static str, + input_order: Option<&'static [(i32, bool)]>, + partition_by: &'static [i32], + order_by: &'static [i32], + /// The keys of the sort the planner adds, as (column, descending, nulls first). + added_sort: Option>, +} + +#[test] +fn window_sorts_input_without_partition_key_order() { + let asc = |index: usize| (index, false, true); + let cases = [ + Case { + name: "no declared order", + input_order: None, + partition_by: &[0], + order_by: &[2], + added_sort: Some(vec![asc(0), asc(2)]), + }, + Case { + name: "ordered on the partition key", + input_order: Some(&[(0, false)]), + partition_by: &[0], + order_by: &[2], + added_sort: Some(vec![asc(0), asc(2)]), + }, + Case { + name: "ordered on the partition key and the ORDER BY key", + input_order: Some(&[(0, false), (2, false)]), + partition_by: &[0], + order_by: &[2], + added_sort: None, + }, + Case { + name: "ordered on the first of two partition keys", + input_order: Some(&[(0, false)]), + partition_by: &[0, 1], + order_by: &[2], + added_sort: Some(vec![asc(0), asc(1), asc(2)]), + }, + Case { + name: "ordered on the partition keys in another order", + input_order: Some(&[(1, false), (0, false), (2, false)]), + partition_by: &[0, 1], + order_by: &[2], + added_sort: Some(vec![asc(0), asc(1), asc(2)]), + }, + Case { + name: "ordered descending on the partition key", + input_order: Some(&[(0, true), (2, false)]), + partition_by: &[0], + order_by: &[2], + added_sort: Some(vec![asc(0), asc(2)]), + }, + Case { + name: "no PARTITION BY", + input_order: None, + partition_by: &[], + order_by: &[2], + added_sort: None, + }, + ]; + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new()), 0); + for (function, exec) in FUNCTIONS { + for case in &cases { + let child = match case.input_order { + Some(order) => sort(scan(), order), + None => scan(), + }; + let operator = window(child, function, case.partition_by, case.order_by); + let (_, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + let window = find_window(&planned.native_plan); + let context = format!("{function}, {}", case.name); + assert_eq!(window.name(), exec, "{context}"); + assert_eq!( + count_sorts(&window), + usize::from(case.input_order.is_some()) + usize::from(case.added_sort.is_some()), + "{context}" + ); + if case.added_sort.is_some() { + assert_eq!(input_sort_keys(&window), case.added_sort, "{context}"); + } + } + } +} + +/// The window `operator` with the orders the serde sends for its input. +fn with_input_orders( + mut operator: Operator, + declared: Vec, + unverified: Vec, +) -> Operator { + if let Some(OpStruct::Window(window)) = operator.op_struct.as_mut() { + window.input_sort_orders = declared; + window.unverified_input_sort_orders = unverified; + } + operator +} + +/// A `row_number` window by `a` and `v` over `child`, with the orders the serde sends for its +/// input. +fn window_with_orders(child: Operator, declared: Vec, unverified: Vec) -> Operator { + with_input_orders( + window(child, "row_number", &[0], &[2]), + declared, + unverified, + ) +} + +/// The message of the error that planning `operator` fails with. +fn plan_error(operator: &Operator) -> String { + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new()), 0); + match planner.create_plan(operator, &mut vec![], 1) { + Ok(_) => panic!("the window was planned"), + Err(error) => error.to_string(), + } +} + +#[test] +fn window_rejects_input_not_ordered_on_its_partition_keys() { + // The unverified order is on b while the window partitions by a, so the input it would sort + // is not ordered on a. + let unverified = vec![sort_order(1, false), sort_order(2, false)]; + let error = plan_error(&window_with_orders(scan(), vec![], unverified)); + assert!( + error.contains("not ordered on every PARTITION BY key [col_0@0]"), + "{error}" + ); + assert!( + error.contains("input ordering [col_1@1 ASC, col_2@2 ASC]"), + "{error}" + ); +} + +#[test] +fn window_rejects_both_input_orders() { + let order = || vec![sort_order(0, false), sort_order(2, false)]; + let error = plan_error(&window_with_orders(scan(), order(), order())); + assert!( + error.contains("input_sort_orders (2) and unverified_input_sort_orders (2) are both set"), + "{error}" + ); +} + +#[test] +fn window_sorts_by_the_unverified_order_it_was_sent() { + // (a ASC NULLS FIRST, b DESC NULLS LAST, v ASC NULLS FIRST), longer than the window's own + // order and with a key it does not use. + let claim = [(0, false), (1, true), (2, false)]; + let unverified = || claim.iter().map(|&(i, desc)| sort_order(i, desc)).collect(); + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new()), 0); + + let operator = window_with_orders(scan(), vec![], unverified()); + let (_, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + let window = find_window(&planned.native_plan); + assert_eq!(count_sorts(&window), 1); + let expected = vec![(0, false, true), (1, true, false), (2, false, true)]; + assert_eq!(input_sort_keys(&window), Some(expected)); + + // An input already in that order is not sorted again. + let operator = window_with_orders(sort(scan(), &claim), vec![], unverified()); + let (_, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + assert_eq!(count_sorts(&find_window(&planned.native_plan)), 1); +} + +#[test] +fn trusted_declaration_checks_partition_keys_only() { + // The serde declares only the PARTITION BY key, and the window relies on Spark for the rest. + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new()), 0); + for (declared, sorts) in [(vec![sort_order(0, false)], 0), (vec![], 1)] { + let context = format!("declared {} keys", declared.len()); + let operator = window_with_orders(scan(), declared, vec![]); + let (_, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + assert_eq!( + count_sorts(&find_window(&planned.native_plan)), + sorts, + "{context}" + ); + } +} + +#[test] +fn window_keeps_input_sorted_on_a_float_partition_key() { + // Both the sort key and the PARTITION BY key go through float normalization, so the window + // sees the input as ordered on its key. + let key = || typed_column(0, DataTypeId::Double); + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new()), 0); + for (function, exec) in FUNCTIONS { + let input = scan_of(vec![data_type(DataTypeId::Double), int32(), int32()]); + let child = sort_by( + input, + vec![sort_order_of(key(), false), sort_order(2, false)], + ); + let order_by = vec![sort_order(2, false)]; + let operator = window_by(child, function, vec![key()], order_by, None); + let (_, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + let window = find_window(&planned.native_plan); + assert_eq!(window.name(), exec, "{function}"); + assert_eq!(count_sorts(&window), 1, "{function}"); + } +} + +/// Rows of (a, b, v) for the execution tests. +fn input_columns() -> (Vec, usize) { + let a = vec![1, 0, 1, 0, 1]; + let v = vec![5, 4, 3, 2, 1]; + let rows = a.len(); + let columns: Vec = vec![ + Arc::new(Int32Array::from(a)), + Arc::new(Int32Array::from(vec![0; rows])), + Arc::new(Int32Array::from(v)), + ]; + (columns, rows) +} + +#[tokio::test] +async fn window_reports_its_input_sort_metrics() { + let session = Arc::new(SessionContext::new()); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + for (cast_result, sorted_input) in [(false, false), (false, true), (true, false), (true, true)] + { + let context = format!("cast_result={cast_result} sorted_input={sorted_input}"); + let child = if sorted_input { + sort(scan(), &[(0, false), (2, false)]) + } else { + scan() + }; + // `row_number` returns UInt64 natively, so an Int32 result type adds a cast. + let result_type = cast_result.then(int32); + let partition_by = vec![column(0)]; + let order_by = vec![sort_order(2, false)]; + let operator = window_by(child, "row_number", partition_by, order_by, result_type); + let (mut scans, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + + let additional: Vec<&str> = planned + .additional_native_plans + .iter() + .map(|p| p.name()) + .collect(); + let expected: &[&str] = match (cast_result, sorted_input) { + (false, false) => &["SortExec"], + (false, true) => &[], + (true, false) => &["BoundedWindowAggExec", "SortExec"], + (true, true) => &["BoundedWindowAggExec"], + }; + assert_eq!(additional, expected, "{context}"); + + let (columns, rows) = input_columns(); + let mut inputs = [vec![InputBatch::Batch(columns, rows), InputBatch::EOF].into_iter()]; + let stream = planned.native_plan.execute(0, session.task_ctx()).unwrap(); + let batches = collect_with_inputs(stream, &mut scans, &mut inputs).await; + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + rows, + "{context}" + ); + + let elapsed = |plan: &Arc| { + plan.metrics() + .and_then(|m| m.elapsed_compute()) + .unwrap_or(0) as i64 + }; + let total_time = std::iter::once(&planned.native_plan) + .chain(&planned.additional_native_plans) + .map(elapsed) + .sum::(); + let metrics = to_native_metric_node(&planned).unwrap(); + // Output rows come from the top operator only, and the time covers all of them. + assert_eq!(metrics.metrics["output_rows"], rows as i64, "{context}"); + assert_eq!(metrics.metrics["elapsed_compute"], total_time, "{context}"); + } +} + +#[tokio::test] +async fn window_over_unordered_input_ranks_each_partition() { + let session = Arc::new(SessionContext::new()); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + for (function, expected) in [ + ("row_number", [1.0, 2.0, 1.0, 2.0, 3.0]), + ("percent_rank", [0.0, 1.0, 0.0, 0.5, 1.0]), + ] { + let operator = window(scan(), function, &[0], &[2]); + let (mut scans, _, planned) = planner.create_plan(&operator, &mut vec![], 1).unwrap(); + let (columns, num_rows) = input_columns(); + let mut inputs = [vec![InputBatch::Batch(columns, num_rows), InputBatch::EOF].into_iter()]; + let stream = planned.native_plan.execute(0, session.task_ctx()).unwrap(); + let batches = collect_with_inputs(stream, &mut scans, &mut inputs).await; + let mut rows = vec![]; + for batch in &batches { + let ints = |i: usize| { + let column = batch.column(i).as_any().downcast_ref::(); + column.unwrap().clone() + }; + let result = cast(batch.column(3), &DataType::Float64).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + let (keys, values) = (ints(0), ints(2)); + for row in 0..batch.num_rows() { + rows.push((keys.value(row), values.value(row), result.value(row))); + } + } + rows.sort_by_key(|&(k, v, _)| (k, v)); + let ranks: Vec = rows.iter().map(|r| r.2).collect(); + assert_eq!(rows.len(), num_rows, "{function}"); + assert_eq!(ranks, expected, "{function}"); + } +} + +/// A window over another window with the same partition key and a longer ORDER BY, (a, b, v) +/// read as (k, v, w): the lower one orders by v, the upper one by v, w. The input declares no +/// order. +fn stacked_windows() -> Operator { + let lower = window(scan(), "row_number", &[0], &[1]); + window(lower, "row_number", &[0], &[1, 2]) +} + +#[test] +fn upper_window_input_satisfies_its_own_order_by() { + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new()), 0); + let (_, _, planned) = planner + .create_plan(&stacked_windows(), &mut vec![], 1) + .unwrap(); + let upper = find_window(&planned.native_plan); + let input = Arc::clone(upper.children()[0]); + let schema = input.schema(); + let required = (0..3).map(|i| { + PhysicalSortExpr::new( + Arc::new(Column::new(schema.field(i).name(), i)), + SortOptions { + descending: false, + nulls_first: true, + }, + ) + }); + let satisfied = input + .properties() + .equivalence_properties() + .ordering_satisfy(required) + .unwrap(); + assert!( + satisfied, + "the upper window's input is not ordered on (k, v, w):\n{}", + datafusion::physical_plan::displayable(planned.native_plan.as_ref()).indent(true) + ); +} + +#[tokio::test] +async fn upper_window_numbers_rows_by_its_own_order_by() { + // The expected numbers are a rank over (v, w) within k, so they hold for any correct plan. + // Ties on (k, v) hold w in descending order, and a sort by (k, v) alone keeps them so, as the + // sort of this small input keeps ties in input order. That is what lets the test fail when + // the upper window's input is not sorted by (k, v, w). + let k = vec![0, 0, 0, 0, 1, 1, 1, 1]; + let v = vec![1, 1, 0, 0, 1, 1, 0, 0]; + let w = vec![8, 7, 6, 5, 4, 3, 2, 1]; + let mut expected: Vec<(i32, u64)> = vec![]; + for (i, &wi) in w.iter().enumerate() { + let rank = (0..w.len()) + .filter(|&j| k[j] == k[i] && (v[j], w[j]) <= (v[i], wi)) + .count(); + expected.push((wi, rank as u64)); + } + expected.sort(); + + let session = Arc::new(SessionContext::new()); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + let (mut scans, _, planned) = planner + .create_plan(&stacked_windows(), &mut vec![], 1) + .unwrap(); + let rows = w.len(); + let columns: Vec = vec![ + Arc::new(Int32Array::from(k)), + Arc::new(Int32Array::from(v)), + Arc::new(Int32Array::from(w)), + ]; + let mut inputs = [vec![InputBatch::Batch(columns, rows), InputBatch::EOF].into_iter()]; + let stream = planned.native_plan.execute(0, session.task_ctx()).unwrap(); + let batches = collect_with_inputs(stream, &mut scans, &mut inputs).await; + let mut actual: Vec<(i32, u64)> = vec![]; + for batch in &batches { + let ws = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let numbers = cast(batch.column(4), &DataType::UInt64).unwrap(); + let numbers = numbers.as_any().downcast_ref::().unwrap(); + for row in 0..batch.num_rows() { + actual.push((ws.value(row), numbers.value(row))); + } + } + actual.sort(); + assert_eq!(actual, expected, "(w, row_number by v, w)"); +} diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs index 9d24cea2f8b..ee7a396d360 100644 --- a/native/core/src/execution/planner/window.rs +++ b/native/core/src/execution/planner/window.rs @@ -33,11 +33,13 @@ use datafusion::{ window::WindowExpr, LexOrdering, PhysicalExpr, PhysicalSortExpr, }, + physical_expr_common::physical_expr::format_physical_expr_list, physical_plan::{ limit::LocalLimitExec, projection::ProjectionExec, + sorts::sort::SortExec, windows::{BoundedWindowAggExec, WindowAggExec}, - ExecutionPlan, InputOrderMode, + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, }, }; use datafusion_comet_operators::{PartitionedRankLimitExec, WindowFnKind}; @@ -66,6 +68,17 @@ use crate::{ use ExecutionError::GeneralError; +/// Whether `plan` is ordered on every one of `keys`, in any order of them. +fn orders_every_key( + plan: &Arc, + keys: &[Arc], +) -> Result { + let (_, ordered) = plan + .equivalence_properties() + .find_longest_permutation(keys)?; + Ok(ordered.len() == keys.len()) +} + /// Builder for `Window`. pub struct WindowBuilder; @@ -98,6 +111,70 @@ impl OperatorBuilder for WindowBuilder { let sort_exprs = &sort_exprs?; let partition_exprs = &partition_exprs?; + // Declare the order Spark sent for the input; rows from the JVM declare no order unless + // Spark sent one. A declared input is sorted unless it is ordered on every PARTITION BY + // key. Any other input with PARTITION BY keys is sorted unless it is already in the + // order to restore: the unverified order Spark relies on when sent, otherwise the + // PARTITION BY keys ascending with nulls first, then the ORDER BY keys. + if !wnd.input_sort_orders.is_empty() && !wnd.unverified_input_sort_orders.is_empty() { + return Err(GeneralError(format!( + "Window: input_sort_orders ({}) and unverified_input_sort_orders ({}) \ + are both set", + wnd.input_sort_orders.len(), + wnd.unverified_input_sort_orders.len() + ))); + } + let declared = + planner.with_declared_order(Arc::clone(&child.native_plan), &wnd.input_sort_orders)?; + let restore: Vec = if wnd.unverified_input_sort_orders.is_empty() { + partition_exprs + .iter() + .map(|e| { + PhysicalSortExpr::new( + Arc::clone(e), + SortOptions { + descending: false, + nulls_first: true, + }, + ) + }) + .chain(sort_exprs.iter().cloned()) + .collect() + } else { + wnd.unverified_input_sort_orders + .iter() + .map(|expr| planner.create_sort_expr(expr, Arc::clone(&input_schema))) + .collect::>()? + }; + let needs_sort = if partition_exprs.is_empty() { + false + } else if !wnd.input_sort_orders.is_empty() { + !orders_every_key(&declared, partition_exprs)? + } else { + !declared + .equivalence_properties() + .ordering_satisfy(restore.iter().cloned())? + }; + let input_sort: Option> = match LexOrdering::new(restore) { + Some(ordering) if needs_sort => { + Some(Arc::new(SortExec::new(ordering, Arc::clone(&declared)))) + } + _ => None, + }; + let window_input = input_sort.as_ref().map_or(declared, Arc::clone); + // The window finds its partitions by its own keys, so an input that is not ordered on + // all of them would be read as fewer partitions than it holds. + if !orders_every_key(&window_input, partition_exprs)? { + let ordering = window_input + .output_ordering() + .map_or_else(|| "none".to_string(), |o| format!("[{o}]")); + return Err(GeneralError(format!( + "Window: the input is not ordered on every PARTITION BY key {}, input ordering {}", + format_physical_expr_list(partition_exprs.iter()), + ordering + ))); + } + let window_expr: Result>, ExecutionError> = wnd .window_expr .iter() @@ -132,20 +209,17 @@ impl OperatorBuilder for WindowBuilder { let window_expr = window_expr?; let all_bounded = window_expr.iter().all(|e| e.uses_bounded_memory()); - // Both window operators reject an input that does not declare its partition keys' order. - let input = - planner.with_declared_order(Arc::clone(&child.native_plan), &wnd.input_sort_orders)?; let window_agg: Arc = if all_bounded { Arc::new(BoundedWindowAggExec::try_new( window_expr, - input, + window_input, InputOrderMode::Sorted, !partition_exprs.is_empty(), )?) } else { Arc::new(WindowAggExec::try_new( window_expr, - input, + window_input, !partition_exprs.is_empty(), )?) }; @@ -187,16 +261,21 @@ impl OperatorBuilder for WindowBuilder { proj_exprs, Arc::clone(&window_agg), )?); - // The window is an additional plan so that its time is reported with the - // projection's rather than left out. + // The window and any input sort are additional plans so that their time is + // reported with the projection's rather than left out. SparkPlan::new_with_additional( spark_plan.plan_id, projection, vec![child], - vec![window_agg], + std::iter::once(window_agg).chain(input_sort).collect(), ) } else { - SparkPlan::new(spark_plan.plan_id, window_agg, vec![child]) + SparkPlan::new_with_additional( + spark_plan.plan_id, + window_agg, + vec![child], + input_sort.into_iter().collect(), + ) } }; diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 90c01a7511d..868cc965506 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1255,6 +1255,11 @@ message Window { // The partition keys ascending, when Spark planned no sort below the window because its input // reports that order. Empty otherwise. repeated spark.spark_expression.Expr input_sort_orders = 5; + // The order the input reports and Spark relies on, when Comet cannot verify it, led by the + // window's PARTITION BY keys ascending and its ORDER BY keys in place of the reported keys they + // match. The native window sorts its input by it unless the input is already in that order. + // Never sent together with input_sort_orders. + repeated spark.spark_expression.Expr unverified_input_sort_orders = 6; } // Top-K rows per partition group. Corresponds to Spark's WindowGroupLimitExec diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 81ed81a7028..84c0e94ad5a 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf._ import org.apache.comet.iceberg.IcebergWriteStrategy -import org.apache.comet.rules.{CometCoalesceShufflePartitions, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule} +import org.apache.comet.rules.{CometCoalesceShufflePartitions, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometShuffledHashJoinStrategy, CometSpark34AqeDppFallbackRule} import org.apache.comet.shims.ShimCometSparkSessionExtensions /** @@ -116,6 +116,7 @@ class CometSparkSessionExtensions injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery) injectQueryStageOptimizerRuleShim(extensions, CometCoalesceShufflePartitions) extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) } + extensions.injectPlannerStrategy { session => CometShuffledHashJoinStrategy(session) } } case class CometColumnar(session: SparkSession) extends ColumnarRule { diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 70d108f2506..1d0b49477ff 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -1087,8 +1087,8 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) * Sorts the input of a window that runs in Spark, a Python one included, when its order may be * false. Spark plans no sort below a window whose child reports the order it needs, and its * window trusts that order. A Comet operator that may not keep the order it reports, such as a - * native nested loop join (#6898), then hands the window rows out of order, and it groups them - * wrongly. + * right outer hash join whose probe input declared no order, then hands the window rows out of + * order, and it groups them wrongly. */ private def sortDoubtedWindowInputs(plan: SparkPlan): SparkPlan = plan.transformUp { case window: WindowExecBase @@ -1120,34 +1120,24 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } else { val planWithDateEvaluation = MakeDateEvaluation.preserve(plan) preserveCollectMetricsRowConsumption(planWithDateEvaluation) - // The rewrite removes the sorts of each join it converts, so put back any sort a kept - // operator above still needs. This runs before transform(), which converts the sorts it - // adds like any other SortExec. - val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) { - RewriteJoin.restoreRequiredOrdering(planWithDateEvaluation.transformUp { case p => - if (MakeDateEvaluation.reason(p).isDefined) p else RewriteJoin.rewrite(p, conf) - }) - } else { - planWithDateEvaluation - } // Tag Partial aggregates that must not be converted to Comet because a // corresponding Final or PartialMerge cannot be converted and the intermediate buffer // formats are incompatible. This runs before transform() so the tags are checked // during the bottom-up conversion. Tags persist through AQE stage creation. - tagUnsafePartialAggregates(planWithJoinRewritten) + tagUnsafePartialAggregates(planWithDateEvaluation) // Tag leaves on broadcast join build sides so shouldApplySparkToColumnar can bridge an // unsupported build-side scan (e.g. ORC) to Arrow and let the join run natively (#6008). - tagBroadcastBuildSideLeaves(planWithJoinRewritten) + tagBroadcastBuildSideLeaves(planWithDateEvaluation) - var transformed = transform(planWithJoinRewritten) + var transformed = transform(planWithDateEvaluation) // A build-side bridge is only worth keeping when its broadcast join went native. Where the // join stayed on Spark, clearWastedBuildSideBridges clears the tag from those scans and we // transform once more so the build side rebuilds exactly as it would without the feature. // See issue #6008. if (clearWastedBuildSideBridges(transformed)) { - transformed = transform(planWithJoinRewritten) + transformed = transform(planWithDateEvaluation) } var newPlan = sortDoubtedWindowInputs(revertUnsafePartialAggregates(transformed)) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala new file mode 100644 index 00000000000..3a55bd8e85b --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometShuffledHashJoinStrategy.scala @@ -0,0 +1,182 @@ +/* + * 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. + */ + +package org.apache.comet.rules + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.RowOrdering +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} +import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys +import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, LeftSemi} +import org.apache.spark.sql.catalyst.plans.logical.{Join, JoinHint, JoinStrategyHint, LogicalPlan} +import org.apache.spark.sql.execution.{SparkPlan, SparkStrategy} +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, LogicalQueryStage} +import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.{isCometLoaded, withFallbackReason, withInfo} +import org.apache.comet.shims.ShimJoinSelection + +/** + * Plans a ShuffledHashJoinExec in place of a SortMergeJoinExec, before EnsureRequirements and + * RemoveRedundantSorts, so the sorts above the join survive. Adapted from Apache Gluten. + */ +case class CometShuffledHashJoinStrategy(session: SparkSession) + extends SparkStrategy + with ShimJoinSelection { + + override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match { + case join: Join if isEnabled(join) => planJoin(join) + case _ => Nil + } + + private def sqlConf: SQLConf = session.sessionState.conf + + // Planner strategies run with Comet off and before CometRule, which applies plan-only mode. + private def isEnabled(join: Join): Boolean = + CometConf.COMET_FORCE_SHJ.get(sqlConf) && + CometConf.COMET_EXEC_ENABLED.get(sqlConf) && + !CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(sqlConf) && + !join.isStreaming && + isCometLoaded(sqlConf) + + private def planJoin(join: Join): Seq[SparkPlan] = join match { + case ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, condition, _, left, right, hint) + if !isLeftToSpark(join, hint) && canSortMergeJoin(joinType) && + RowOrdering.isOrderable(leftKeys) && hashJoinSupportedShim(leftKeys, rightKeys) => + buildSide(join) match { + case Some(BuildRight) if joinType == LeftSemi || joinType.isInstanceOf[ExistenceJoin] => + // LeftSemi https://github.com/apache/datafusion-comet/issues/2667 + // ExistenceJoin https://github.com/apache/datafusion-comet/issues/2697 + planBySpark(join) { smj => + withFallbackReason(smj, declineReason(s"BuildRight with $joinType is not supported")) + } + case Some(side) => + buildSideOverLimit(join, side) match { + case Some(reason) => + // Not a fallback: the kept sort-merge join still runs natively. + planBySpark(join)(smj => withInfo(smj, declineReason(reason))) + case None => + Seq( + ShuffledHashJoinExec( + leftKeys, + rightKeys, + joinType, + side, + condition, + planLater(left), + planLater(right))) + } + case None => Nil + } + case _ => Nil + } + + private def isLeftToSpark(join: Join, hint: JoinHint): Boolean = + isBroadcastStage(join.left) || isBroadcastStage(join.right) || + canPlanAsBroadcastHashJoin(join, sqlConf) || hasStrategyHint(hint) + + // canPlanAsBroadcastHashJoin misses a broadcast stage when the adaptive threshold is lower or + // off, and AQE discards a re-plan that hash joins one. + private def isBroadcastStage(plan: LogicalPlan): Boolean = plan match { + case LogicalQueryStage(_, _: BroadcastQueryStageExec) => true + case _ => false + } + + // AQE's NO_BROADCAST_HASH and PREFER_SHUFFLE_HASH are outside JoinStrategyHint.strategies, so + // the join stays with this strategy, which builds on the side PREFER_SHUFFLE_HASH names. + private def hasStrategyHint(hint: JoinHint): Boolean = + Seq(hint.leftHint, hint.rightHint).exists( + _.exists(_.strategy.exists(JoinStrategyHint.strategies.contains))) + + private def declineReason(reason: String): String = + s"Cannot rewrite SortMergeJoin to HashJoin: $reason" + + private def planBySpark(join: Join)( + record: SortMergeJoinExec => SortMergeJoinExec): Seq[SparkPlan] = { + val planned = session.sessionState.planner.JoinSelection(join) + planned.flatMap(_.collect { case smj: SortMergeJoinExec => smj }).foreach(record) + planned + } + + // When AQE hints only one side to build on, that side wins, much as in + // getShuffleHashJoinBuildSide. + private def buildSide(join: Join): Option[BuildSide] = { + val leftBuildable = canBuildShuffledHashJoinLeft(join.joinType) + val rightBuildable = canBuildShuffledHashJoinRight(join.joinType) + val leftHinted = leftBuildable && hintToPreferShuffleHashJoinLeft(join.hint) + val rightHinted = rightBuildable && hintToPreferShuffleHashJoinRight(join.hint) + if (leftHinted != rightHinted) { + Some(if (leftHinted) BuildLeft else BuildRight) + } else if (leftBuildable && rightBuildable) { + Some(optimalBuildSide(join)) + } else if (leftBuildable) { + Some(BuildLeft) + } else if (rightBuildable) { + Some(BuildRight) + } else { + None + } + } + + private def optimalBuildSide(join: Join): BuildSide = { + val leftSize = join.left.stats.sizeInBytes + val rightSize = join.right.stats.sizeInBytes + val leftRowCount = join.left.stats.rowCount + val rightRowCount = join.right.stats.rowCount + if (leftSize == rightSize && rightRowCount.isDefined && leftRowCount.isDefined) { + if (rightRowCount.get <= leftRowCount.get) BuildRight else BuildLeft + } else if (rightSize <= leftSize) { + BuildRight + } else { + BuildLeft + } + } + + /** The largest build side to hash join, or None when the limit is off. */ + private def maxBuildSize: Option[BigInt] = + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.get(sqlConf) match { + case Some(limit) if limit <= 0 => None + case Some(limit) => Some(BigInt(limit)) + case None => + val threshold = sqlConf.autoBroadcastJoinThreshold + val perPartition = + if (threshold > 0) threshold else SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get + Some(BigInt(perPartition) * sqlConf.numShufflePartitions) + } + + /** + * Why the build side is not strictly under the limit, as in canBuildLocalHashMapBySize. An + * unknown size is spark.sql.defaultSizeInBytes, so it never fits. + */ + private def buildSideOverLimit(join: Join, side: BuildSide): Option[String] = + maxBuildSize.flatMap { limit => + val buildSize = side match { + case BuildLeft => join.left.stats.sizeInBytes + case BuildRight => join.right.stats.sizeInBytes + } + if (buildSize >= limit) { + Some( + s"build side size estimate of $buildSize bytes is not under the limit of $limit bytes") + } else { + None + } + } +} diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 3a6f0d6c196..5a1b45c5917 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -197,10 +197,10 @@ private[rules] object MakeDateEvaluation { .getOrElse(plan) /** - * Run before bottom-up conversion and join rewriting. Reasons propagate back to the consuming - * operator, keeping the row pipeline together instead of inserting an eager Arrow batch between - * a throwing projection and its lazy consumer. Exchanges materialize their input independently; - * local sorts and aggregates do so only after their consumer requests input. + * Run before bottom-up conversion. Reasons propagate back to the consuming operator, keeping + * the row pipeline together instead of inserting an eager Arrow batch between a throwing + * projection and its lazy consumer. Exchanges materialize their input independently; local + * sorts and aggregates do so only after their consumer requests input. */ def preserve(plan: SparkPlan): SparkPlan = { if (!enabled(plan)) return plan diff --git a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala b/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala deleted file mode 100644 index 5f543cc84aa..00000000000 --- a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala +++ /dev/null @@ -1,183 +0,0 @@ -/* - * 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. - */ - -package org.apache.comet.rules - -import org.apache.spark.sql.catalyst.expressions.SortOrder -import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, JoinSelectionHelper} -import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, LeftSemi} -import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.execution.{SortExec, SparkPlan} -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} -import org.apache.spark.sql.internal.SQLConf - -import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.{withFallbackReason, withInfo} - -/** - * Adapted from equivalent rule in Apache Gluten. - * - * This rule replaces [[SortMergeJoinExec]] with [[ShuffledHashJoinExec]]. - */ -object RewriteJoin extends JoinSelectionHelper { - - private def getSmjBuildSide(join: SortMergeJoinExec): Option[BuildSide] = { - val leftBuildable = canBuildShuffledHashJoinLeft(join.joinType) - val rightBuildable = canBuildShuffledHashJoinRight(join.joinType) - if (!leftBuildable && !rightBuildable) { - return None - } - if (!leftBuildable) { - return Some(BuildRight) - } - if (!rightBuildable) { - return Some(BuildLeft) - } - val side = join.logicalLink - .flatMap { - case join: Join => Some(getOptimalBuildSide(join)) - case _ => None - } - .getOrElse { - // If smj has no logical link, or its logical link is not a join, - // then we always choose left as build side. - BuildLeft - } - Some(side) - } - - private def removeSort(plan: SparkPlan) = plan match { - case _: SortExec => plan.children.head - case _ => plan - } - - /** - * Adds a local sort under any operator whose child no longer has the ordering the operator - * requires, using the same check as EnsureRequirements. EnsureRequirements puts no sort above a - * sort-merge join whose output ordering already satisfies its parent, so when rewrite turns - * that join into a hash join and removes its sorts, a kept parent such as another sort-merge - * join on the same key would read unsorted input. Plans that already satisfy their required - * orderings are returned unchanged. - */ - def restoreRequiredOrdering(plan: SparkPlan): SparkPlan = plan.transformUp { case p => - p.withNewChildren(p.children.zip(p.requiredChildOrdering).map { - case (child, required) if !SortOrder.orderingSatisfies(child.outputOrdering, required) => - SortExec(required, global = false, child = child) - case (child, _) => child - }) - } - - /** - * The largest build side the rewrite converts, or None when the limit is switched off. Unset - * means Spark's own rule for choosing a shuffled hash join by size: the broadcast threshold - * times the initial shuffle partition count. A non-positive threshold only disables broadcast - * joins and says nothing about the hash table an executor can hold, so Spark's default - * threshold stands in for it. - */ - private def maxBuildSize(conf: SQLConf): Option[BigInt] = - CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.get(conf) match { - case Some(limit) if limit <= 0 => None - case Some(limit) => Some(BigInt(limit)) - case None => - val threshold = conf.autoBroadcastJoinThreshold - val perPartition = - if (threshold > 0) threshold else SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get - Some(BigInt(perPartition) * conf.numShufflePartitions) - } - - /** - * Why the build side is too large to rewrite, or None when it fits. Like Spark's - * canBuildLocalHashMapBySize, the size must be strictly under the limit. The hash join cannot - * spill, so a build side without statistics is refused as well. Spark reports an unknown size - * as Long.MaxValue, which fails the comparison on its own. - */ - private def buildSideOverLimit( - smj: SortMergeJoinExec, - buildSide: BuildSide, - conf: SQLConf): Option[String] = - maxBuildSize(conf).flatMap { limit => - smj.logicalLink match { - case Some(join: Join) => - val buildSize = buildSide match { - case BuildLeft => join.left.stats.sizeInBytes - case BuildRight => join.right.stats.sizeInBytes - } - if (buildSize >= limit) { - Some( - s"build side size estimate of $buildSize bytes is not under the limit of " + - s"$limit bytes") - } else { - None - } - case _ => Some("no statistics are available for the build side") - } - } - - def rewrite(plan: SparkPlan, conf: SQLConf): SparkPlan = plan match { - case smj: SortMergeJoinExec => - getSmjBuildSide(smj) match { - case Some(BuildRight) - if smj.joinType == LeftSemi || smj.joinType.isInstanceOf[ExistenceJoin] => - // LeftSemi https://github.com/apache/datafusion-comet/issues/2667 - // ExistenceJoin https://github.com/apache/datafusion-comet/issues/2697 - withFallbackReason( - smj, - "Cannot rewrite SortMergeJoin to HashJoin: " + - s"BuildRight with ${smj.joinType} is not supported") - plan - case Some(buildSide) => - buildSideOverLimit(smj, buildSide, conf) match { - case Some(reason) => - // The sort-merge join still runs natively, so this is information for extended - // explain rather than a fallback reason. - withInfo(smj, s"Cannot rewrite SortMergeJoin to HashJoin: $reason") - plan - case None => - ShuffledHashJoinExec( - smj.leftKeys, - smj.rightKeys, - smj.joinType, - buildSide, - smj.condition, - removeSort(smj.left), - removeSort(smj.right), - smj.isSkewJoin) - } - case _ => plan - } - case _ => plan - } - - def getOptimalBuildSide(join: Join): BuildSide = { - val leftSize = join.left.stats.sizeInBytes - val rightSize = join.right.stats.sizeInBytes - val leftRowCount = join.left.stats.rowCount - val rightRowCount = join.right.stats.rowCount - if (leftSize == rightSize && rightRowCount.isDefined && leftRowCount.isDefined) { - if (rightRowCount.get <= leftRowCount.get) { - return BuildRight - } - return BuildLeft - } - if (rightSize <= leftSize) { - return BuildRight - } - BuildLeft - } -} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala index 3c3b2df03d5..8be707c741b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala @@ -78,11 +78,13 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { // Spark plans no sort below the window when its input already reports the order the window // needs. The native window relies on that order only when mayReportFalseOrder trusts it. A // window that runs in Spark over a doubted order gets its input sorted again by - // CometExecRule, since Spark's window would trust that order too. + // CometExecRule, since Spark's window would trust that order too. An order that does not + // cover the window's own keys cannot be sorted by alone, as a Comet operator can report less + // than the Spark operator it replaced. val requiredOrdering = op.requiredChildOrdering.head + val satisfiesRequired = SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) val trustsInputOrder = - SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) && - !CometHashJoinExec.mayReportFalseInputOrder(op.child) + satisfiesRequired && !CometHashJoinExec.mayReportFalseInputOrder(op.child) // The native window finds partitions and peers by comparing keys byte by byte, so it would // tell apart collated strings that Spark treats as equal, as CometWindowGroupLimitExec notes. @@ -98,8 +100,8 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { } // Without partition keys the native window checks no order at all, so over a doubted order it - // would compute over rows that may be out of order. With partition keys it rejects an input - // whose order it cannot see. + // would compute over rows that may be out of order. With partition keys the native planner + // sorts an input whose order it cannot see. if (op.partitionSpec.isEmpty && op.orderSpec.nonEmpty && !trustsInputOrder) { withFallbackReason( op, @@ -131,12 +133,38 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { } else { Nil } + // Spark relies on the whole order the input reports, also above the window, which reports + // it too. When Comet cannot verify it, the native window sorts its input by all of it, so + // all of it has to be expressible. A reported key may match the window's key through an + // equal column only, an alias or the other key of an inner join, so the window's own keys + // replace the prefix they satisfy, which sorts the rows the same way. + val unverifiedInputSortOrders = + if (op.partitionSpec.nonEmpty && satisfiesRequired && !trustsInputOrder) { + val claim = requiredOrdering ++ op.child.outputOrdering.drop(requiredOrdering.length) + val orders = claim.iterator + .takeWhile(order => !hasNonDefaultStringCollation(order.dataType)) + .map(exprToProto(_, op.child.output)) + .takeWhile(_.isDefined) + .flatten + .toSeq + if (orders.length < claim.length) { + withFallbackReason( + op, + "the order of the window's input, which Spark relies on, cannot be fully " + + "expressed natively") + return None + } + orders + } else { + Nil + } val windowBuilder = OperatorOuterClass.Window.newBuilder() windowBuilder.addAllWindowExpr(windowExprProto.map(_.get).asJava) windowBuilder.addAllPartitionByList(partitionExprs.map(_.get).asJava) windowBuilder.addAllOrderByList(sortOrders.map(_.get).asJava) windowBuilder.addAllInputSortOrders(inputSortOrders.asJava) + windowBuilder.addAllUnverifiedInputSortOrders(unverifiedInputSortOrders.asJava) Some(builder.setWindow(windowBuilder).build()) } else { // Whichever of the window / partition / order expressions failed has already recorded its diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index d07490ee914..a242fd3ee1b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -2966,26 +2966,29 @@ trait CometHashJoin { // and an input sorted outside the native plan, a relation sorted before it was cached for // instance, declares no order to DataFusion. Send the order so the native join declares it // on the probe input; see SortedInputExec. Every other join type keeps the probe order as - // it is. The null-aware anti join is the exception, as Comet builds that one on its streamed - // side, but it runs natively only when that side reports no order, so it reports none. + // it is, but an operator above it in the native plan, a window for one, sees that order only + // if the probe input declares it, so the order is sent for those joins too. The null-aware + // anti join is the exception, as Comet builds that one on its streamed side, but it runs + // natively only when that side reports no order, so it reports none. // - // The order Spark reports can be false to begin with. A Comet join that may move its - // streamed rows reports it all the same: a right outer join like this one whose probe input - // declared no order, and the nested loop join, which keeps the order of neither input. A - // cache reports the order of the plan that wrote it, so the rows of such a join can reach - // this one through a cache, as a JVM input that the native planner cannot look behind. - // Declaring the order there would let an operator above rely on it, an - // aggregate in sorted mode for instance, which then splits groups. So nothing is declared - // when such a join feeds the streamed side with no sort in between, or when a leaf reports - // an order that cannot be checked; see mayReportFalseOrder. When it comes from the JVM or a - // nested loop join, the join stays in Spark, as Spark may have dropped a sort above it. + // The order Spark reports can be false to begin with. A right outer join like this one whose + // probe input declared no order reports it all the same. A cache reports the order of the + // plan that wrote it, so the rows of such a join can reach this one through a cache, as a JVM + // input that the native planner cannot look behind. Declaring the order there would let an + // operator above rely on it, an aggregate in sorted mode for instance, which then splits + // groups. So nothing is declared when such a join feeds the streamed side with no sort in + // between, or when a leaf reports an order that cannot be checked; see mayReportFalseOrder. + // When it comes from the JVM, an outer join of the two shapes stays in Spark, as Spark may + // have dropped a sort above it. // // Declaring any order makes DataFusion keep the probe rows in place, so the longest prefix // of the order that Comet can express is enough. A key that exprToProto takes counts as // expressible, unless it holds a collated string, at the top level or nested, which Spark // orders by its collation but the native side would declare in raw byte order. That is the // check `supportedSortType` makes for a sort. Nothing here can tell whether an operator above - // relies on the order, so a join whose first key Comet cannot express stays in Spark. + // relies on the order, so an outer join of the two shapes whose first key Comet cannot express + // stays in Spark. Any other join keeps its probe rows in place regardless, so it stays native + // and declares what it can. val (streamed, _) = streamedSide(join.buildSide, join.left, join.right) val keepsStreamedOrder = runsAsRightOuterJoin(join.joinType, join.buildSide) && join.outputOrdering.nonEmpty @@ -2994,14 +2997,14 @@ trait CometHashJoin { return None } val streamedSortOrders = - if (keepsStreamedOrder && !mayReportFalseOrder(streamed, countInPlan = true)) { + if (join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed, countInPlan = true)) { val orders = join.outputOrdering.iterator .takeWhile(order => !hasNonDefaultStringCollation(order.dataType)) .map(exprToProto(_, streamed.output)) .takeWhile(_.isDefined) .flatten .toSeq - if (orders.isEmpty) { + if (orders.isEmpty && keepsStreamedOrder) { withFallbackReason( join, "the streamed side's sort order, which this outer join has to keep, cannot be " + @@ -3040,10 +3043,9 @@ trait CometHashJoin { // Whether the order that `plan` reports may be false, because something that feeds it with no // sort in between cannot be relied on to keep it. Comet's hash joins report the order of their - // streamed side, which Spark's own joins keep, but two Comet joins do not always: an outer join - // that runs as DataFusion's right outer join, which keeps its probe rows in place only if the - // planner declared an order on its probe input; and the nested loop join, which DataFusion runs - // build row by build row. So does a leaf that reports an order from state this walk cannot + // streamed side, which Spark's own joins keep, but an outer join that runs as DataFusion's right + // outer join does not always, as it keeps its probe rows in place only if the planner declared + // an order on its probe input. So does a leaf that reports an order from state this walk cannot // check, an RDD checkpointed from a plan that is gone for one. // // The walk follows the order, so a hash join, which passes on only its streamed side's, is @@ -3060,9 +3062,6 @@ trait CometHashJoin { mayReportFalseOrder(adaptive.executedPlan, countInPlan) case _ if plan.outputOrdering.isEmpty => false case _: SortExec | _: CometSortExec => false - // Its rows stay in order only over a one-row build side, and a native join above would move - // them. Over more build rows they are already out of order (#6898); drop this once fixed. - case _: CometBroadcastNestedLoopJoinExec => true case join: CometBroadcastHashJoinExec => hashJoinMayReportFalseOrder( join.nativeOp, @@ -3206,11 +3205,23 @@ object CometBroadcastNestedLoopJoinExec extends CometOperatorSerde[BroadcastNest " DataFusion's NestedLoopJoinExec does not provide. Affects: LeftOuter+BuildLeft," + " RightOuter+BuildRight, FullOuter, LeftSemi+BuildLeft, LeftAnti+BuildLeft." + // Spark's nested loop join keeps its streamed side's order and reports it, so a sort above it + // may be gone, but DataFusion's emits its rows build row by build row. Remove this fallback once + // the native join keeps that order. + private[apache] val streamedOrderReason: String = + "nested loop join over an ordered streamed side is not supported natively, because the " + + "native join does not keep that order " + + "(https://github.com/apache/datafusion-comet/issues/6898)" + override def getSupportLevel(op: BroadcastNestedLoopJoinExec): SupportLevel = (op.joinType, op.buildSide) match { case (Inner, _) | (LeftOuter, BuildRight) | (RightOuter, BuildLeft) | (LeftSemi, BuildRight) | (LeftAnti, BuildRight) => - Compatible(None) + if (op.outputOrdering.nonEmpty) { + Unsupported(Some(streamedOrderReason)) + } else { + Compatible(None) + } case (LeftOuter | RightOuter | FullOuter | LeftSemi | LeftAnti, _) => Unsupported(Some(broadcastBuildReplicationReason)) case _ => Unsupported(Some(s"Join type ${op.joinType} is not supported")) diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimJoinSelection.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimJoinSelection.scala new file mode 100644 index 00000000000..c26ffe295f0 --- /dev/null +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimJoinSelection.scala @@ -0,0 +1,35 @@ +/* + * 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. + */ + +package org.apache.comet.shims + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.optimizer.JoinSelectionHelper +import org.apache.spark.sql.catalyst.plans.JoinType + +trait ShimJoinSelection extends JoinSelectionHelper { + + /** Spark 3.x has no collations. */ + protected def hashJoinSupportedShim( + leftKeys: Seq[Expression], + rightKeys: Seq[Expression]): Boolean = true + + /** Spark 3.x has no LeftSingle. */ + protected def canSortMergeJoin(joinType: JoinType): Boolean = true +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimJoinSelection.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimJoinSelection.scala new file mode 100644 index 00000000000..f975aa5ddbb --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimJoinSelection.scala @@ -0,0 +1,35 @@ +/* + * 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. + */ + +package org.apache.comet.shims + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.optimizer.JoinSelectionHelper +import org.apache.spark.sql.catalyst.plans.{JoinType, LeftSingle} + +trait ShimJoinSelection extends JoinSelectionHelper { + + /** False for keys whose collation is not binary stable, which Spark never hash joins. */ + protected def hashJoinSupportedShim( + leftKeys: Seq[Expression], + rightKeys: Seq[Expression]): Boolean = hashJoinSupported(leftKeys, rightKeys) + + /** Spark plans a LeftSingle join as a hash or nested loop join, never a sort-merge join. */ + protected def canSortMergeJoin(joinType: JoinType): Boolean = joinType != LeftSingle +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 0ce9241635e..b6eda56162d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -43,11 +43,11 @@ import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometBroadcastNes import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FileSourceScanExec, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FileSourceScanExec, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.functions.{col, count, lit, max, min, sum} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} @@ -3920,12 +3920,13 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("a window that stays in Spark over a nested loop join sorts the join's rows again") { - // A native nested loop join emits its rows build row by build row, out of the streamed side's - // order that it reports (#6898), and Spark plans no sort above it. A window that runs in Spark - // there, one without partition keys or one with a collated key, gets its input sorted again. - // One partition, so no exchange comes between the join and the window, with or without AQE. - // With AQE, the stage's own pass can make the added sort native and then the window above it. + test("a window over a nested loop join over a sorted cache keeps the join in Spark") { + // Spark plans no sort above a nested loop join that reports its streamed side's order, but + // the native join emits its rows build row by build row (#6898), so the join stays in Spark. + // A window above it, one without partition keys or one with a collated key, then sees the + // rows in order, with no sort of its own. It runs in Spark as well, since the join's rows + // have no way into a native plan. One partition, so no exchange comes between the join and + // the window, with or without AQE. Seq(false, true).foreach { aqe => withConversions(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { try { @@ -3963,16 +3964,12 @@ class CometInMemoryCacheSuite extends CometTestBase { withClue(s"$plan\n") { assert(plan.isInstanceOf[AdaptiveSparkPlanExec] == aqe) val finalPlan = stripAQEPlan(plan) - val joins = collect(finalPlan) { case j: CometBroadcastNestedLoopJoinExec => j } - assert(joins.nonEmpty) - val windowInputs = collect(finalPlan) { - case w: WindowExec => w.child - case w: CometWindowExec => w.child - } - assert(windowInputs.nonEmpty && windowInputs.forall { - case _: SortExec | _: CometSortExec | WholeStageCodegenExec(_: SortExec) => true - case _ => false - }) + assert(collect(finalPlan) { case j: BroadcastNestedLoopJoinExec => j }.nonEmpty) + assert(collect(finalPlan) { case j: CometBroadcastNestedLoopJoinExec => j }.isEmpty) + assert(collect(finalPlan) { case w: WindowExec => w }.size == 1) + assert(collect(finalPlan) { case w: CometWindowExec => w }.isEmpty) + val sorts = collect(finalPlan) { case s @ (_: CometSortExec | _: SortExec) => s } + assert(sorts.isEmpty) } } } finally { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 7cd9923cc45..e12e0829937 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -25,6 +25,8 @@ import org.scalactic.source.Position import org.scalatest.Tag import org.apache.hadoop.fs.Path +import org.apache.logging.log4j.{Level, LogManager} +import org.apache.logging.log4j.core.LoggerContext import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.spark.{SparkException, TaskContext} @@ -35,23 +37,23 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeReference, BoundReference, Cast, Descending, DynamicPruningExpression, Expression, IsNotNull, Literal, SortOrder} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.{Inner, JoinType, LeftOuter, LeftSemi} -import org.apache.spark.sql.catalyst.plans.logical.Join +import org.apache.spark.sql.catalyst.plans.logical.{HintInfo, Join, JoinHint, LogicalPlan, NO_BROADCAST_HASH, PREFER_SHUFFLE_HASH} import org.apache.spark.sql.catalyst.plans.physical.SinglePartition -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometCartesianProductExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometCartesianProductExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec, CometWindowGroupLimitExec} import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec -import org.apache.spark.sql.execution.{CoalesceExec, ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, SQLExecution, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{CoalesceExec, ColumnarToRowTransition, CommandResultExec, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, SparkStrategy, SQLExecution, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.aggregate.SortAggregateExec import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions.{array, count, lit, struct, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, DoubleType, IntegerType, LongType, MetadataBuilder, StructField, StructType} import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, ExtendedExplainInfo} -import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus} -import org.apache.comet.rules.RewriteJoin +import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus, isSpark40Plus} +import org.apache.comet.rules.CometShuffledHashJoinStrategy import org.apache.comet.serde.OperatorOuterClass class CometJoinSuite extends CometTestBase { @@ -1099,11 +1101,8 @@ class CometJoinSuite extends CometTestBase { } } - // The MERGE hint makes Spark plan a SortMergeJoin, which forceShuffledHashJoin then rewrites. - private val forcedSmjQuery = - "SELECT /*+ MERGE(tbl_a) */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1" + private val forcedSmjQuery = "SELECT * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1" - // tbl_b is the smaller side, so the rewrite builds on the right. private def withForcedSmjTables(f: => Unit): Unit = { withParquetTable((0 until 1000).map(i => (i, i % 5)), "tbl_a") { withParquetTable((0 until 10).map(i => (i % 10, i + 2)), "tbl_b") { @@ -1115,9 +1114,7 @@ class CometJoinSuite extends CometTestBase { private def rewriteReason(reason: String): String = s"Cannot rewrite SortMergeJoin to HashJoin: $reason" - // Checks that the SortMergeJoin and its sorts survived and that the kept join, which still - // runs natively, is not reported as a fallback. Returns Spark's size estimate of the build - // side, read through the same logical link the rewrite used. + // Returns Spark's size estimate of the build side. private def assertSortMergeJoinKept(plan: SparkPlan): BigInt = { val joins = collect(plan) { case j: CometSortMergeJoinExec => j } assert(joins.size == 1, plan) @@ -1161,6 +1158,8 @@ class CometJoinSuite extends CometTestBase { test(s"forceShuffledHashJoin keeps SortMergeJoin over maxBuildSize, AQE=$adaptive") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") { withForcedSmjTables { @@ -1173,6 +1172,8 @@ class CometJoinSuite extends CometTestBase { test(s"forceShuffledHashJoin rewrites SortMergeJoin under maxBuildSize, AQE=$adaptive") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1g") { withForcedSmjTables { @@ -1189,6 +1190,7 @@ class CometJoinSuite extends CometTestBase { var estimate = BigInt(0) withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) @@ -1197,6 +1199,7 @@ class CometJoinSuite extends CometTestBase { // A build side exactly at the limit is not under it, as in canBuildLocalHashMapBySize. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> estimate.toString) { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) @@ -1204,6 +1207,7 @@ class CometJoinSuite extends CometTestBase { } withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> (estimate + 1).toString) { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) @@ -1216,6 +1220,7 @@ class CometJoinSuite extends CometTestBase { withSQLConf( CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "0", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { withForcedSmjTables { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) @@ -1225,21 +1230,22 @@ class CometJoinSuite extends CometTestBase { } test("forceShuffledHashJoin without maxBuildSize uses Spark's shuffled hash join size rule") { - // The limit is autoBroadcastJoinThreshold times the shuffle partition count. - withSQLConf( - CometConf.COMET_FORCE_SHJ.key -> "true", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { - withForcedSmjTables { + withForcedSmjTables { + var estimate = BigInt(0) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + estimate = assertSortMergeJoinKept(cometPlan) assertKeptWithReason(cometPlan, 2) } - } - withSQLConf( - CometConf.COMET_FORCE_SHJ.key -> "true", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1g") { - withForcedSmjTables { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> (estimate - 1).toString) { val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) assertRewrittenToHashJoin(cometPlan) } @@ -1271,30 +1277,47 @@ class CometJoinSuite extends CometTestBase { } } - test("forceShuffledHashJoin keeps a SortMergeJoin that has no statistics") { - val left = spark.range(1).toDF("a").queryExecution.sparkPlan - val right = spark.range(1).toDF("b").queryExecution.sparkPlan - def smj: SortMergeJoinExec = - SortMergeJoinExec(left.output, right.output, Inner, None, left, right) - assert(smj.logicalLink.isEmpty) - - val kept = smj - assert(RewriteJoin.rewrite(kept, spark.sessionState.conf) eq kept) - assert(!hasFallbackReason(kept)) - val info = kept.getTagValue(CometExplainInfo.EXTENSION_INFO) - assert( - info.exists(_.contains(rewriteReason("no statistics are available for the build side"))), - info) - - withSQLConf(CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { - val rewritten = RewriteJoin.rewrite(smj, spark.sessionState.conf) - assert(rewritten.isInstanceOf[ShuffledHashJoinExec], rewritten) + test("forceShuffledHashJoin keeps a SortMergeJoin whose build side has no size estimate") { + val schema = StructType(Seq(StructField("k", IntegerType), StructField("v", IntegerType))) + def unsized(name: String): Unit = spark + .createDataFrame(spark.sparkContext.parallelize((0 until 10).map(i => Row(i, i))), schema) + .createOrReplaceTempView(name) + val query = "SELECT * FROM rdd_a JOIN rdd_b ON rdd_a.k = rdd_b.k" + withTempView("rdd_a", "rdd_b") { + unsized("rdd_a") + unsized("rdd_b") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_FORCE_SHJ.key -> "true") { + val (_, cometPlan) = checkSparkAnswer(query) + val kept = collect(cometPlan) { + case j: SortMergeJoinExec => j + case j: CometSortMergeJoinExec => j.originalPlan + } + assert(kept.size == 1, cometPlan) + val estimate = kept.head.logicalLink match { + case Some(join: Join) => join.right.stats.sizeInBytes + case other => fail(s"Expected a Join logical link, got $other") + } + assert(estimate == spark.sessionState.conf.defaultSizeInBytes, estimate) + val info = kept.head.getTagValue(CometExplainInfo.EXTENSION_INFO).toSeq.flatten + val expected = rewriteReason(s"build side size estimate of $estimate bytes") + assert(info.exists(_.contains(expected)), info) + + withSQLConf(CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + val (_, cometPlan) = checkSparkAnswer(query) + val hashJoins = collect(cometPlan) { + case j: ShuffledHashJoinExec => j + case j: CometHashJoinExec => j + } + assert(hashJoins.size == 1, cometPlan) + } + } } } - // Both joins are on the same key, so EnsureRequirements puts no sort between them: the inner - // sort-merge join's output ordering satisfies the outer one. When the inner join is rewritten - // to a hash join and the outer one is kept, the outer join needs that ordering restored. + // Both joins are on the same key, so EnsureRequirements puts no sort between them. private def withChainedJoinTables(midRows: Int, midKeys: Int)(f: => Unit): Unit = { withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { @@ -1305,6 +1328,16 @@ class CometJoinSuite extends CometTestBase { } } + private def withForcedShjTables(adaptive: Boolean)(f: => Unit): Unit = withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + withChainedJoinTables(midRows = 3000, midKeys = 100)(f) + } + // Whether a hash join is reachable without leaving the stage. private def reachesHashJoinInStage(plan: SparkPlan): Boolean = plan match { case _: CometHashJoinExec => true @@ -1323,8 +1356,6 @@ class CometJoinSuite extends CometTestBase { case _ => false } - // Checks the answer first, then that a kept sort-merge join reads the rewritten hash join - // through a sort. private def checkSortRestoredOverHashJoin(query: String): Unit = { val (_, cometPlan) = checkSparkAnswer(query) val inputs = collect(cometPlan) { @@ -1366,23 +1397,10 @@ class CometJoinSuite extends CometTestBase { } } - // The rewrite never converts a LeftSemi or ExistenceJoin that builds on the right, whatever - // the build size, so these joins are kept above the rewritten inner join with no size limit. - // Each predicate also refers to small, so neither the filter nor the semi join is pushed below - // the inner join. + // Each predicate refers to small, so neither the filter nor the semi join is pushed down. for (adaptive <- Seq(false, true)) { - def withKeptSemiJoinConf(f: => Unit): Unit = withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, - SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - CometConf.COMET_FORCE_SHJ.key -> "true", - CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { - withChainedJoinTables(midRows = 3000, midKeys = 100)(f) - } - test(s"forceShuffledHashJoin restores ordering for a kept ExistenceJoin, AQE=$adaptive") { - withKeptSemiJoinConf { + withForcedShjTables(adaptive) { checkSortRestoredOverHashJoin( "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " + "WHERE small._2 < 0 OR EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1)") @@ -1393,7 +1411,7 @@ class CometJoinSuite extends CometTestBase { withSQLConf( SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> "org.apache.spark.sql.catalyst.optimizer.PushLeftSemiLeftAntiThroughJoin") { - withKeptSemiJoinConf { + withForcedShjTables(adaptive) { checkSortRestoredOverHashJoin( "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " + "WHERE EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1 AND mid._2 > small._2)") @@ -1406,19 +1424,491 @@ class CometJoinSuite extends CometTestBase { // output ordering satisfies it, so the window needs that ordering restored too. for (adaptive <- Seq(false, true)) { test(s"forceShuffledHashJoin restores ordering for a window over the join, AQE=$adaptive") { + withForcedShjTables(adaptive) { + val (_, cometPlan) = checkSparkAnswer( + "SELECT big._1, big._2, small._2, count(*) OVER (PARTITION BY big._1) " + + "FROM big JOIN small ON big._1 = small._1") + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) + assert(collect(cometPlan) { case w: CometWindowExec => w }.nonEmpty, cometPlan) + } + } + } + + test("forceShuffledHashJoin chain keeps early join dynamic filtering") { + withForcedShjTables(adaptive = false) { + val query = "SELECT big._1, big._2, mid._2, small._2 FROM big " + + "JOIN mid ON big._1 = mid._1 JOIN small ON big._1 = small._1" + for (enabled <- Seq(false, true)) { + withSQLConf(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> enabled.toString) { + val (_, plan) = checkSparkAnswerAndOperator(sql(query)) + val joins = collect(plan) { case join: CometHashJoinExec => join } + assert(joins.map(_.buildSide) == Seq(BuildRight, BuildRight), plan) + // big has 100 rows per key, mid 30, and small one row for keys 0 to 90 by 10. + val outputs = joins.map(_.metrics("output_rows").value).sorted + assert(outputs == (if (enabled) Seq(30000L, 30000L) else Seq(30000L, 300000L))) + if (enabled) { + assert( + joins.map(_.metrics("dynamic_filter_early_rows_evaluated").value).sum == 10000L) + assert(joins.map(_.metrics("dynamic_filter_early_rows_pruned").value).sum == 9000L) + } else { + assert(joins.forall(!_.metrics.contains("dynamic_filter_early_rows_pruned"))) + } + } + } + } + } + + private def checkPartitionsSortedOverHashJoin(df: => DataFrame): Unit = { + // withSQLConf returns Unit on Spark 3.x, so the expected rows are assigned inside it. + var expected = Seq.empty[String] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = df.collect().map(_.toString).sorted.toSeq + } + val cometDf = df + assert(cometDf.collect().map(_.toString).sorted.toSeq == expected) + val cometPlan = cometDf.queryExecution.executedPlan + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) + val sorted = cometDf.queryExecution.toRdd + .mapPartitions { it => + val keys = it.map(_.getInt(0)).toArray + Iterator(keys.sameElements(keys.sorted)) + } + .collect() + assert(sorted.forall(identity), sorted.mkString(", ")) + } + + private val sortLossJoin = + "SELECT big._1 AS k, big._2 AS v FROM big JOIN small ON big._1 = small._1" + + // Spark drops a sort that the sort-merge join's ordering satisfies (#6770). + for (adaptive <- Seq(false, true)) { + test(s"forceShuffledHashJoin keeps sortWithinPartitions on the join key, AQE=$adaptive") { + withForcedShjTables(adaptive) { + checkPartitionsSortedOverHashJoin(sql(sortLossJoin).sortWithinPartitions("k")) + } + } + + // Unsorted partition values make the single dynamic partition writer reopen a closed file. + test(s"forceShuffledHashJoin keeps the sort of a partitioned insert, AQE=$adaptive") { + withForcedShjTables(adaptive) { + withTable("out") { + sql("CREATE TABLE out (v INT, w INT, k INT) USING parquet PARTITIONED BY (k)") + val insert = sql( + "INSERT OVERWRITE TABLE out " + + "SELECT big._2, small._2, big._1 FROM big JOIN small ON big._1 = small._1") + val written = insert.queryExecution.executedPlan match { + case c: CommandResultExec => c.commandPhysicalPlan + case p => p + } + assert(collect(written) { case j: CometHashJoinExec => j }.nonEmpty, written) + assert(spark.table("out").count() == 1000) + val files = spark.table("out").inputFiles + val perPartition = files.groupBy(_.split("/").find(_.startsWith("k="))) + assert(perPartition.size == 10, files.mkString("\n")) + assert(perPartition.values.forall(_.length == 1), files.mkString("\n")) + } + } + } + + for ((joinType, key) <- Seq("RIGHT" -> "big._1", "FULL" -> "coalesce(big._1, small._1)")) { + test( + s"forceShuffledHashJoin keeps sortWithinPartitions over a $joinType join, " + + s"AQE=$adaptive") { + withForcedShjTables(adaptive) { + checkPartitionsSortedOverHashJoin( + sql(s"SELECT $key AS k, big._2 AS v, small._2 AS w " + + s"FROM small $joinType JOIN big ON small._1 = big._1") + .sortWithinPartitions("k")) + } + } + } + } + + test("forceShuffledHashJoin keeps probe side order when no exchange is added") { + withForcedShjTables(adaptive = false) { + val streamed = spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1") + val build = spark.table("small") + checkPartitionsSortedOverHashJoin( + streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_1").as("k"), streamed("_2"), build("_2")) + .sortWithinPartitions("k")) + } + } + + // The cache reaches the native join with no declared order, so the join must declare it. + test("forceShuffledHashJoin keeps the order of a sorted cache through an outer join") { + withForcedShjTables(adaptive = false) { + withSQLConf( + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + val streamed = spark + .table("big") + .repartition(2, $"_1") + .sortWithinPartitions("_1", "_2") + .cache() + try { + streamed.count() + val build = spark.table("small") + val valueOrder = SortOrder(BoundReference(1, IntegerType, nullable = true), Ascending) + for { + hint <- Seq(None, Some("shuffle_hash")) + (joinType, buildSide) <- Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft) + } { + val hinted = hint.fold(build)(build.hint(_)) + val joined = if (joinType == "left_outer") { + streamed.join(hinted, streamed("_1") === build("_1"), joinType) + } else { + hinted.join(streamed, build("_1") === streamed("_1"), joinType) + } + def df: DataFrame = joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k", "v") + val plan = df.queryExecution.executedPlan + val joins = collect(plan) { case j: CometHashJoinExec => j } + assert(joins.map(_.buildSide) == Seq(buildSide), plan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert(collect(plan) { case c: CometSparkToColumnarExec => c }.size == 1, plan) + assert(collect(plan) { case s @ (_: SortExec | _: CometSortExec) => s }.isEmpty, plan) + sameRowsAsSpark(df) + assert( + partitionsInOrder(df, keyOrder(Ascending), valueOrder), + s"$joinType with hint $hint is out of order\n$plan") + } + } finally { + streamed.unpersist() + } + } + } + } + + // Spark drops the final sort for the checkpoint's order, which Comet cannot check. + test("forceShuffledHashJoin keeps the order of a sorted checkpoint through an outer join") { + withOrderedCheckpoint { streamed => + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> "true") { + val build = spark.range(10).selectExpr("id * 10 AS k") + val order = + (0 to 1).map(i => SortOrder(BoundReference(i, LongType, nullable = true), Ascending)) + for { + (joinType, buildSide) <- Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft) + } { + def df: DataFrame = { + val on = streamed("k") === build("k") + val joined = if (joinType == "left_outer") { + streamed.join(build, on, joinType) + } else { + build.join(streamed, on, joinType) + } + joined.select(streamed("k"), streamed("v")).sortWithinPartitions("k", "v") + } + val plan = df.queryExecution.executedPlan + val context = s"$joinType\n$plan" + assert(partitionsInOrder(df, order: _*), context) + checkSparkAnswerAndFallbackReason(df, CometHashJoin.uncheckedStreamedOrderReason) + val joins = collect(plan) { case j: ShuffledHashJoinExec => j.buildSide } + assert(joins == Seq(buildSide), context) + assert(collect(plan) { case j: CometHashJoinExec => j }.isEmpty, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + } + } + } + } + + test("forceShuffledHashJoin respects join strategy hints") { + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withForcedSmjTables { + def hinted(hint: String): DataFrame = + sql(s"SELECT /*+ $hint */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1") + val (_, merge) = checkSparkAnswerAndOperator(hinted("MERGE(tbl_a)")) + assert(collect(merge) { case j: CometSortMergeJoinExec => j }.size == 1, merge) + assert(collect(merge) { case j: CometHashJoinExec => j }.isEmpty, merge) + // BuildLeft proves Spark planned it: the forced conversion would build the smaller tbl_b. + val (_, shuffleHash) = checkSparkAnswer(hinted("SHUFFLE_HASH(tbl_a)")) + assert( + collect(shuffleHash) { case j: CometHashJoinExec => j.buildSide } == Seq(BuildLeft), + shuffleHash) + val (_, broadcast) = checkSparkAnswer(hinted("BROADCAST(tbl_b)")) + assert(collect(broadcast) { case j: CometBroadcastHashJoinExec => j }.nonEmpty, broadcast) + } + } + } + + test("forceShuffledHashJoin ignores AQE's NO_BROADCAST_HASH hint") { + var hints = Seq.empty[JoinHint] + val recordJoinHints = new SparkStrategy { + override def apply(plan: LogicalPlan): Seq[SparkPlan] = { + plan match { + case j: Join if CometConf.COMET_ENABLED.get(SQLConf.get) => hints :+= j.hint + case _ => + } + Nil + } + } + withForcedShjTables(adaptive = true) { + // oneKey fills 1 of 10 shuffle partitions, so AQE adds NO_BROADCAST_HASH on re-plan. + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "10") { + withParquetTable((0 until 100).map(i => (10, i)), "oneKey") { + val strategies = spark.experimental.extraStrategies + spark.experimental.extraStrategies = strategies :+ recordJoinHints + try { + val (_, cometPlan) = checkSparkAnswer( + "SELECT big._2, oneKey._2 FROM big JOIN oneKey ON big._1 = oneKey._1") + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) + val sortMergeJoins = collect(cometPlan) { + case j: SortMergeJoinExec => j + case j: CometSortMergeJoinExec => j + } + assert(sortMergeJoins.isEmpty, cometPlan) + val strategyHints = + hints.flatMap(h => Seq(h.leftHint, h.rightHint)).flatten.flatMap(_.strategy) + assert(strategyHints.contains(NO_BROADCAST_HASH), hints) + } finally { + spark.experimental.extraStrategies = strategies + } + } + } + } + } + + test("forceShuffledHashJoin builds on the side AQE's PREFER_SHUFFLE_HASH hint names") { + // Every shuffle partition of the right side fits a local hash map, so AQE hints that side on + // re-plan. The skewed left side is smaller in all, so the strategy would build on it otherwise. + var hints = Seq.empty[JoinHint] + val recordJoinHints = new SparkStrategy { + override def apply(plan: LogicalPlan): Seq[SparkPlan] = { + plan match { + case j: Join if CometConf.COMET_ENABLED.get(SQLConf.get) => hints :+= j.hint + case _ => + } + Nil + } + } + def isPreferred(hint: Option[HintInfo]): Boolean = + hint.exists(_.strategy.contains(PREFER_SHUFFLE_HASH)) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SKEW_JOIN_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "16", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "131072", + SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "1572864", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true", + CometConf.COMET_FORCE_SHJ.key -> "true") { + val left = spark + .range(90000) + .selectExpr("IF(id < 60000, 0L, id) AS lk", "sha2(cast(id AS string), 256) AS ls") + val right = + spark.range(200000).selectExpr("id AS rk", "sha2(cast(id AS string), 256) AS rs") + val strategies = spark.experimental.extraStrategies + spark.experimental.extraStrategies = strategies :+ recordJoinHints + try { + val (_, cometPlan) = checkSparkAnswer( + left + .join(right, $"lk" === $"rk") + .selectExpr("sum(length(ls))", "sum(length(rs))")) + assert(hints.exists(h => isPreferred(h.rightHint) && !isPreferred(h.leftHint)), hints) + assert( + collect(cometPlan) { case j: CometHashJoinExec => j.buildSide } == Seq(BuildRight), + cometPlan) + } finally { + spark.experimental.extraStrategies = strategies + } + } + } + + test("forceShuffledHashJoin builds on the one buildable side PREFER_SHUFFLE_HASH names") { + // tbl_b is the smaller table, so with no hint, or both sides hinted, the strategy builds on + // it. A LEFT ANTI join cannot build on its left side, so a hint there changes nothing. + val prefer = Some(HintInfo(strategy = Some(PREFER_SHUFFLE_HASH))) + val left = JoinHint(prefer, None) + val right = JoinHint(None, prefer) + val both = JoinHint(prefer, prefer) + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withForcedSmjTables { + def buildSide(query: String, hint: JoinHint): Seq[BuildSide] = { + val join = sql(query).queryExecution.optimizedPlan.collectFirst { case j: Join => j } + assert(join.isDefined, query) + CometShuffledHashJoinStrategy(spark)(join.get.copy(hint = hint)).collect { + case j: ShuffledHashJoinExec => j.buildSide + } + } + val bigFirst = "SELECT * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1" + val smallFirst = "SELECT * FROM tbl_b JOIN tbl_a ON tbl_a._2 = tbl_b._1" + val anti = "SELECT * FROM tbl_a LEFT ANTI JOIN tbl_b ON tbl_a._2 = tbl_b._1" + for ((query, hint, expected) <- Seq( + (bigFirst, JoinHint.NONE, BuildRight), + (bigFirst, left, BuildLeft), + (bigFirst, both, BuildRight), + (smallFirst, JoinHint.NONE, BuildLeft), + (smallFirst, right, BuildRight), + (smallFirst, both, BuildLeft), + (anti, left, BuildRight))) { + assert(buildSide(query, hint) == Seq(expected), s"$query $hint") + } + } + } + } + + test("forceShuffledHashJoin leaves a runtime broadcast to AQE") { + withForcedShjTables(adaptive = true) { + withSQLConf(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { + val (_, cometPlan) = checkSparkAnswer(sql(sortLossJoin)) + val broadcasts = collect(cometPlan) { + case j: CometBroadcastHashJoinExec => j + case j: BroadcastHashJoinExec => j + } + assert(broadcasts.nonEmpty, cometPlan) + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.isEmpty, cometPlan) + assert(collect(cometPlan) { case j: ShuffledHashJoinExec => j }.isEmpty, cometPlan) + } + } + } + + test("forceShuffledHashJoin handles a skewed join under AQE") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "10", + SQLConf.SKEW_JOIN_ENABLED.key -> "true", + SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100", + CometConf.COMET_FORCE_SHJ.key -> "true") { + withParquetTable((0 until 10000).map(i => (if (i < 9000) 0 else i % 100, i)), "skewed") { + withParquetTable((0 until 100).map(i => (i, i)), "keys") { + val (_, cometPlan) = checkSparkAnswer( + "SELECT skewed._2, keys._2 FROM skewed JOIN keys ON skewed._1 = keys._1") + val hashJoins = collect(cometPlan) { + case j: ShuffledHashJoinExec => j + case j: CometHashJoinExec => j.originalPlan + } + val skewJoins = hashJoins.collect { case j: ShuffledHashJoinExec if j.isSkewJoin => j } + assert(skewJoins.nonEmpty, cometPlan) + } + } + } + } + + // A computed key stops Spark from rewriting the join to compare collation keys. + test("forceShuffledHashJoin keeps a SortMergeJoin on collated keys") { + assume(isSpark40Plus) + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withTable("c1", "c2") { + sql("CREATE TABLE c1 (s STRING COLLATE UTF8_LCASE, v INT) USING parquet") + sql("INSERT INTO c1 VALUES ('a', 1), ('A', 2), ('b', 3)") + sql("CREATE TABLE c2 (s STRING COLLATE UTF8_LCASE, w INT) USING parquet") + sql("INSERT INTO c2 VALUES ('a', 10), ('B', 20)") + val (_, cometPlan) = + checkSparkAnswer("SELECT v, w FROM c1 JOIN c2 ON trim(c1.s) = c2.s") + assert(collect(cometPlan) { case j: ShuffledHashJoinExec => j }.isEmpty, cometPlan) + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.isEmpty, cometPlan) + } + } + } + + // The condition references both sides, so it stays on the join, which make_date keeps in Spark. + for (adaptive <- Seq(false, true)) { + test(s"forceShuffledHashJoin hash joins a make_date condition in Spark, AQE=$adaptive") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", CometConf.COMET_FORCE_SHJ.key -> "true", CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { - withChainedJoinTables(midRows = 3000, midKeys = 100) { - val (_, cometPlan) = checkSparkAnswer( - "SELECT big._1, big._2, small._2, count(*) OVER (PARTITION BY big._1) " + - "FROM big JOIN small ON big._1 = small._1") - assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) - assert(collect(cometPlan) { case w: CometWindowExec => w }.nonEmpty, cometPlan) + withParquetTable(Seq((1, 2024), (2, 2025), (3, 2024), (4, 2023)), "date_years") { + withParquetTable(Seq((1, 1), (2, 6), (3, 12), (5, 3)), "date_months") { + val (_, cometPlan) = checkSparkAnswer( + "SELECT y._1, y._2, m._2 FROM date_years y JOIN date_months m " + + "ON y._1 = m._1 AND make_date(y._2, m._2, 1) >= DATE '2024-06-01'") + assert(collect(cometPlan) { case j: ShuffledHashJoinExec => j }.size == 1, cometPlan) + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.isEmpty, cometPlan) + assert( + collect(cometPlan) { + case j: SortMergeJoinExec => j + case j: CometSortMergeJoinExec => j + }.isEmpty, + cometPlan) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(cometPlan) + .exists(_.contains("make_date requires Spark evaluation in ShuffledHashJoin")), + cometPlan) + } + } + } + } + } + + // AQE discards a re-plan that hash joins the broadcast stage, so only its warning shows one. + test("forceShuffledHashJoin leaves a planned broadcast stage to AQE") { + val appender = new LogAppender("AQE re-plans") + val logger = classOf[AdaptiveSparkPlanExec].getName + val context = LogManager.getContext(false).asInstanceOf[LoggerContext] + val isUnconfigured = !context.getConfiguration.getLoggers.containsKey(logger) + val threshold = 10L * 1024 + withForcedShjTables(adaptive = true) { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> threshold.toString, + SQLConf.ADAPTIVE_EXECUTION_LOG_LEVEL.key -> "WARN") { + val bigSize = spark.table("big").queryExecution.optimizedPlan.stats.sizeInBytes + val smallSize = spark.table("small").queryExecution.optimizedPlan.stats.sizeInBytes + assert(bigSize > threshold, s"big is $bigSize bytes, not over $threshold") + assert(smallSize <= threshold, s"small is $smallSize bytes, over $threshold") + try { + withLogAppender(appender, Seq(logger), Some(Level.WARN)) { + val (_, cometPlan) = checkSparkAnswer(sql(sortLossJoin)) + val broadcasts = collect(cometPlan) { + case j: CometBroadcastHashJoinExec => j + case j: BroadcastHashJoinExec => j + } + assert(broadcasts.nonEmpty, cometPlan) + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.isEmpty, cometPlan) + assert(collect(cometPlan) { case j: ShuffledHashJoinExec => j }.isEmpty, cometPlan) + } + } finally { + // withLogAppender never removes the config it adds for an unconfigured logger. + if (isUnconfigured) { + context.getConfiguration.removeLogger(logger) + context.updateLoggers() + } + } + } + } + val messages = appender.loggingEvents.map(_.getMessage.getFormattedMessage) + assert(messages.exists(_.contains("Final plan")), messages.mkString("\n")) + val invalidStages = messages.filter(_.contains("Invalid broadcast query stage")) + assert(invalidStages.isEmpty, invalidStages.mkString("\n")) + } + + for ((key, value) <- Seq( + CometConf.COMET_ENABLED.key -> "false", + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true")) { + test(s"forceShuffledHashJoin does nothing with $key=$value") { + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + key -> value, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withForcedSmjTables { + val df = sql(forcedSmjQuery) + df.collect() + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case j: ShuffledHashJoinExec => j }.isEmpty, plan) + assert(collect(plan) { case j: SortMergeJoinExec => j }.nonEmpty, plan) } } } @@ -2495,6 +2985,46 @@ class CometJoinSuite extends CometTestBase { } } + test("BroadcastNestedLoopJoin over an ordered streamed side falls back and keeps its order") { + // Spark's nested loop join keeps the native range's order, and Spark drops the sort for it, + // either directly above the join or above a hash join over it. The native nested loop join + // emits its rows build row by build row, so it stays in Spark. + val nestedLoop = + "SELECT /*+ BROADCAST(n) */ p.id AS x, n.id AS z FROM range(0, 20, 1, 1) p " + + "JOIN range(0, 2, 1, 1) n ON p.id % 2 <= n.id" + def overHashJoin(hint: String): String = + s"SELECT $hint j.x, j.z, b.id AS w FROM ($nestedLoop) j JOIN range(0, 2, 1, 1) b " + + "ON j.x % 4 = b.id * 2 ORDER BY x" + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true") { + for ((query, force, hashJoins) <- Seq( + (s"SELECT * FROM ($nestedLoop) j ORDER BY x", false, 0), + (overHashJoin(""), true, 1), + (overHashJoin("/*+ SHUFFLE_HASH(b) */"), false, 1), + (overHashJoin("/*+ BROADCAST(b) */"), false, 1))) { + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> force.toString) { + val (_, plan) = checkSparkAnswerAndFallbackReason( + sql(query), + CometBroadcastNestedLoopJoinExec.streamedOrderReason) + val context = s"$query\n$plan" + assert(collect(plan) { case j: BroadcastNestedLoopJoinExec => j }.size == 1, context) + assert(collect(plan) { case j: CometBroadcastNestedLoopJoinExec => j }.isEmpty, context) + // The hash join over the nested loop join stays in Spark as well. + val joins = collect(plan) { + case j: ShuffledHashJoinExec => j + case j: BroadcastHashJoinExec => j + } + assert(joins.size == hashJoins, context) + val xs = sql(query).collect().map(_.getLong(0)).toSeq + assert(xs == xs.sorted, s"$xs\n$context") + } + } + } + } + test("ExistenceJoin via BroadcastHashJoin (EXISTS combined with OR)") { withSQLConf( CometConf.COMET_EXEC_EXISTENCE_JOIN_ENABLED.key -> "true", @@ -2650,10 +3180,7 @@ class CometJoinSuite extends CometTestBase { } test("ExistenceJoin via SortMergeJoin stays on Spark even with forceShuffledHashJoin") { - // RewriteJoin refuses to rewrite an existence SortMergeJoin into a BuildRight ShuffledHashJoin - // (https://github.com/apache/datafusion-comet/issues/2697), so even with forced SHJ the join - // must stay on Spark rather than run as a native hash join. Assert on the plan (no native join - // form) rather than the fallback reason, which is the robust signal for this rewrite guard. + // Never a BuildRight hash join: https://github.com/apache/datafusion-comet/issues/2697 withSQLConf( CometConf.COMET_EXEC_EXISTENCE_JOIN_ENABLED.key -> "true", CometConf.COMET_FORCE_SHJ.key -> "true", @@ -2873,6 +3400,85 @@ class CometJoinSuite extends CometTestBase { } } + test("an inner hash join declares a sorted cache's order for a native window above it") { + // Spark drops the window's sort, as the join reports the cache's order, so the native window + // needs that order declared on the join's probe input, which reads the cache from the JVM. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + val cached = spark + .range(2000) + .selectExpr("id % 20 AS k", "id AS v") + .repartition(2, $"k") + .sortWithinPartitions("k", "v") + .cache() + try { + cached.count() + val keys = spark.range(20).toDF("id") + for ((build, force) <- Seq( + keys -> true, + keys.hint("shuffle_hash") -> false, + keys.hint("broadcast") -> false)) { + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> force.toString) { + def df: DataFrame = cached + .join(build, $"k" === $"id") + .selectExpr("k", "v", "dense_rank() OVER (PARTITION BY k ORDER BY v) AS r") + sameRowsAsSpark(df) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case w: CometWindowExec => w }.size == 1, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + val declared = collect(plan) { + case j: CometHashJoinExec => j.nativeOp.getHashJoin.getStreamedSortOrdersCount + case j: CometBroadcastHashJoinExec => + j.nativeOp.getHashJoin.getStreamedSortOrdersCount + } + assert(declared == Seq(2), plan) + } + } + } finally { + cached.unpersist() + } + } + } + + test("an inner hash join declares the part of the streamed order Comet can express") { + // An inner join keeps its probe rows in place without a declared order, so one whose first + // key Comet cannot express declares nothing and still runs natively. With Add disabled, Comet + // cannot express `_1 + 1`. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.getExprEnabledConfigKey("Add") -> "false") { + withParquetTable((0 until 1000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + val build = spark.table("small") + try { + for ((order, declared) <- Seq(Seq($"_1" + 1) -> 0, Seq($"_1", $"_1" + 1) -> 1)) { + val cached = + spark.table("big").repartition(2, $"_1").sortWithinPartitions(order: _*).cache() + cached.count() + val (_, plan) = checkSparkAnswer( + cached.join(build.hint("shuffle_hash"), cached("_1") === build("_1"))) + val joins = collect(plan) { case j: CometHashJoinExec => j } + assert(joins.size == 1, plan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount == declared, plan) + cached.unpersist() + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + test("outer hash joins keep a descending, nulls last streamed order with null keys") { // Null keys match nothing, so all of them take the unmatched path, at the end of each // partition. @@ -2938,6 +3544,15 @@ class CometJoinSuite extends CometTestBase { .localCheckpoint(eager = true)) } + /** + * Build sides for an inner join with a checkpoint, each with whether to force a shuffled hash + * join: unhinted, then hinted for a shuffled hash join, then for a broadcast. + */ + private def checkpointJoinBuildSides: Seq[(DataFrame, Boolean)] = { + val keys = spark.range(100).toDF("id") + Seq(keys -> true, keys.hint("shuffle_hash") -> false, keys.hint("broadcast") -> false) + } + test("outer hash joins over a checkpointed streamed side fall back and keep its order") { // Spark drops the sort above the join for the checkpoint's order, which Comet cannot check. withOrderedCheckpoint { streamed => @@ -3005,6 +3620,126 @@ class CometJoinSuite extends CometTestBase { } } + test("an inner hash join over a checkpointed streamed side declares no order") { + // Comet cannot check the checkpoint's order, so it declares none, and the inner join keeps its + // probe rows in place without one, so it stays native. + withOrderedCheckpoint { streamed => + val build = spark.range(10).selectExpr("id * 10 AS k", "id AS w") + val order = + (0 to 1).map(i => SortOrder(BoundReference(i, LongType, nullable = true), Ascending)) + for (hint <- Seq("shuffle_hash", "broadcast")) { + def df: DataFrame = streamed + .join(build.hint(hint), streamed("k") === build("k")) + .select(streamed("k"), streamed("v"), build("w")) + .sortWithinPartitions("k", "v") + val (_, plan) = checkSparkAnswer(df) + val declared = collect(plan) { + case j: CometHashJoinExec => j.nativeOp.getHashJoin.getStreamedSortOrdersCount + case j: CometBroadcastHashJoinExec => j.nativeOp.getHashJoin.getStreamedSortOrdersCount + } + assert(declared == Seq(0), s"$hint\n$plan") + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + s"$hint\n$plan") + assert(partitionsInOrder(df, order: _*), s"$hint\n$plan") + } + } + } + + test("a native window above an inner hash join over a checkpointed streamed side runs") { + // Spark drops the window's sort for the checkpoint's order, which the join reports, but Comet + // cannot check that order, so the join declares none and the native window sorts its input. + withOrderedCheckpoint { streamed => + for { + (build, force) <- checkpointJoinBuildSides + window <- Seq( + "dense_rank() OVER (PARTITION BY k ORDER BY v)", + "sum(v) OVER (PARTITION BY k)") + } { + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> force.toString) { + def df: DataFrame = streamed + .join(build, $"k" === $"id") + .selectExpr("k", "v", s"$window AS r") + val plan = df.queryExecution.executedPlan + val context = s"force=$force $window\n$plan" + val joins = collect(plan) { + case j: CometHashJoinExec => j + case j: CometBroadcastHashJoinExec => j + } + assert(joins.size == 1, context) + val windows = collect(plan) { case w: CometWindowExec => w } + assert( + windows.map(_.nativeOp.getWindow.getUnverifiedInputSortOrdersCount) == Seq(2), + context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + sameRowsAsSpark(df) + } + } + } + } + + test("a native window group limit above an inner hash join over a checkpoint ranks it") { + assume(isSpark35Plus, "WindowGroupLimit was added in Spark 3.5") + // The native limit declares the order Spark requires of its input without checking it, which + // holds for the checkpoint's rows. + withOrderedCheckpoint { streamed => + for ((build, force) <- checkpointJoinBuildSides) { + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> force.toString) { + def df: DataFrame = streamed + .join(build, $"k" === $"id") + .selectExpr("k", "v", "rank() OVER (PARTITION BY k ORDER BY v) AS rk") + .where("rk <= 2") + val plan = df.queryExecution.executedPlan + val context = s"force=$force\n$plan" + assert(collect(plan) { case l: CometWindowGroupLimitExec => l }.nonEmpty, context) + assert(collect(plan) { case w: CometWindowExec => w }.size == 1, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + sameRowsAsSpark(df) + } + } + } + } + + test("a native window over an outer hash join that may move rows restores the whole order") { + // The left outer join runs as DataFusion's right outer join over a native window, which + // declares no order on the join's probe input, so the order the join reports is doubted. The + // window above it partitions by k alone and the sort above it, which Spark drops, needs v too. + withParquetTable((0 until 2000).map(i => (i % 40, i)), "t") { + withParquetTable((0 until 40 by 2).map(k => (k, k * 10)), "d1") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + def df: DataFrame = sql("""SELECT k, v, rn, sum(x) OVER (PARTITION BY k) AS s + |FROM (SELECT w.k, w.v, w.rn, d1._2 AS x + | FROM (SELECT k, v, row_number() OVER (PARTITION BY k ORDER BY v) AS rn + | FROM (SELECT _1 AS k, _2 AS v FROM t)) w + | LEFT JOIN d1 ON w.k = d1._1)""".stripMargin) + .sortWithinPartitions("k", "v") + val (_, plan) = checkSparkAnswer(df) + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j } + assert(joins.size == 1 && joins.head.joinType == LeftOuter, plan) + val upper = collect(plan) { case w: CometWindowExec if w.orderSpec.isEmpty => w } + assert( + upper.map(_.nativeOp.getWindow.getUnverifiedInputSortOrdersCount) == Seq(2), + plan) + assert( + collect(plan) { + case s @ (_: CometSortExec | _: SortExec) + if s.find(_.isInstanceOf[CometBroadcastHashJoinExec]).isDefined => + s + }.isEmpty, + plan) + val order = (0 to 1).map { i => + SortOrder(BoundReference(i, IntegerType, nullable = true), Ascending) + } + assert(partitionsInOrder(df, order: _*), plan) + } + } + } + } + test("chained outer hash joins over a native window keep its order natively") { // The window and both joins share one native plan, where DataFusion sees the order itself. withParquetTable((0 until 2000).map(i => (i % 40, i)), "t") { @@ -3542,36 +4277,27 @@ class CometJoinSuite extends CometTestBase { } } - test("an outer hash join over a native nested loop join falls back") { - // Over two build rows the nested loop join emits rows out of order (#6898), so they are - // compared as a multiset. Over one build row they stay in order, and Spark's join keeps it. + test("an outer hash join over a nested loop join over a sorted cache keeps its order") { + // The nested loop join stays in Spark, which keeps the cache's order over any number of build + // rows, and so does the outer join over it, so the order holds without the sort Spark drops. withSortedCache() { withParquetTable(Seq(Tuple1(0), Tuple1(1)), "two") { - withParquetTable(Seq(Tuple1(1)), "one") { - def query(nlBuild: String): String = - s"""SELECT /*+ BROADCAST(n), BROADCAST(b) */ a.k, a.v, n._1 AS z, b._2 AS w - |FROM cached a JOIN $nlBuild n ON a.v % 2 <= n._1 - |LEFT JOIN small b ON a.k = b._1""".stripMargin - val reason = CometHashJoin.uncheckedStreamedOrderReason - def checkFallback(df: => DataFrame): SparkPlan = { - sameRowsAsSpark(df) - val plan = df.queryExecution.executedPlan - val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) - assert(reasons.exists(_.contains(reason)), plan) - val nestedLoopJoins = collect(plan) { case j: CometBroadcastNestedLoopJoinExec => j } - assert(nestedLoopJoins.size == 1, plan) - val outer = collect(plan) { case j: BroadcastHashJoinExec => j.joinType } - assert(outer == Seq(LeftOuter), plan) - plan - } - - checkFallback(sql(query("two"))) - - def ordered: DataFrame = sql(query("one")).sortWithinPartitions("k") - val plan = checkFallback(ordered) - assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) - assert(partitionsInOrder(ordered, keyOrder(Ascending)), plan) - } + def df: DataFrame = + sql("""SELECT /*+ BROADCAST(n), BROADCAST(b) */ a.k, a.v, n._1 AS z, b._2 AS w + |FROM cached a JOIN two n ON a.v % 2 <= n._1 + |LEFT JOIN small b ON a.k = b._1""".stripMargin) + .sortWithinPartitions("k") + sameRowsAsSpark(df) + val plan = df.queryExecution.executedPlan + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert( + reasons.exists(_.contains(CometBroadcastNestedLoopJoinExec.streamedOrderReason)), + plan) + assert(collect(plan) { case j: BroadcastNestedLoopJoinExec => j }.size == 1, plan) + val outer = collect(plan) { case j: BroadcastHashJoinExec => j.joinType } + assert(outer == Seq(LeftOuter), plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 5d09646d61e..3716f2aad5d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -20,23 +20,25 @@ package org.apache.comet.exec import scala.util.Random +import scala.util.control.NonFatal import org.scalactic.source.Position import org.scalatest.Tag import org.apache.hadoop.fs.Path -import org.apache.spark.sql.{CometTestBase, DataFrame, Row} -import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Divide, Expression, MakeDecimal, WindowExpression} +import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, BoundReference, Cast, Descending, Divide, Expression, MakeDecimal, NullsFirst, NullsLast, SortOrder, WindowExpression} import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeExec, CometNativeScanExec, CometSortExec, CometWindowExec, CometWindowGroupLimitExec} -import org.apache.spark.sql.execution.{InputAdapter, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{InputAdapter, SortExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.window.{WindowExec => SparkWindowExec} import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions.{count, expr, lead, sum} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.DecimalType +import org.apache.spark.sql.types.{DecimalType, DoubleType, LongType} import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} +import org.apache.comet.serde.OperatorOuterClass class CometWindowExecSuite extends CometTestBase { @@ -117,6 +119,18 @@ class CometWindowExecSuite extends CometTestBase { digits.toString() } + // Rows of (k, v, w), w unique, in two partitions by k, each sorted by k, `v` and w. + private def sortedCheckpoint(rows: Long, kMod: Int, vMod: Int, v: Column): DataFrame = + spark + .range(rows) + .selectExpr( + s"id % $kMod AS k", + s"CASE WHEN id % 7 = 0 THEN NULL ELSE id % $vMod END AS v", + "id AS w") + .repartition(2, $"k") + .sortWithinPartitions($"k", v, $"w") + .localCheckpoint(eager = true) + for { orderColumn <- Seq("f", "d") // Scalar floating-point order keys are admitted under strict floating point (#5506), so the @@ -543,6 +557,272 @@ class CometWindowExecSuite extends CometTestBase { } } + test("a native window over a sorted checkpoint runs without Spark's sort") { + // Spark drops the window's sort for the checkpoint's order, but those rows enter the native + // plan with no declared order, so the native window sorts them itself. + // AQE is off, as an adaptive plan reports no order. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { + val checkpoint = spark + .range(2000) + .selectExpr("id % 20 AS k", "id AS v") + .repartition(2, $"k") + .sortWithinPartitions("k", "v") + .localCheckpoint(eager = true) + for (window <- Seq( + "dense_rank() OVER (PARTITION BY k ORDER BY v)", + "sum(v) OVER (PARTITION BY k)")) { + val (_, plan) = checkSparkAnswer(checkpoint.selectExpr("k", "v", s"$window AS r")) + val context = s"$window\n$plan" + val windows = collect(plan) { case w: CometWindowExec => w } + assert(windows.size == 1, context) + assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 0, context) + assert(windows.head.nativeOp.getWindow.getUnverifiedInputSortOrdersCount == 2, context) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, context) + } + } + } + + test("stacked native windows over a sorted checkpoint order each by its own ORDER BY") { + // Spark plans no sort for either window, as the checkpoint reports (k, v, w). The native + // planner sorts the lower window's input, and the upper window relies on that sort. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { + val byV = "dense_rank() OVER (PARTITION BY k ORDER BY v) AS r1" + val byVW = "row_number() OVER (PARTITION BY k ORDER BY v, w) AS r2" + val failures = scala.collection.mutable.ArrayBuffer.empty[String] + for { + (rows, kMod, vMod) <- Seq((2000L, 20, 50), (200000L, 4, 3)) + windows <- Seq(Seq(byV, byVW), Seq(byVW, byV)) + } { + val checkpoint = sortedCheckpoint(rows, kMod, vMod, $"v") + val context = s"rows=$rows k%$kMod v%$vMod windows=${windows.mkString(", ")}" + val df = checkpoint.selectExpr(Seq("k", "v", "w") ++ windows: _*) + val plan = df.queryExecution.executedPlan + val nativeWindows = collect(plan) { case w: CometWindowExec => w } + assert( + nativeWindows.map(_.nativeOp.getWindow.getUnverifiedInputSortOrdersCount) == Seq(3, 3), + s"$context\n$plan") + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + s"$context\n$plan") + try { + checkSparkAnswer(df) + } catch { + case NonFatal(e) => + val lines = String.valueOf(e.getMessage).linesIterator.toSeq + val diff = + lines.filter(l => l.startsWith("!") || l.contains("Expects PARTITION BY")) + val shown = if (diff.nonEmpty) diff else lines + failures += s"$context: ${e.getClass.getName}, ${diff.size} matching lines\n" + + shown.take(12).mkString("\n") + } + } + assert(failures.isEmpty, failures.mkString("\n\n")) + } + } + + test("a native window over a sorted checkpoint keeps the order Spark dropped a sort for") { + // The window reports the checkpoint's order, so Spark plans no sort for the one above it, and + // the native window has to restore all of that order, not only its own keys. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { + for { + (rows, kMod, vMod) <- Seq((2000L, 20, 50), (200000L, 4, 3)) + (v, vOrder, direction, nulls) <- Seq( + ($"v", "v", Ascending, NullsFirst), + ($"v".desc_nulls_last, "v DESC NULLS LAST", Descending, NullsLast)) + } { + val checkpoint = sortedCheckpoint(rows, kMod, vMod, v) + val window = s"dense_rank() OVER (PARTITION BY k ORDER BY $vOrder) AS r" + val df = checkpoint.selectExpr("k", "v", "w", window).sortWithinPartitions($"k", v, $"w") + val (_, plan) = checkSparkAnswer(df) + val context = s"rows=$rows $vOrder\n$plan" + val windows = collect(plan) { case w: CometWindowExec => w } + assert( + windows.map(_.nativeOp.getWindow.getUnverifiedInputSortOrdersCount) == Seq(3), + context) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, context) + val column = (i: Int) => BoundReference(i, LongType, nullable = true) + val order = Seq( + SortOrder(column(0), Ascending), + SortOrder(column(1), direction, nulls, Seq.empty), + SortOrder(column(2), Ascending)) + assert(partitionsInOrder(df, order: _*), context) + } + } + } + + test("a native window keeps a checkpoint's order on a float key holding -0.0 and 0.0") { + // Spark sorts -0.0 and 0.0 as equal, then by w. So does the native sort, which normalizes its + // keys, so a float key before the last one stays native. f is -0.0 for a third of the rows + // and 0.0 for the rest. Spark writes the checkpoint. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true", + CometConf.COMET_ENABLED.key -> "false") { + val checkpoint = spark + .range(2000) + .selectExpr("id % 4 AS k", "(id % 3 - 1) * 0.0D AS f", "id AS w") + .repartition(2, $"k") + .sortWithinPartitions("k", "f", "w") + .localCheckpoint(eager = true) + withSQLConf(CometConf.COMET_ENABLED.key -> "true") { + val df = checkpoint + .selectExpr("k", "f", "w", "dense_rank() OVER (PARTITION BY k ORDER BY f) AS r") + .sortWithinPartitions("k", "f", "w") + val (_, plan) = checkSparkAnswer(df) + val context = s"${checkpoint.queryExecution.executedPlan.outputOrdering}\n$plan" + val windows = collect(plan) { case w: CometWindowExec => w } + assert( + windows.map(_.nativeOp.getWindow.getUnverifiedInputSortOrdersCount) == Seq(3), + context) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, context) + val order = Seq(LongType, DoubleType, LongType).zipWithIndex.map { case (t, i) => + SortOrder(BoundReference(i, t, nullable = true), Ascending) + } + assert(partitionsInOrder(df, order: _*), context) + } + } + } + + test("a native window sorts by its own key where the checkpoint reports an equal one") { + // The checkpoint reports the order of kk also for k, which kk is an alias of, and that of lk + // also for rk, the key it was joined on. A window by k or rk finds its partitions by that + // column, so its input has to be sorted by it, not by the twin. Each checkpoint is one + // partition, so the windows need no exchange. Spark writes the checkpoints. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true", + CometConf.COMET_ENABLED.key -> "false") { + val sorted = spark + .range(2000) + .selectExpr( + "id % 20 AS k", + "CASE WHEN id % 7 = 0 THEN NULL ELSE id % 50 END AS v", + "id AS w") + .coalesce(1) + .sortWithinPartitions("k", "v", "w") + val aliased = sorted.selectExpr("k", "k AS kk", "v", "w").localCheckpoint(eager = true) + val right = spark + .range(20) + .selectExpr("id AS rk", "id * 3 AS x") + .coalesce(1) + .sortWithinPartitions("rk") + val joined = sorted + .selectExpr("k AS lk", "v", "w") + .join(right, $"lk" === $"rk") + .localCheckpoint(eager = true) + val cases = Seq( + ( + aliased, + "kk", + "k", + Seq( + "percent_rank() OVER (PARTITION BY k ORDER BY v)", + "cume_dist() OVER (PARTITION BY k ORDER BY v)", + "ntile(3) OVER (PARTITION BY k ORDER BY v)", + "row_number() OVER (PARTITION BY k ORDER BY v)"), + 3), + ( + joined, + "lk", + "rk", + Seq("count(*) OVER (PARTITION BY rk)", "sum(v) OVER (PARTITION BY rk)"), + 1)) + withSQLConf(CometConf.COMET_ENABLED.key -> "true") { + for ((checkpoint, twin, key, windows, claimed) <- cases; window <- windows) { + val reported = checkpoint.queryExecution.optimizedPlan.outputOrdering + val names = (es: Seq[Expression]) => es.flatMap(_.references.map(_.name)) + val context = s"$window\n" + reported + .map(o => s"${o.sql} also ${names(o.sameOrderExpressions).mkString(", ")}") + .mkString("; ") + assert(names(Seq(reported.head.child)) == Seq(twin), context) + assert(names(reported.head.sameOrderExpressions) == Seq(key), context) + val (_, plan) = checkSparkAnswer(checkpoint.selectExpr("*", s"$window AS r")) + val nativeWindows = collect(plan) { case w: CometWindowExec => w } + assert( + nativeWindows.map(_.nativeOp.getWindow.getUnverifiedInputSortOrdersCount) == + Seq(claimed), + s"$context\n$plan") + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + s"$context\n$plan") + } + } + } + } + + test("a window falls back when the order Spark relies on cannot be fully expressed") { + // The checkpoint's order ends in a key Comet cannot express, so the native window could not + // restore it. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true", + CometConf.getExprEnabledConfigKey("Abs") -> "false") { + val checkpoint = spark + .range(2000) + .selectExpr("id % 20 AS k", "id % 50 AS v", "id AS w") + .repartition(2, $"k") + .sortWithinPartitions($"k", $"v", expr("abs(w)")) + .localCheckpoint(eager = true) + val window = "dense_rank() OVER (PARTITION BY k ORDER BY v) AS r" + val (_, plan) = checkSparkAnswerAndFallbackReason( + checkpoint.selectExpr("k", "v", "w", window), + "cannot be fully expressed natively") + assert(collect(plan) { case w: CometWindowExec => w }.isEmpty, plan) + } + } + + test("a window falls back when the order Spark relies on ends in a collated string") { + assume(isSpark40Plus, "COLLATE requires Spark 4.0") + // The native sort compares strings byte by byte, so it could not restore an order that + // UTF8_LCASE breaks ties in, even on a key the window does not use. No scan hands a collated + // column to a native plan, so the serde is called on the window Spark plans over the + // checkpoint. One with a plain string key instead stays native. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val reason = "the order of the window's input, which Spark relies on, cannot be fully " + + "expressed natively" + for (collation <- Seq("UTF8_BINARY", "UTF8_LCASE")) { + val checkpoint = spark + .range(2000) + .selectExpr( + "id % 20 AS k", + "id % 50 AS v", + s"concat(IF(id % 2 = 0, 'A', 'a'), id % 3) COLLATE $collation AS s") + .repartition(2, $"k") + .sortWithinPartitions("k", "v", "s") + .localCheckpoint(eager = true) + val df = checkpoint.selectExpr( + "k", + "v", + "s", + "dense_rank() OVER (PARTITION BY k ORDER BY v) AS r") + val plan = df.queryExecution.sparkPlan + val window = collect(plan) { case w: SparkWindowExec => w }.head + val context = s"$collation\n${window.child.outputOrdering.map(_.sql)}\n$plan" + assert(window.child.outputOrdering.size == 3, context) + val proto = CometWindowExec.convert( + window, + OperatorOuterClass.Operator.newBuilder(), + OperatorOuterClass.Operator.getDefaultInstance) + val reasons = window.getTagValue(CometExplainInfo.FALLBACK_REASONS).getOrElse(Set.empty) + if (collation == "UTF8_LCASE") { + assert(proto.isEmpty, context) + assert(reasons.contains(reason), s"$reasons\n$context") + } else { + assert(proto.map(_.getWindow.getUnverifiedInputSortOrdersCount).contains(3), context) + assert(reasons.isEmpty, s"$reasons\n$context") + } + } + } + } + test("Window range frame with long boundary should not fail") { val df = Seq((1L, "1"), (1L, "1"), (2147483650L, "1"), (3L, "2"), (2L, "1"), (2147483650L, "2")) diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index ffbe1d0e1b0..35972ba5495 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -413,12 +413,10 @@ class CometExecRuleSuite extends CometTestBase { } } - test("make_date join input demand reaches local sorts before hash join rewriting") { + test("make_date join input demand reaches local sorts") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", - CometConf.COMET_FORCE_SHJ.key -> "true", - CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { val left = createSparkPlan(spark, "SELECT * FROM VALUES (2024), (2025) AS t(y)") val right = createSparkPlan(spark, "SELECT * FROM VALUES (2024) AS t(y)") val date = MakeDate(left.output.head, Literal(1), Literal(1), failOnError = false)