andygrove commented on code in PR #6835: URL: https://github.com/apache/datafusion-comet/pull/6835#discussion_r4285111316
########## native/core/src/execution/operators/sorted_input.rs: ########## @@ -0,0 +1,124 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! An input that declares the order its rows arrive in. + +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_err, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::{ + apply_expression_roots, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SendableRecordBatchStream, +}; + +/// Passes `input` through unchanged and declares to DataFusion the order its rows arrive in. +/// +/// Spark's plan knows the order of an input that enters the native plan from the JVM, a relation +/// 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. +#[derive(Debug)] +pub struct SortedInputExec { + input: Arc<dyn ExecutionPlan>, + ordering: LexOrdering, + cache: Arc<PlanProperties>, +} + +impl SortedInputExec { + pub fn new(input: Arc<dyn ExecutionPlan>, ordering: LexOrdering) -> Self { + let input_properties = input.properties(); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new_with_orderings(input.schema(), [ordering.clone()]), Review Comment: Agreed, it was dropping the constants and equivalence classes of the input. `SortedInputExec` now starts from `input.equivalence_properties()` and adds the ordering to it. Fixed in a2ad967e98. ########## spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala: ########## @@ -2254,4 +2255,636 @@ class CometJoinSuite extends CometTestBase { } } } + + // Whether every partition of `df` is in `order`, which is bound to its output. + private def partitionsInOrder(df: DataFrame, order: SortOrder*): Boolean = { + val ordering = new InterpretedOrdering(order) + df.queryExecution.toRdd.map(_.copy()).glom().collect().forall { part => + part.zip(part.drop(1)).forall { case (a, b) => !ordering.gt(a, b) } + } + } + + // The order of the integer key in the first output column. + private def keyOrder(direction: SortDirection): SortOrder = Review Comment: Added `outer hash joins keep the streamed order for decimal, string and timestamp keys` in ddce25a6fe. It runs both outer shapes on a `decimal(20,2)` key with values at both ends of the range, a string key with accented, Chinese, Japanese, Korean and emoji values, and a timestamp key before 1970 in `America/Los_Angeles`. Each case checks that the serde sent the order, that no sort is left in the plan, that the partitions are in order, and that the rows and a sorted aggregate on the key match Spark's. The test comment says why struct, array and map keys are out of scope. ########## native/core/src/execution/planner.rs: ########## @@ -2937,6 +2951,63 @@ 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, but the `ScanExec` it arrives through declares none. Only that order is declared: + /// when an operator in between may have moved rows, `plan` is left as it is. + 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_jvm_input_order(&plan) { Review Comment: Thanks, that was the shape. `CometSampleExec` now reports no order, as Spark's does, so the serde sends nothing for a join over a sample and the two checks agree there. I kept the `debug_assert` out and wrote the invariant down instead, see my reply on parthchandra's question. Fixed in a2ad967e98. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
