fix: declare the input order of a native window over a sorted cache

docs/source/contributor-guide/debugging.md

@@ -260,7 +260,8 @@ so `spark.comet.debug.enabled` and `COMET_LOG_LEVEL` do not affect it.
 #### Example: dumping window operator output
 
 In `WindowBuilder::build` in `native/core/src/execution/planner/window.rs`,
-the window operator builds `window_agg`, either a `BoundedWindowAggExec` or a
+the window operator builds `window_agg` over `input`, the child plan with
+Spark's input order declared, either a `BoundedWindowAggExec` or a
 `WindowAggExec` depending on whether every window expression runs with bounded
 memory. It then optionally wraps `window_agg` in a casting `ProjectionExec` and
 builds the `SparkPlan` from either the projection or `window_agg`. To see the raw
@@ -271,14 +272,14 @@ projection:
 let window_agg: Arc<dyn ExecutionPlan> = 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(),
     )?)
 };

native/core/src/execution/planner.rs

@@ -4270,6 +4270,48 @@ mod tests {
         assert!(!PhysicalPlanner::keeps_source_order(&scan(2)));
     }
 
+    #[test]
+    fn declared_order_extends_an_order_the_input_declares() {
+        // An order is declared over an empty bucket, and next to the order the input declares,
+        // unless an operator in between may have moved rows.
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("col_0", DataType::Int32, true),
+            Field::new("col_1", DataType::Int32, true),
+        ]));
+        let planner = PhysicalPlanner::default();
+        let declare = |plan: &Arc<dyn ExecutionPlan>, index: i32| {
+            planner
+                .with_declared_order(
+                    Arc::clone(plan),
+                    &[create_sort_order(index, 3, false, true)],
+                )
+                .unwrap()
+        };
+        let empty: Arc<dyn ExecutionPlan> =
+            Arc::new(datafusion::physical_plan::empty::EmptyExec::new(schema));
+        let by_col_0 = declare(&empty, 0);
+        assert_eq!("SortedInputExec", by_col_0.name());
+        assert!(Arc::ptr_eq(&declare(&by_col_0, 0), &by_col_0));
+
+        let both = declare(&by_col_0, 1);
+        assert!(Arc::ptr_eq(both.children()[0], &by_col_0));
+        for index in 0..2 {
+            let partition_by: Vec<Arc<dyn PhysicalExpr>> =
+                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<dyn ExecutionPlan> = Arc::new(SortExec::new(
+            LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(Column::new(
+                "col_0", 0,
+            )))])
+            .unwrap(),
+            empty,
+        ));
+        assert!(Arc::ptr_eq(&declare(&sort, 1), &sort));
+    }
+
     #[tokio::test]
     async fn spark_plan_dynamic_filter_preserves_metrics() {
         use crate::execution::metrics::utils::to_native_metric_node;
@@ -4337,24 +4379,22 @@ mod tests {
         }
     }
 
