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/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 8785874c6f6..25deaf3ab05 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -133,19 +133,26 @@ express falls back to Spark and records a fallback reason (see an operator above the join relies on the order. An outer join falls back to Spark when Comet cannot check the order Spark reports and that order -reaches the join through an input from the JVM or through a nested loop join. Spark keeps the rows -in the order they arrive. One such input is an RDD checkpoint. It reaches a native join only with +reaches the join through an input from the JVM. Spark keeps the rows in the order they arrive. One +such input is an RDD checkpoint. It reaches a native join only with `spark.comet.convert.rdd.enabled=true`, which defaults to `false`. Without that setting the checkpoint stays a Spark input and the join stays in Spark anyway. Another is a cache written by an outer hash join whose own probe input was not a scan -([#6922](https://github.com/apache/datafusion-comet/issues/6922)). The output of a nested loop -join can already be out of order when it reaches the join +([#6922](https://github.com/apache/datafusion-comet/issues/6922)). + +The native nested loop join falls back to Spark when Spark reports an output order for it, since +it emits its rows build row by build row ([#6898](https://github.com/apache/datafusion-comet/issues/6898)). When the order Comet cannot check comes only from another native hash join in the same plan, the join does not fall back. It runs natively and declares no order, since DataFusion keeps the order it can see there. +Hash joins of the other types that report an order, such as inner joins, keep their streamed rows +in place either way. They declare the order when Comet can check and express it, so a native +operator above them, a window for instance, sees it. Otherwise they declare nothing and still run +natively. + ## Window Functions Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of @@ -203,6 +210,10 @@ supported range of 0001-01-01 to 9999-12-31, such as one produced by `date_add` native frame can treat the bound as the edge of the partition where Spark compares against it ([#5022](https://github.com/apache/datafusion-comet/issues/5022)). +Spark drops the sort below a window when it believes the window's input is already ordered. For a window with +`PARTITION BY` keys, when Comet cannot verify that order, as with rows read from a checkpoint, the native window sorts +its input by the `PARTITION BY` keys and then the `ORDER BY` keys. The results match Spark, at the cost of a sort. + `WindowGroupLimitExec` (window-based limit pushdown for `ROW_NUMBER`, `RANK`, and `DENSE_RANK`) runs natively; it is controlled by `spark.comet.exec.windowGroupLimit.enabled` (default: true). diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 93d39973b0e..b0a9ff1756a 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -96,12 +96,12 @@ omitted from the tables below and may be reconsidered based on demand: ## Joins -| Operator | Status | Notes | -| ----------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `BroadcastHashJoinExec` | ✅ | Falls back to Spark for a `NOT IN` (null-aware anti join) whose streamed side reports an ordering. | -| `ShuffledHashJoinExec` | ✅ | | -| `SortMergeJoinExec` | ✅ | Supports `BINARY` join keys. Nested-type (struct, array, map) and collated-string join keys fall back to Spark. | -| `BroadcastNestedLoopJoinExec` | ✅ | Falls back to Spark when the preserved side is broadcast (for example LEFT OUTER with BROADCAST on the left) ([#4429](https://github.com/apache/datafusion-comet/pull/4429)). | +| Operator | Status | Notes | +| ----------------------------- | ------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `BroadcastHashJoinExec` | ✅ | Falls back to Spark for a `NOT IN` (null-aware anti join) whose streamed side reports an ordering. | +| `ShuffledHashJoinExec` | ✅ | | +| `SortMergeJoinExec` | ✅ | Supports `BINARY` join keys. Nested-type (struct, array, map) and collated-string join keys fall back to Spark. | +| `BroadcastNestedLoopJoinExec` | ✅ | Falls back to Spark when the preserved side is broadcast (for example LEFT OUTER with BROADCAST on the left) ([#4429](https://github.com/apache/datafusion-comet/pull/4429)), and when Spark reports an output ordering for the join, which the native join does not keep ([#6898](https://github.com/apache/datafusion-comet/issues/6898)). | `ExistenceJoin` (produced when `EXISTS` / `IN` is combined with another predicate via `OR`) runs natively on `BroadcastHashJoinExec` and `ShuffledHashJoinExec`. Existence sort-merge joins, residual diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 6aab2c0d36f..572961f08e7 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 7789e8eacc2..aa8b6ba2a62 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -224,6 +224,8 @@ Keep the following in mind when reading the reports: native plan still counts as accelerated, so treat the percentage as an upper bound. - Comet's split Iceberg V2 write (`spark.comet.write.iceberg.enabled`) is declined in plan-only mode, so such writes run on, and are reported as, Spark. +- `spark.comet.exec.forceShuffledHashJoin` is not applied in plan-only mode, so the report shows + the sort-merge joins Spark plans rather than the hash joins a real run would use. - Under AQE the report describes the plan before any adaptive re-planning, so coverage of the plan that finally executes can differ. In particular, AQE plans subqueries into the outer query only after the report is produced, so the outer report counts their operators as 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 a31053bd6e8..94acefc7c32 100644 --- a/native/core/src/execution/operators/sorted_input.rs +++ b/native/core/src/execution/operators/sorted_input.rs @@ -37,13 +37,14 @@ use datafusion::physical_plan::{ /// sorted before it was cached for instance, but the `ScanExec` it arrives through declares none, /// and DataFusion's hash join keeps the unmatched probe rows of a right outer join in place only /// when the probe input declares an order. Otherwise it moves them after the matched rows of each -/// batch, which loses an order Spark has already dropped a sort for. The planner puts this only -/// on that join's probe input, with the order Spark reports for it, and only when the rows reach -/// that input from the JVM through operators that keep their order. The join keeps the order in -/// its own output, so the operators above it in the native plan see it too: an aggregate grouped -/// by the key runs in sorted, streaming mode, a sort aggregate takes the order as already -/// satisfied, and a top-k sort whose keys start with it stops reading early. A wrong order would -/// therefore give wrong groups or wrong rows, not only rows out of order. +/// batch, which loses an order Spark has already dropped a sort for. The planner puts this on a +/// hash join's probe input, with the order Spark reports for it, when the rows reach that input +/// from the JVM or a file scan through operators that keep their order, and on the native range, +/// whose order is real, so that the join and the operators above it in the native plan see the +/// order: an aggregate grouped by the key runs in sorted, streaming mode, a window finds the order +/// it requires, a sort aggregate takes the order as already satisfied, and a top-k sort whose keys +/// start with it stops reading early. A wrong order would therefore give wrong groups or wrong +/// rows, not only rows out of order. #[derive(Debug)] pub struct SortedInputExec { input: Arc, @@ -70,6 +71,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 a21c61cfe1d..8d5871b09e1 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -3115,6 +3115,7 @@ fn needs_fields_coercion(sig: &TypeSignature) -> bool { mod tests { mod empty_native_scan; mod multi_store_scan; + mod window_input_order; use futures::{poll, StreamExt}; use std::{ @@ -4070,7 +4071,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, @@ -4104,7 +4105,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(); @@ -4125,6 +4126,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 { @@ -4134,7 +4168,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], @@ -4183,7 +4217,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![ @@ -4234,7 +4268,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(); @@ -4255,7 +4289,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], 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..6dbd3bb8685 --- /dev/null +++ b/native/core/src/execution/planner/tests/window_input_order.rs @@ -0,0 +1,406 @@ +// 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 every PARTITION BY key is ordered. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, Float64Array, Int32Array}; +use arrow::compute::cast; +use arrow::datatypes::DataType; +use datafusion::physical_expr::expressions::Column; +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)]), + }, + // The ORDER BY key is not declared either. Only the PARTITION BY keys are checked. + Case { + name: "ordered on the partition key", + input_order: Some(&[(0, 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)]), + partition_by: &[0, 1], + order_by: &[2], + added_sort: None, + }, + Case { + name: "ordered descending on the partition key", + input_order: Some(&[(0, true)]), + partition_by: &[0], + order_by: &[2], + added_sort: None, + }, + 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}"); + } + } + } +} + +#[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)]); + 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)]) + } 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}"); + } +} diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs index 2d15fdc86e1..0cbfe186f7b 100644 --- a/native/core/src/execution/planner/window.rs +++ b/native/core/src/execution/planner/window.rs @@ -36,8 +36,9 @@ use datafusion::{ physical_plan::{ limit::LocalLimitExec, projection::ProjectionExec, + sorts::sort::SortExec, windows::{BoundedWindowAggExec, WindowAggExec}, - ExecutionPlan, InputOrderMode, + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, }, }; use datafusion_comet_operators::{PartitionedRankLimitExec, WindowFnKind}; @@ -98,6 +99,36 @@ impl OperatorBuilder for WindowBuilder { let sort_exprs = &sort_exprs?; let partition_exprs = &partition_exprs?; + // Sort the input unless it is already ordered on every PARTITION BY key; rows from the + // JVM declare no order even when Spark reports them as sorted. The check looks only at + // the PARTITION BY keys; the sort orders by them ascending with nulls first, then by the + // ORDER BY keys. + let (_, ordered) = child + .native_plan + .equivalence_properties() + .find_longest_permutation(partition_exprs)?; + let required = partition_exprs + .iter() + .map(|e| { + PhysicalSortExpr::new( + Arc::clone(e), + SortOptions { + descending: false, + nulls_first: true, + }, + ) + }) + .chain(sort_exprs.iter().cloned()); + let input_sort: Option> = match LexOrdering::new(required) { + Some(ordering) if ordered.len() < partition_exprs.len() => Some(Arc::new( + SortExec::new(ordering, Arc::clone(&child.native_plan)), + )), + _ => None, + }; + let window_input = input_sort + .as_ref() + .map_or_else(|| Arc::clone(&child.native_plan), Arc::clone); + let window_expr: Result>, ExecutionError> = wnd .window_expr .iter() @@ -134,14 +165,14 @@ impl OperatorBuilder for WindowBuilder { let window_agg: Arc = if all_bounded { Arc::new(BoundedWindowAggExec::try_new( window_expr, - Arc::clone(&child.native_plan), + window_input, InputOrderMode::Sorted, !partition_exprs.is_empty(), )?) } else { Arc::new(WindowAggExec::try_new( window_expr, - Arc::clone(&child.native_plan), + window_input, !partition_exprs.is_empty(), )?) }; @@ -183,16 +214,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/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 d79927f0c04..c52398b3d76 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -1041,34 +1041,23 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) plan } } else { - // 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(plan.transformUp { case p => - RewriteJoin.rewrite(p, conf) - }) - } else { - plan - } - // 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(plan) // 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(plan) - var transformed = transform(planWithJoinRewritten) + var transformed = transform(plan) // 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(plan) } var newPlan = 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/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/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 3f331fc5b9c..ee0abc0df03 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 @@ -2950,26 +2950,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 @@ -2978,14 +2981,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 " + @@ -3019,10 +3022,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 @@ -3039,9 +3041,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, @@ -3185,11 +3184,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/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index b39d16f1235..270006d6127 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 @@ -34,21 +36,21 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, Ascending, AttributeReference, BoundReference, Descending, DynamicPruningExpression, IsNotNull, SortOrder} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter, LeftSemi} -import org.apache.spark.sql.catalyst.plans.logical.Join -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, SQLExecution, WholeStageCodegenExec} +import org.apache.spark.sql.catalyst.plans.logical.{HintInfo, Join, JoinHint, LogicalPlan, NO_BROADCAST_HASH, PREFER_SHUFFLE_HASH} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoin, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometRangeExec, CometSampleExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometUnionExec, CometWindowExec, CometWindowGroupLimitExec} +import org.apache.spark.sql.execution.{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, 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, 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 class CometJoinSuite extends CometTestBase { @@ -1076,11 +1078,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") { @@ -1092,9 +1091,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) @@ -1138,6 +1135,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 { @@ -1150,6 +1149,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 { @@ -1166,6 +1167,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)) @@ -1174,6 +1176,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)) @@ -1181,6 +1184,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)) @@ -1193,6 +1197,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)) @@ -1202,21 +1207,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) } @@ -1248,30 +1254,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") { @@ -1282,6 +1305,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 @@ -1300,8 +1333,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) { @@ -1343,23 +1374,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)") @@ -1370,7 +1388,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)") @@ -1383,19 +1401,458 @@ 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( - 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) { + 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 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) + "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) + } + } + } + + // 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) } } } @@ -2286,6 +2743,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", @@ -2441,10 +2938,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", @@ -2664,6 +3158,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. @@ -2729,6 +3302,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 => @@ -2796,6 +3378,87 @@ 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) + 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 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("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") { @@ -3333,36 +3996,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 642a73e337a..2c7651913c3 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -28,7 +28,7 @@ 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.comet.{CometColumnarToRowExec, CometNativeExec, 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} @@ -543,6 +543,30 @@ 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" + assert(collect(plan) { case w: CometWindowExec => w }.size == 1, context) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, 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"))