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/planner.rs b/native/core/src/execution/planner.rs index 0241e1f1ce8..e6b451933cf 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4270,6 +4270,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; @@ -4337,24 +4379,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, @@ -4376,9 +4416,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 { @@ -4433,6 +4484,53 @@ mod tests { ); } + /// `func` over a window partitioned and ordered by `col_0` over `input`. + fn partitioned_window(func: &str, input: Operator, input_sort_orders: Vec) -> Operator { + Operator { + children: vec![input], + 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() + } + } + + #[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 (_, _, planned) = PhysicalPlanner::default() + .create_plan( + &partitioned_window(func, create_scan(), input_sort_orders), + &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 0b91f27187d..6de4f3821ac 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, @@ -402,27 +403,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 @@ -430,8 +435,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), } } @@ -439,11 +450,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::() { 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 406592cbce8..1fb6fd02dc1 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1233,6 +1233,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, then the order keys, 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 69fa6e371ae..1b0caae361d 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] { @@ -81,10 +81,37 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { if (windowExprProto.forall(_.isDefined) && partitionExprs.forall(_.isDefined) && sortOrders.forall(_.isDefined)) { + // Spark plans no sort below the window when its input already reports the order the window + // needs, but rows from the JVM arrive through a ScanExec that declares none. Send that order + // for the planner to declare (see SortedInputExec), unless a native sort below declares it + // or it may be false. A key Comet cannot express, or a collated string key, sends the window + // back to Spark. + val requiredOrdering = op.requiredChildOrdering.head + val inputSortOrders = + if (requiredOrdering.nonEmpty && !op.child.isInstanceOf[CometSortExec] && + SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) && + !CometHashJoinExec.mayReportFalseInputOrder(op.child)) { + val orders = requiredOrdering.map { order => + if (hasNonDefaultStringCollation(order.dataType)) None + else exprToProto(order, op.child.output) + } + if (!orders.forall(_.isDefined)) { + withFallbackReason( + op, + "the order of the window's input, which Spark relies on, cannot be expressed " + + "natively") + return None + } + orders.flatten + } 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 e7c8140bf5f..f6cf8697c8f 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 @@ -3017,6 +3017,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 2e1f4d56d33..3887f64891f 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} @@ -3692,6 +3692,157 @@ 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 order (#7223). A nullable key and a descending order key check the direction and null + // placement. 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", "1 Descending NullsLast"))) + } + } + } finally { + spark.catalog.clearCache() + } + } + } + } + + test("stacked native windows over a sorted cache") { + // Spark plans no sort between two windows that the cache's order satisfies, so the upper + // window's input, which declares only the lower window's order, gets the rest declared next + // to it (#7223). Alias a of sort key k shares its order. One partition, so no exchange + // separates the windows. + 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"))) + } 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 side would declare a collated key in raw byte order, not Spark's collation + // order, so the window falls back. + withConversions( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + try { + val cached = 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 AS v") + .repartition(2, col("k")) + .sortWithinPartitions("k", "v") + .cache() + cached.count() + + val df = cached.selectExpr("k", "v", "dense_rank() OVER (PARTITION BY k ORDER BY v) AS r") + // UTF8_LCASE puts 'a' and 'A' in one partition, so each of the four keys holds 50 rows. + val rows = (0L until 200L).map { v => + Row(s"${v % 4}${if (v % 8 < 4) "a" else "A"}", v, (v / 4 + 1).toInt) + } + QueryTest.sameRows(rows, 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(_.contains("the order of the window's input")), 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..d1ed4f4600c 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 == 2, plan) + } + } + } }