-    #[tokio::test]
-    async fn window_cast_projection_reports_window_time() {
-        use crate::execution::metrics::utils::to_native_metric_node;
+    /// `func() OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)`, returning
+    /// `result_type`.
+    fn window_function_expr(
+        func: &str,
+        result_type: Option<spark_expression::DataType>,
+    ) -> spark_operator::WindowExpr {
         use spark_operator::{
             lower_window_frame_bound::LowerFrameBoundStruct,
             upper_window_frame_bound::UpperFrameBoundStruct, CurrentRow, LowerWindowFrameBound,
             UnboundedPreceding, UpperWindowFrameBound, WindowExpr, WindowFrame, WindowFrameType,
             WindowSpecDefinition,
         };
-
-        let session = Arc::new(SessionContext::new());
-        let planner = PhysicalPlanner::new(Arc::clone(&session), 0);
-        // `row_number` returns UInt64 natively and Spark expects Int32, so the planner adds a
-        // casting projection above the window.
-        let row_number = WindowExpr {
+        WindowExpr {
             built_in_window_function: Some(Expr {
                 expr_struct: Some(ScalarFunc(spark_expression::ScalarFunc {
-                    func: "row_number".to_string(),
+                    func: func.to_string(),
                     ..Default::default()
                 })),
                 query_context: None,
@@ -4376,9 +4416,20 @@ mod tests {
                 }),
                 ..Default::default()
             }),
-            result_type: Some(create_proto_datatype()),
+            result_type,
             ..Default::default()
-        };
+        }
+    }
+
+    #[tokio::test]
+    async fn window_cast_projection_reports_window_time() {
+        use crate::execution::metrics::utils::to_native_metric_node;
+
+        let session = Arc::new(SessionContext::new());
+        let planner = PhysicalPlanner::new(Arc::clone(&session), 0);
+        // `row_number` returns UInt64 natively and Spark expects Int32, so the planner adds a
+        // casting projection above the window.
+        let row_number = window_function_expr("row_number", Some(create_proto_datatype()));
         let operator = Operator {
             children: vec![create_scan()],
             op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window {
@@ -4433,6 +4484,53 @@ mod tests {
         );
     }
 
+    /// `func` over a window partitioned and ordered by `col_0` over `input`.
+    fn partitioned_window(func: &str, input: Operator, input_sort_orders: Vec<Expr>) -> Operator {
+        Operator {
+            children: vec![input],
+            op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window {
+                window_expr: vec![window_function_expr(func, None)],
+                partition_by_list: vec![create_bound_reference(0)],
+                order_by_list: vec![create_sort_order(0, 3, false, true)],
+                input_sort_orders,
+                ..Default::default()
+            }))),
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn window_declares_input_order_on_a_scan() {
+        // Both window operators need their partition keys' order, which a `ScanExec` declares
+        // only through the order Spark sent. `percent_rank` runs in `WindowAggExec`.
+        let partition_by: Vec<Arc<dyn PhysicalExpr>> = vec![Arc::new(Column::new("col_0", 0))];
+        for (func, operator) in [
+            ("row_number", "BoundedWindowAggExec"),
+            ("percent_rank", "WindowAggExec"),
+        ] {
+            let window_input = |input_sort_orders| {
+                let (_, _, planned) = PhysicalPlanner::default()
+                    .create_plan(
+                        &partitioned_window(func, create_scan(), input_sort_orders),
+                        &mut vec![],
+                        1,
+                    )
+                    .unwrap();
+                assert_eq!(operator, planned.native_plan.name());
+                Arc::clone(planned.native_plan.children()[0])
+            };
+            let input = window_input(vec![create_sort_order(0, 3, false, true)]);
+            assert_eq!("SortedInputExec", input.name());
+            let ordered = get_ordered_partition_by_indices(&partition_by, &input).unwrap();
+            assert_eq!(1, ordered.len(), "{func}");
+
+            let input = window_input(vec![]);
+            assert_eq!("ScanExec", input.name());
+            let ordered = get_ordered_partition_by_indices(&partition_by, &input).unwrap();
+            assert!(ordered.is_empty(), "{func}");
+        }
+    }
+
     #[test]
     fn projection_rolls_up_same_plan_id_child_metrics() {
         let op_scan = Operator {

native/core/src/execution/planner/join.rs

@@ -35,6 +35,7 @@ use datafusion::{
         LexOrdering, PhysicalExpr,
     },
     physical_plan::{
+        empty::EmptyExec,
         filter::FilterExec,
         joins::{
             utils::JoinFilter, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec,
@@ -402,48 +403,59 @@ impl PhysicalPlanner {
         Ok(Arc::new(prepared))
     }
 
-    /// `plan`, declaring `sort_orders` as the order its rows arrive in, unless it already
-    /// declares one. Spark's plan knows the order of an input that enters the native plan from
-    /// the JVM, and of a scan of sorted files, but the `ScanExec` or the file scan the rows
-    /// arrive through declares none. Only that order is declared: when an operator in between may
-    /// have moved rows, `plan` is left as it is.
+    /// `plan`, declaring `sort_orders` as the order its rows arrive in, unless the order it
+    /// declares already satisfies them. Spark's plan knows the order of an input that enters the
+    /// native plan from the JVM, and of a scan of sorted files, but the `ScanExec` or the file
+    /// scan the rows arrive through declares none. Only that order is declared: when an operator
+    /// in between may have moved rows, `plan` is left as it is.
     ///
-    /// The serde sends an order only when the streamed side's Comet operators report it, and it
-    /// checks that nothing under them moves rows (`mayReportFalseOrder`), while this walk asks
-    /// the native operators the same thing through `maintains_input_order`. The two agree as
-    /// long as a Comet operator that reports an order keeps it natively and says so, and a Comet
-    /// leaf that reports one is a source this walk knows or declares its order itself, as the
-    /// native range does. One that does not lands here with `plan` unchanged, and the join then
-    /// moves its unmatched rows after Spark has dropped the sort above it. Nothing here can fall
-    /// back at that point, so the invariant is kept at the operator: see "Adding a New Operator"
-    /// in the contributor guide.
-    fn with_declared_order(
+    /// An input that declares only part of `sort_orders`, such as a window over another window,
+    /// gets them declared next to its own order.
+    ///
+    /// The serde sends an order only when the Comet operators of a join's streamed side or of a
+    /// window's input report it, and it checks that nothing under them moves rows
+    /// (`mayReportFalseOrder`), while this walk asks the native operators the same thing through
+    /// `maintains_input_order`. The two agree as long as a Comet operator that reports an order
+    /// keeps it natively and says so, and a Comet leaf that reports one is a source this walk
+    /// knows or declares its order itself, as the native range does. One that does not lands
+    /// here with `plan` unchanged, and the join then moves its unmatched rows after Spark has
+    /// dropped the sort above it, or the window rejects its input. Nothing here can fall back at
+    /// that point, so the invariant is kept at the operator: see "Adding a New Operator" in the
+    /// contributor guide.
+    pub(super) fn with_declared_order(
         &self,
         plan: Arc<dyn ExecutionPlan>,
         sort_orders: &[Expr],
     ) -> Result<Arc<dyn ExecutionPlan>, ExecutionError> {
-        if plan.output_ordering().is_some() || !Self::keeps_source_order(&plan) {
+        if sort_orders.is_empty() || !Self::keeps_source_order(&plan) {
             return Ok(plan);
         }
         let exprs = sort_orders
             .iter()
             .map(|expr| self.create_sort_expr(expr, plan.schema()))
             .collect::<Result<Vec<_>, _>>()?;
         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),
         }
     }
 
     /// Whether `plan`'s rows come from a source whose order Spark knows through operators that
     /// keep their input's order, so that the order Spark reports for them still holds. A source is
     /// an input from the JVM, or a scan of files, which reads the files of its one file group one
     /// after the other: Spark reports an order for a file scan only when every partition reads one
-    /// sorted file, as for a sorted bucketed table with one file per bucket. Any other operator may
-    /// have moved rows, such as the null-aware anti join, whose streamed side Comet builds on, and
-    /// declaring Spark's order over it would be wrong.
+    /// 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<dyn ExecutionPlan>) -> bool {
-        if plan.is::<ScanExec>() || plan.is::<ShuffleScanExec>() {
+        if plan.is::<ScanExec>() || plan.is::<ShuffleScanExec>() || plan.is::<EmptyExec>() {
             return true;
         }
         if let Some(scan) = plan.downcast_ref::<DataSourceExec>() {

native/core/src/execution/planner/window.rs

@@ -131,17 +131,21 @@ impl OperatorBuilder for WindowBuilder {
         // trigger a retract call.
         let window_expr = window_expr?;
         let all_bounded = window_expr.iter().all(|e| e.uses_bounded_memory());
+
+        // Both window operators reject an input that does not declare its partition keys' order.
+        let input =
+            planner.with_declared_order(Arc::clone(&child.native_plan), &wnd.input_sort_orders)?;
         let window_agg: Arc<dyn ExecutionPlan> = 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(),
             )?)
         };

native/proto/src/proto/operator.proto

@@ -1233,6 +1233,9 @@ message Window {
   repeated spark.spark_expression.Expr order_by_list = 2;
   repeated spark.spark_expression.Expr partition_by_list = 3;
   Operator child = 4;
+  // The partition keys ascending, then the order keys, when Spark planned no sort below the
+  // window because its input reports that order. Empty otherwise.
+  repeated spark.spark_expression.Expr input_sort_orders = 5;
 }
 
 // Top-K rows per partition group. Corresponds to Spark's WindowGroupLimitExec

spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala

@@ -36,7 +36,7 @@ import org.apache.comet.{CometConf, ConfigEntry}
 import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
 import org.apache.comet.serde.{AggSerde, CometOperatorSerde, CometSortOrder, LiteralOuterClass, OperatorOuterClass}
 import org.apache.comet.serde.OperatorOuterClass.Operator
-import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, liftFallbackReasons, scalarFunctionExprToProto, serializeDataType}
+import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, liftFallbackReasons, scalarFunctionExprToProto, serializeDataType}
 
 object CometWindowExec extends CometOperatorSerde[WindowExec] {
 
@@ -81,10 +81,37 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] {
 
     if (windowExprProto.forall(_.isDefined) && partitionExprs.forall(_.isDefined)
       && sortOrders.forall(_.isDefined)) {
+      // Spark plans no sort below the window when its input already reports the order the window
+      // needs, but rows from the JVM arrive through a ScanExec that declares none. Send that order
+      // for the planner to declare (see SortedInputExec), unless a native sort below declares it
+      // or it may be false. A key Comet cannot express, or a collated string key, sends the window
+      // back to Spark.
+      val requiredOrdering = op.requiredChildOrdering.head
+      val inputSortOrders =
+        if (requiredOrdering.nonEmpty && !op.child.isInstanceOf[CometSortExec] &&
+          SortOrder.orderingSatisfies(op.child.outputOrdering, requiredOrdering) &&
+          !CometHashJoinExec.mayReportFalseInputOrder(op.child)) {
+          val orders = requiredOrdering.map { order =>
+            if (hasNonDefaultStringCollation(order.dataType)) None
+            else exprToProto(order, op.child.output)
+          }
+          if (!orders.forall(_.isDefined)) {
+            withFallbackReason(
+              op,
+              "the order of the window's input, which Spark relies on, cannot be expressed " +
+                "natively")
+            return None
+          }
+          orders.flatten
+        } else {
+          Nil
+        }
+
       val windowBuilder = OperatorOuterClass.Window.newBuilder()
       windowBuilder.addAllWindowExpr(windowExprProto.map(_.get).asJava)
       windowBuilder.addAllPartitionByList(partitionExprs.map(_.get).asJava)
       windowBuilder.addAllOrderByList(sortOrders.map(_.get).asJava)
+      windowBuilder.addAllInputSortOrders(inputSortOrders.asJava)
       Some(builder.setWindow(windowBuilder).build())
     } else {
       // Whichever of the window / partition / order expressions failed has already recorded its

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

@@ -3017,6 +3017,10 @@ trait CometHashJoin {
     }
   }
 
+  // mayReportFalseOrder for a window's input, checked before CometWindowExec sends its order.
+  private[comet] def mayReportFalseInputOrder(plan: SparkPlan): Boolean =
+    mayReportFalseOrder(plan, countInPlan = true)
+
   // Whether the order that `plan` reports may be false, because something that feeds it with no
   // sort in between cannot be relied on to keep it. Comet's hash joins report the order of their
   // streamed side, which Spark's own joins keep, but two Comet joins do not always: an outer join