diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index caf45198336..a7f7d457f4a 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -357,6 +357,29 @@ case class CometProjectExec( } ``` +#### Reporting an Output Ordering + +An operator that overrides `outputOrdering`, as `CometFilterExec` does above with +`child.outputOrdering`, makes a claim that Spark acts on: it drops a sort above the operator that +the order satisfies. Report what the Spark operator reports and nothing more. `SampleExec`, for +example, reports no order, so `CometSampleExec` must not report its child's. + +The claim has to hold at run time as well. A native hash join declares the order Spark reports for +its streamed side to DataFusion, and an aggregate above the join may then run in sorted mode, or a +top-k sort may stop reading early, on the strength of that order. If the order is false, the +result is wrong groups or wrong rows, not only rows out of order. So an operator that reports an +ordering must keep it natively, and the native operator must say so by returning `true` from +`maintains_input_order` for that child, as DataFusion's filter and projection do. The planner +declares the order through an operator only when the operator says it keeps it. An operator that +cannot keep the order natively must not report it, or has to be added to `mayReportFalseOrder` in +`operators.scala`, which makes the hash join serde send no order over it. + +A leaf that reports an ordering is where the order starts, so the planner has to know that the +native operator produces its rows in that order. It knows that of an input from the JVM and of a +scan of files, which reads the files of each partition one after the other, and the native range +declares its order itself through `SortedInputExec`. A new leaf that reports an ordering must +declare it natively the same way. + #### Plan Identity and Exchange Reuse Spark's `ReuseExchangeAndSubquery` identifies equivalent plans through canonicalization and may diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 8312fb520ad..aaa95d859b5 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -117,6 +117,27 @@ for example the output of a window, a sort-merge join, or a cached table sorted partitions, Comet falls back to Spark for the join, so that a sort Spark removed above it on the strength of that order is not lost. +### Output Order + +Spark reports the sort order of a hash join's streamed side as the join's own, and drops a sort +above the join that this order satisfies. Comet's native hash joins keep that order where Comet can +show that it holds. That includes a left or right outer join over a streamed side that was sorted +outside the native plan, such as a sorted relation read from the [in-memory cache](#in-memory-cache). + +Comet passes the order to the native join as sort keys, and an outer join needs at least the first +key to be one that Comet can express natively. Declaring any leading part of the order keeps the +rows in place, so the join does not need the rest. An outer join whose first sort key Comet cannot +express falls back to Spark and records a fallback reason (see +[Understanding Comet Plans](../understanding-comet-plans.md)), because Comet cannot tell whether +an operator above the join relies on the order. + +Where the order Spark reports might not hold, Comet declares nothing and an outer join can return +its unmatched rows out of order. The output of a nested loop join +([#6898](https://github.com/apache/datafusion-comet/issues/6898)) is one such case, and so is a +cache or checkpoint that one wrote. Comet also declares nothing when it cannot follow the +order back to its source, as with an RDD checkpoint, and over some caches written by an outer hash +join whose order does hold ([#6922](https://github.com/apache/datafusion-comet/issues/6922)). + ## Window Functions Comet runs `WindowExec` natively and it is enabled by default (`spark.comet.exec.window.enabled`). A broad set of diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index a5c30ed79fe..cc20a4746ad 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -44,5 +44,7 @@ mod csv_scan; pub mod projection; mod scan; mod shuffle_scan; +mod sorted_input; pub use csv_scan::init_csv_datasource_exec; pub use shuffle_scan::ShuffleScanExec; +pub use sorted_input::SortedInputExec; diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs new file mode 100644 index 00000000000..a31053bd6e8 --- /dev/null +++ b/native/core/src/execution/operators/sorted_input.rs @@ -0,0 +1,135 @@ +// 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. + +//! An input that declares the order its rows arrive in. + +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_err, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{LexOrdering, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::{ + apply_expression_roots, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, SendableRecordBatchStream, +}; + +/// Passes `input` through unchanged and declares to DataFusion the order its rows arrive in. +/// +/// Spark's plan knows the order of an input that enters the native plan from the JVM, a relation +/// 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. +#[derive(Debug)] +pub struct SortedInputExec { + input: Arc, + ordering: LexOrdering, + cache: Arc, +} + +impl SortedInputExec { + pub fn new(input: Arc, ordering: LexOrdering) -> Self { + let input_properties = input.properties(); + // Start from the input's own properties, so the constants and equivalence classes it + // knows of stay visible to the operators above. + let mut equivalence_properties = input.equivalence_properties().clone(); + equivalence_properties.add_ordering(ordering.iter().cloned()); + let cache = Arc::new(PlanProperties::new( + equivalence_properties, + input_properties.output_partitioning().clone(), + input_properties.emission_type, + input_properties.boundedness, + )); + Self { + input, + ordering, + cache, + } + } +} + +impl DisplayAs for SortedInputExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "SortedInputExec: ordering=[{}]", self.ordering) + } +} + +impl ExecutionPlan for SortedInputExec { + fn name(&self) -> &str { + "SortedInputExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + apply_expression_roots(self.ordering.iter().map(|sort| Arc::clone(&sort.expr)), f) + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + /// The rows pass through, so the input's metrics are this operator's, and a range whose root + /// this wraps still reports its output rows. + fn metrics(&self) -> Option { + self.input.metrics() + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!("SortedInputExec requires one child"); + } + Ok(Arc::new(Self::new( + children.remove(0), + self.ordering.clone(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(partition, context) + } +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 26a746d0710..2367c4f4089 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -40,6 +40,7 @@ mod lance_scan; use crate::execution::operators::init_csv_datasource_exec; use crate::execution::operators::DynamicFilterJoinExec; use crate::execution::operators::IcebergScanExec; +use crate::execution::operators::SortedInputExec; use crate::execution::operators::TopKReaderFilterExec; use crate::execution::{ operators::{ @@ -88,7 +89,7 @@ use datafusion::{ prelude::SessionContext, }; use datafusion_comet_operators::{ - range_exec, CometFilterExec, CometProjectionExec, ExpandExec, ExplodeExec, + range_exec, range_sort_expr, CometFilterExec, CometProjectionExec, ExpandExec, ExplodeExec, PartitionedRankLimitExec, SampleExec, WindowFnKind, }; use datafusion_comet_spark_expr::{ @@ -97,6 +98,8 @@ use datafusion_comet_spark_expr::{ CsvWriteOptions, EvalMode, ListPositionsExpr, SparkArraysZipFunc, SparkBloomFilterVersion, SparkListAgg, SparkPercentile, SparkPivotFirst, Subquery, SumInteger, ToCsv, }; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::TableSchema; use iceberg::expr::Bind; @@ -2090,7 +2093,7 @@ impl PhysicalPlanner { } OpStruct::RangeScan(range) => { // A leaf with no JVM input: each task produces its own partition of the range. - let range: Arc = Arc::new(range_exec( + let source: Arc = Arc::new(range_exec( range.start, range.step, range.num_elements, @@ -2098,10 +2101,17 @@ impl PhysicalPlanner { self.partition, self.session_ctx.copied_config().batch_size(), )?); + // The range's order is real, and Spark's `RangeExec` reports it, so declare it + // for the operators above, a hash join in particular: see `SortedInputExec`. + let ordering = LexOrdering::new([range_sort_expr(range.step)]); + let source: Arc = match ordering { + Some(ordering) => Arc::new(SortedInputExec::new(source, ordering)), + None => source, + }; Ok(( vec![], vec![], - Arc::new(SparkPlan::new(spark_plan.plan_id, range, vec![])), + Arc::new(SparkPlan::new(spark_plan.plan_id, source, vec![])), )) } OpStruct::Sample(sample) => { @@ -2782,8 +2792,21 @@ impl PhysicalPlanner { *probe = Self::prepare_probe_for_runtime_filters(Arc::clone(probe))?; } - let left = Arc::clone(&join_params.left.native_plan); - let right = Arc::clone(&join_params.right.native_plan); + let mut left = Arc::clone(&join_params.left.native_plan); + let mut right = Arc::clone(&join_params.right.native_plan); + + // Spark dropped any sort above this join that the streamed side's order + // satisfies, and the native join keeps that order only if the probe input + // declares it; see `SortedInputExec`. + if !join.streamed_sort_orders.is_empty() { + let probe = if join.build_side == BuildSide::BuildLeft as i32 { + &mut right + } else { + &mut left + }; + *probe = + self.with_declared_order(Arc::clone(probe), &join.streamed_sort_orders)?; + } // Null-aware anti-join must run in CollectLeft mode. In Partitioned mode // each partition only sees per-partition null/emptiness state, which can @@ -3092,6 +3115,83 @@ impl PhysicalPlanner { Ok(Arc::new(prepared)) } + /// `plan`, declaring `sort_orders` as the order its rows arrive in, unless it already + /// declares one. Spark's plan knows the order of an input that enters the native plan from + /// the JVM, and of a scan of sorted files, but the `ScanExec` or the file scan the rows + /// arrive through declares none. Only that order is declared: when an operator in between may + /// have moved rows, `plan` is left as it is. + /// + /// The serde sends an order only when the streamed side's Comet operators report it, and it + /// checks that nothing under them moves rows (`mayReportFalseOrder`), while this walk asks + /// the native operators the same thing through `maintains_input_order`. The two agree as + /// long as a Comet operator that reports an order keeps it natively and says so, and a Comet + /// leaf that reports one is a source this walk knows or declares its order itself, as the + /// native range does. One that does not lands here with `plan` unchanged, and the join then + /// moves its unmatched rows after Spark has dropped the sort above it. Nothing here can fall + /// back at that point, so the invariant is kept at the operator: see "Adding a New Operator" + /// in the contributor guide. + fn with_declared_order( + &self, + plan: Arc, + sort_orders: &[Expr], + ) -> Result, ExecutionError> { + if plan.output_ordering().is_some() || !Self::keeps_source_order(&plan) { + return Ok(plan); + } + let exprs = sort_orders + .iter() + .map(|expr| self.create_sort_expr(expr, plan.schema())) + .collect::, _>>()?; + match LexOrdering::new(exprs) { + Some(ordering) => Ok(Arc::new(SortedInputExec::new(plan, ordering))), + None => Ok(plan), + } + } + + /// Whether `plan`'s rows come from a source whose order Spark knows through operators that + /// keep their input's order, so that the order Spark reports for them still holds. A source is + /// an input from the JVM, or a scan of files, which reads the files of its one file group one + /// after the other: Spark reports an order for a file scan only when every partition reads one + /// sorted file, as for a sorted bucketed table with one file per bucket. Any other operator may + /// have moved rows, such as the null-aware anti join, whose streamed side Comet builds on, and + /// declaring Spark's order over it would be wrong. + fn keeps_source_order(plan: &Arc) -> bool { + if plan.is::() || plan.is::() { + return true; + } + if let Some(scan) = plan.downcast_ref::() { + return scan + .data_source() + .downcast_ref::() + .is_some_and(|config| config.file_groups.len() == 1); + } + // DataFusion reports every right hash join as keeping its probe order, but a right outer + // join keeps it only when the probe declares one, and no input on this path declares an + // order. An inner join emits its matches probe row by probe row, and a right semi, anti or + // mark join builds its probe indices in ascending order, so none of them depends on a + // declaration. Any other join type keeps no probe order at all. + if let Some(join) = plan.downcast_ref::() { + if !matches!( + join.join_type(), + DFJoinType::Inner + | DFJoinType::RightSemi + | DFJoinType::RightAnti + | DFJoinType::RightMark + ) { + return false; + } + } + let mut kept = plan + .children() + .into_iter() + .zip(plan.maintains_input_order()) + .filter_map(|(child, keeps)| keeps.then_some(child)); + match (kept.next(), kept.next()) { + (Some(child), None) => Self::keeps_source_order(child), + _ => false, + } + } + /// Attach after choosing the final build side, including the projection emitted /// by `swap_inputs`, without running DataFusion's physical optimizer. pub(crate) fn apply_join_dynamic_filter( @@ -6369,6 +6469,7 @@ mod tests { build_side: 0, null_aware_anti_join: false, dynamic_filter_enabled: false, + streamed_sort_orders: vec![], })), }; @@ -6383,6 +6484,239 @@ mod tests { assert_eq!("ScanExec", hash_join_exec.children[1].native_plan.name()); } + /// The hash join in `plan`, under the projection that `swap_inputs` adds when Spark builds + /// on the right. + fn planned_hash_join( + plan: &Arc, + ) -> &datafusion::physical_plan::joins::HashJoinExec { + let join = if plan.is::() { + plan.children()[0] + } else { + plan + }; + join.downcast_ref::() + .expect("a hash join") + } + + fn sorted_outer_hash_join( + join_type: spark_operator::JoinType, + build_side: spark_operator::BuildSide, + children: Vec, + ) -> Operator { + Operator { + children, + op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin { + left_join_keys: vec![create_bound_reference(0)], + right_join_keys: vec![create_bound_reference(0)], + join_type: join_type as i32, + build_side: build_side as i32, + streamed_sort_orders: vec![create_sort_order(0, 3, false, true)], + ..Default::default() + })), + ..Default::default() + } + } + + #[test] + fn hash_join_declares_streamed_order_on_probe_input() { + // The two outer joins Comet sends the streamed side's order for both run as DataFusion's + // right join, whose probe input is its right one, whichever side Spark builds on. + for (join_type, build_side) in [ + ( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + ), + ( + spark_operator::JoinType::LeftOuter, + spark_operator::BuildSide::BuildRight, + ), + ] { + let operator = + sorted_outer_hash_join(join_type, 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::Right, join.join_type()); + assert_eq!("ScanExec", join.left().name()); + let probe = join.right(); + let sorted = probe + .downcast_ref::() + .expect("the probe input declares the streamed side's order"); + assert_eq!("ScanExec", sorted.children()[0].name()); + assert_eq!( + "col_0@0 ASC", + probe.properties().output_ordering().unwrap().to_string() + ); + // The join keeps the order in its own output, so the operators above it see it too. + assert!(planned.native_plan.properties().output_ordering().is_some()); + } + } + + #[test] + fn hash_join_declares_order_through_operators_that_keep_it() { + let projection = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::Projection(spark_operator::Projection { + project_list: vec![create_bound_reference(0)], + })), + ..Default::default() + }; + let operator = sorted_outer_hash_join( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + vec![create_scan(), projection], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + let probe = planned_hash_join(&planned.native_plan).right(); + assert_eq!("SortedInputExec", probe.name()); + assert_eq!("ProjectionExec", probe.children()[0].name()); + } + + /// A hash join over two scans, for use as the probe input of the join under test. + fn hash_join_over_scans( + join_type: spark_operator::JoinType, + build_side: spark_operator::BuildSide, + null_aware_anti_join: bool, + ) -> Operator { + Operator { + children: vec![create_scan(), create_scan()], + op_struct: Some(OpStruct::HashJoin(spark_operator::HashJoin { + left_join_keys: vec![create_bound_reference(0)], + right_join_keys: vec![create_bound_reference(0)], + join_type: join_type as i32, + build_side: build_side as i32, + null_aware_anti_join, + ..Default::default() + })), + ..Default::default() + } + } + + #[test] + fn hash_join_declares_order_through_joins_that_keep_probe_order() { + // An inner join emits its matches probe row by probe row, and a right semi, anti or mark + // join emits its probe indices in ascending order. None of them depends on the probe + // declaring an order, so the order Spark reports for the outer join's streamed side + // still holds above them. + use datafusion::common::JoinType as DF; + use spark_operator::BuildSide::{BuildLeft, BuildRight}; + use spark_operator::JoinType::{Existence, Inner, LeftAnti, LeftSemi}; + for (join_type, build_side, planned_type) in [ + (Inner, BuildLeft, DF::Inner), + (Inner, BuildRight, DF::Inner), + (LeftSemi, BuildRight, DF::RightSemi), + (LeftAnti, BuildRight, DF::RightAnti), + (Existence, BuildRight, DF::RightMark), + ] { + let operator = sorted_outer_hash_join( + spark_operator::JoinType::RightOuter, + BuildLeft, + vec![ + create_scan(), + hash_join_over_scans(join_type, build_side, false), + ], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + let probe = planned_hash_join(&planned.native_plan).right(); + assert_eq!( + "SortedInputExec", + probe.name(), + "{join_type:?} {build_side:?}" + ); + let declared = planned_hash_join(probe.children()[0]); + assert_eq!(&planned_type, declared.join_type()); + } + } + + #[test] + fn hash_join_declares_no_order_over_an_operator_that_may_move_rows() { + use spark_operator::BuildSide::{BuildLeft, BuildRight}; + use spark_operator::JoinType::{FullOuter, LeftAnti, LeftSemi, RightOuter}; + for (probe_input, reason) in [ + // Comet builds a null-aware anti join on Spark's streamed side, so its output is not + // in the order Spark reports for it. + ( + hash_join_over_scans(LeftAnti, BuildRight, true), + "a null-aware anti join", + ), + // A right join whose probe declares no order moves its unmatched rows, although + // DataFusion reports it as keeping its probe order. + ( + hash_join_over_scans(RightOuter, BuildLeft, false), + "a right outer join", + ), + // The remaining types keep no probe order at all: a left semi join built on the left + // emits build rows, and a full outer join emits its unmatched build rows last. + ( + hash_join_over_scans(LeftSemi, BuildLeft, false), + "a left semi join", + ), + ( + hash_join_over_scans(FullOuter, BuildLeft, false), + "a full outer join", + ), + ] { + let operator = + sorted_outer_hash_join(RightOuter, BuildLeft, vec![create_scan(), probe_input]); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + let probe = planned_hash_join(&planned.native_plan).right(); + assert_eq!("HashJoinExec", probe.name(), "{reason}"); + } + } + + #[test] + fn hash_join_keeps_probe_input_that_declares_an_order() { + // A probe input that already declares an order, such as a sort in the native plan, is + // left as it is. + let sort = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::Sort(spark_operator::Sort { + sort_orders: vec![create_sort_order(0, 3, false, true)], + ..Default::default() + })), + ..Default::default() + }; + let operator = sorted_outer_hash_join( + spark_operator::JoinType::RightOuter, + spark_operator::BuildSide::BuildLeft, + vec![create_scan(), sort], + ); + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&operator, &mut vec![], 1) + .unwrap(); + assert_eq!( + "SortExec", + planned_hash_join(&planned.native_plan).right().name() + ); + } + + #[test] + fn declared_order_reaches_a_scan_of_one_file_group() { + // A scan of files reads the files of its one file group one after the other, so the order + // Spark reports for it, a sorted bucketed table's, holds. With more than one file group, + // DataFusion may hand a file to another partition. + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let scan = |groups: usize| -> Arc { + let source = Arc::new(ParquetSource::new(Arc::clone(&schema))) as Arc; + let file_groups = (0..groups) + .map(|i| FileGroup::new(vec![PartitionedFile::new(format!("f{i}.parquet"), 8)])) + .collect(); + let config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file_groups(file_groups) + .build(); + Arc::new(DataSourceExec::new(Arc::new(config))) + }; + assert!(PhysicalPlanner::keeps_source_order(&scan(1))); + assert!(!PhysicalPlanner::keeps_source_order(&scan(2))); + } + #[tokio::test] async fn spark_plan_dynamic_filter_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; diff --git a/native/operators/src/lib.rs b/native/operators/src/lib.rs index b5f762025ad..97c80a97895 100644 --- a/native/operators/src/lib.rs +++ b/native/operators/src/lib.rs @@ -37,6 +37,6 @@ pub use expand::ExpandExec; pub use explode::ExplodeExec; pub use filter::CometFilterExec; pub use projection::CometProjectionExec; -pub use range::range_exec; +pub use range::{range_exec, range_sort_expr}; pub use rank_limit::{PartitionedRankLimitExec, WindowFnKind}; pub use sample::SampleExec; diff --git a/native/operators/src/range.rs b/native/operators/src/range.rs index 3766c82b1a6..4ba4f6d6429 100644 --- a/native/operators/src/range.rs +++ b/native/operators/src/range.rs @@ -16,8 +16,11 @@ // under the License. use arrow::array::{Int64Array, RecordBatch}; +use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::common::Result; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalSortExpr; use datafusion::physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec}; use datafusion::physical_plan::Partitioning; use parking_lot::RwLock; @@ -58,6 +61,23 @@ pub fn range_exec( Ok(exec) } +/// The order a range's values come out in. They are `partition_start + i * step` and never +/// overflow, so they ascend for a positive step and descend for a negative one, which is the order +/// Spark's `RangeExec` reports. +/// +/// `LazyMemoryExec::add_ordering` does not refresh the plan properties' cached output ordering in +/// DataFusion 55, so `output_ordering()` would still be `None` after it. The planner declares this +/// order by wrapping the range in a `SortedInputExec` instead. +pub fn range_sort_expr(step: i64) -> PhysicalSortExpr { + PhysicalSortExpr::new( + Arc::new(Column::new("id", 0)), + SortOptions { + descending: step < 0, + nulls_first: step > 0, + }, + ) +} + /// The first value of partition `partition` and the number of values in it, computed the way /// Spark's `RangeExec` does (`initRange` in its generated code): in `BigInteger` arithmetic, with /// the partition's start and end clamped to the `Long` range. @@ -247,6 +267,20 @@ mod tests { assert_eq!(collect(i64::MIN, step, n, 1, 1000).await, expected); } + #[tokio::test] + async fn test_sort_expr_matches_the_values() { + for step in [1i64, 3, -1, -7] { + let n = num_elements(0, 100 * step, step); + let values = collect(0, step, n, 2, 8).await; + let expr = range_sort_expr(step); + assert_eq!(expr.options.descending, step < 0, "step {step}"); + let mut sorted = values.clone(); + sorted.sort_by_key(|v| if expr.options.descending { -v } else { *v }); + // Both partitions are in order, and the partitions follow each other. + assert_eq!(values, sorted, "step {step}"); + } + } + #[test] fn test_partition_bounds() { // range(0, 10, 1, 3): Spark puts 3, 3 and 4 values in the three partitions. diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index ac04e9e2320..48e67ec7c11 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1099,6 +1099,11 @@ message HashJoin { // Opt-in completed-build filtering of probe batches and eligible native Parquet readers. // Older serialized plans leave it false. bool dynamic_filter_enabled = 10; + // The streamed side's sort order, or the longest prefix of it Comet can express, as + // `SortOrder` expressions over its output, when Spark reports that order for the join and the + // native join keeps it only if the probe input declares it: a left outer join built on the + // right, or a right outer join built on the left. Empty otherwise. + repeated spark.spark_expression.Expr streamed_sort_orders = 11; } message SortMergeJoin { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometSampleExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometSampleExec.scala index dcdae2c779d..54391580a44 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometSampleExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometSampleExec.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.comet -import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} +import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.execution.{SampleExec, SparkPlan} @@ -97,8 +97,6 @@ case class CometSampleExec( override def outputPartitioning: Partitioning = child.outputPartitioning - override def outputOrdering: Seq[SortOrder] = child.outputOrdering - override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = this.copy(child = newChild) 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 87268b5fc69..733f56aaded 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 @@ -40,8 +40,10 @@ import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNative import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, BroadcastQueryStageExec, QueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanExecBase import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} @@ -62,7 +64,7 @@ import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde -import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, isStringCollationType, supportedSortType} +import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, supportedSortType} import org.apache.comet.serde.operator.CometSink import org.apache.comet.shims.MergeRowsMetricsShim @@ -2918,6 +2920,55 @@ trait CometHashJoin { val leftKeys = join.leftKeys.map(exprToProto(_, join.left.output)) val rightKeys = join.rightKeys.map(exprToProto(_, join.right.output)) + // Spark reports the streamed side's order as the join's, and drops a sort above the join + // that it satisfies. DataFusion's hash join keeps the unmatched probe rows of a right outer + // join in place only when the probe input declares an order, and otherwise moves them after + // the matched rows of each batch. Comet builds on the build side, so a left outer join built + // on the right and a right outer join built on the left both run as that right outer join, + // 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. + // + // 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. + // + // 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. + val (streamed, _) = streamedSide(join.buildSide, join.left, join.right) + val streamedSortOrders = + if (runsAsRightOuterJoin(join.joinType, join.buildSide) && + join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed)) { + val orders = join.outputOrdering.iterator + .takeWhile(order => !hasNonDefaultStringCollation(order.dataType)) + .map(exprToProto(_, streamed.output)) + .takeWhile(_.isDefined) + .flatten + .toSeq + if (orders.isEmpty) { + withFallbackReason( + join, + "the streamed side's sort order, which this outer join has to keep, cannot be " + + "expressed natively") + return None + } + orders + } else { + Nil + } + if (leftKeys.forall(_.isDefined) && rightKeys.forall(_.isDefined) && childOp.nonEmpty) { @@ -2930,12 +2981,103 @@ trait CometHashJoin { else OperatorOuterClass.BuildSide.BuildRight) .setNullAwareAntiJoin(isNullAwareAntiJoin) .setDynamicFilterEnabled(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get(join.conf)) + .addAllStreamedSortOrders(streamedSortOrders.asJava) condition.foreach(joinBuilder.setCondition) Some(builder.setHashJoin(joinBuilder).build()) } else { None } } + + // 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 + // 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 + // walked through that side alone. A cache reports the order of the plan that wrote it, and the + // walk goes on into that plan, picked as Spark's scan picks it: an adaptive plan reports no + // order, but the scan reads the final plan behind it once there is one. While CometExecRule + // converts a plan, a native plan's JVM inputs sit behind a CometScanWrapper, a leaf, so the + // walk looks through that too. + private def mayReportFalseOrder(plan: SparkPlan): Boolean = plan match { + case adaptive: AdaptiveSparkPlanExec if adaptive.isFinalPlan => + mayReportFalseOrder(adaptive.executedPlan) + case _ if plan.outputOrdering.isEmpty => false + case _: SortExec | _: CometSortExec => false + case _: CometBroadcastNestedLoopJoinExec => true + case join: CometBroadcastHashJoinExec => + hashJoinMayReportFalseOrder( + join.nativeOp, + join.joinType, + join.buildSide, + join.left, + join.right) + case join: CometHashJoinExec => + hashJoinMayReportFalseOrder( + join.nativeOp, + join.joinType, + join.buildSide, + join.left, + join.right) + case join: HashJoin => + mayReportFalseOrder(streamedSide(join.buildSide, join.left, join.right)._1) + case wrapper: CometScanWrapper => mayReportFalseOrder(wrapper.originalPlan) + case scan: CometInMemoryTableScanExec => + mayReportFalseOrder(scan.originalPlan.relation.cachedPlan) + case scan: InMemoryTableScanExec => mayReportFalseOrder(scan.relation.cachedPlan) + case stage: QueryStageExec => mayReportFalseOrder(stage.plan) + // A scan reports an order that Spark already relies on. + case _: DataSourceScanExec | _: DataSourceV2ScanExecBase | _: RangeExec | _: CometLeafExec => + false + case _ if plan.children.isEmpty => true + case _ => plan.children.exists(mayReportFalseOrder) + } + + // Whether a Comet hash join that reports an order may not keep it. Only the streamed side's + // order reaches the join's output, so a false order on the build side is of no concern, and an + // outer join of the two shapes has to have declared the order on its probe input as well. + private def hashJoinMayReportFalseOrder( + nativeOp: Operator, + joinType: JoinType, + buildSide: BuildSide, + left: SparkPlan, + right: SparkPlan): Boolean = { + val (streamed, streamedIndex) = streamedSide(buildSide, left, right) + mayReportFalseOrder(streamed) || + (runsAsRightOuterJoin(joinType, buildSide) && !declaresStreamedOrder(nativeOp, streamedIndex)) + } + + // Whether the native planner declares the streamed order on the probe input of the join that + // `nativeOp` describes. It does so when the serde sent the order and the probe input reads from + // the JVM, which is the case when the join's plan has a scan there. A native operator over that + // scan is declared as well, as far as the native walk goes, but that is for the planner to say, + // so it counts as not declared. + private def declaresStreamedOrder(nativeOp: Operator, streamedIndex: Int): Boolean = + nativeOp.getHashJoin.getStreamedSortOrdersCount > 0 && + nativeOp.getChildrenCount > streamedIndex && { + val probe = nativeOp.getChildren(streamedIndex) + probe.hasScan || probe.hasShuffleScan + } + + // The streamed side of a hash join, which is the side it is not built on, and its child index. + private def streamedSide( + buildSide: BuildSide, + left: SparkPlan, + right: SparkPlan): (SparkPlan, Int) = + if (buildSide == BuildRight) (left, 0) else (right, 1) + + // A left outer join built on the right and a right outer join built on the left both run as + // DataFusion's right outer join, which keeps its unmatched probe rows in place only when the + // probe input declares an order. + private def runsAsRightOuterJoin(joinType: JoinType, buildSide: BuildSide): Boolean = + (joinType, buildSide) match { + case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) => true + case _ => false + } } case class CometBroadcastNestedLoopJoinExec( diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 4dd68e6b5f3..c300365f536 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -36,13 +36,15 @@ import org.apache.arrow.vector.types.pojo.ArrowType import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.CometDriverPlugin import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row} -import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal} +import org.apache.spark.sql.catalyst.expressions.{And, Ascending, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} +import org.apache.spark.sql.catalyst.plans.LeftOuter import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} -import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoinExec, CometInMemoryTableScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlanInfo} +import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} @@ -3288,4 +3290,210 @@ class CometInMemoryCacheSuite extends CometTestBase { }.nonEmpty) } } + + test("hash joins over a sorted cache keep the streamed side's order") { + // CometInMemoryTableScan reports the cached plan's sort order, so Spark drops the sort it + // would otherwise plan above a join that keeps that order. DataFusion's hash join keeps the + // unmatched probe rows of a left outer join built on the right, or of a right outer join built + // on the left, in place only when the probe input declares an order, which the native scan + // over the cache does not, so the native join declares the cache's order on its probe input. + // An inner join never moves probe rows. + 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 -> "true") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + + def query(hint: String, joinType: String): DataFrame = { + val joined = if (joinType == "right_outer") { + build.hint(hint).join(streamed, build("_1") === streamed("_1"), joinType) + } else { + streamed.join(build.hint(hint), streamed("_1") === build("_1"), joinType) + } + joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + } + + def joins(plan: SparkPlan): Seq[(Class[_], BuildSide)] = collect(plan) { + case j: CometHashJoinExec => (j.getClass, j.buildSide) + case j: CometBroadcastHashJoinExec => (j.getClass, j.buildSide) + } + + def sorts(plan: SparkPlan): Seq[SparkPlan] = + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s } + + val outerRows = (0 until 10000).map { i => + val k = i % 100 + Row(k, i, if (k % 10 == 0) k / 10 else null) + } + val innerRows = outerRows.filter(_.get(2) != null) + + Seq( + "shuffle_hash" -> classOf[CometHashJoinExec], + "broadcast" -> classOf[CometBroadcastHashJoinExec]).foreach { + case (hint, joinClass) => + Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft).foreach { + case (joinType, buildSide) => + val outer = query(hint, joinType) + val outerPlan = outer.queryExecution.executedPlan + assert(joins(outerPlan) == Seq((joinClass, buildSide)), outerPlan) + assert(collect(outerPlan) { case s: CometInMemoryTableScanExec => + s + }.size == 1) + assert(sorts(outerPlan).isEmpty, outerPlan) + QueryTest.sameRows(outerRows, outer.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(outer, keyOrder(Ascending)), outerPlan) + } + + val inner = query(hint, "inner") + val innerPlan = inner.queryExecution.executedPlan + assert(joins(innerPlan) == Seq((joinClass, BuildRight)), innerPlan) + assert(sorts(innerPlan).isEmpty, innerPlan) + QueryTest.sameRows(innerRows, inner.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(inner, keyOrder(Ascending)), innerPlan) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + + test("a broadcast outer hash join over a sorted cache keeps its order with default settings") { + // A cache written with no exchange in its plan is not adaptive, so with AQE on, as it is by + // default, the scan reports the cached plan's order just as it does with AQE off. Only the + // native cache scan is switched on here, and the broadcast join is the one Spark plans for + // a small build side. + withSQLConf(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + try { + val streamed = spark.table("big").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + val df = streamed + .join(build, streamed("_1") === build("_1"), "left_outer") + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val expected = (0 until 10000).map { i => + val k = i % 100 + Row(k, i, if (k % 10 == 0) k / 10 else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + // The plan AQE chose is final once the query has run. + val plan = df.queryExecution.executedPlan + assert(plan.isInstanceOf[AdaptiveSparkPlanExec], plan) + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j } + assert(joins.size == 1 && joins.head.joinType == LeftOuter, plan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, plan) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) + } finally { + spark.catalog.clearCache() + } + } + } + } + } + + test("an aggregate on the key above a hash join over a sorted cache groups each key once") { + // The order the join declares on its probe input reaches past the join, so the aggregate + // above it in the same native plan groups in sorted, streaming mode. That is right while + // Spark's reported order holds, and a wrong order would split a key across groups, which only + // a comparison of the rows as a multiset catches. With 40 keys and 8-row batches each group + // spans several batches. With 1000 keys a batch holds four keys, alternately matched and + // unmatched, so a join that moved its unmatched rows would also change the order in which the + // aggregate first sees the keys, and its output would not come back sorted. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + Seq(40, 1000).foreach { keys => + val data = (0 until 2000).map(i => ((i % keys) * 1000, i)) + withParquetTable(data, "big") { + withParquetTable((0 until keys by 2).map(m => (m * 1000, m)), "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + + // A match in the build side for every even key. + val expected = data.groupBy(_._1).toSeq.map { case (k, rows) => + Row( + k, + rows.size.toLong, + rows.map(_._2.toLong).sum, + if ((k / 1000) % 2 == 0) { + rows.size.toLong + } else 0L) + } + + for { + (hint, joinClass) <- Seq( + "shuffle_hash" -> classOf[CometHashJoinExec], + "broadcast" -> classOf[CometBroadcastHashJoinExec]) + (joinType, buildSide) <- Seq( + "left_outer" -> BuildRight, + "right_outer" -> BuildLeft) + } { + val joined = if (joinType == "right_outer") { + build.hint(hint).join(streamed, build("_1") === streamed("_1"), joinType) + } else { + streamed.join(build.hint(hint), streamed("_1") === build("_1"), joinType) + } + val df = joined + .groupBy(streamed("_1")) + .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) + val plan = df.queryExecution.executedPlan + val joins: Seq[(SparkPlan, BuildSide)] = collect(plan) { + case j: CometHashJoinExec => (j, j.buildSide) + case j: CometBroadcastHashJoinExec => (j, j.buildSide) + } + assert( + joins.map { case (j, side) => (j.getClass, side) } == Seq( + (joinClass, buildSide)), + plan) + // The serde sent the order, which is what lets the aggregates above run in + // sorted mode, and a sorted aggregate emits its groups in key order. + assert( + collect(plan) { + case j: CometHashJoinExec => j.nativeOp + case j: CometBroadcastHashJoinExec => j.nativeOp + }.forall(_.getHashJoin.getStreamedSortOrdersCount > 0), + plan) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) + // A final aggregate after a shuffle would merge a key a partial aggregate split, + // so both aggregates have to sit above the join with no exchange in between. + def aboveJoin(p: SparkPlan): Boolean = p.find(_ eq joins.head._1).isDefined + assert( + collect(plan) { case a: CometHashAggregateExec => a }.count(aboveJoin) == 2, + plan) + assert( + !collect(plan) { case e: ShuffleExchangeLike => e }.exists(aboveJoin), + plan) + assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1, plan) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + } } 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 d0e1438ee4a..6ae7b1081f9 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -28,22 +28,23 @@ import org.apache.hadoop.fs.Path import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.spark.SparkException -import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.{Column, CometTestBase, DataFrame, QueryTest, Row} import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation -import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull} -import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.catalyst.plans.Inner +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, CometHashJoinExec, CometNativeScanExec, CometProjectExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, 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.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.functions.{array, count, lit, struct, sum} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} +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} @@ -2458,4 +2459,689 @@ class CometJoinSuite extends CometTestBase { } } } + + test("outer hash joins keep a streamed side sorted outside the native plan in order") { + // Spark drops the sort above the join because the join reports the streamed side's order, + // and the streamed side arrives sorted through a JVM input, which declares no order to + // DataFusion unless the planner declares it on the join's probe input; see + // CometHashJoin.doConvert. CometInMemoryCacheSuite covers the native cache scan. + 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") { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_1").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + + // A residual condition, which keeps only the build rows past 3, sends more probe rows + // down the join's unmatched path. + def expected(withCondition: Boolean): Seq[Row] = (0 until 10000).map { i => + val k = i % 100 + val matched = k % 10 == 0 && (!withCondition || k / 10 > 3) + Row(k, i, if (matched) k / 10 else null) + } + + // Both shapes run as DataFusion's right join: Comet swaps a join built on the right + // and leaves one built on the left as it is. + def outerJoin(hint: String, joinType: String, withCondition: Boolean): DataFrame = { + val keys = streamed("_1") === build("_1") + val on = if (withCondition) keys && build("_2") > 3 else keys + val joined = if (joinType == "left_outer") { + streamed.join(build.hint(hint), on, joinType) + } else { + build.hint(hint).join(streamed, on, joinType) + } + joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + } + + for { + (hint, joinClass) <- Seq( + "shuffle_hash" -> classOf[CometHashJoinExec], + "broadcast" -> classOf[CometBroadcastHashJoinExec]) + (joinType, buildSide) <- Seq("left_outer" -> BuildRight, "right_outer" -> BuildLeft) + withCondition <- Seq(false, true) + } { + val df = outerJoin(hint, joinType, withCondition) + val plan = df.queryExecution.executedPlan + val joins: Seq[(Class[_], BuildSide)] = collect(plan) { + case j: CometHashJoinExec => (j.getClass, j.buildSide) + case j: CometBroadcastHashJoinExec => (j.getClass, j.buildSide) + } + assert(joins == Seq((joinClass, buildSide)), plan) + assert(collect(plan) { case c: CometSparkToColumnarExec => c }.size == 1, plan) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + plan) + QueryTest.sameRows(expected(withCondition), df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + + test("an outer hash join declares the part of the streamed order Comet can express") { + // Declaring any order keeps the probe rows in place, so a prefix of the order is enough. + // Nothing tells whether an operator above relies on the order, so a join whose first key Comet + // cannot express stays in Spark rather than lose it. 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") { + def sortedCache(order: Column*): DataFrame = { + val cached = + spark.table("big").repartition(2, $"_1").sortWithinPartitions(order: _*).cache() + cached.count() + cached + } + val build = spark.table("small") + try { + val unexpressed = sortedCache($"_1" + 1) + val (_, fallbackPlan) = checkSparkAnswerAndFallbackReason( + unexpressed.join( + build.hint("shuffle_hash"), + unexpressed("_1") === build("_1"), + "left_outer"), + "the streamed side's sort order, which this outer join has to keep, cannot be " + + "expressed natively") + assert( + collect(fallbackPlan) { case j: ShuffledHashJoinExec => j }.size == 1, + fallbackPlan) + assert(collect(fallbackPlan) { case j: CometHashJoinExec => j }.isEmpty, fallbackPlan) + + val prefixed = sortedCache($"_1", $"_1" + 1) + val df = prefixed + .join(build.hint("shuffle_hash"), prefixed("_1") === build("_1"), "left_outer") + .select(prefixed("_1").as("k"), prefixed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case j: CometHashJoinExec => j }.size == 1, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + val expected = (0 until 1000).map { i => + val k = i % 100 + Row(k, i, if (k % 10 == 0) k / 10 else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) + } 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. + 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") { + val data = (0 until 10000).map(i => (if (i % 7 == 0) None else Some(i % 100), i)) + withParquetTable(data, "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + try { + val streamed = spark + .table("big") + .repartition(2, $"_1") + .sortWithinPartitions($"_1".desc_nulls_last) + .cache() + streamed.count() + val build = spark.table("small").hint("shuffle_hash") + val expected = data.map { case (k, i) => + Row(k.getOrElse(null), i, k.filter(_ % 10 == 0).map(_ / 10).getOrElse(null)) + } + Seq("left_outer", "right_outer").foreach { joinType => + val joined = if (joinType == "left_outer") { + streamed.join(build, streamed("_1") === build("_1"), joinType) + } else { + build.join(streamed, build("_1") === streamed("_1"), joinType) + } + val df = joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions($"k".desc_nulls_last) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case j: CometHashJoinExec => j }.size == 1, plan) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + plan) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Descending)), plan) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + + // Compares the rows of a fresh `df` with the ones Spark returns for it, as a multiset unless + // `isSorted`. The rows of a key tie under a sort, so their order within it can differ between + // the two. + private def sameRowsAsSpark(df: => DataFrame, isSorted: Boolean = false): Unit = { + // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. + var sparkRows = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = df.collect().toSeq + } + QueryTest.sameRows(sparkRows, df.collect().toSeq, isSorted).foreach(fail(_)) + } + + test("outer hash joins keep the order of a native range") { + // Spark's RangeExec reports id ascending, or descending for a negative step, and the sort above + // the join is dropped on that basis. The range is a native leaf, not a JVM input, so the + // planner has to declare its order itself or the join moves its unmatched rows. A broadcast + // join, since a shuffled one would put an exchange between the range and the join. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true") { + withParquetTable((0 until 40).map(k => (k.toLong, k)), "small") { + val build = spark.table("small") + for { + (start, end, step) <- Seq((0L, 2000L, 1L), (1999L, -1L, -1L)) + joinType <- Seq("left_outer", "right_outer") + } { + val streamed = spark.range(start, end, step, 1) + val joined = if (joinType == "left_outer") { + streamed.join(build.hint("broadcast"), streamed("id") === build("_1"), joinType) + } else { + build.hint("broadcast").join(streamed, streamed("id") === build("_1"), joinType) + } + val direction = if (step > 0) Ascending else Descending + val df = joined + .select(streamed("id").as("k"), build("_2").as("w")) + .sortWithinPartitions(if (step > 0) $"k" else $"k".desc) + val plan = df.queryExecution.executedPlan + val context = s"$joinType step $step\n$plan" + assert(collect(plan) { case r: CometRangeExec => r }.size == 1, context) + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j.nativeOp } + assert(joins.size == 1, context) + assert(joins.head.getHashJoin.getStreamedSortOrdersCount > 0, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + val expected = + (start until end by step).map(id => Row(id, if (id < 40) id.toInt else null)) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert( + partitionsInOrder( + df, + SortOrder(BoundReference(0, LongType, nullable = true), direction)), + context) + } + } + } + } + + test("outer hash joins keep the order of a sorted bucketed table read by the native scan") { + // With spark.sql.legacy.bucketedTableScan.outputOrdering on, Spark reports a bucketed table's + // sort columns as the scan's order when every bucket has one file, and the sort above the + // join is dropped on that basis. The native scan reads each bucket's file in order, but it is + // a native leaf, not a JVM input, so the planner has to count it as a source whose order + // Spark reports, or the join moves its unmatched rows. A broadcast join, since a shuffled one + // would put an exchange between the scan and the join. + withTable("sorted_bucketed") { + // A single writer task writes one file per bucket. + (0 until 2000) + .map(i => (i % 40, i)) + .toDF("k", "v") + .coalesce(1) + .write + .format("parquet") + .bucketBy(2, "k") + .sortBy("k") + .saveAsTable("sorted_bucketed") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.BUCKETING_ENABLED.key -> "true", + SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "false", + SQLConf.LEGACY_BUCKETED_TABLE_SCAN_OUTPUT_ORDERING.key -> "true") { + withParquetTable((0 until 40 by 2).map(k => (k, k / 2)), "small") { + val streamed = spark.table("sorted_bucketed") + val build = spark.table("small") + val scanOrder = collectFirst(streamed.queryExecution.executedPlan) { + case scan: CometNativeScanExec if scan.bucketedScan => scan.outputOrdering + } + assert(scanOrder.exists(_.nonEmpty), streamed.queryExecution.executedPlan) + Seq("left_outer", "right_outer").foreach { joinType => + val joined = if (joinType == "left_outer") { + streamed.join(build.hint("broadcast"), streamed("k") === build("_1"), joinType) + } else { + build.hint("broadcast").join(streamed, streamed("k") === build("_1"), joinType) + } + val df = joined + .select(streamed("k"), streamed("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + val context = s"$joinType\n$plan" + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j } + assert(joins.size == 1, context) + assert( + collect(joins.head) { + case s: CometNativeScanExec if s.bucketedScan => s + }.size == 1, + context) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + val expected = (0 until 2000).map { i => + val k = i % 40 + Row(k, i, if (k % 2 == 0) k / 2 else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), context) + } + } + } + } + } + + test("outer hash joins keep the streamed order for decimal, string and timestamp keys") { + // SortedInputExec passes the rows through, so the order is physically Spark's, but the serde + // and the native sort expression handle the key, and the sorted aggregate above the join finds + // its group boundaries by comparing it. The timestamps run in a session time zone that is not + // UTC and include values before 1970. Double, array and struct keys have tests of their own, + // and Spark cannot sort by a map. + val keyTypes = Seq( + "decimal(20,2)" -> Seq( + "-999999999999999999.99", + "-0.01", + "0.00", + "0.01", + "999999999999999999.99"), + "string" -> Seq("apple", "Äpfel", "naïve", "日本語", "中文", "한국어", "😀"), + "timestamp" -> Seq( + "1900-01-01 00:00:00", + "1969-12-31 23:59:59.999999", + "1970-01-01 00:00:00", + "2026-11-01 01:30:00", + "2030-06-15 12:00:00")) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Los_Angeles", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + keyTypes.foreach { case (typeName, keys) => + val data = (0 until 2000).map(i => (keys(i % keys.size), i)) + val buildData = keys.zipWithIndex.collect { case (k, i) if i % 2 == 0 => (k, i) } + withParquetTable(data, "big") { + withParquetTable(buildData, "small") { + try { + val streamed = spark + .table("big") + .select($"_1".cast(typeName).as("_1"), $"_2") + .repartition(2, $"_1") + .sortWithinPartitions("_1") + .cache() + streamed.count() + val build = spark.table("small").select($"_1".cast(typeName).as("_1"), $"_2") + val order = SortOrder( + BoundReference(0, streamed.schema("_1").dataType, nullable = true), + Ascending) + + Seq("left_outer", "right_outer").foreach { joinType => + def joined: DataFrame = { + val hinted = build.hint("shuffle_hash") + if (joinType == "left_outer") { + streamed.join(hinted, streamed("_1") === build("_1"), joinType) + } else { + hinted.join(streamed, build("_1") === streamed("_1"), joinType) + } + } + def sorted: DataFrame = joined + .select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("k") + val df = sorted + val plan = df.queryExecution.executedPlan + val joins = collect(plan) { case j: CometHashJoinExec => j } + assert(joins.size == 1, s"$typeName $joinType\n$plan") + assert( + joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, + s"$typeName $joinType\n$plan") + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + s"$typeName $joinType\n$plan") + sameRowsAsSpark(sorted) + assert(partitionsInOrder(df, order), s"$typeName $joinType\n$plan") + + // The aggregate runs in sorted mode over the declared order, and groups a key once + // only if the join keeps it. + def agg: DataFrame = joined + .groupBy(streamed("_1").as("k")) + .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) + sameRowsAsSpark(agg) + assert(agg.count() == keys.size, s"$typeName $joinType groups a key twice") + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + } + + test("outer hash joins keep a streamed double order with signed zeros and NaN") { + // Spark sorts -0.0 and 0.0 as equal and NaN above every other value. The planner declares the + // order on the key normalized as the native sort normalizes it, so an operator above the join + // compares keys as Spark does. A top-k above the join runs in a native plan of its own, over + // the join's output, so it cannot stop early on the declared order, and one sorting by the key + // alone reads only the first rows of each partition, which holds only if the join keeps them in + // order. A broadcast join, since Spark joins on the normalized key and a shuffled join would + // shuffle the streamed side again. + val keys = + Seq(Double.NegativeInfinity, -1.5, -0.0, 0.0, 1.5, Double.PositiveInfinity, Double.NaN) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + // The build side holds -0.0 and not 0.0, so a streamed 0.0 matches only on the normalized + // key. + val buildData = Seq((Double.NegativeInfinity, 0), (-0.0, 1), (1.5, 2), (Double.NaN, 3)) + withParquetTable((0 until 2000).map(i => (keys(i % keys.size), i)), "big") { + withParquetTable(buildData, "small") { + try { + val streamed = + spark.table("big").repartition(2, $"_2").sortWithinPartitions("_1").cache() + streamed.count() + val build = spark.table("small") + val order = SortOrder(BoundReference(0, DoubleType, nullable = true), Ascending) + Seq("left_outer", "right_outer").foreach { joinType => + def joined: DataFrame = { + val hinted = build.hint("broadcast") + if (joinType == "left_outer") { + streamed.join(hinted, streamed("_1") === build("_1"), joinType) + } else { + hinted.join(streamed, build("_1") === streamed("_1"), joinType) + } + } + def rows: DataFrame = + joined.select(streamed("_1").as("k"), streamed("_2").as("v"), build("_2").as("w")) + val df = rows.sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + val context = s"$joinType\n$plan" + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j.nativeOp } + assert(joins.size == 1, context) + assert(joins.head.getHashJoin.getStreamedSortOrdersCount > 0, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + sameRowsAsSpark(rows.sortWithinPartitions("k")) + assert(partitionsInOrder(df, order), context) + + def agg: DataFrame = joined + .groupBy(streamed("_1").as("k")) + .agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))) + sameRowsAsSpark(agg) + val groupKeys = agg.collect().map(_.getDouble(0)).toSeq + assert(groupKeys.distinct.size == groupKeys.size, s"$joinType groups a key twice") + + // The first 600 rows by the key end among the zeros, which tie, so which of them + // are taken depends on v alone. The first 300 descending take every NaN, which tie + // too, and the rest from the infinities. A limit on the key alone ends where the + // zeros start, since which of them it took would be arbitrary. + sameRowsAsSpark(rows.orderBy("k", "v").limit(600), isSorted = true) + sameRowsAsSpark(rows.orderBy($"k".desc, $"v").limit(300), isSorted = true) + val belowZero = (0 until 2000).count(i => keys(i % keys.size) < -1.0) + sameRowsAsSpark(rows.select("k").orderBy("k").limit(belowZero), isSorted = true) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + + test("outer hash joins keep a streamed order that starts with an array or a struct") { + // exprToProto takes an array or a struct sort key, so a join whose streamed order starts with + // one sends it, stays native, and the planner declares it on the JVM input. Spark drops a sort + // right below a join, so the sorted rows come from a cache. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => (i % 40, i)), "big") { + withParquetTable((0 until 40 by 2).map(k => (k, k / 2)), "small") { + val build = spark.table("small") + // An array of strings, the one array type that Comet reads from Spark's cache. + Seq( + "array" -> array(($"_1" % 7).cast("string"), $"_1".cast("string")), + "struct" -> struct($"_1" % 7, $"_1")).foreach { case (keyName, key) => + try { + val streamed = spark + .table("big") + .select(key.as("s"), $"_1", $"_2") + .repartition(2, $"_1") + .sortWithinPartitions("s") + .cache() + streamed.count() + val order = SortOrder( + BoundReference(0, streamed.schema("s").dataType, nullable = true), + Ascending) + Seq("left_outer", "right_outer").foreach { joinType => + def joined: DataFrame = { + val hinted = build.hint("broadcast") + 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("s"), streamed("_2").as("v"), build("_2").as("w")) + .sortWithinPartitions("s") + val plan = df.queryExecution.executedPlan + val context = s"$keyName $joinType\n$plan" + val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j.nativeOp } + assert(joins.size == 1, context) + assert(joins.head.getHashJoin.getStreamedSortOrdersCount > 0, context) + assert( + collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, + context) + sameRowsAsSpark(df) + assert(partitionsInOrder(df, order), context) + sameRowsAsSpark( + joined + .groupBy(streamed("s")) + .agg(count(lit(1)), sum(streamed("_2")), count(build("_2")))) + } + } finally { + spark.catalog.clearCache() + } + } + } + } + } + } + + // The setup several cache tests share: `big(_1, _2)` with 2000 rows over 40 sparse keys and + // `small` with a row for every even key, and `big` as `cached(k, v)`, partitioned on the key, + // sorted within each partition and cached, with 8-row batches and the native cache on. `views` + // names the temporary views that `f` creates, which are dropped afterwards. + private def withSortedCache(views: String*)(f: => Unit): Unit = { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COLUMN_BATCH_SIZE.key -> "8", + CometConf.COMET_BATCH_SIZE.key -> "8", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)), "small") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + f + } finally { + spark.catalog.clearCache() + (views :+ "cached").foreach(spark.catalog.dropTempView) + } + } + } + } + } + + test("NOT IN below an outer hash join over a sorted cache keeps each group once") { + // Spark reports the streamed side's order for a null-aware anti join, but a native one does + // not keep it, so Comet leaves a NOT IN over an ordered side in Spark. Its order then holds + // and the outer join above may declare it. Declaring it over a native NOT IN would put the + // aggregates in sorted mode over unsorted rows and split groups. Sparse keys keep DataFusion + // off its perfect hash build path, which keeps the build batches in order and would hide + // that, and small batches spread each partition over many. + withSortedCache() { + withParquetTable(Seq(1000, 10000, 20000, 100000).map(x => Tuple1(Option(x))), "excl") { + val query = + """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v), count(b._2) + |FROM (SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM excl)) a + |LEFT JOIN small b ON a.k = b._1 + |GROUP BY a.k""".stripMargin + checkEachKeyGroupedOnce(query) + } + } + } + + test("a cache of an outer hash join over a native sample reports no order, as Spark does") { + // Spark's SampleExec reports no output ordering, so the sort above a join over a sample stays, + // and CometSampleExec reports none either. The native sample does not say it keeps its input's + // order, so the planner could not declare an order through it, and a cache that reported the + // order would carry a false claim to whatever reads it next. + withSortedCache("joined") { + sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + |FROM (SELECT * FROM cached TABLESAMPLE (90 PERCENT) REPEATABLE (7)) a + |LEFT JOIN small b ON a.k = b._1""".stripMargin) + .cache() + .createOrReplaceTempView("joined") + spark.table("joined").count() + val cachedPlan = cachedPlanOf(spark.table("joined")) + assert( + find(cachedPlan) { + case join: CometBroadcastHashJoinExec => + join.joinType == LeftOuter && join.nativeOp.getHashJoin.getStreamedSortOrdersCount == 0 + case _ => false + }.isDefined, + s"the cache was not written by a native outer join that sent no order:\n$cachedPlan") + assert( + find(cachedPlan)(_.isInstanceOf[CometSampleExec]).isDefined, + s"the outer join does not read a native sample:\n$cachedPlan") + assert( + cachedPlan.outputOrdering.isEmpty, + s"the cache reports an order that nothing keeps:\n$cachedPlan") + // The sort above the cache stays, so the partitions come back sorted. + assert( + partitionsInOrder(spark.table("joined").sortWithinPartitions("k"), keyOrder(Ascending))) + val query = + """SELECT /*+ BROADCAST(c) */ j.k, count(*), sum(j.v), count(c._2) + |FROM joined j LEFT JOIN small c ON j.k = c._1 + |GROUP BY j.k""".stripMargin + checkEachKeyGroupedOnce(query) + } + } + + test("an outer hash join over a cache of a declared outer hash join declares the order too") { + // The first outer join's probe input is the sorted cache, a JVM input that the planner + // declares the order on, so the join keeps its unmatched rows in place and the cache it + // writes is in the order it reports. The second outer join can rely on that. + withSortedCache("joined") { + sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + |FROM cached a LEFT JOIN small b ON a.k = b._1""".stripMargin) + .cache() + .createOrReplaceTempView("joined") + spark.table("joined").count() + val cachedPlan = cachedPlanOf(spark.table("joined")) + assert( + find(cachedPlan) { + case join: CometBroadcastHashJoinExec => + join.joinType == LeftOuter && + join.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0 + case _ => false + }.isDefined, + s"the cache was not written by a native outer join that sent the order:\n$cachedPlan") + assert( + partitionsInOrder(spark.table("joined"), keyOrder(Ascending)), + "the cached outer join rows are out of order, so the first join did not keep them") + val df = sql("""SELECT /*+ BROADCAST(c) */ j.k, j.v, c._2 + |FROM joined j LEFT JOIN small c ON j.k = c._1""".stripMargin) + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + val joins = collect(plan) { case join: CometBroadcastHashJoinExec => join } + assert(joins.size == 1, plan) + assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + val expected = (0 until 2000).map { i => + val m = i % 40 + Row(m * 1000, i, if (m % 2 == 0) m else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) + } + } + + test("an outer hash join declares the streamed order over a semi join") { + // `k IN (SELECT ...)` runs as a left semi join, which emits its probe rows in ascending + // order whether or not the probe input declares an order. The outer join above it therefore + // declares the cache's order on the semi join's output, and keeps it. + withSortedCache() { + withParquetTable((0 until 40).filter(_ % 7 != 3).map(k => Tuple1(k * 1000)), "incl") { + val df = sql("""SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w + |FROM (SELECT * FROM cached WHERE k IN (SELECT _1 FROM incl)) a + |LEFT JOIN small b ON a.k = b._1""".stripMargin) + .sortWithinPartitions("k") + val plan = df.queryExecution.executedPlan + // Spark also pushes the IN to the other input of the outer join, so the outer join + // is found by its type and has to read the semi join over the cache. + val outer = collect(plan) { + case join: CometBroadcastHashJoinExec if join.joinType == LeftOuter => join + } + assert(outer.size == 1, plan) + assert( + outer.head.left match { + case semi: CometBroadcastHashJoinExec => semi.joinType == LeftSemi + case _ => false + }, + plan) + assert(outer.head.nativeOp.getHashJoin.getStreamedSortOrdersCount > 0, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + val expected = (0 until 2000).filter(i => (i % 40) % 7 != 3).map { i => + val m = i % 40 + Row(m * 1000, i, if (m % 2 == 0) m else null) + } + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + assert(partitionsInOrder(df, keyOrder(Ascending)), plan) + } + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 2f7ba3896d4..a264cc416d1 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -38,17 +38,19 @@ import org.apache.parquet.schema.{MessageType, MessageTypeParser} import org.apache.spark._ import org.apache.spark.internal.config.{MEMORY_OFFHEAP_ENABLED, MEMORY_OFFHEAP_SIZE, SHUFFLE_MANAGER} import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.expressions.{BoundReference, InterpretedOrdering, SortDirection, SortOrder} import org.apache.spark.sql.catalyst.plans.logical import org.apache.spark.sql.catalyst.util.sideBySide -import org.apache.spark.sql.comet.CometPlanChecker +import org.apache.spark.sql.comet.{CometInMemoryTableScanExec, CometPlanChecker} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, QueryStageExec} +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.internal._ import org.apache.spark.sql.sources.{BaseRelation, TableScan} import org.apache.spark.sql.test._ -import org.apache.spark.sql.types.{DecimalType, StructType} +import org.apache.spark.sql.types.{DecimalType, IntegerType, StructType} import org.apache.comet._ import org.apache.comet.shims.ShimCometSparkSessionExtensions @@ -1662,4 +1664,52 @@ abstract class CometTestBase (expected.failed.toOption, actual.failed.toOption) } } + + // Whether every partition of `df` is in `order`, which is bound to its output. + protected def partitionsInOrder(df: DataFrame, order: SortOrder*): Boolean = { + val ordering = new InterpretedOrdering(order) + df.queryExecution.toRdd.map(_.copy()).glom().collect().forall { part => + part.zip(part.drop(1)).forall { case (a, b) => !ordering.gt(a, b) } + } + } + + // The order of the integer key in the first output column. + protected def keyOrder(direction: SortDirection): SortOrder = + SortOrder(BoundReference(0, IntegerType, nullable = true), direction) + + // Runs `query`, which groups by the integer key in its first column and returns three longs, + // with Comet and with Spark, and checks that Comet returns each key once and Spark's rows. + // `checkPlan` sees the plan Comet ran. + protected def checkEachKeyGroupedOnce( + query: String, + checkPlan: SparkPlan => Unit = _ => ()): Unit = { + def rows(): Seq[(Int, Long, Long, Long)] = + sql(query) + .collect() + .map(r => (r.getInt(0), r.getLong(1), r.getLong(2), r.getLong(3))) + .toSeq + .sorted + val cometRows = rows() + checkPlan(sql(query).queryExecution.executedPlan) + // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned inside it. + var sparkRows = Seq.empty[(Int, Long, Long, Long)] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkRows = rows() + } + assert(cometRows.map(_._1).distinct.size == cometRows.size, "a key is grouped twice") + assert(cometRows == sparkRows) + } + + // The physical plan that wrote the cache `df` reads, as Spark's scan reports its order: for an + // adaptive cached plan that is the final plan behind it. + protected def cachedPlanOf(df: DataFrame): SparkPlan = { + val cachedPlan = collectFirst(df.queryExecution.executedPlan) { + case scan: CometInMemoryTableScanExec => scan.originalPlan.relation.cachedPlan + case scan: InMemoryTableScanExec => scan.relation.cachedPlan + }.getOrElse(fail(s"no cache scan in\n${df.queryExecution.executedPlan}")) + cachedPlan match { + case adaptive: AdaptiveSparkPlanExec if adaptive.isFinalPlan => adaptive.executedPlan + case other => other + } + } }