diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index 0c36921a94a..ed48f4fba0f 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -365,17 +365,17 @@ the order satisfies. Report what the Spark operator reports and nothing more. `S 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`. The hash join serde then sends no order over it, and an outer hash join above it -that has to keep that order falls back to Spark. Comet's hash joins are the exception within one -native plan: an outer hash join over one of them stays native, even over one that declared no -order, since DataFusion keeps the order it can see there. +its streamed side to DataFusion, and a native window the order of its partition keys. An aggregate +above 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`. The hash join and window serdes then send no order over +it, and an outer hash join above it that has to keep that order falls back to Spark. Comet's hash +joins are the exception within one native plan: an outer hash join over one of them stays native, +even over one that declared no order, since DataFusion keeps the order it can see there. A leaf that reports an ordering is where the order starts, so the planner has to know that the native operator produces its rows in that order. It knows that of an input from the JVM and of a diff --git a/docs/source/contributor-guide/debugging.md b/docs/source/contributor-guide/debugging.md index 8795854193d..51d74a1ea94 100644 --- a/docs/source/contributor-guide/debugging.md +++ b/docs/source/contributor-guide/debugging.md @@ -260,7 +260,8 @@ so `spark.comet.debug.enabled` and `COMET_LOG_LEVEL` do not affect it. #### Example: dumping window operator output In `WindowBuilder::build` in `native/core/src/execution/planner/window.rs`, -the window operator builds `window_agg`, either a `BoundedWindowAggExec` or a +the window operator builds `window_agg` over `input`, the child plan with +Spark's input order declared, either a `BoundedWindowAggExec` or a `WindowAggExec` depending on whether every window expression runs with bounded memory. It then optionally wraps `window_agg` in a casting `ProjectionExec` and builds the `SparkPlan` from either the projection or `window_agg`. To see the raw @@ -271,14 +272,14 @@ projection: let window_agg: Arc = if all_bounded { Arc::new(BoundedWindowAggExec::try_new( window_expr, - Arc::clone(&child.native_plan), + input, InputOrderMode::Sorted, !partition_exprs.is_empty(), )?) } else { Arc::new(WindowAggExec::try_new( window_expr, - Arc::clone(&child.native_plan), + input, !partition_exprs.is_empty(), )?) }; diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs index a31053bd6e8..ad561cf49e3 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. A window rejects an input +/// that does not declare its partition keys' order. The planner puts this only on that join's +/// probe input or that window's input, with the order Spark reports for it, and only when the +/// rows reach that input from the JVM through operators that keep their order. Both keep the +/// order in their own output, so the operators above them in the native plan see it too: an +/// aggregate grouped by the key runs in sorted, streaming mode, a sort aggregate takes the order +/// as already satisfied, and a top-k sort whose keys start with it stops reading early. A wrong +/// order would therefore give wrong groups or wrong rows, not only rows out of order. #[derive(Debug)] pub struct SortedInputExec { input: Arc, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 603bba872bc..8ea033d11da 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4163,8 +4163,10 @@ mod tests { "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. + // The join keeps the order in its own output, so the operators above it see it too, + // and an order declared above it holds, such as a window's for an alias of the key. assert!(planned.native_plan.properties().output_ordering().is_some()); + assert!(PhysicalPlanner::keeps_source_order(&planned.native_plan)); } } @@ -4332,6 +4334,48 @@ mod tests { assert!(!PhysicalPlanner::keeps_source_order(&scan(2))); } + #[test] + fn declared_order_extends_an_order_the_input_declares() { + // An order is declared over an empty bucket, and next to the order the input declares, + // unless an operator in between may have moved rows. + let schema = Arc::new(Schema::new(vec![ + Field::new("col_0", DataType::Int32, true), + Field::new("col_1", DataType::Int32, true), + ])); + let planner = PhysicalPlanner::default(); + let declare = |plan: &Arc, index: i32| { + planner + .with_declared_order( + Arc::clone(plan), + &[create_sort_order(index, 3, false, true)], + ) + .unwrap() + }; + let empty: Arc = + Arc::new(datafusion::physical_plan::empty::EmptyExec::new(schema)); + let by_col_0 = declare(&empty, 0); + assert_eq!("SortedInputExec", by_col_0.name()); + assert!(Arc::ptr_eq(&declare(&by_col_0, 0), &by_col_0)); + + let both = declare(&by_col_0, 1); + assert!(Arc::ptr_eq(both.children()[0], &by_col_0)); + for index in 0..2 { + let partition_by: Vec> = + vec![Arc::new(Column::new(&format!("col_{index}"), index))]; + let ordered = get_ordered_partition_by_indices(&partition_by, &both).unwrap(); + assert_eq!(1, ordered.len(), "col_{index}"); + } + + let sort: Arc = Arc::new(SortExec::new( + LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(Column::new( + "col_0", 0, + )))]) + .unwrap(), + empty, + )); + assert!(Arc::ptr_eq(&declare(&sort, 1), &sort)); + } + #[tokio::test] async fn spark_plan_dynamic_filter_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; @@ -4399,24 +4443,22 @@ mod tests { } } - #[tokio::test] - async fn window_cast_projection_reports_window_time() { - use crate::execution::metrics::utils::to_native_metric_node; + /// `func() OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)`, returning + /// `result_type`. + fn window_function_expr( + func: &str, + result_type: Option, + ) -> spark_operator::WindowExpr { use spark_operator::{ lower_window_frame_bound::LowerFrameBoundStruct, upper_window_frame_bound::UpperFrameBoundStruct, CurrentRow, LowerWindowFrameBound, UnboundedPreceding, UpperWindowFrameBound, WindowExpr, WindowFrame, WindowFrameType, WindowSpecDefinition, }; - - let session = Arc::new(SessionContext::new()); - let planner = PhysicalPlanner::new(Arc::clone(&session), 0); - // `row_number` returns UInt64 natively and Spark expects Int32, so the planner adds a - // casting projection above the window. - let row_number = WindowExpr { + WindowExpr { built_in_window_function: Some(Expr { expr_struct: Some(ScalarFunc(spark_expression::ScalarFunc { - func: "row_number".to_string(), + func: func.to_string(), ..Default::default() })), query_context: None, @@ -4438,9 +4480,20 @@ mod tests { }), ..Default::default() }), - result_type: Some(create_proto_datatype()), + result_type, ..Default::default() - }; + } + } + + #[tokio::test] + async fn window_cast_projection_reports_window_time() { + use crate::execution::metrics::utils::to_native_metric_node; + + let session = Arc::new(SessionContext::new()); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + // `row_number` returns UInt64 natively and Spark expects Int32, so the planner adds a + // casting projection above the window. + let row_number = window_function_expr("row_number", Some(create_proto_datatype())); let operator = Operator { children: vec![create_scan()], op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window { @@ -4495,6 +4548,45 @@ mod tests { ); } + #[test] + fn window_declares_input_order_on_a_scan() { + // Both window operators need their partition keys' order, which a `ScanExec` declares + // only through the order Spark sent. `percent_rank` runs in `WindowAggExec`. + let partition_by: Vec> = vec![Arc::new(Column::new("col_0", 0))]; + for (func, operator) in [ + ("row_number", "BoundedWindowAggExec"), + ("percent_rank", "WindowAggExec"), + ] { + let window_input = |input_sort_orders| { + let window = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window { + window_expr: vec![window_function_expr(func, None)], + partition_by_list: vec![create_bound_reference(0)], + order_by_list: vec![create_sort_order(0, 3, false, true)], + input_sort_orders, + ..Default::default() + }))), + ..Default::default() + }; + let (_, _, planned) = PhysicalPlanner::default() + .create_plan(&window, &mut vec![], 1) + .unwrap(); + assert_eq!(operator, planned.native_plan.name()); + Arc::clone(planned.native_plan.children()[0]) + }; + let input = window_input(vec![create_sort_order(0, 3, false, true)]); + assert_eq!("SortedInputExec", input.name()); + let ordered = get_ordered_partition_by_indices(&partition_by, &input).unwrap(); + assert_eq!(1, ordered.len(), "{func}"); + + let input = window_input(vec![]); + assert_eq!("ScanExec", input.name()); + let ordered = get_ordered_partition_by_indices(&partition_by, &input).unwrap(); + assert!(ordered.is_empty(), "{func}"); + } + } + #[test] fn projection_rolls_up_same_plan_id_child_metrics() { let op_scan = Operator { diff --git a/native/core/src/execution/planner/join.rs b/native/core/src/execution/planner/join.rs index b4ed291eaf2..ffefd158e92 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -35,6 +35,7 @@ use datafusion::{ LexOrdering, PhysicalExpr, }, physical_plan::{ + empty::EmptyExec, filter::FilterExec, joins::{ utils::JoinFilter, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, @@ -478,27 +479,31 @@ 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. + /// `plan`, declaring `sort_orders` as the order its rows arrive in, unless the order it + /// declares already satisfies them. 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( + /// An input that declares only part of `sort_orders`, such as a window over another window, + /// gets them declared next to its own order. + /// + /// The serde sends an order only when the Comet operators of a join's streamed side or of a + /// window's input 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, or the window rejects its input. Nothing here can fall back at + /// that point, so the invariant is kept at the operator: see "Adding a New Operator" in the + /// contributor guide. + pub(super) fn with_declared_order( &self, plan: Arc, sort_orders: &[Expr], ) -> Result, ExecutionError> { - if plan.output_ordering().is_some() || !Self::keeps_source_order(&plan) { + if sort_orders.is_empty() || !Self::keeps_source_order(&plan) { return Ok(plan); } let exprs = sort_orders @@ -506,8 +511,14 @@ impl PhysicalPlanner { .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), + Some(ordering) + if !plan + .equivalence_properties() + .ordering_satisfy(ordering.iter().cloned())? => + { + Ok(Arc::new(SortedInputExec::new(plan, ordering))) + } + _ => Ok(plan), } } @@ -515,11 +526,12 @@ impl PhysicalPlanner { /// 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. + /// sorted file, as for a sorted bucketed table with one file per bucket. An `EmptyExec`, a + /// bucket with no file, has no rows to be out of order. 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. pub(super) fn keeps_source_order(plan: &Arc) -> bool { - if plan.is::() || plan.is::() { + if plan.is::() || plan.is::() || plan.is::() { return true; } if let Some(scan) = plan.downcast_ref::() { @@ -529,18 +541,21 @@ impl PhysicalPlanner { .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. + // join keeps its unmatched probe rows in place only when the probe declares an order, as + // the planner does for the join's streamed side. 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(), + let keeps_probe_order = match join.join_type() { DFJoinType::Inner - | DFJoinType::RightSemi - | DFJoinType::RightAnti - | DFJoinType::RightMark - ) { + | DFJoinType::RightSemi + | DFJoinType::RightAnti + | DFJoinType::RightMark => true, + DFJoinType::Right => join.right().output_ordering().is_some(), + _ => false, + }; + if !keeps_probe_order { return false; } } diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs index 2d15fdc86e1..9d24cea2f8b 100644 --- a/native/core/src/execution/planner/window.rs +++ b/native/core/src/execution/planner/window.rs @@ -131,17 +131,21 @@ impl OperatorBuilder for WindowBuilder { // trigger a retract call. let window_expr = window_expr?; let all_bounded = window_expr.iter().all(|e| e.uses_bounded_memory()); + + // Both window operators reject an input that does not declare its partition keys' order. + let input = + planner.with_declared_order(Arc::clone(&child.native_plan), &wnd.input_sort_orders)?; let window_agg: Arc = if all_bounded { Arc::new(BoundedWindowAggExec::try_new( window_expr, - Arc::clone(&child.native_plan), + input, InputOrderMode::Sorted, !partition_exprs.is_empty(), )?) } else { Arc::new(WindowAggExec::try_new( window_expr, - Arc::clone(&child.native_plan), + input, !partition_exprs.is_empty(), )?) }; diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 28d94e78bf3..e99b91e6cc9 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1242,6 +1242,9 @@ message Window { repeated spark.spark_expression.Expr order_by_list = 2; repeated spark.spark_expression.Expr partition_by_list = 3; Operator child = 4; + // The partition keys ascending, when Spark planned no sort below the window because its input + // reports that order. Empty otherwise. + repeated spark.spark_expression.Expr input_sort_orders = 5; } // Top-K rows per partition group. Corresponds to Spark's WindowGroupLimitExec diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala index a162d93a098..192bdf444bf 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala @@ -36,7 +36,7 @@ import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometSparkSessionExtensions.withFallbackReason import org.apache.comet.serde.{AggSerde, CometOperatorSerde, CometSortOrder, LiteralOuterClass, OperatorOuterClass} import org.apache.comet.serde.OperatorOuterClass.Operator -import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, liftFallbackReasons, scalarFunctionExprToProto, serializeDataType} +import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, liftFallbackReasons, scalarFunctionExprToProto, serializeDataType} object CometWindowExec extends CometOperatorSerde[WindowExec] { @@ -75,16 +75,58 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { } } + // Spark plans no sort below the window when its input already reports the order the window + // needs. That order is relied on, by declaring it natively or by handing the rows to Spark's + // window, only when mayReportFalseOrder trusts it: Spark's window would group the rows of a + // false order wrongly, while the native one rejects an input whose order it cannot see. + val requiredOrdering = op.requiredChildOrdering.head + val trustsInputOrder = + SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) && + !CometHashJoinExec.mayReportFalseInputOrder(op.child) + + // The native window finds partitions and peers by comparing keys byte by byte, so it would + // tell apart collated strings that Spark treats as equal, as CometWindowGroupLimitExec notes. + val collated = (op.partitionSpec ++ op.orderSpec.map(_.child)) + .filter(key => hasNonDefaultStringCollation(key.dataType)) + if (collated.nonEmpty && trustsInputOrder) { + withFallbackReason( + op, + collated + .map(_.sql) + .mkString("Window: non-default string collation on key(s): ", ", ", "")) + return None + } + val partitionExprs = op.partitionSpec.map(exprToProto(_, op.child.output)) val sortOrders = op.orderSpec.map(exprToProto(_, op.child.output)) if (windowExprProto.forall(_.isDefined) && partitionExprs.forall(_.isDefined) && sortOrders.forall(_.isDefined)) { + // Rows from the JVM arrive through a ScanExec that declares no order. Send the partition + // keys' order, the only one the native window checks, for the planner to declare (see + // SortedInputExec). + val inputSortOrders = + if (op.partitionSpec.nonEmpty && trustsInputOrder) { + val partitionKeys = requiredOrdering.take(op.partitionSpec.length) + val orders = partitionKeys.flatMap(exprToProto(_, op.child.output)) + if (orders.length < partitionKeys.length) { + withFallbackReason( + op, + "the order of the window's input, which Spark relies on, cannot be expressed " + + "natively") + return None + } + orders + } else { + Nil + } + val windowBuilder = OperatorOuterClass.Window.newBuilder() windowBuilder.addAllWindowExpr(windowExprProto.map(_.get).asJava) windowBuilder.addAllPartitionByList(partitionExprs.map(_.get).asJava) windowBuilder.addAllOrderByList(sortOrders.map(_.get).asJava) + windowBuilder.addAllInputSortOrders(inputSortOrders.asJava) Some(builder.setWindow(windowBuilder).build()) } else { // Whichever of the window / partition / order expressions failed has already recorded its diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index dd7c59f6556..a043910e82c 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 @@ -3033,6 +3033,10 @@ trait CometHashJoin { } } + // mayReportFalseOrder for a window's input, checked before CometWindowExec sends its order. + private[comet] def mayReportFalseInputOrder(plan: SparkPlan): Boolean = + mayReportFalseOrder(plan, countInPlan = true) + // 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 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 59cede8d87e..8c9d68d01eb 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -40,7 +40,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Ascending, Attribute, Att import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.{LeftOuter, RightOuter} import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} -import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoin, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoin, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometWindowExec} 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 @@ -49,7 +49,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffl import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} -import org.apache.spark.sql.functions.{count, lit, max, min, sum} +import org.apache.spark.sql.functions.{col, count, lit, max, min, sum} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} @@ -3749,6 +3749,177 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + // The order a native window declares on its input, one "index direction nulls" per key. + private def windowInputOrder(window: CometWindowExec): Seq[String] = + window.nativeOp.getWindow.getInputSortOrdersList.asScala.toSeq.map { expr => + val order = expr.getSortOrder + s"${order.getChild.getBound.getIndex} ${order.getDirection} ${order.getNullOrdering}" + } + + // Checks `df`'s rows, the order each native window declares, the top one first, and that no + // sort is planned. + private def checkWindows(df: DataFrame, rows: Seq[Row], orders: Seq[Seq[String]]) = { + QueryTest.sameRows(rows, df.collect().toSeq).foreach(fail(_)) + val plan = df.queryExecution.executedPlan + withClue(s"$plan\n") { + assert(collect(plan) { case w: CometWindowExec => w }.map(windowInputOrder) == orders) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) + } + } + + test("a native window over a sorted cache declares the cache's order") { + // Spark plans no sort below a window over the cache's sort keys, so the scan has to declare + // the partition key's order (#7223). The key is nullable, so the null placement has to match, + // and the descending order key is not declared. dense_rank runs in BoundedWindowAggExec and + // the UNBOUNDED FOLLOWING frame in WindowAggExec. Within a key, v = k + 20 * j for j in + // 0 until 100, and a NULL key holds the multiples of 20. + val key = (v: Long) => if (v % 20 == 0) null else v % 20 + val queries = Seq( + "dense_rank() OVER (PARTITION BY k ORDER BY v DESC)" -> + ((v: Long) => (100 - v / 20).toInt), + "sum(v) OVER (PARTITION BY k ORDER BY v DESC " + + "ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING)" -> { (v: Long) => + val j = v / 20 + (j + 1) * (v % 20) + 10 * j * (j + 1) + }) + Seq(true, false).foreach { nativeCache => + withConversions( + 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 -> nativeCache.toString) { + try { + val cached = spark + .range(2000) + .selectExpr("CASE WHEN id % 20 = 0 THEN NULL ELSE id % 20 END AS k", "id AS v") + .repartition(2, col("k")) + .sortWithinPartitions(col("k"), col("v").desc) + .cache() + cached.count() + + queries.foreach { case (window, expected) => + withClue(s"nativeCache=$nativeCache, $window: ") { + checkWindows( + cached.selectExpr("k", "v", s"$window AS w"), + (0L until 2000L).map(v => Row(key(v), v, expected(v))), + Seq(Seq("0 Ascending NullsFirst"))) + } + } + } finally { + spark.catalog.clearCache() + } + } + } + } + + test("a native window declares its order next to the order its input declares") { + // Spark plans no sort above an operator that keeps the cache's order, so a window's input may + // declare only the order that operator needed: an upper window's input the lower window's, + // and an outer join's output its streamed side's. The window's order is declared next to it + // (#7223). Alias a of sort key k shares its order. One partition, so no exchange comes in + // between. + val ids = 0L until 24L + val key = (id: Long) => if (id % 4 == 0) None else Some(id % 3) + val nullable = (k: Option[Long]) => k.map(Long.box).orNull + val keySums = ids.groupBy(key).map { case (k, group) => k -> group.sum } + val keyCounts = ids.groupBy(key).map { case (k, group) => k -> group.size } + val pairCounts = + ids.groupBy(id => (key(id), id % 2)).map { case (p, group) => p -> group.size } + withConversions( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + try { + val cached = spark + .range(0, 24, 1, 1) + .selectExpr( + "CASE WHEN id % 4 = 0 THEN NULL ELSE id % 3 END AS k", + "id % 2 AS x", + "id AS v") + .sortWithinPartitions("k", "x") + .selectExpr("k", "k AS a", "x", "v") + .cache() + cached.count() + + checkWindows( + cached + .selectExpr("k", "x", "v", "sum(v) OVER (PARTITION BY k) AS s") + .selectExpr("k", "x", "v", "sum(s) OVER (PARTITION BY k, x) AS t"), + ids.map { id => + val k = key(id) + Row(nullable(k), id % 2, id, keySums(k) * pairCounts((k, id % 2))) + }, + Seq( + Seq("0 Ascending NullsFirst", "1 Ascending NullsFirst"), + Seq("0 Ascending NullsFirst"))) + checkWindows( + cached + .selectExpr("k", "a", "v", "sum(v) OVER (PARTITION BY k) AS s") + .selectExpr("k", "a", "v", "sum(s) OVER (PARTITION BY a) AS t"), + ids.map { id => + val k = key(id) + Row(nullable(k), nullable(k), id, keySums(k) * keyCounts(k)) + }, + Seq(Seq("1 Ascending NullsFirst"), Seq("0 Ascending NullsFirst"))) + // The join declares the order of a, which Spark reports in place of k, so the window by k + // needs its order declared next to the join's. The residual condition leaves rows of key 0 + // unmatched among matched ones, so the key stays together only if the join keeps its + // unmatched probe rows in place. + val one = spark.range(1).toDF("j").hint("broadcast") + checkWindows( + cached + .join(one, col("k") === col("j") && col("v") % 2 === 0, "left_outer") + .selectExpr("k", "a", "v", "sum(v) OVER (PARTITION BY k) AS s"), + ids.map { id => + val k = key(id) + Row(nullable(k), nullable(k), id, keySums(k)) + }, + Seq(Seq("0 Ascending NullsFirst"))) + } finally { + spark.catalog.clearCache() + } + } + } + + test("a window over a cache sorted by a collated string key stays in Spark") { + assume(isSpark40Plus, "collated string types require Spark 4.0+") + // The native window compares keys byte by byte, so it would split the UTF8_LCASE partition + // of '0a' and '0A' and rank them apart. + val key = (v: Long) => s"${v % 4}${if (v % 8 < 4) "a" else "A"}" + withConversions( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + try { + val source = spark + .range(200) + .selectExpr( + "CAST(id % 4 AS STRING) || CASE WHEN id % 8 < 4 THEN 'a' ELSE 'A' END " + + "COLLATE UTF8_LCASE AS k", + "id % 2 AS n", + "id AS v") + val byKey = source.repartition(2, col("k")).sortWithinPartitions("k", "v").cache() + val byGroup = source.repartition(2, col("n")).sortWithinPartitions("n", "k").cache() + byKey.count() + byGroup.count() + + Seq( + // Each of the four keys holds 50 rows. + byKey.selectExpr("k", "v", "dense_rank() OVER (PARTITION BY k ORDER BY v) AS r") -> + ((v: Long) => Row(key(v), v, (v / 4 + 1).toInt)), + // Each group holds two keys, '0' and '2' or '1' and '3'. + byGroup.selectExpr("n", "k", "dense_rank() OVER (PARTITION BY n ORDER BY k) AS r") -> + ((v: Long) => Row(v % 2, key(v), (v % 4 / 2 + 1).toInt))).foreach { case (df, row) => + QueryTest.sameRows((0L until 200L).map(row), df.collect().toSeq).foreach(fail(_)) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case w: CometWindowExec => w }.isEmpty, plan) + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert(reasons.exists(_.startsWith("Window: non-default string collation")), reasons) + } + } 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 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..5d09646d61e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -27,7 +27,7 @@ import org.scalatest.Tag import org.apache.hadoop.fs.Path import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Divide, Expression, MakeDecimal, WindowExpression} -import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeExec, CometSortExec, CometWindowExec, CometWindowGroupLimitExec} +import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeExec, CometNativeScanExec, CometSortExec, CometWindowExec, CometWindowGroupLimitExec} import org.apache.spark.sql.execution.{InputAdapter, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.window.{WindowExec => SparkWindowExec} import org.apache.spark.sql.expressions.Window @@ -2073,4 +2073,40 @@ class CometWindowExecSuite extends CometTestBase { } } } + + test("window over a sorted bucketed table with empty buckets runs natively") { + // With legacy bucketed output ordering on, Spark plans no sort below the window. Two keys fill + // at most two of the eight buckets, and each empty one has to declare the order too (#7223). + withTable("sorted_bucketed") { + // A single writer task writes one file per bucket. + (0 until 40) + .map(i => (i % 2, i)) + .toDF("k", "v") + .coalesce(1) + .write + .format("parquet") + .bucketBy(8, "k") + .sortBy("k", "v") + .saveAsTable("sorted_bucketed") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.BUCKETING_ENABLED.key -> "true", + SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "false", + SQLConf.LEGACY_BUCKETED_TABLE_SCAN_OUTPUT_ORDERING.key -> "true") { + val df = sql( + "SELECT k, v, dense_rank() OVER (PARTITION BY k ORDER BY v) AS r " + + "FROM sorted_bucketed") + val (_, plan) = checkSparkAnswerAndOperator(df) + val windows = collect(plan) { case w: CometWindowExec => w } + assert(windows.size == 1, plan) + assert( + windows.head.child match { + case scan: CometNativeScanExec => scan.bucketedScan + case _ => false + }, + plan) + assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 1, plan) + } + } + } }