sunchao commented on code in PR #7256:
URL: https://github.com/apache/datafusion-comet/pull/7256#discussion_r4426184807
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala:
##########
@@ -81,10 +88,45 @@ 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). 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 (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)
Review Comment:
The partition keys are sufficient for the ordering declaration, but a window
still compares its `ORDER BY` keys to find peers. With the native cache reader,
I cached `(k, s)` sorted by both columns, where `k` is an integer and `s` uses
`UTF8_LCASE`. For `A, a, B, b` within each `k`, Spark 4.1.3 gives `dense_rank()
OVER (PARTITION BY k ORDER BY s)` ranks `1, 1, 2, 2`. This commit runs
`CometWindow` directly over `CometInMemoryTableScan` and gives `1, 2, 3, 4`.
The previous revision fell back and matched Spark. Checking collation on the
order keys before conversion would still allow the declaration itself to
contain only partition keys.
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala:
##########
@@ -81,10 +88,45 @@ 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). 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 (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
Review Comment:
Falling back to Spark here leaves the window above a join whose output is
out of order. I cached 12 rows in one partition, sorted by `k = v % 3`, then
broadcast two build rows (`z = 0, 1`) and joined on `v % 2 <= z`. The native
nested-loop join emits keys `0, 1, 2` twice, once for each build row. Spark has
already omitted the sort, so `count(*) OVER (PARTITION BY k)` counts the
separate runs. Six rows get 2 and twelve get 4. Vanilla Spark and a control
with a physical sort after the same native join both return 6 on all 18 rows.
Caching the join output also reproduces it. The previous revision raised an
ordering error on this query. The fallback needs to restore the required order
before the Spark window consumes those rows.
##########
native/core/src/execution/planner/join.rs:
##########
@@ -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) {
Review Comment:
This early return still drops a valid ordering above a left outer broadcast
hash join. I sorted `range(12)` by `k = id % 3`, projected `k AS a, k AS b, id
AS v`, and cached the result in one partition. I then left-joined a broadcast
one-row table on `a`, kept both aliases in the output, and ran `sum(v) OVER
(PARTITION BY a)`. Spark knows `a` and `b` have the same order and returns all
12 rows (group sums `18, 22, 26`). The native join declares its probe ordered
by `b`, and the window sends `a`, but `keeps_source_order` rejects the
DataFusion right outer join even with that ordered probe. The task fails with
`All partition by columns should have an ordering`. Changing only the window to
`PARTITION BY b` returns all 12 expected rows with the same native join and no
sort.
--
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]
