From 68bff2e2c48d575d3515807c00ad89b29624841a Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 27 Oct 2026 08:35:29 -0700 Subject: [PATCH 1/7] fix: declare the input order of a native window over a sorted cache Spark puts no sort below a window when its input reports the order the window needs, a relation sorted before it was cached for instance. When those rows enter the native plan from the JVM, the ScanExec they arrive through declares no order, and the native window failed the task with "Expects PARTITION BY expression to be ordered". Send the window's required ordering in a new input_sort_orders field on the Window proto when the child reports an order that satisfies it and mayReportFalseOrder trusts that order. The planner declares it on the window's input with with_declared_order, as the hash join does for its probe input. A key that Comet cannot express, or that holds a collated string, sends the window back to Spark. Closes #7223. --- native/core/src/execution/planner.rs | 106 ++++++++++++++++-- native/core/src/execution/planner/join.rs | 21 ++-- native/core/src/execution/planner/window.rs | 13 ++- native/proto/src/proto/operator.proto | 4 + .../spark/sql/comet/CometWindowExec.scala | 33 +++++- .../apache/spark/sql/comet/operators.scala | 5 + .../comet/exec/CometInMemoryCacheSuite.scala | 42 ++++++- 7 files changed, 198 insertions(+), 26 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 0241e1f1ce8..1fd03d0bf7d 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4337,21 +4337,18 @@ mod tests { } } - #[tokio::test] - async fn window_cast_projection_reports_window_time() { - use crate::execution::metrics::utils::to_native_metric_node; + /// `row_number() OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)`, returning + /// `result_type`. + fn row_number_window_expr( + 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(), @@ -4376,9 +4373,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 = row_number_window_expr(Some(create_proto_datatype())); let operator = Operator { children: vec![create_scan()], op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window { @@ -4433,6 +4441,82 @@ mod tests { ); } + /// A window partitioned and ordered by `col_0` over `input`, declaring `input_sort_orders` + /// for it. + fn partitioned_window(input: Operator, input_sort_orders: Vec) -> Operator { + Operator { + children: vec![input], + op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window { + window_expr: vec![row_number_window_expr(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() { + // Spark puts no sort below a window whose input reports the order it needs, and the rows + // then arrive through a `ScanExec`, which declares none. The window finds its partition + // keys' order only once the planner declares the order Spark sent. + let window_over = |input_sort_orders| { + let (_, _, planned) = PhysicalPlanner::default() + .create_plan( + &partitioned_window(create_scan(), input_sort_orders), + &mut vec![], + 1, + ) + .unwrap(); + let window = planned + .native_plan + .downcast_ref::() + .expect("a bounded window"); + (Arc::clone(window.input()), window.partition_by_sort_keys()) + }; + let (input, keys) = window_over(vec![create_sort_order(0, 3, false, true)]); + let sorted = input + .downcast_ref::() + .expect("the window's input declares the order Spark sent"); + assert_eq!("ScanExec", sorted.children()[0].name()); + assert_eq!( + "col_0@0 ASC", + input.properties().output_ordering().unwrap().to_string() + ); + assert_eq!(1, keys.unwrap().len()); + + let (input, keys) = window_over(vec![]); + assert_eq!("ScanExec", input.name()); + assert!(keys + .unwrap_err() + .to_string() + .contains("Expects PARTITION BY expression to be ordered")); + } + + #[test] + fn window_keeps_input_that_declares_an_order() { + // An input that already declares an order, a sort in the native plan for instance, 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 (_, _, planned) = PhysicalPlanner::default() + .create_plan( + &partitioned_window(sort, vec![create_sort_order(0, 3, false, true)]), + &mut vec![], + 1, + ) + .unwrap(); + assert_eq!("SortExec", planned.native_plan.children()[0].name()); + } + #[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..7ed0fd8459a 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -408,16 +408,17 @@ impl PhysicalPlanner { /// 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( + /// 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], diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs index 2d15fdc86e1..69c4301dc18 100644 --- a/native/core/src/execution/planner/window.rs +++ b/native/core/src/execution/planner/window.rs @@ -131,17 +131,26 @@ 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()); + + // Spark put no sort below this window when its input reports the order the window + // needs, and both window operators reject an input that does not declare its partition + // keys' order; see `with_declared_order`. + let input = if wnd.input_sort_orders.is_empty() { + Arc::clone(&child.native_plan) + } else { + 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..c21395721f7 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1233,6 +1233,10 @@ message Window { repeated spark.spark_expression.Expr order_by_list = 2; repeated spark.spark_expression.Expr partition_by_list = 3; Operator child = 4; + // The order the window needs from its input, the partition keys ascending followed by the + // order keys, as `SortOrder` expressions over the input's output, when Spark reports that the + // input arrives in it and so put no sort below the window. 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..70edb537901 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] { @@ -79,12 +79,43 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { val sortOrders = op.orderSpec.map(exprToProto(_, op.child.output)) + // Spark puts no sort below the window when its input reports the order the window needs, a + // relation sorted before it was cached for instance. When those rows enter the native plan + // from the JVM, the ScanExec they arrive through declares no order, and the native window + // then rejects its input. Send the order so the planner declares it on the window's input; + // see SortedInputExec. Nothing is declared when the order the input reports may be false; see + // CometHashJoin.mayReportFalseOrder. A key that Comet cannot express, or that holds a + // collated string, which Spark orders by its collation but the native side would declare in + // raw byte order, sends the window back to Spark. + val requiredOrdering = op.requiredChildOrdering.head + val inputSortOrders = + if (requiredOrdering.nonEmpty && + SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) && + !CometHashJoinExec.mayReportFalseInputOrder(op.child)) { + val orders = + if (requiredOrdering.exists(order => hasNonDefaultStringCollation(order.dataType))) { + Nil + } else { + requiredOrdering.map(exprToProto(_, op.child.output)) + } + if (orders.isEmpty || !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 + } + if (windowExprProto.forall(_.isDefined) && partitionExprs.forall(_.isDefined) && sortOrders.forall(_.isDefined)) { 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..3e94a14b5e6 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,11 @@ trait CometHashJoin { } } + // Whether the order that the input of a window reports may be false; see mayReportFalseOrder. + // CometWindowExec asks before it sends that order for the native window to declare. + 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..e433a40506d 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,44 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("a native window over a sorted cache declares the cache's order") { + // The cache reports the order it was written in, so Spark puts no sort below a window + // partitioned and ordered by the cache's sort keys. The native window then reads the cache + // through a scan that declares no order, and rejected it with "Expects PARTITION BY + // expression to be ordered" until the window declared the order Spark sent (#7223). + 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("id % 20 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") + val expected = (0L until 2000L).map(v => Row(v % 20, v, (v / 20 + 1).toInt)) + QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) + val plan = df.queryExecution.executedPlan + withClue(s"nativeCache=$nativeCache\n$plan\n") { + val windows = collect(plan) { case w: CometWindowExec => w } + assert(windows.size == 1) + assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 2) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) + } + } 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 From eb0a999bd06c825dfe9652d0b19a4cffb942ca0f Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 27 Oct 2026 08:53:30 -0700 Subject: [PATCH 2/7] refactor: build the window's input order only where it is used - Send no input order when a native sort below the window declares its own. - Build the order only once the window's own expressions have serialized, so a window that falls back for another reason does not walk the plan or record an extra reason. - Map a collated key straight to None instead of signalling it with an empty list. - Drop the empty-list branch in WindowBuilder, as with_declared_order already leaves the plan unchanged for no orders. --- native/core/src/execution/planner/window.rs | 7 +-- .../spark/sql/comet/CometWindowExec.scala | 58 +++++++++---------- 2 files changed, 31 insertions(+), 34 deletions(-) diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs index 69c4301dc18..ca0c64c4fb9 100644 --- a/native/core/src/execution/planner/window.rs +++ b/native/core/src/execution/planner/window.rs @@ -135,11 +135,8 @@ impl OperatorBuilder for WindowBuilder { // Spark put no sort below this window when its input reports the order the window // needs, and both window operators reject an input that does not declare its partition // keys' order; see `with_declared_order`. - let input = if wnd.input_sort_orders.is_empty() { - Arc::clone(&child.native_plan) - } else { - planner.with_declared_order(Arc::clone(&child.native_plan), &wnd.input_sort_orders)? - }; + 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, 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 70edb537901..9e866de71b3 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 @@ -79,38 +79,38 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { val sortOrders = op.orderSpec.map(exprToProto(_, op.child.output)) - // Spark puts no sort below the window when its input reports the order the window needs, a - // relation sorted before it was cached for instance. When those rows enter the native plan - // from the JVM, the ScanExec they arrive through declares no order, and the native window - // then rejects its input. Send the order so the planner declares it on the window's input; - // see SortedInputExec. Nothing is declared when the order the input reports may be false; see - // CometHashJoin.mayReportFalseOrder. A key that Comet cannot express, or that holds a - // collated string, which Spark orders by its collation but the native side would declare in - // raw byte order, sends the window back to Spark. - val requiredOrdering = op.requiredChildOrdering.head - val inputSortOrders = - if (requiredOrdering.nonEmpty && - SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) && - !CometHashJoinExec.mayReportFalseInputOrder(op.child)) { - val orders = - if (requiredOrdering.exists(order => hasNonDefaultStringCollation(order.dataType))) { - Nil - } else { - requiredOrdering.map(exprToProto(_, op.child.output)) + if (windowExprProto.forall(_.isDefined) && partitionExprs.forall(_.isDefined) + && sortOrders.forall(_.isDefined)) { + // Spark puts no sort below the window when its input reports the order the window needs, + // a relation sorted before it was cached for instance. When those rows enter the native + // plan from the JVM, the ScanExec they arrive through declares no order, and the native + // window then rejects its input. Send the order so the planner declares it on the window's + // input; see SortedInputExec. A native sort below declares its own order, and nothing is + // declared when the order the input reports may be false; see + // CometHashJoin.mayReportFalseOrder. A key that Comet cannot express, or that holds a + // collated string, which Spark orders by its collation but the native side would declare + // in raw byte order, 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.isEmpty || !orders.forall(_.isDefined)) { - withFallbackReason( - op, - "the order of the window's input, which Spark relies on, cannot be expressed natively") - return None + 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 } - orders.flatten - } else { - Nil - } - if (windowExprProto.forall(_.isDefined) && partitionExprs.forall(_.isDefined) - && sortOrders.forall(_.isDefined)) { val windowBuilder = OperatorOuterClass.Window.newBuilder() windowBuilder.addAllWindowExpr(windowExprProto.map(_.get).asJava) windowBuilder.addAllPartitionByList(partitionExprs.map(_.get).asJava) From 8bc8908611a22e996491e5dff0a2c9f2e5e85ac3 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 27 Oct 2026 09:07:27 -0700 Subject: [PATCH 3/7] chore: trim comments and drop a window test the hash join test already covers --- native/core/src/execution/planner.rs | 30 ++----------------- native/core/src/execution/planner/window.rs | 4 +-- native/proto/src/proto/operator.proto | 5 ++-- .../spark/sql/comet/CometWindowExec.scala | 15 ++++------ .../apache/spark/sql/comet/operators.scala | 3 +- .../comet/exec/CometInMemoryCacheSuite.scala | 7 ++--- 6 files changed, 16 insertions(+), 48 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 1fd03d0bf7d..68f8c25a660 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4441,8 +4441,7 @@ mod tests { ); } - /// A window partitioned and ordered by `col_0` over `input`, declaring `input_sort_orders` - /// for it. + /// A window partitioned and ordered by `col_0` over `input`. fn partitioned_window(input: Operator, input_sort_orders: Vec) -> Operator { Operator { children: vec![input], @@ -4459,9 +4458,8 @@ mod tests { #[test] fn window_declares_input_order_on_a_scan() { - // Spark puts no sort below a window whose input reports the order it needs, and the rows - // then arrive through a `ScanExec`, which declares none. The window finds its partition - // keys' order only once the planner declares the order Spark sent. + // The window finds its partition keys' order over a `ScanExec` only once the planner + // declares the order Spark sent. let window_over = |input_sort_orders| { let (_, _, planned) = PhysicalPlanner::default() .create_plan( @@ -4495,28 +4493,6 @@ mod tests { .contains("Expects PARTITION BY expression to be ordered")); } - #[test] - fn window_keeps_input_that_declares_an_order() { - // An input that already declares an order, a sort in the native plan for instance, 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 (_, _, planned) = PhysicalPlanner::default() - .create_plan( - &partitioned_window(sort, vec![create_sort_order(0, 3, false, true)]), - &mut vec![], - 1, - ) - .unwrap(); - assert_eq!("SortExec", planned.native_plan.children()[0].name()); - } - #[test] fn projection_rolls_up_same_plan_id_child_metrics() { let op_scan = Operator { diff --git a/native/core/src/execution/planner/window.rs b/native/core/src/execution/planner/window.rs index ca0c64c4fb9..9d24cea2f8b 100644 --- a/native/core/src/execution/planner/window.rs +++ b/native/core/src/execution/planner/window.rs @@ -132,9 +132,7 @@ impl OperatorBuilder for WindowBuilder { let window_expr = window_expr?; let all_bounded = window_expr.iter().all(|e| e.uses_bounded_memory()); - // Spark put no sort below this window when its input reports the order the window - // needs, and both window operators reject an input that does not declare its partition - // keys' order; see `with_declared_order`. + // 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 { diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index c21395721f7..1fb6fd02dc1 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1233,9 +1233,8 @@ message Window { repeated spark.spark_expression.Expr order_by_list = 2; repeated spark.spark_expression.Expr partition_by_list = 3; Operator child = 4; - // The order the window needs from its input, the partition keys ascending followed by the - // order keys, as `SortOrder` expressions over the input's output, when Spark reports that the - // input arrives in it and so put no sort below the window. Empty otherwise. + // 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; } 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 9e866de71b3..b38396de096 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 @@ -81,15 +81,12 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { if (windowExprProto.forall(_.isDefined) && partitionExprs.forall(_.isDefined) && sortOrders.forall(_.isDefined)) { - // Spark puts no sort below the window when its input reports the order the window needs, - // a relation sorted before it was cached for instance. When those rows enter the native - // plan from the JVM, the ScanExec they arrive through declares no order, and the native - // window then rejects its input. Send the order so the planner declares it on the window's - // input; see SortedInputExec. A native sort below declares its own order, and nothing is - // declared when the order the input reports may be false; see - // CometHashJoin.mayReportFalseOrder. A key that Comet cannot express, or that holds a - // collated string, which Spark orders by its collation but the native side would declare - // in raw byte order, sends the window back to Spark. + // Spark plans no sort below the window when its input reports the order the window needs, + // a sorted cache for instance, but rows from the JVM arrive through a ScanExec that + // declares none, and the native window rejects them. Send the order for the planner to + // declare (see SortedInputExec), unless a native sort below declares it or the reported + // order may be false. A key Comet cannot express or a collated string key, which the + // native side would order by raw bytes, sends the window back to Spark. val requiredOrdering = op.requiredChildOrdering.head val inputSortOrders = if (requiredOrdering.nonEmpty && !op.child.isInstanceOf[CometSortExec] && 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 3e94a14b5e6..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,8 +3017,7 @@ trait CometHashJoin { } } - // Whether the order that the input of a window reports may be false; see mayReportFalseOrder. - // CometWindowExec asks before it sends that order for the native window to declare. + // mayReportFalseOrder for a window's input, checked before CometWindowExec sends its order. private[comet] def mayReportFalseInputOrder(plan: SparkPlan): Boolean = mayReportFalseOrder(plan, countInPlan = true) 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 e433a40506d..e532495c12e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3693,10 +3693,9 @@ class CometInMemoryCacheSuite extends CometTestBase { } test("a native window over a sorted cache declares the cache's order") { - // The cache reports the order it was written in, so Spark puts no sort below a window - // partitioned and ordered by the cache's sort keys. The native window then reads the cache - // through a scan that declares no order, and rejected it with "Expects PARTITION BY - // expression to be ordered" until the window declared the order Spark sent (#7223). + // Spark plans no sort below a window over the cache's sort keys, and the native window + // rejected the scan, which declares no order, with "Expects PARTITION BY expression to be + // ordered" (#7223). Seq(true, false).foreach { nativeCache => withConversions( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", From 94aaa194dc19251cdccbf9463662f44898748a8b Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 27 Oct 2026 13:16:56 -0700 Subject: [PATCH 4/7] test: cover WindowAggExec, nulls, descending keys and the collated fallback From review: - The cache test also runs a frame up to UNBOUNDED FOLLOWING, which runs in WindowAggExec, over a nullable partition key and a descending order key, so the declared order has to match the cache's. - The planner test covers both window operators through the partition order lookup they share, since WindowAggExec fails only when its stream starts, not in partition_by_sort_keys. - A Spark 4.x test checks that a window over a cache sorted by a collated key stays in Spark for the new reason and returns Spark's rows. - The debugging guide's window snippet builds the operators over `input`. --- docs/source/contributor-guide/debugging.md | 7 +- native/core/src/execution/planner.rs | 84 ++++++++++--------- .../comet/exec/CometInMemoryCacheSuite.scala | 74 +++++++++++++--- 3 files changed, 109 insertions(+), 56 deletions(-) diff --git a/docs/source/contributor-guide/debugging.md b/docs/source/contributor-guide/debugging.md index 8795854193d..f1f94c4238a 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 any +order Spark sent declared on it, 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 68f8c25a660..ceaa7dde47a 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4337,9 +4337,10 @@ mod tests { } } - /// `row_number() OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)`, returning + /// `func() OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)`, returning /// `result_type`. - fn row_number_window_expr( + fn window_function_expr( + func: &str, result_type: Option, ) -> spark_operator::WindowExpr { use spark_operator::{ @@ -4351,7 +4352,7 @@ mod tests { 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, @@ -4386,7 +4387,7 @@ mod tests { 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 = row_number_window_expr(Some(create_proto_datatype())); + 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 { @@ -4441,12 +4442,12 @@ mod tests { ); } - /// A window partitioned and ordered by `col_0` over `input`. - fn partitioned_window(input: Operator, input_sort_orders: Vec) -> Operator { + /// `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![row_number_window_expr(None)], + 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, @@ -4458,39 +4459,42 @@ mod tests { #[test] fn window_declares_input_order_on_a_scan() { - // The window finds its partition keys' order over a `ScanExec` only once the planner - // declares the order Spark sent. - let window_over = |input_sort_orders| { - let (_, _, planned) = PhysicalPlanner::default() - .create_plan( - &partitioned_window(create_scan(), input_sort_orders), - &mut vec![], - 1, - ) - .unwrap(); - let window = planned - .native_plan - .downcast_ref::() - .expect("a bounded window"); - (Arc::clone(window.input()), window.partition_by_sort_keys()) - }; - let (input, keys) = window_over(vec![create_sort_order(0, 3, false, true)]); - let sorted = input - .downcast_ref::() - .expect("the window's input declares the order Spark sent"); - assert_eq!("ScanExec", sorted.children()[0].name()); - assert_eq!( - "col_0@0 ASC", - input.properties().output_ordering().unwrap().to_string() - ); - assert_eq!(1, keys.unwrap().len()); - - let (input, keys) = window_over(vec![]); - assert_eq!("ScanExec", input.name()); - assert!(keys - .unwrap_err() - .to_string() - .contains("Expects PARTITION BY expression to be ordered")); + // Both window operators need their partition keys' order declared, and over a `ScanExec` + // only the order Spark sent declares it. `percent_rank` does not run with bounded memory. + use datafusion::physical_plan::windows::get_ordered_partition_by_indices; + 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)]); + let sorted = input + .downcast_ref::() + .expect("the window's input declares the order Spark sent"); + assert_eq!("ScanExec", sorted.children()[0].name()); + assert_eq!( + "col_0@0 ASC", + input.properties().output_ordering().unwrap().to_string() + ); + 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] 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 e532495c12e..6f1e57313eb 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3695,7 +3695,19 @@ class CometInMemoryCacheSuite extends CometTestBase { 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, and the native window // rejected the scan, which declares no order, with "Expects PARTITION BY expression to be - // ordered" (#7223). + // ordered" (#7223). The partition key is nullable and the order key descending, so the + // declared order has to match the cache's direction and null placement. dense_rank runs in + // DataFusion's BoundedWindowAggExec, and the frame up to UNBOUNDED FOLLOWING 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", @@ -3705,22 +3717,23 @@ class CometInMemoryCacheSuite extends CometTestBase { try { val cached = spark .range(2000) - .selectExpr("id % 20 AS k", "id AS v") + .selectExpr("CASE WHEN id % 20 = 0 THEN NULL ELSE id % 20 END AS k", "id AS v") .repartition(2, col("k")) - .sortWithinPartitions("k", "v") + .sortWithinPartitions(col("k"), col("v").desc) .cache() cached.count() - val df = - cached.selectExpr("k", "v", "dense_rank() OVER (PARTITION BY k ORDER BY v) AS r") - val expected = (0L until 2000L).map(v => Row(v % 20, v, (v / 20 + 1).toInt)) - QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_)) - val plan = df.queryExecution.executedPlan - withClue(s"nativeCache=$nativeCache\n$plan\n") { - val windows = collect(plan) { case w: CometWindowExec => w } - assert(windows.size == 1) - assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 2) - assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) + queries.foreach { case (window, expected) => + val df = cached.selectExpr("k", "v", s"$window AS w") + val rows = (0L until 2000L).map(v => Row(key(v), v, expected(v))) + QueryTest.sameRows(rows, df.collect().toSeq).foreach(fail(_)) + val plan = df.queryExecution.executedPlan + withClue(s"nativeCache=$nativeCache, $window\n$plan\n") { + val windows = collect(plan) { case w: CometWindowExec => w } + assert(windows.size == 1) + assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 2) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) + } } } finally { spark.catalog.clearCache() @@ -3729,6 +3742,41 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + 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 the order of a collated key in raw byte order, which is not + // the order Spark sorted the cache in, so the window falls back instead. + 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 From 16b1add97e2f9e203944d882e6d40343827c8e81 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 29 Oct 2026 13:21:34 -0700 Subject: [PATCH 5/7] fix: keep a stacked window's longer input order and declare it on empty buckets The window now sends the rest of the order its child reports past the required keys, up to the first that cannot be expressed, so a window above it in the same native plan sees the longer order. An empty bucket of a sorted bucketed scan is planned as an EmptyExec, which now counts as a source whose order holds. The cache test checks the declared columns, directions and null placements. --- native/core/src/execution/planner.rs | 25 +++++++- native/core/src/execution/planner/join.rs | 6 +- native/proto/src/proto/operator.proto | 5 +- .../spark/sql/comet/CometWindowExec.scala | 21 ++++++- .../comet/exec/CometInMemoryCacheSuite.scala | 58 ++++++++++++++++++- .../comet/exec/CometWindowExecSuite.scala | 40 ++++++++++++- 6 files changed, 145 insertions(+), 10 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index ceaa7dde47a..376db839228 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4270,6 +4270,30 @@ mod tests { assert!(!PhysicalPlanner::keeps_source_order(&scan(2))); } + #[test] + fn declared_order_reaches_an_empty_bucket() { + // A bucket with no file is planned as an `EmptyExec`. A window over it checks its input's + // order before polling it, so the order Spark reports has to be declared there as well. + let schema = Arc::new(Schema::new(vec![Field::new( + "col_0", + DataType::Int32, + true, + )])); + let empty: Arc = + Arc::new(datafusion::physical_plan::empty::EmptyExec::new(schema)); + let input = PhysicalPlanner::default() + .with_declared_order(empty, &[create_sort_order(0, 3, false, true)]) + .unwrap(); + assert_eq!("SortedInputExec", input.name()); + let partition_by: Vec> = vec![Arc::new(Column::new("col_0", 0))]; + assert_eq!( + 1, + get_ordered_partition_by_indices(&partition_by, &input) + .unwrap() + .len() + ); + } + #[tokio::test] async fn spark_plan_dynamic_filter_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; @@ -4461,7 +4485,6 @@ mod tests { fn window_declares_input_order_on_a_scan() { // Both window operators need their partition keys' order declared, and over a `ScanExec` // only the order Spark sent declares it. `percent_rank` does not run with bounded memory. - use datafusion::physical_plan::windows::get_ordered_partition_by_indices; let partition_by: Vec> = vec![Arc::new(Column::new("col_0", 0))]; for (func, operator) in [ ("row_number", "BoundedWindowAggExec"), diff --git a/native/core/src/execution/planner/join.rs b/native/core/src/execution/planner/join.rs index 7ed0fd8459a..94fd9523139 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, @@ -440,11 +441,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 + /// sorted file, as for a sorted bucketed table with one file per bucket. A bucket with no file + /// is planned as an `EmptyExec`, which 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/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 1fb6fd02dc1..338eaca1615 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1233,8 +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. + // The partition keys ascending, then the order keys, then any further keys the input reports, + // 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; } 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 b38396de096..552d4c6d546 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, hasNonDefaultStringCollation, liftFallbackReasons, scalarFunctionExprToProto, serializeDataType} +import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, exprToProtoWithFallbackReasons, hasNonDefaultStringCollation, liftFallbackReasons, scalarFunctionExprToProto, serializeDataType} object CometWindowExec extends CometOperatorSerde[WindowExec] { @@ -103,7 +103,24 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { "natively") return None } - orders.flatten + // The keys the child reports past the required ones hold too, and this window reports + // them, so a window above in the same native plan may rely on them. Send them up to the + // first one that cannot be expressed, in a scope of their own, so that a key left out + // records no fallback reason. + val rest = op.child.outputOrdering + .drop(requiredOrdering.length) + .iterator + .map { order => + if (hasNonDefaultStringCollation(order.dataType) || + !order.references.subsetOf(op.child.outputSet)) { + None + } else { + exprToProtoWithFallbackReasons(order, op.child.output)._1 + } + } + .takeWhile(_.isDefined) + .flatten + orders.flatten ++ rest } else { Nil } 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 6f1e57313eb..1c24bd8d342 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3692,6 +3692,13 @@ 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}" + } + 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, and the native window // rejected the scan, which declares no order, with "Expects PARTITION BY expression to be @@ -3731,7 +3738,10 @@ class CometInMemoryCacheSuite extends CometTestBase { withClue(s"nativeCache=$nativeCache, $window\n$plan\n") { val windows = collect(plan) { case w: CometWindowExec => w } assert(windows.size == 1) - assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 2) + assert( + windowInputOrder(windows.head) == Seq( + "0 Ascending NullsFirst", + "1 Descending NullsLast")) assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) } } @@ -3742,6 +3752,52 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("stacked native windows over a sorted cache keep the cache's whole order") { + // The lower window needs only k ordered and the upper one k and v. Spark plans no sort + // between them, since the lower window reports the cache's order, so the lower window has to + // declare all of it for the upper one to start (#7223). Within a key, v = k + 20 * j for j + // in 0 until 100, and a NULL key holds the multiples of 20, so each key sums to + // 100 * k + 99000. + val key = (v: Long) => if (v % 20 == 0) null else v % 20 + 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("k", "v") + .cache() + cached.count() + + val df = cached + .selectExpr("k", "v", "sum(v) OVER (PARTITION BY k) AS s") + .selectExpr("k", "v", "sum(s) OVER (PARTITION BY k, v) AS t") + val rows = (0L until 2000L).map(v => Row(key(v), v, 100 * (v % 20) + 99000)) + QueryTest.sameRows(rows, df.collect().toSeq).foreach(fail(_)) + val plan = df.queryExecution.executedPlan + withClue(s"nativeCache=$nativeCache\n$plan\n") { + val windows = collect(plan) { case w: CometWindowExec => w } + assert(windows.size == 2) + windows.foreach { window => + assert( + windowInputOrder(window) == Seq( + "0 Ascending NullsFirst", + "1 Ascending NullsFirst")) + } + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) + } + } 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 the order of a collated key in raw byte order, which is not 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..8e3a76cd71f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -27,8 +27,8 @@ 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.execution.{InputAdapter, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeExec, CometNativeScanExec, CometSortExec, CometWindowExec, CometWindowGroupLimitExec} +import org.apache.spark.sql.execution.{InputAdapter, SortExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.window.{WindowExec => SparkWindowExec} import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions.{count, expr, lead, sum} @@ -2073,4 +2073,40 @@ class CometWindowExecSuite extends CometTestBase { } } } + + test("window over a sorted bucketed table with empty buckets runs natively") { + // With spark.sql.legacy.bucketedTableScan.outputOrdering on, Spark reports a bucketed table's + // sort columns as the scan's order and plans no sort below the window. Two keys fill at most + // two of the eight buckets, and the native scan plans each of the others as an empty input, + // which has to declare the order as well, or the window rejects it before reading (#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.isInstanceOf[CometNativeScanExec], plan) + assert(windows.head.child.asInstanceOf[CometNativeScanExec].bucketedScan, plan) + assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 2, plan) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) + assert(df.collect().length == 40) + } + } + } } From dc94ccf3d3bac54e8ff7ca981a92b0801a49b693 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 29 Oct 2026 18:23:13 -0700 Subject: [PATCH 6/7] fix: declare a window's order next to a partial one its input declares Two aliases of one sorted column share an order in Spark, so a window over the window for one alias may need the other alias's order with no sort in between. with_declared_order now returns early only when the input's declared order already satisfies the one sent, and otherwise declares it next to the input's when keeps_source_order trusts the input. The serde checks only a suffix key's child against the child's output, since its same-order expressions are not sent and may name pruned columns. --- native/core/src/execution/planner.rs | 42 +++++++++++ native/core/src/execution/planner/join.rs | 27 ++++--- .../spark/sql/comet/CometWindowExec.scala | 5 +- .../comet/exec/CometInMemoryCacheSuite.scala | 70 +++++++++++++++++++ 4 files changed, 134 insertions(+), 10 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 376db839228..5727d7f55d5 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4294,6 +4294,48 @@ mod tests { ); } + #[test] + fn declared_order_extends_an_order_the_input_declares() { + // Two aliases of one column share an order, so a window over the window for one alias can + // need the other's order while its input declares only the first. That order is declared + // next to the input's, 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!(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; diff --git a/native/core/src/execution/planner/join.rs b/native/core/src/execution/planner/join.rs index 94fd9523139..4d7bbe4bb84 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -403,11 +403,16 @@ 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. + /// + /// An input can declare part of the order already, as a window does over an input declared + /// for the window below it. Spark may know more than that declaration: two aliases of one + /// column share an order, so the window above may need the other alias's order. The rest is + /// declared next to what `plan` declares. /// /// 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 @@ -424,7 +429,7 @@ impl PhysicalPlanner { plan: Arc, sort_orders: &[Expr], ) -> Result, ExecutionError> { - if plan.output_ordering().is_some() || !Self::keeps_source_order(&plan) { + if !Self::keeps_source_order(&plan) { return Ok(plan); } let exprs = sort_orders @@ -432,8 +437,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), } } 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 552d4c6d546..2e197da2671 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 @@ -106,13 +106,14 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { // The keys the child reports past the required ones hold too, and this window reports // them, so a window above in the same native plan may rely on them. Send them up to the // first one that cannot be expressed, in a scope of their own, so that a key left out - // records no fallback reason. + // records no fallback reason. Only a key's child is sent, and its same-order + // expressions may name columns the child no longer outputs. val rest = op.child.outputOrdering .drop(requiredOrdering.length) .iterator .map { order => if (hasNonDefaultStringCollation(order.dataType) || - !order.references.subsetOf(op.child.outputSet)) { + !order.child.references.subsetOf(op.child.outputSet)) { None } else { exprToProtoWithFallbackReasons(order, op.child.output)._1 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 1c24bd8d342..47e0d6586c1 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3798,6 +3798,76 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("stacked native windows over a sorted cache with aliases of a sort key") { + // Spark knows that two aliases of one sorted column share its order, so a window over the + // window for one alias needs no sort for the other. The upper window's input declares the + // order the lower one needed, so the planner declares the rest next to it (#7223). In the + // second query the scan drops alias a, which only the order's same-order expressions name. + // One partition, so no exchange comes between 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 } + val source = () => + 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") + Seq(true, false).foreach { nativeCache => + withConversions( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> nativeCache.toString) { + val check = (df: DataFrame, rows: Seq[Row], orders: Seq[Seq[String]]) => { + QueryTest.sameRows(rows, df.collect().toSeq).foreach(fail(_)) + val plan = df.queryExecution.executedPlan + withClue(s"nativeCache=$nativeCache\n$plan\n") { + val windows = collect(plan) { case w: CometWindowExec => w } + assert(windows.map(windowInputOrder) == orders) + assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) + } + } + try { + val byKey = + source().sortWithinPartitions("k").selectExpr("k AS a", "k AS b", "v").cache() + byKey.count() + check( + byKey + .selectExpr("a", "b", "v", "sum(v) OVER (PARTITION BY a) AS s") + .selectExpr("a", "b", "v", "sum(s) OVER (PARTITION BY b) 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"))) + + val byPair = source() + .sortWithinPartitions("k", "x") + .selectExpr("k", "x AS a", "x AS b", "v") + .cache() + byPair.count() + val both = Seq("0 Ascending NullsFirst", "1 Ascending NullsFirst") + check( + byPair + .selectExpr("k", "b", "v", "sum(v) OVER (PARTITION BY k) AS s") + .selectExpr("k", "b", "v", "sum(s) OVER (PARTITION BY k, b) AS t"), + ids.map { id => + val k = key(id) + Row(nullable(k), id % 2, id, keySums(k) * pairCounts((k, id % 2))) + }, + Seq(both, both)) + } 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 the order of a collated key in raw byte order, which is not From 9c5dcd0d843ddd70db76af92cf3a56d63adb3f2a Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 29 Oct 2026 18:38:50 -0700 Subject: [PATCH 7/7] refactor: send only a window's required order and merge overlapping tests with_declared_order extends a partial order at the upper window, so the serde no longer sends the keys a window's child reports past the ones the window requires. The two stacked-window tests become one over a single cache, the EmptyExec unit test folds into the extension test, and redundant assertions and comments are trimmed. --- docs/source/contributor-guide/debugging.md | 4 +- native/core/src/execution/planner.rs | 43 +--- native/core/src/execution/planner/join.rs | 16 +- native/proto/src/proto/operator.proto | 5 +- .../spark/sql/comet/CometWindowExec.scala | 33 +-- .../comet/exec/CometInMemoryCacheSuite.scala | 190 ++++++------------ .../comet/exec/CometWindowExecSuite.scala | 18 +- 7 files changed, 98 insertions(+), 211 deletions(-) diff --git a/docs/source/contributor-guide/debugging.md b/docs/source/contributor-guide/debugging.md index f1f94c4238a..51d74a1ea94 100644 --- a/docs/source/contributor-guide/debugging.md +++ b/docs/source/contributor-guide/debugging.md @@ -260,8 +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` over `input`, the child plan with any -order Spark sent declared on it, 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 diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 5727d7f55d5..e6b451933cf 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4270,35 +4270,10 @@ mod tests { assert!(!PhysicalPlanner::keeps_source_order(&scan(2))); } - #[test] - fn declared_order_reaches_an_empty_bucket() { - // A bucket with no file is planned as an `EmptyExec`. A window over it checks its input's - // order before polling it, so the order Spark reports has to be declared there as well. - let schema = Arc::new(Schema::new(vec![Field::new( - "col_0", - DataType::Int32, - true, - )])); - let empty: Arc = - Arc::new(datafusion::physical_plan::empty::EmptyExec::new(schema)); - let input = PhysicalPlanner::default() - .with_declared_order(empty, &[create_sort_order(0, 3, false, true)]) - .unwrap(); - assert_eq!("SortedInputExec", input.name()); - let partition_by: Vec> = vec![Arc::new(Column::new("col_0", 0))]; - assert_eq!( - 1, - get_ordered_partition_by_indices(&partition_by, &input) - .unwrap() - .len() - ); - } - #[test] fn declared_order_extends_an_order_the_input_declares() { - // Two aliases of one column share an order, so a window over the window for one alias can - // need the other's order while its input declares only the first. That order is declared - // next to the input's, unless an operator in between may have moved rows. + // 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), @@ -4315,6 +4290,7 @@ mod tests { 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); @@ -4525,8 +4501,8 @@ mod tests { #[test] fn window_declares_input_order_on_a_scan() { - // Both window operators need their partition keys' order declared, and over a `ScanExec` - // only the order Spark sent declares it. `percent_rank` does not run with bounded memory. + // 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"), @@ -4544,14 +4520,7 @@ mod tests { Arc::clone(planned.native_plan.children()[0]) }; let input = window_input(vec![create_sort_order(0, 3, false, true)]); - let sorted = input - .downcast_ref::() - .expect("the window's input declares the order Spark sent"); - assert_eq!("ScanExec", sorted.children()[0].name()); - assert_eq!( - "col_0@0 ASC", - input.properties().output_ordering().unwrap().to_string() - ); + assert_eq!("SortedInputExec", input.name()); let ordered = get_ordered_partition_by_indices(&partition_by, &input).unwrap(); assert_eq!(1, ordered.len(), "{func}"); diff --git a/native/core/src/execution/planner/join.rs b/native/core/src/execution/planner/join.rs index 4d7bbe4bb84..6de4f3821ac 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -409,10 +409,8 @@ impl PhysicalPlanner { /// 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. /// - /// An input can declare part of the order already, as a window does over an input declared - /// for the window below it. Spark may know more than that declaration: two aliases of one - /// column share an order, so the window above may need the other alias's order. The rest is - /// declared next to what `plan` declares. + /// 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 @@ -429,7 +427,7 @@ impl PhysicalPlanner { plan: Arc, sort_orders: &[Expr], ) -> Result, ExecutionError> { - if !Self::keeps_source_order(&plan) { + if sort_orders.is_empty() || !Self::keeps_source_order(&plan) { return Ok(plan); } let exprs = sort_orders @@ -452,10 +450,10 @@ 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. A bucket with no file - /// is planned as an `EmptyExec`, which 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. + /// 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::() || plan.is::() { return true; diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 338eaca1615..1fb6fd02dc1 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1233,9 +1233,8 @@ 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, then any further keys the input reports, - // when Spark planned no sort below the window because its input reports that order. Empty - // otherwise. + // 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; } 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 2e197da2671..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, exprToProtoWithFallbackReasons, hasNonDefaultStringCollation, liftFallbackReasons, scalarFunctionExprToProto, serializeDataType} +import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, liftFallbackReasons, scalarFunctionExprToProto, serializeDataType} object CometWindowExec extends CometOperatorSerde[WindowExec] { @@ -81,12 +81,11 @@ 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 reports the order the window needs, - // a sorted cache for instance, but rows from the JVM arrive through a ScanExec that - // declares none, and the native window rejects them. Send the order for the planner to - // declare (see SortedInputExec), unless a native sort below declares it or the reported - // order may be false. A key Comet cannot express or a collated string key, which the - // native side would order by raw bytes, sends the window back to Spark. + // 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] && @@ -103,25 +102,7 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { "natively") return None } - // The keys the child reports past the required ones hold too, and this window reports - // them, so a window above in the same native plan may rely on them. Send them up to the - // first one that cannot be expressed, in a scope of their own, so that a key left out - // records no fallback reason. Only a key's child is sent, and its same-order - // expressions may name columns the child no longer outputs. - val rest = op.child.outputOrdering - .drop(requiredOrdering.length) - .iterator - .map { order => - if (hasNonDefaultStringCollation(order.dataType) || - !order.child.references.subsetOf(op.child.outputSet)) { - None - } else { - exprToProtoWithFallbackReasons(order, op.child.output)._1 - } - } - .takeWhile(_.isDefined) - .flatten - orders.flatten ++ rest + orders.flatten } else { Nil } 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 47e0d6586c1..3887f64891f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3699,13 +3699,23 @@ class CometInMemoryCacheSuite extends CometTestBase { 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, and the native window - // rejected the scan, which declares no order, with "Expects PARTITION BY expression to be - // ordered" (#7223). The partition key is nullable and the order key descending, so the - // declared order has to match the cache's direction and null placement. dense_rank runs in - // DataFusion's BoundedWindowAggExec, and the frame up to UNBOUNDED FOLLOWING in WindowAggExec. - // Within a key, v = k + 20 * j for j in 0 until 100, and a NULL key holds the multiples of 20. + // 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)" -> @@ -3731,65 +3741,12 @@ class CometInMemoryCacheSuite extends CometTestBase { cached.count() queries.foreach { case (window, expected) => - val df = cached.selectExpr("k", "v", s"$window AS w") - val rows = (0L until 2000L).map(v => Row(key(v), v, expected(v))) - QueryTest.sameRows(rows, df.collect().toSeq).foreach(fail(_)) - val plan = df.queryExecution.executedPlan - withClue(s"nativeCache=$nativeCache, $window\n$plan\n") { - val windows = collect(plan) { case w: CometWindowExec => w } - assert(windows.size == 1) - assert( - windowInputOrder(windows.head) == Seq( - "0 Ascending NullsFirst", - "1 Descending NullsLast")) - assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) - } - } - } finally { - spark.catalog.clearCache() - } - } - } - } - - test("stacked native windows over a sorted cache keep the cache's whole order") { - // The lower window needs only k ordered and the upper one k and v. Spark plans no sort - // between them, since the lower window reports the cache's order, so the lower window has to - // declare all of it for the upper one to start (#7223). Within a key, v = k + 20 * j for j - // in 0 until 100, and a NULL key holds the multiples of 20, so each key sums to - // 100 * k + 99000. - val key = (v: Long) => if (v % 20 == 0) null else v % 20 - 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("k", "v") - .cache() - cached.count() - - val df = cached - .selectExpr("k", "v", "sum(v) OVER (PARTITION BY k) AS s") - .selectExpr("k", "v", "sum(s) OVER (PARTITION BY k, v) AS t") - val rows = (0L until 2000L).map(v => Row(key(v), v, 100 * (v % 20) + 99000)) - QueryTest.sameRows(rows, df.collect().toSeq).foreach(fail(_)) - val plan = df.queryExecution.executedPlan - withClue(s"nativeCache=$nativeCache\n$plan\n") { - val windows = collect(plan) { case w: CometWindowExec => w } - assert(windows.size == 2) - windows.foreach { window => - assert( - windowInputOrder(window) == Seq( - "0 Ascending NullsFirst", - "1 Ascending NullsFirst")) + 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"))) } - assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) } } finally { spark.catalog.clearCache() @@ -3798,12 +3755,11 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("stacked native windows over a sorted cache with aliases of a sort key") { - // Spark knows that two aliases of one sorted column share its order, so a window over the - // window for one alias needs no sort for the other. The upper window's input declares the - // order the lower one needed, so the planner declares the rest next to it (#7223). In the - // second query the scan drops alias a, which only the order's same-order expressions name. - // One partition, so no exchange comes between the windows. + 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 @@ -3811,67 +3767,51 @@ class CometInMemoryCacheSuite extends CometTestBase { 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 } - val source = () => - 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") - Seq(true, false).foreach { nativeCache => - withConversions( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> nativeCache.toString) { - val check = (df: DataFrame, rows: Seq[Row], orders: Seq[Seq[String]]) => { - QueryTest.sameRows(rows, df.collect().toSeq).foreach(fail(_)) - val plan = df.queryExecution.executedPlan - withClue(s"nativeCache=$nativeCache\n$plan\n") { - val windows = collect(plan) { case w: CometWindowExec => w } - assert(windows.map(windowInputOrder) == orders) - assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty) - } - } - try { - val byKey = - source().sortWithinPartitions("k").selectExpr("k AS a", "k AS b", "v").cache() - byKey.count() - check( - byKey - .selectExpr("a", "b", "v", "sum(v) OVER (PARTITION BY a) AS s") - .selectExpr("a", "b", "v", "sum(s) OVER (PARTITION BY b) 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"))) - - val byPair = source() - .sortWithinPartitions("k", "x") - .selectExpr("k", "x AS a", "x AS b", "v") - .cache() - byPair.count() - val both = Seq("0 Ascending NullsFirst", "1 Ascending NullsFirst") - check( - byPair - .selectExpr("k", "b", "v", "sum(v) OVER (PARTITION BY k) AS s") - .selectExpr("k", "b", "v", "sum(s) OVER (PARTITION BY k, b) AS t"), - ids.map { id => - val k = key(id) - Row(nullable(k), id % 2, id, keySums(k) * pairCounts((k, id % 2))) - }, - Seq(both, both)) - } finally { - spark.catalog.clearCache() - } + 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 the order of a collated key in raw byte order, which is not - // the order Spark sorted the cache in, so the window falls back instead. + // 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") { 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 8e3a76cd71f..d1ed4f4600c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -28,7 +28,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, Divide, Expression, MakeDecimal, WindowExpression} import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeExec, CometNativeScanExec, CometSortExec, CometWindowExec, CometWindowGroupLimitExec} -import org.apache.spark.sql.execution.{InputAdapter, SortExec, SparkPlan, WholeStageCodegenExec} +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 import org.apache.spark.sql.functions.{count, expr, lead, sum} @@ -2075,10 +2075,8 @@ class CometWindowExecSuite extends CometTestBase { } test("window over a sorted bucketed table with empty buckets runs natively") { - // With spark.sql.legacy.bucketedTableScan.outputOrdering on, Spark reports a bucketed table's - // sort columns as the scan's order and plans no sort below the window. Two keys fill at most - // two of the eight buckets, and the native scan plans each of the others as an empty input, - // which has to declare the order as well, or the window rejects it before reading (#7223). + // 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) @@ -2101,11 +2099,13 @@ class CometWindowExecSuite extends CometTestBase { val (_, plan) = checkSparkAnswerAndOperator(df) val windows = collect(plan) { case w: CometWindowExec => w } assert(windows.size == 1, plan) - assert(windows.head.child.isInstanceOf[CometNativeScanExec], plan) - assert(windows.head.child.asInstanceOf[CometNativeScanExec].bucketedScan, plan) + assert( + windows.head.child match { + case scan: CometNativeScanExec => scan.bucketedScan + case _ => false + }, + plan) assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 2, plan) - assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty, plan) - assert(df.collect().length == 40) } } }