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

docs/source/contributor-guide/adding_a_new_operator.md

@@ -365,17 +365,17 @@ the order satisfies. Report what the Spark operator reports and nothing more. `S
 example, reports no order, so `CometSampleExec` must not report its child's.
 
 The claim has to hold at run time as well. A native hash join declares the order Spark reports for
-its streamed side to DataFusion, and an aggregate above the join may then run in sorted mode, or a
-top-k sort may stop reading early, on the strength of that order. If the order is false, the
-result is wrong groups or wrong rows, not only rows out of order. So an operator that reports an
-ordering must keep it natively, and the native operator must say so by returning `true` from
-`maintains_input_order` for that child, as DataFusion's filter and projection do. The planner
-declares the order through an operator only when the operator says it keeps it. An operator that
-cannot keep the order natively must not report it, or has to be added to `mayReportFalseOrder` in
-`operators.scala`. The hash join serde then sends no order over it, and an outer hash join above it
-that has to keep that order falls back to Spark. Comet's hash joins are the exception within one
-native plan: an outer hash join over one of them stays native, even over one that declared no
-order, since DataFusion keeps the order it can see there.
+its streamed side to DataFusion, and a native window the order of its partition keys. An aggregate
+above may then run in sorted mode, or a top-k sort may stop reading early, on the strength of that
+order. If the order is false, the result is wrong groups or wrong rows, not only rows out of order.
+So an operator that reports an ordering must keep it natively, and the native operator must say so
+by returning `true` from `maintains_input_order` for that child, as DataFusion's filter and
+projection do. The planner declares the order through an operator only when the operator says it
+keeps it. An operator that cannot keep the order natively must not report it, or has to be added to
+`mayReportFalseOrder` in `operators.scala`. The hash join and window serdes then send no order over
+it, and an outer hash join above it that has to keep that order falls back to Spark. Comet's hash
+joins are the exception within one native plan: an outer hash join over one of them stays native,
+even over one that declared no order, since DataFusion keeps the order it can see there.
 
 A leaf that reports an ordering is where the order starts, so the planner has to know that the
 native operator produces its rows in that order. It knows that of an input from the JVM and of a

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/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<dyn ExecutionPlan>,

native/core/src/execution/planner.rs

@@ -4163,8 +4163,10 @@ mod tests {
                 "col_0@0 ASC",
                 probe.properties().output_ordering().unwrap().to_string()
             );
-            // The join keeps the order in its own output, so the operators above it see it too.
+            // The join keeps the order in its own output, so the operators above it see it too,
+            // and an order declared above it holds, such as a window's for an alias of the key.
             assert!(planned.native_plan.properties().output_ordering().is_some());
+            assert!(PhysicalPlanner::keeps_source_order(&planned.native_plan));
         }
     }
 
@@ -4332,6 +4334,48 @@ mod tests {
         assert!(!PhysicalPlanner::keeps_source_order(&scan(2)));
     }
 
+    #[test]
+    fn declared_order_extends_an_order_the_input_declares() {
+        // An order is declared over an empty bucket, and next to the order the input declares,
+        // unless an operator in between may have moved rows.
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("col_0", DataType::Int32, true),
+            Field::new("col_1", DataType::Int32, true),
+        ]));
+        let planner = PhysicalPlanner::default();
+        let declare = |plan: &Arc<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;
@@ -4399,24 +4443,22 @@ mod tests {
         }
     }
 
-    #[tokio::test]
-    async fn window_cast_projection_reports_window_time() {
-        use crate::execution::metrics::utils::to_native_metric_node;
+    /// `func() OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)`, returning
+    /// `result_type`.
+    fn window_function_expr(
+        func: &str,
+        result_type: Option<spark_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,
@@ -4438,9 +4480,20 @@ mod tests {
                 }),
                 ..Default::default()
             }),
-            result_type: Some(create_proto_datatype()),
+            result_type,
             ..Default::default()
-        };
+        }
+    }
+
+    #[tokio::test]
+    async fn window_cast_projection_reports_window_time() {
+        use crate::execution::metrics::utils::to_native_metric_node;
+
+        let session = Arc::new(SessionContext::new());
+        let planner = PhysicalPlanner::new(Arc::clone(&session), 0);
+        // `row_number` returns UInt64 natively and Spark expects Int32, so the planner adds a
+        // casting projection above the window.
+        let row_number = window_function_expr("row_number", Some(create_proto_datatype()));
         let operator = Operator {
             children: vec![create_scan()],
             op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window {
@@ -4495,6 +4548,45 @@ mod tests {
         );
     }
 
+    #[test]
+    fn window_declares_input_order_on_a_scan() {
+        // Both window operators need their partition keys' order, which a `ScanExec` declares
+        // only through the order Spark sent. `percent_rank` runs in `WindowAggExec`.
+        let partition_by: Vec<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 window = Operator {
+                    children: vec![create_scan()],
+                    op_struct: Some(OpStruct::Window(Box::new(spark_operator::Window {
+                        window_expr: vec![window_function_expr(func, None)],
+                        partition_by_list: vec![create_bound_reference(0)],
+                        order_by_list: vec![create_sort_order(0, 3, false, true)],
+                        input_sort_orders,
+                        ..Default::default()
+                    }))),
+                    ..Default::default()
+                };
+                let (_, _, planned) = PhysicalPlanner::default()
+                    .create_plan(&window, &mut vec![], 1)
+                    .unwrap();
+                assert_eq!(operator, planned.native_plan.name());
+                Arc::clone(planned.native_plan.children()[0])
+            };
+            let input = window_input(vec![create_sort_order(0, 3, false, true)]);
+            assert_eq!("SortedInputExec", input.name());
+            let ordered = get_ordered_partition_by_indices(&partition_by, &input).unwrap();
+            assert_eq!(1, ordered.len(), "{func}");
+
+            let input = window_input(vec![]);
+            assert_eq!("ScanExec", input.name());
+            let ordered = get_ordered_partition_by_indices(&partition_by, &input).unwrap();
+            assert!(ordered.is_empty(), "{func}");
+        }
+    }
+
     #[test]
     fn projection_rolls_up_same_plan_id_child_metrics() {
         let op_scan = Operator {

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

@@ -35,6 +35,7 @@ use datafusion::{
         LexOrdering, PhysicalExpr,
     },
     physical_plan::{
+        empty::EmptyExec,
         filter::FilterExec,
         joins::{
             utils::JoinFilter, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec,
@@ -478,48 +479,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>() {
@@ -529,18 +541,21 @@ impl PhysicalPlanner {
                 .is_some_and(|config| config.file_groups.len() == 1);
         }
         // DataFusion reports every right hash join as keeping its probe order, but a right outer
-        // join keeps it only when the probe declares one, and no input on this path declares an
-        // order. An inner join emits its matches probe row by probe row, and a right semi, anti or
-        // mark join builds its probe indices in ascending order, so none of them depends on a
-        // declaration. Any other join type keeps no probe order at all.
+        // join keeps its unmatched probe rows in place only when the probe declares an order, as
+        // the planner does for the join's streamed side. An inner join emits its matches probe
+        // row by probe row, and a right semi, anti or mark join builds its probe indices in
+        // ascending order, so none of them depends on a declaration. Any other join type keeps no
+        // probe order at all.
         if let Some(join) = plan.downcast_ref::<HashJoinExec>() {
-            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;
             }
         }