From dc6009771ddedc708432969612896a9225ccad2d Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 27 Oct 2026 08:35:29 -0700 Subject: [PATCH 01/12] 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 3c92c16e54f..46ff8ec6c96 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4282,21 +4282,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(), @@ -4321,9 +4318,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 { @@ -4378,6 +4386,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 b4ed291eaf2..3389ae47427 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -484,16 +484,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 28d94e78bf3..8a7b6e0c166 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1242,6 +1242,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 a162d93a098..f5695db8082 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 dd7c59f6556..e31f5df66ef 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -3033,6 +3033,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 59cede8d87e..d3e4c7dfb3e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -40,7 +40,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Ascending, Attribute, Att import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.{LeftOuter, RightOuter} import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} -import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoin, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometHashJoin, CometHashJoinExec, CometInMemoryTableScanExec, CometNativeScanExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec, CometSparkToColumnarExec, CometWindowExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager import org.apache.spark.sql.comet.util.Utils @@ -49,7 +49,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffl import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} -import org.apache.spark.sql.functions.{count, lit, max, min, sum} +import org.apache.spark.sql.functions.{col, count, lit, max, min, sum} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} @@ -3749,6 +3749,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 5a3af5bbbc203c149a37b3f78939945bd32bfba8 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 27 Oct 2026 08:53:30 -0700 Subject: [PATCH 02/12] 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 f5695db8082..83dd383526a 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 c231e48808b6b3f5fb9045599234e5989ba65ef1 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 27 Oct 2026 09:07:27 -0700 Subject: [PATCH 03/12] 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 46ff8ec6c96..ce14c954062 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4386,8 +4386,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], @@ -4404,9 +4403,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( @@ -4440,28 +4438,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 8a7b6e0c166..48c0ba151c6 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1242,9 +1242,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 83dd383526a..f5ef6a9d12e 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 e31f5df66ef..a043910e82c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -3033,8 +3033,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 d3e4c7dfb3e..292a7076a07 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3750,10 +3750,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 4508a55fd293d2b450ae2e362614649105d73a6b Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 27 Oct 2026 13:16:56 -0700 Subject: [PATCH 04/12] 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 ce14c954062..600208a4afe 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4282,9 +4282,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::{ @@ -4296,7 +4297,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, @@ -4331,7 +4332,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 { @@ -4386,12 +4387,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, @@ -4403,39 +4404,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 292a7076a07..95645efe525 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3752,7 +3752,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", @@ -3762,22 +3774,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() @@ -3786,6 +3799,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 ca1a53a979c80909ef05321f38058627a1ad71ba Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 29 Oct 2026 13:21:34 -0700 Subject: [PATCH 05/12] 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 600208a4afe..5ec4c1ad890 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4215,6 +4215,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; @@ -4406,7 +4430,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 3389ae47427..9139b9e1125 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, @@ -516,11 +517,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 48c0ba151c6..02a27f68311 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1242,8 +1242,9 @@ message Window { repeated spark.spark_expression.Expr order_by_list = 2; repeated spark.spark_expression.Expr partition_by_list = 3; Operator child = 4; - // The partition keys ascending, 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 f5ef6a9d12e..472638cbe6d 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 95645efe525..9b28a1fdbff 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3749,6 +3749,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 @@ -3788,7 +3795,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) } } @@ -3799,6 +3809,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 0438a5927d2192ccfdc17be7e6084285e8025158 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 29 Oct 2026 18:23:13 -0700 Subject: [PATCH 06/12] 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 5ec4c1ad890..03d6dea1c3d 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4239,6 +4239,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 9139b9e1125..8a4bcadb498 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -479,11 +479,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 @@ -500,7 +505,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 @@ -508,8 +513,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 472638cbe6d..6d41df307da 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 9b28a1fdbff..e0fcef91723 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3855,6 +3855,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 fce10f06dc9548a6d93227a382c86fdd28fd6dd4 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 29 Oct 2026 18:38:50 -0700 Subject: [PATCH 07/12] 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 03d6dea1c3d..cd8b3c1ff6f 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4215,35 +4215,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), @@ -4260,6 +4235,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); @@ -4470,8 +4446,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"), @@ -4489,14 +4465,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 8a4bcadb498..37806293de5 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -485,10 +485,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 @@ -505,7 +503,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 @@ -528,10 +526,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 02a27f68311..48c0ba151c6 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1242,9 +1242,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 6d41df307da..f603a44d92a 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 e0fcef91723..5c9cd94698b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3756,13 +3756,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)" -> @@ -3788,65 +3798,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() @@ -3855,12 +3812,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 @@ -3868,67 +3824,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) } } } From e87d39ad6dea363c4fec2637a28d4b257537d183 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 30 Oct 2026 09:43:44 -0700 Subject: [PATCH 08/12] fix: declare only a window's partition keys and fall back on an unchecked order The native window checks only its partition keys' order, so the serde sends only those, and a float order key that Spark sorts with -0.0 and 0.0 as equal is never declared. When the order the child reports may be false past the native plan, the window stays in Spark, as the outer hash join does, instead of failing the task. The SortedInputExec doc and the operator guide now cover the window. --- .../adding_a_new_operator.md | 13 ++--- .../src/execution/operators/sorted_input.rs | 15 +++--- native/core/src/execution/planner.rs | 32 +++++------- native/proto/src/proto/operator.proto | 4 +- .../spark/sql/comet/CometWindowExec.scala | 49 +++++++++++------- .../apache/spark/sql/comet/operators.scala | 6 +-- .../comet/exec/CometInMemoryCacheSuite.scala | 50 ++++++++++++++++--- .../comet/exec/CometWindowExecSuite.scala | 2 +- 8 files changed, 109 insertions(+), 62 deletions(-) diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index 0c36921a94a..5d18ff311c9 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -365,17 +365,18 @@ the order satisfies. Report what the Spark operator reports and nothing more. `S example, reports no order, so `CometSampleExec` must not report its child's. The claim has to hold at run time as well. A native hash join declares the order Spark reports for -its streamed side to DataFusion, and an aggregate above the join may then run in sorted mode, or a -top-k sort may stop reading early, on the strength of that order. If the order is false, the +its streamed side to DataFusion, and a native window the order of its partition keys. An aggregate +above may then run in sorted mode, or a top-k sort may stop reading early, on the strength of that +order. If the order is false, the result is wrong groups or wrong rows, not only rows out of order. So an operator that reports an ordering must keep it natively, and the native operator must say so by returning `true` from `maintains_input_order` for that child, as DataFusion's filter and projection do. The planner declares the order through an operator only when the operator says it keeps it. An operator that cannot keep the order natively must not report it, or has to be added to `mayReportFalseOrder` in -`operators.scala`. The hash join serde then sends no order over it, and an outer hash join above it -that has to keep that order falls back to Spark. Comet's hash joins are the exception within one -native plan: an outer hash join over one of them stays native, even over one that declared no -order, since DataFusion keeps the order it can see there. +`operators.scala`. The hash join and window serdes then send no order over it, and an outer hash +join or a window above it that needs that order falls back to Spark. Comet's hash joins are the +exception within one native plan: a join or window over one of them stays native, even over one +that declared no order, since DataFusion keeps the order it can see there. A leaf that reports an ordering is where the order starts, so the planner has to know that the native operator produces its rows in that order. It knows that of an input from the JVM and of a diff --git a/native/core/src/execution/operators/sorted_input.rs b/native/core/src/execution/operators/sorted_input.rs index a31053bd6e8..ad561cf49e3 100644 --- a/native/core/src/execution/operators/sorted_input.rs +++ b/native/core/src/execution/operators/sorted_input.rs @@ -37,13 +37,14 @@ use datafusion::physical_plan::{ /// sorted before it was cached for instance, but the `ScanExec` it arrives through declares none, /// and DataFusion's hash join keeps the unmatched probe rows of a right outer join in place only /// when the probe input declares an order. Otherwise it moves them after the matched rows of each -/// batch, which loses an order Spark has already dropped a sort for. The planner puts this only -/// on that join's probe input, with the order Spark reports for it, and only when the rows reach -/// that input from the JVM through operators that keep their order. The join keeps the order in -/// its own output, so the operators above it in the native plan see it too: an aggregate grouped -/// by the key runs in sorted, streaming mode, a sort aggregate takes the order as already -/// satisfied, and a top-k sort whose keys start with it stops reading early. A wrong order would -/// therefore give wrong groups or wrong rows, not only rows out of order. +/// batch, which loses an order Spark has already dropped a sort for. A window rejects an input +/// that does not declare its partition keys' order. The planner puts this only on that join's +/// probe input or that window's input, with the order Spark reports for it, and only when the +/// rows reach that input from the JVM through operators that keep their order. Both keep the +/// order in their own output, so the operators above them in the native plan see it too: an +/// aggregate grouped by the key runs in sorted, streaming mode, a sort aggregate takes the order +/// as already satisfied, and a top-k sort whose keys start with it stops reading early. A wrong +/// order would therefore give wrong groups or wrong rows, not only rows out of order. #[derive(Debug)] pub struct SortedInputExec { input: Arc, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index cd8b3c1ff6f..4aa88b90377 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4429,21 +4429,6 @@ 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 @@ -4454,12 +4439,19 @@ mod tests { ("percent_rank", "WindowAggExec"), ] { let window_input = |input_sort_orders| { + let window = Operator { + children: vec![create_scan()], + op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window { + window_expr: vec![window_function_expr(func, None)], + partition_by_list: vec![create_bound_reference(0)], + order_by_list: vec![create_sort_order(0, 3, false, true)], + input_sort_orders, + ..Default::default() + }))), + ..Default::default() + }; let (_, _, planned) = PhysicalPlanner::default() - .create_plan( - &partitioned_window(func, create_scan(), input_sort_orders), - &mut vec![], - 1, - ) + .create_plan(&window, &mut vec![], 1) .unwrap(); assert_eq!(operator, planned.native_plan.name()); Arc::clone(planned.native_plan.children()[0]) diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 48c0ba151c6..e99b91e6cc9 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1242,8 +1242,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, when Spark planned no sort below the - // window because its input reports that order. Empty otherwise. + // The partition keys ascending, when Spark planned no sort below the window because its input + // reports that order. Empty otherwise. repeated spark.spark_expression.Expr input_sort_orders = 5; } 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 f603a44d92a..add6eb85e5f 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 @@ -40,6 +40,13 @@ import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNo object CometWindowExec extends CometOperatorSerde[WindowExec] { + private[apache] val inexpressibleInputOrderReason: String = + "the order of the window's input, which Spark relies on, cannot be expressed natively" + + private[apache] val uncheckedInputOrderReason: String = + "the order of the window's input, which Spark relies on, cannot be checked, so the window " + + "stays in Spark" + override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( CometConf.COMET_EXEC_WINDOW_ENABLED) @@ -82,27 +89,35 @@ 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. + // needs, but rows from the JVM arrive through a ScanExec that declares none. Send the + // partition keys' order, the only one the native window checks, for the planner to declare + // (see SortedInputExec). When the order may be false, nothing native can declare it, so the + // window stays in Spark, unless the doubt is a Comet join in the same native plan, which + // DataFusion sees itself. A key Comet cannot express, or a collated string key, sends the + // window back to Spark too. 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") + if (op.partitionSpec.nonEmpty && + SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering)) { + // The walk that counts Comet joins in the same native plan flags everything the other + // one does, so the other runs only when this one flags the order. + if (!CometHashJoinExec.mayReportFalseInputOrder(op.child, countInPlan = true)) { + val partitionKeys = requiredOrdering.take(op.partitionSpec.length) + val orders = partitionKeys.flatMap { order => + if (hasNonDefaultStringCollation(order.dataType)) None + else exprToProto(order, op.child.output) + } + if (orders.length < partitionKeys.length) { + withFallbackReason(op, inexpressibleInputOrderReason) + return None + } + orders + } else if (CometHashJoinExec.mayReportFalseInputOrder(op.child, countInPlan = false)) { + withFallbackReason(op, uncheckedInputOrderReason) return None + } else { + Nil } - orders.flatten } else { Nil } 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 a043910e82c..61db25c13ec 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -3033,9 +3033,9 @@ 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) + // mayReportFalseOrder for a window's input, which CometWindowExec checks like a join does. + private[comet] def mayReportFalseInputOrder(plan: SparkPlan, countInPlan: Boolean): Boolean = + mayReportFalseOrder(plan, countInPlan) // 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 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 5c9cd94698b..3136a79580e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3769,10 +3769,10 @@ 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, 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. + // the partition key's order (#7223). The key is nullable, so the null placement has to match, + // and the descending order key is not declared. dense_rank runs in BoundedWindowAggExec and + // the UNBOUNDED FOLLOWING frame in WindowAggExec. Within a key, v = k + 20 * j for j in + // 0 until 100, and a NULL key holds the multiples of 20. val key = (v: Long) => if (v % 20 == 0) null else v % 20 val queries = Seq( "dense_rank() OVER (PARTITION BY k ORDER BY v DESC)" -> @@ -3802,7 +3802,7 @@ class CometInMemoryCacheSuite extends CometTestBase { 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"))) + Seq(Seq("0 Ascending NullsFirst"))) } } } finally { @@ -3893,13 +3893,51 @@ class CometInMemoryCacheSuite extends CometTestBase { 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) + assert(reasons.contains(CometWindowExec.inexpressibleInputOrderReason), reasons) } finally { spark.catalog.clearCache() } } } + test("a window over a cache whose order cannot be checked stays in Spark") { + // The filter keeps the serde from seeing that the outer join declared its order, so the + // window, reading that join's cache from the JVM, cannot check it. Declaring nothing would + // fail the native window. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "false") { + withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { + withParquetTable((0 until 40 by 2).map(m => (m * 1000, m)), "small") { + withTempView("cached", "joined") { + try { + spark + .table("big") + .toDF("k", "v") + .repartition(2, $"k") + .sortWithinPartitions("k") + .cache() + .createOrReplaceTempView("cached") + spark.table("cached").count() + sql("SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w FROM " + + "(SELECT * FROM cached WHERE v >= 0) a LEFT JOIN small b ON a.k = b._1") + .cache() + .createOrReplaceTempView("joined") + spark.table("joined").count() + checkSparkAnswerAndFallbackReason( + sql("SELECT k, v, sum(v) OVER (PARTITION BY k) AS s FROM joined"), + CometWindowExec.uncheckedInputOrderReason) + } 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 d1ed4f4600c..5d09646d61e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -2105,7 +2105,7 @@ class CometWindowExecSuite extends CometTestBase { case _ => false }, plan) - assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 2, plan) + assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount == 1, plan) } } } From ee88fd96aa716dfb7405ca4a5195ed09be419f23 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 30 Oct 2026 17:45:13 -0700 Subject: [PATCH 09/12] fix: keep a window's doubted order native, fall back on collated keys Falling back when the window's input order may be false handed the same rows to Spark's window, which trusts that order too and grouped them wrongly. The window again sends no order there, so the native window rejects its input as on main. A window whose partition or order key holds a collated string now stays in Spark, as WindowGroupLimit does, since the native window compares keys byte by byte. keeps_source_order accepts a right outer hash join whose probe declares an order, which is when DataFusion keeps its unmatched probe rows in place. --- .../adding_a_new_operator.md | 19 ++-- native/core/src/execution/planner.rs | 4 +- native/core/src/execution/planner/join.rs | 23 +++-- .../spark/sql/comet/CometWindowExec.scala | 54 +++++----- .../apache/spark/sql/comet/operators.scala | 6 +- .../comet/exec/CometInMemoryCacheSuite.scala | 98 ++++++++----------- 6 files changed, 92 insertions(+), 112 deletions(-) diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index 5d18ff311c9..ed48f4fba0f 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -367,16 +367,15 @@ example, reports no order, so `CometSampleExec` must not report its child's. The claim has to hold at run time as well. A native hash join declares the order Spark reports for its streamed side to DataFusion, and a native window the order of its partition keys. An aggregate above may then run in sorted mode, or a top-k sort may stop reading early, on the strength of that -order. If the order is false, the -result is wrong groups or wrong rows, not only rows out of order. So an operator that reports an -ordering must keep it natively, and the native operator must say so by returning `true` from -`maintains_input_order` for that child, as DataFusion's filter and projection do. The planner -declares the order through an operator only when the operator says it keeps it. An operator that -cannot keep the order natively must not report it, or has to be added to `mayReportFalseOrder` in -`operators.scala`. The hash join and window serdes then send no order over it, and an outer hash -join or a window above it that needs that order falls back to Spark. Comet's hash joins are the -exception within one native plan: a join or window over one of them stays native, even over one -that declared no order, since DataFusion keeps the order it can see there. +order. If the order is false, the result is wrong groups or wrong rows, not only rows out of order. +So an operator that reports an ordering must keep it natively, and the native operator must say so +by returning `true` from `maintains_input_order` for that child, as DataFusion's filter and +projection do. The planner declares the order through an operator only when the operator says it +keeps it. An operator that cannot keep the order natively must not report it, or has to be added to +`mayReportFalseOrder` in `operators.scala`. The hash join and window serdes then send no order over +it, and an outer hash join above it that has to keep that order falls back to Spark. Comet's hash +joins are the exception within one native plan: an outer hash join over one of them stays native, +even over one that declared no order, since DataFusion keeps the order it can see there. A leaf that reports an ordering is where the order starts, so the planner has to know that the native operator produces its rows in that order. It knows that of an input from the JVM and of a diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 4aa88b90377..70648070282 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4046,8 +4046,10 @@ mod tests { "col_0@0 ASC", probe.properties().output_ordering().unwrap().to_string() ); - // The join keeps the order in its own output, so the operators above it see it too. + // The join keeps the order in its own output, so the operators above it see it too, + // and an order declared above it holds, such as a window's for an alias of the key. assert!(planned.native_plan.properties().output_ordering().is_some()); + assert!(PhysicalPlanner::keeps_source_order(&planned.native_plan)); } } diff --git a/native/core/src/execution/planner/join.rs b/native/core/src/execution/planner/join.rs index 37806293de5..ffefd158e92 100644 --- a/native/core/src/execution/planner/join.rs +++ b/native/core/src/execution/planner/join.rs @@ -541,18 +541,21 @@ impl PhysicalPlanner { .is_some_and(|config| config.file_groups.len() == 1); } // DataFusion reports every right hash join as keeping its probe order, but a right outer - // join keeps it only when the probe declares one, and no input on this path declares an - // order. An inner join emits its matches probe row by probe row, and a right semi, anti or - // mark join builds its probe indices in ascending order, so none of them depends on a - // declaration. Any other join type keeps no probe order at all. + // join keeps its unmatched probe rows in place only when the probe declares an order, as + // the planner does for the join's streamed side. An inner join emits its matches probe + // row by probe row, and a right semi, anti or mark join builds its probe indices in + // ascending order, so none of them depends on a declaration. Any other join type keeps no + // probe order at all. if let Some(join) = plan.downcast_ref::() { - if !matches!( - join.join_type(), + let keeps_probe_order = match join.join_type() { DFJoinType::Inner - | DFJoinType::RightSemi - | DFJoinType::RightAnti - | DFJoinType::RightMark - ) { + | DFJoinType::RightSemi + | DFJoinType::RightAnti + | DFJoinType::RightMark => true, + DFJoinType::Right => join.right().output_ordering().is_some(), + _ => false, + }; + if !keeps_probe_order { return false; } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala index add6eb85e5f..f25f9225b6a 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 @@ -40,13 +40,6 @@ import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNo object CometWindowExec extends CometOperatorSerde[WindowExec] { - private[apache] val inexpressibleInputOrderReason: String = - "the order of the window's input, which Spark relies on, cannot be expressed natively" - - private[apache] val uncheckedInputOrderReason: String = - "the order of the window's input, which Spark relies on, cannot be checked, so the window " + - "stays in Spark" - override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( CometConf.COMET_EXEC_WINDOW_ENABLED) @@ -82,6 +75,19 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { } } + // The native window finds partitions and peers by comparing keys byte by byte, so it would + // tell apart collated strings that Spark treats as equal, as CometWindowGroupLimitExec notes. + val collated = (op.partitionSpec ++ op.orderSpec.map(_.child)) + .filter(key => hasNonDefaultStringCollation(key.dataType)) + if (collated.nonEmpty) { + withFallbackReason( + op, + collated + .map(_.sql) + .mkString("Window: non-default string collation on key(s): ", ", ", "")) + return None + } + val partitionExprs = op.partitionSpec.map(exprToProto(_, op.child.output)) val sortOrders = op.orderSpec.map(exprToProto(_, op.child.output)) @@ -91,33 +97,23 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { // 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 the // partition keys' order, the only one the native window checks, for the planner to declare - // (see SortedInputExec). When the order may be false, nothing native can declare it, so the - // window stays in Spark, unless the doubt is a Comet join in the same native plan, which - // DataFusion sees itself. A key Comet cannot express, or a collated string key, sends the - // window back to Spark too. + // (see SortedInputExec). Nothing is sent when that order may be false. Falling back would + // not help, since Spark's window trusts the order too. val requiredOrdering = op.requiredChildOrdering.head val inputSortOrders = if (op.partitionSpec.nonEmpty && - SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering)) { - // The walk that counts Comet joins in the same native plan flags everything the other - // one does, so the other runs only when this one flags the order. - if (!CometHashJoinExec.mayReportFalseInputOrder(op.child, countInPlan = true)) { - val partitionKeys = requiredOrdering.take(op.partitionSpec.length) - val orders = partitionKeys.flatMap { order => - if (hasNonDefaultStringCollation(order.dataType)) None - else exprToProto(order, op.child.output) - } - if (orders.length < partitionKeys.length) { - withFallbackReason(op, inexpressibleInputOrderReason) - return None - } - orders - } else if (CometHashJoinExec.mayReportFalseInputOrder(op.child, countInPlan = false)) { - withFallbackReason(op, uncheckedInputOrderReason) + SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) && + !CometHashJoinExec.mayReportFalseInputOrder(op.child)) { + val partitionKeys = requiredOrdering.take(op.partitionSpec.length) + val orders = partitionKeys.flatMap(exprToProto(_, op.child.output)) + if (orders.length < partitionKeys.length) { + withFallbackReason( + op, + "the order of the window's input, which Spark relies on, cannot be expressed " + + "natively") return None - } else { - Nil } + orders } else { Nil } 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 61db25c13ec..a043910e82c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -3033,9 +3033,9 @@ trait CometHashJoin { } } - // mayReportFalseOrder for a window's input, which CometWindowExec checks like a join does. - private[comet] def mayReportFalseInputOrder(plan: SparkPlan, countInPlan: Boolean): Boolean = - mayReportFalseOrder(plan, countInPlan) + // 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 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 3136a79580e..bcba4cd8a6d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3812,11 +3812,12 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - 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. + test("a native window declares its order next to the order its input declares") { + // Spark plans no sort above an operator that keeps the cache's order, so a window's input may + // declare only the order that operator needed: an upper window's input the lower window's, + // and an outer join's output its streamed side's. The window's order is declared next to it + // (#7223). Alias a of sort key k shares its order. One partition, so no exchange comes in + // between. val ids = 0L until 24L val key = (id: Long) => if (id % 4 == 0) None else Some(id % 3) val nullable = (k: Option[Long]) => k.map(Long.box).orNull @@ -3859,6 +3860,18 @@ class CometInMemoryCacheSuite extends CometTestBase { Row(nullable(k), nullable(k), id, keySums(k) * keyCounts(k)) }, Seq(Seq("1 Ascending NullsFirst"), Seq("0 Ascending NullsFirst"))) + // The residual condition leaves rows of key 0 unmatched among matched ones, so the key + // stays together only if the join keeps its unmatched probe rows in place. + val one = spark.range(1).toDF("j").hint("broadcast") + checkWindows( + cached + .join(one, col("k") === col("j") && col("v") % 2 === 0, "left_outer") + .selectExpr("k", "a", "v", "sum(v) OVER (PARTITION BY a) AS s"), + ids.map { id => + val k = key(id) + Row(nullable(k), nullable(k), id, keySums(k)) + }, + Seq(Seq("1 Ascending NullsFirst"))) } finally { spark.catalog.clearCache() } @@ -3867,77 +3880,44 @@ 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 a collated key in raw byte order, not Spark's collation - // order, so the window falls back. + // The native window compares keys byte by byte, so it would split the UTF8_LCASE partition + // of '0a' and '0A' and rank them apart. + val key = (v: Long) => s"${v % 4}${if (v % 8 < 4) "a" else "A"}" withConversions( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "2") { try { - val cached = spark + val source = spark .range(200) .selectExpr( "CAST(id % 4 AS STRING) || CASE WHEN id % 8 < 4 THEN 'a' ELSE 'A' END " + "COLLATE UTF8_LCASE AS k", + "id % 2 AS n", "id AS v") - .repartition(2, col("k")) - .sortWithinPartitions("k", "v") - .cache() - cached.count() + val byKey = source.repartition(2, col("k")).sortWithinPartitions("k", "v").cache() + val byGroup = source.repartition(2, col("n")).sortWithinPartitions("n", "k").cache() + byKey.count() + byGroup.count() - 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) + Seq( + // Each of the four keys holds 50 rows. + byKey.selectExpr("k", "v", "dense_rank() OVER (PARTITION BY k ORDER BY v) AS r") -> + ((v: Long) => Row(key(v), v, (v / 4 + 1).toInt)), + // Each group holds two keys, '0' and '2' or '1' and '3'. + byGroup.selectExpr("n", "k", "dense_rank() OVER (PARTITION BY n ORDER BY k) AS r") -> + ((v: Long) => Row(v % 2, key(v), (v % 4 / 2 + 1).toInt))).foreach { case (df, row) => + QueryTest.sameRows((0L until 200L).map(row), df.collect().toSeq).foreach(fail(_)) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case w: CometWindowExec => w }.isEmpty, plan) + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert(reasons.exists(_.startsWith("Window: non-default string collation")), reasons) } - 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.contains(CometWindowExec.inexpressibleInputOrderReason), reasons) } finally { spark.catalog.clearCache() } } } - test("a window over a cache whose order cannot be checked stays in Spark") { - // The filter keeps the serde from seeing that the outer join declared its order, so the - // window, reading that join's cache from the JVM, cannot check it. Declaring nothing would - // fail the native window. - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", - CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "false") { - withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") { - withParquetTable((0 until 40 by 2).map(m => (m * 1000, m)), "small") { - withTempView("cached", "joined") { - try { - spark - .table("big") - .toDF("k", "v") - .repartition(2, $"k") - .sortWithinPartitions("k") - .cache() - .createOrReplaceTempView("cached") - spark.table("cached").count() - sql("SELECT /*+ BROADCAST(b) */ a.k, a.v, b._2 AS w FROM " + - "(SELECT * FROM cached WHERE v >= 0) a LEFT JOIN small b ON a.k = b._1") - .cache() - .createOrReplaceTempView("joined") - spark.table("joined").count() - checkSparkAnswerAndFallbackReason( - sql("SELECT k, v, sum(v) OVER (PARTITION BY k) AS s FROM joined"), - CometWindowExec.uncheckedInputOrderReason) - } 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 1bb942a33aa7815e9817792414a5b6091291d3e1 Mon Sep 17 00:00:00 2001 From: comphead Date: Sun, 1 Nov 2026 14:47:16 -0800 Subject: [PATCH 10/12] fix: fall back on a collated window key only over a trusted order Falling back hands the rows to Spark's window, which trusts the order the input reports. Over an order mayReportFalseOrder doubts, such as a native nested loop join's, that grouped the rows wrongly, so the window now stays native there, as on main, and the native window rejects the input. The outer join regression test partitions by k, which Spark reports through the alias a, so it needs the order declared next to the join's. --- .../spark/sql/comet/CometWindowExec.scala | 24 +++++++++++-------- .../comet/exec/CometInMemoryCacheSuite.scala | 10 ++++---- 2 files changed, 20 insertions(+), 14 deletions(-) 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 f25f9225b6a..192bdf444bf 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala @@ -75,11 +75,20 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { } } + // Spark plans no sort below the window when its input already reports the order the window + // needs. That order is relied on, by declaring it natively or by handing the rows to Spark's + // window, only when mayReportFalseOrder trusts it: Spark's window would group the rows of a + // false order wrongly, while the native one rejects an input whose order it cannot see. + val requiredOrdering = op.requiredChildOrdering.head + val trustsInputOrder = + SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) && + !CometHashJoinExec.mayReportFalseInputOrder(op.child) + // The native window finds partitions and peers by comparing keys byte by byte, so it would // tell apart collated strings that Spark treats as equal, as CometWindowGroupLimitExec notes. val collated = (op.partitionSpec ++ op.orderSpec.map(_.child)) .filter(key => hasNonDefaultStringCollation(key.dataType)) - if (collated.nonEmpty) { + if (collated.nonEmpty && trustsInputOrder) { withFallbackReason( op, collated @@ -94,16 +103,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 already reports the order the window - // needs, but rows from the JVM arrive through a ScanExec that declares none. Send the - // partition keys' order, the only one the native window checks, for the planner to declare - // (see SortedInputExec). Nothing is sent when that order may be false. Falling back would - // not help, since Spark's window trusts the order too. - val requiredOrdering = op.requiredChildOrdering.head + // Rows from the JVM arrive through a ScanExec that declares no order. Send the partition + // keys' order, the only one the native window checks, for the planner to declare (see + // SortedInputExec). val inputSortOrders = - if (op.partitionSpec.nonEmpty && - SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) && - !CometHashJoinExec.mayReportFalseInputOrder(op.child)) { + if (op.partitionSpec.nonEmpty && trustsInputOrder) { val partitionKeys = requiredOrdering.take(op.partitionSpec.length) val orders = partitionKeys.flatMap(exprToProto(_, op.child.output)) if (orders.length < partitionKeys.length) { 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 bcba4cd8a6d..8c9d68d01eb 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3860,18 +3860,20 @@ class CometInMemoryCacheSuite extends CometTestBase { Row(nullable(k), nullable(k), id, keySums(k) * keyCounts(k)) }, Seq(Seq("1 Ascending NullsFirst"), Seq("0 Ascending NullsFirst"))) - // The residual condition leaves rows of key 0 unmatched among matched ones, so the key - // stays together only if the join keeps its unmatched probe rows in place. + // The join declares the order of a, which Spark reports in place of k, so the window by k + // needs its order declared next to the join's. The residual condition leaves rows of key 0 + // unmatched among matched ones, so the key stays together only if the join keeps its + // unmatched probe rows in place. val one = spark.range(1).toDF("j").hint("broadcast") checkWindows( cached .join(one, col("k") === col("j") && col("v") % 2 === 0, "left_outer") - .selectExpr("k", "a", "v", "sum(v) OVER (PARTITION BY a) AS s"), + .selectExpr("k", "a", "v", "sum(v) OVER (PARTITION BY k) AS s"), ids.map { id => val k = key(id) Row(nullable(k), nullable(k), id, keySums(k)) }, - Seq(Seq("1 Ascending NullsFirst"))) + Seq(Seq("0 Ascending NullsFirst"))) } finally { spark.catalog.clearCache() } From 3ff09398ec81656f972a6bc7f0b3601ee034b5a8 Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 3 Nov 2026 11:07:13 -0800 Subject: [PATCH 11/12] fix: sort the input of a Spark window over a doubted order Spark plans no sort below a window whose child reports the order it needs, and its window trusts that order. Over a Comet operator that may not keep it, such as a native nested loop join (#6898), a window that stays in Spark grouped rows out of order. CometExecRule now sorts that input again. With that, a window with a collated key falls back in every case, and so does a window with ORDER BY but no PARTITION BY over a doubted order, since the native window checks no order there. The compatibility guide lists collated keys among the window's fallbacks. --- .../adding_a_new_operator.md | 7 ++- .../latest/compatibility/operators.md | 3 + .../apache/comet/rules/CometExecRule.scala | 16 ++++- .../spark/sql/comet/CometWindowExec.scala | 18 ++++-- .../apache/spark/sql/comet/operators.scala | 5 +- .../comet/exec/CometInMemoryCacheSuite.scala | 58 ++++++++++++++++++- 6 files changed, 95 insertions(+), 12 deletions(-) diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index ed48f4fba0f..b2beae1740b 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -373,9 +373,10 @@ by returning `true` from `maintains_input_order` for that child, as DataFusion's projection do. The planner declares the order through an operator only when the operator says it keeps it. An operator that cannot keep the order natively must not report it, or has to be added to `mayReportFalseOrder` in `operators.scala`. The hash join and window serdes then send no order over -it, and an outer hash join above it that has to keep that order falls back to Spark. Comet's hash -joins are the exception within one native plan: an outer hash join over one of them stays native, -even over one that declared no order, since DataFusion keeps the order it can see there. +it, an outer hash join above it that has to keep that order falls back to Spark, and a window that +runs in Spark above it gets its input sorted again. Comet's hash joins are the exception within one +native plan: an outer hash join over one of them stays native, even over one that declared no order, +since DataFusion keeps the order it can see there. A leaf that reports an ordering is where the order starts, so the planner has to know that the native operator produces its rows in that order. It knows that of an input from the JVM and of a diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 43f118942d1..8723cd58d39 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -219,6 +219,9 @@ incorrect result. When any single window expression in a `WindowExec` falls back - `lag` / `lead` with a non-literal default value ([#4268](https://github.com/apache/datafusion-comet/issues/4268)). - A `ROWS` offset that is not an integer or long, or a `RANGE` offset that is not numeric. - Any `PARTITION BY` or `ORDER BY` expression that Comet cannot serialize. +- Any `PARTITION BY` or `ORDER BY` key whose type carries a non-default `StringType` collation (e.g. + `UTF8_LCASE`). The native window finds partitions and peers by comparing keys byte by byte, which splits values + that the collation ties. A `RANGE` frame with an explicit offset on a `DATE` `ORDER BY` runs natively, but over a date outside Spark's supported range of 0001-01-01 to 9999-12-31, such as one produced by `date_add` with a very large day count, the diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 390d2954161..46beb814733 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -1082,6 +1082,20 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) visit(plan, requiresRowConsumption = true) } + /** + * Sorts the input of a window that runs in Spark when its order may be false. Spark plans no + * sort below a window whose child reports the order it needs, and its window trusts that order. + * A Comet operator that may not keep the order it reports, such as a native nested loop join + * (#6898), then hands the window rows out of order, and it groups them wrongly. + */ + private def sortDoubtedWindowInputs(plan: SparkPlan): SparkPlan = plan.transformUp { + case window: WindowExec + if window.requiredChildOrdering.head.nonEmpty && + CometHashJoinExec.mayReportFalseInputOrder(window.child) => + window.withNewChildren( + Seq(SortExec(window.requiredChildOrdering.head, global = false, child = window.child))) + } + private def _apply(plan: SparkPlan): SparkPlan = { // We shouldn't transform Spark query plan if Comet is not loaded. if (!isCometLoaded(conf)) { @@ -1133,7 +1147,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) if (clearWastedBuildSideBridges(transformed)) { transformed = transform(planWithJoinRewritten) } - var newPlan = revertUnsafePartialAggregates(transformed) + var newPlan = sortDoubtedWindowInputs(revertUnsafePartialAggregates(transformed)) // if the plan cannot be run fully natively then explain why (when appropriate // config is enabled) 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 192bdf444bf..3c3b2df03d5 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 @@ -76,9 +76,9 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { } // Spark plans no sort below the window when its input already reports the order the window - // needs. That order is relied on, by declaring it natively or by handing the rows to Spark's - // window, only when mayReportFalseOrder trusts it: Spark's window would group the rows of a - // false order wrongly, while the native one rejects an input whose order it cannot see. + // needs. The native window relies on that order only when mayReportFalseOrder trusts it. A + // window that runs in Spark over a doubted order gets its input sorted again by + // CometExecRule, since Spark's window would trust that order too. val requiredOrdering = op.requiredChildOrdering.head val trustsInputOrder = SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) && @@ -88,7 +88,7 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { // tell apart collated strings that Spark treats as equal, as CometWindowGroupLimitExec notes. val collated = (op.partitionSpec ++ op.orderSpec.map(_.child)) .filter(key => hasNonDefaultStringCollation(key.dataType)) - if (collated.nonEmpty && trustsInputOrder) { + if (collated.nonEmpty) { withFallbackReason( op, collated @@ -97,6 +97,16 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { return None } + // Without partition keys the native window checks no order at all, so over a doubted order it + // would compute over rows that may be out of order. With partition keys it rejects an input + // whose order it cannot see. + if (op.partitionSpec.isEmpty && op.orderSpec.nonEmpty && !trustsInputOrder) { + withFallbackReason( + op, + "the order of the window's input, which Spark relies on, may be false") + return None + } + val partitionExprs = op.partitionSpec.map(exprToProto(_, op.child.output)) val sortOrders = op.orderSpec.map(exprToProto(_, op.child.output)) 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 a043910e82c..2d3378388f6 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -3033,8 +3033,9 @@ trait CometHashJoin { } } - // mayReportFalseOrder for a window's input, checked before CometWindowExec sends its order. - private[comet] def mayReportFalseInputOrder(plan: SparkPlan): Boolean = + // mayReportFalseOrder for a window's input, checked by CometWindowExec before it relies on that + // order and by CometExecRule before a window that stays in Spark does. + private[apache] 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 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 8c9d68d01eb..b72bc6a65df 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -40,15 +40,16 @@ 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, CometWindowExec} +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, 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 -import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FileSourceScanExec, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo} +import org.apache.spark.sql.execution.{ColumnarToRowExec, CometSparkPlanInfoHelper, FileSourceScanExec, FilterExec, FormattedMode, RowToColumnarExec, SortExec, SparkPlan, SparkPlanInfo, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.window.WindowExec 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._ @@ -3920,6 +3921,59 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("a window that stays in Spark over a nested loop join sorts the join's rows again") { + // A native nested loop join emits its rows build row by build row, out of the streamed side's + // order that it reports (#6898), and Spark plans no sort above it. A window that runs in Spark + // there, one without partition keys or one with a collated key, gets its input sorted again. + // One partition, so no exchange comes between the join and the window. + withConversions(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + try { + val pairs = spark.range(2).toDF("z").hint("broadcast") + val sorted = (columns: Seq[String], keys: Seq[String]) => + spark + .range(0, 12, 1, 1) + .selectExpr(columns: _*) + .sortWithinPartitions(keys.map(col): _*) + .cache() + .join(pairs, col("v") % 2 <= col("z")) + val queries = Seq( + sorted(Seq("id AS v"), Seq("v")) + .selectExpr("v", "dense_rank() OVER (ORDER BY v) AS r")) ++ + (if (isSpark40Plus) { + val collated = Seq( + "id % 3 AS k", + "id AS v", + "CASE id % 4 WHEN 0 THEN 'A' WHEN 1 THEN 'a' WHEN 2 THEN 'B' ELSE 'b' END " + + "COLLATE UTF8_LCASE AS s") + Seq( + // lower(s) leaves the case-insensitive peers equal in the output, whose order among + // themselves the comparison would otherwise see. + sorted(collated, Seq("s")) + .selectExpr("lower(s) AS s", "dense_rank() OVER (ORDER BY s) AS r"), + sorted(collated, Seq("k", "s")).selectExpr( + "k", + "count(*) OVER (PARTITION BY k ORDER BY s " + + "ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS c")) + } else { + Nil + }) + queries.foreach { df => + val (_, plan) = checkSparkAnswer(df) + withClue(s"$plan\n") { + assert(collect(plan) { case j: CometBroadcastNestedLoopJoinExec => j }.nonEmpty) + val windows = collect(plan) { case w: WindowExec => w } + assert(windows.nonEmpty && windows.forall(_.child match { + case _: SortExec | WholeStageCodegenExec(_: SortExec) => true + case _ => false + })) + } + } + } 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 0f730d5b5cae2566c38c3019cf4a55367771e1b6 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 5 Nov 2026 09:04:47 -0800 Subject: [PATCH 12/12] address review: run the nested loop join window test with AQE on, sort Python windows too --- .../apache/comet/rules/CometExecRule.scala | 13 +-- .../comet/exec/CometInMemoryCacheSuite.scala | 95 ++++++++++--------- 2 files changed, 59 insertions(+), 49 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 46beb814733..b65636f7bb8 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -53,7 +53,7 @@ import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan import org.apache.spark.sql.execution.datasources.v2.text.TextScan import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, Exchange, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} -import org.apache.spark.sql.execution.window.WindowExec +import org.apache.spark.sql.execution.window.{WindowExec, WindowExecBase} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -1083,13 +1083,14 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } /** - * Sorts the input of a window that runs in Spark when its order may be false. Spark plans no - * sort below a window whose child reports the order it needs, and its window trusts that order. - * A Comet operator that may not keep the order it reports, such as a native nested loop join - * (#6898), then hands the window rows out of order, and it groups them wrongly. + * Sorts the input of a window that runs in Spark, a Python one included, when its order may be + * false. Spark plans no sort below a window whose child reports the order it needs, and its + * window trusts that order. A Comet operator that may not keep the order it reports, such as a + * native nested loop join (#6898), then hands the window rows out of order, and it groups them + * wrongly. */ private def sortDoubtedWindowInputs(plan: SparkPlan): SparkPlan = plan.transformUp { - case window: WindowExec + case window: WindowExecBase if window.requiredChildOrdering.head.nonEmpty && CometHashJoinExec.mayReportFalseInputOrder(window.child) => window.withNewChildren( 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 b72bc6a65df..2071ff41edb 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -3925,51 +3925,60 @@ class CometInMemoryCacheSuite extends CometTestBase { // A native nested loop join emits its rows build row by build row, out of the streamed side's // order that it reports (#6898), and Spark plans no sort above it. A window that runs in Spark // there, one without partition keys or one with a collated key, gets its input sorted again. - // One partition, so no exchange comes between the join and the window. - withConversions(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { - try { - val pairs = spark.range(2).toDF("z").hint("broadcast") - val sorted = (columns: Seq[String], keys: Seq[String]) => - spark - .range(0, 12, 1, 1) - .selectExpr(columns: _*) - .sortWithinPartitions(keys.map(col): _*) - .cache() - .join(pairs, col("v") % 2 <= col("z")) - val queries = Seq( - sorted(Seq("id AS v"), Seq("v")) - .selectExpr("v", "dense_rank() OVER (ORDER BY v) AS r")) ++ - (if (isSpark40Plus) { - val collated = Seq( - "id % 3 AS k", - "id AS v", - "CASE id % 4 WHEN 0 THEN 'A' WHEN 1 THEN 'a' WHEN 2 THEN 'B' ELSE 'b' END " + - "COLLATE UTF8_LCASE AS s") - Seq( - // lower(s) leaves the case-insensitive peers equal in the output, whose order among - // themselves the comparison would otherwise see. - sorted(collated, Seq("s")) - .selectExpr("lower(s) AS s", "dense_rank() OVER (ORDER BY s) AS r"), - sorted(collated, Seq("k", "s")).selectExpr( - "k", - "count(*) OVER (PARTITION BY k ORDER BY s " + - "ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS c")) - } else { - Nil - }) - queries.foreach { df => - val (_, plan) = checkSparkAnswer(df) - withClue(s"$plan\n") { - assert(collect(plan) { case j: CometBroadcastNestedLoopJoinExec => j }.nonEmpty) - val windows = collect(plan) { case w: WindowExec => w } - assert(windows.nonEmpty && windows.forall(_.child match { - case _: SortExec | WholeStageCodegenExec(_: SortExec) => true - case _ => false - })) + // One partition, so no exchange comes between the join and the window, with or without AQE. + // With AQE, the stage's own pass can make the added sort native and then the window above it. + Seq(false, true).foreach { aqe => + withConversions(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + try { + val pairs = spark.range(2).toDF("z").hint("broadcast") + val sorted = (columns: Seq[String], keys: Seq[String]) => + spark + .range(0, 12, 1, 1) + .selectExpr(columns: _*) + .sortWithinPartitions(keys.map(col): _*) + .cache() + .join(pairs, col("v") % 2 <= col("z")) + val queries = Seq( + sorted(Seq("id AS v"), Seq("v")) + .selectExpr("v", "dense_rank() OVER (ORDER BY v) AS r")) ++ + (if (isSpark40Plus) { + val collated = Seq( + "id % 3 AS k", + "id AS v", + "CASE id % 4 WHEN 0 THEN 'A' WHEN 1 THEN 'a' WHEN 2 THEN 'B' ELSE 'b' END " + + "COLLATE UTF8_LCASE AS s") + Seq( + // lower(s) leaves the case-insensitive peers equal in the output, whose order + // among themselves the comparison would otherwise see. + sorted(collated, Seq("s")) + .selectExpr("lower(s) AS s", "dense_rank() OVER (ORDER BY s) AS r"), + sorted(collated, Seq("k", "s")).selectExpr( + "k", + "count(*) OVER (PARTITION BY k ORDER BY s " + + "ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS c")) + } else { + Nil + }) + queries.foreach { df => + val (_, plan) = checkSparkAnswer(df) + withClue(s"$plan\n") { + assert(plan.isInstanceOf[AdaptiveSparkPlanExec] == aqe) + val finalPlan = stripAQEPlan(plan) + val joins = collect(finalPlan) { case j: CometBroadcastNestedLoopJoinExec => j } + assert(joins.nonEmpty) + val windowInputs = collect(finalPlan) { + case w: WindowExec => w.child + case w: CometWindowExec => w.child + } + assert(windowInputs.nonEmpty && windowInputs.forall { + case _: SortExec | _: CometSortExec | WholeStageCodegenExec(_: SortExec) => true + case _ => false + }) + } } + } finally { + spark.catalog.clearCache() } - } finally { - spark.catalog.clearCache() } } }