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,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
-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, 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

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(),
     )?)
 };

docs/source/user-guide/latest/compatibility/operators.md

@@ -222,6 +222,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

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 {