dongjoon-hyun commented on code in PR #59306:
URL: https://github.com/apache/spark/pull/59306#discussion_r4235381978
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -253,23 +256,33 @@ case class GroupPartitionsExec(
case _ => true
}
+ /** Whether the k-way merge was planned for this node. See
`sortedMergeOrdering`. */
+ def enableSortedMerge: Boolean = sortedMergeOrdering.isDefined
+
+ /**
+ * The ordering the k-way merge merges by. Before the merge is planned it is
the child's. After
+ * that it is the one the merge was planned with, see `sortedMergeOrdering`.
+ */
+ private def mergeOrdering: Seq[SortOrder] =
sortedMergeOrdering.getOrElse(child.outputOrdering)
+
/** Whether a k-way merge would work at all, leaving aside whether it is
switched on. */
@transient private lazy val kWayMergeIsFeasible: Boolean =
- child.outputOrdering.nonEmpty && childIsSafeForKWayMerge
+ mergeOrdering.nonEmpty && childIsSafeForKWayMerge
- /** Whether this node performs the k-way merge. No config term, see
`enableSortedMerge`. */
+ /** Whether this node performs the k-way merge. No config term, see
`sortedMergeOrdering`. */
@transient private lazy val usesSortedMerge: Boolean =
Review Comment:
This is the same class of issue as SPARK-59279 and this PR, just through a
different input. The scaladoc of `sortedMergeOrdering` says it is "the
decision, not a hint", but `usesSortedMerge` still re-derives
`childIsSafeForKWayMerge` against whatever child the current instance holds,
and every `withNewChildInternal` / `copy` mints a new instance.
So if a rule after `EnsureRequirements` ever rebuilds a planned merge over a
child that holds a node outside the `SafeForKWayMerge` whitelist, for example a
transition or wrapper inserted by a columnar plugin, the new instance
concatenates and `outputOrdering` drops to the key-only branch. Outside the
`AQEShuffleReadRule` path nothing re-checks the parent's requirement, so a
window or sort-merge join above it reads unsorted rows with no error.
I know the comment above `childIsSafeForKWayMerge` says the correctness
relies on all insertable wrappers being whitelisted, and no rule is known to
break that today. But for a node where `enableSortedMerge` is true, would it be
safer to fail explicitly, like `checkChildStillMatches` does for the
partitioning, instead of silently falling back?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -253,23 +256,33 @@ case class GroupPartitionsExec(
case _ => true
}
+ /** Whether the k-way merge was planned for this node. See
`sortedMergeOrdering`. */
+ def enableSortedMerge: Boolean = sortedMergeOrdering.isDefined
+
+ /**
+ * The ordering the k-way merge merges by. Before the merge is planned it is
the child's. After
+ * that it is the one the merge was planned with, see `sortedMergeOrdering`.
+ */
+ private def mergeOrdering: Seq[SortOrder] =
sortedMergeOrdering.getOrElse(child.outputOrdering)
+
/** Whether a k-way merge would work at all, leaving aside whether it is
switched on. */
@transient private lazy val kWayMergeIsFeasible: Boolean =
- child.outputOrdering.nonEmpty && childIsSafeForKWayMerge
+ mergeOrdering.nonEmpty && childIsSafeForKWayMerge
- /** Whether this node performs the k-way merge. No config term, see
`enableSortedMerge`. */
+ /** Whether this node performs the k-way merge. No config term, see
`sortedMergeOrdering`. */
@transient private lazy val usesSortedMerge: Boolean =
enableSortedMerge && hasCoalescing && kWayMergeIsFeasible
/**
* Returns a copy of this node with k-way merge enabled, when the config is
on, this node
* coalesces partitions and the merge is feasible. The only read of
- * `preserveOrderingOnCoalesce`.
+ * `preserveOrderingOnCoalesce`. The copy keeps the ordering checked here,
see
+ * `sortedMergeOrdering`.
*/
def tryEnableSortedMerge(): Option[GroupPartitionsExec] = {
Option.when(conf.v2BucketingPreserveOrderingOnCoalesceEnabled &&
hasCoalescing &&
kWayMergeIsFeasible) {
- val newGroupPartitions = copy(enableSortedMerge = true)
+ val newGroupPartitions = copy(sortedMergeOrdering = Some(mergeOrdering))
Review Comment:
Here the feasibility check and the captured ordering come from two different
reads. `kWayMergeIsFeasible` is a `lazy val`, so it may hold
`mergeOrdering.nonEmpty` from an earlier call on this instance, while
`Some(mergeOrdering)` reads `child.outputOrdering` again.
If `tryEnableSortedMerge` runs a second time on the same instance after
`partitionKeyOrdering.enabled` was turned off, for example from an
`ensureRequirements` re-run inside AQE's `OptimizeSkewedJoin`, the cached
`true` passes and the capture is `Some(Seq())`. That node reports `SortedMerge:
true` and is refused by `withKeyPositionsFor`, but it never merges.
It is hard to reach today, but the doc above says "The copy keeps the
ordering checked here". How about reading the ordering once, so that the check
and the capture are guaranteed to use the same value?
```scala
val ordering = child.outputOrdering
Option.when(conf.v2BucketingPreserveOrderingOnCoalesceEnabled &&
hasCoalescing &&
ordering.nonEmpty && childIsSafeForKWayMerge) {
val newGroupPartitions = copy(sortedMergeOrdering = Some(ordering))
...
}
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -490,8 +512,8 @@ case class GroupPartitionsExec(
s"Reducers: ${truncatedString(names, "[", ", ", "]", joinKeyMaxFields)}"
}
val distributeStr = Iterator(s"DistributePartitions:
$distributePartitions")
- // Rendered from the constructor field, as `DistributePartitions` above
is, and not from
- // `usesSortedMerge`, which reads the child's ordering. This method feeds
`simpleString`, which
+ // Rendered from the constructor field `sortedMergeOrdering`, as
`DistributePartitions` above
+ // is. `usesSortedMerge` would walk the child. This method feeds
`simpleString`, which
// `treeString` calls on error paths, so it stays on what the node was
built with.
val sortedMergeStr = Iterator(s"SortedMerge: $enableSortedMerge")
Review Comment:
Since the frozen ordering can now differ from what the child reports, and
that is exactly the case this PR handles, could EXPLAIN show it? For example
`SortedMerge: [id ASC NULLS FIRST, name ASC NULLS FIRST]` instead of
`SortedMerge: true`.
It is still rendered from the constructor field, so it keeps the rule in the
comment above that error-path rendering must not walk the child. Today a plan
where the scan reports no ordering but the merge sorts by `[id, name]` shows
only `SortedMerge: true`, which makes a wrong-ordering report hard to debug
from EXPLAIN alone.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -428,16 +443,24 @@ case class GroupPartitionsExec(
* `plannedPartitioning` are decided for one child's partitioning, so a node
handed another has to
* be decided again over it.
*
- * Rebuilding is all this does: what the operator above reads, the ordering
it was planned
- * against included, is the caller's to hold, since nothing here knows what
that operator
- * requires.
+ * A merging node is turned away. Its frozen `sortedMergeOrdering` was
checked against this node's
+ * child, not against `newChild`. No merging node reaches this path today.
The only caller,
+ * `CombineAdjacentAggregation`, rebuilds a node over a subtree that holds
an aggregate, which is
+ * not `SafeForKWayMerge`.
+ *
+ * Rebuilding is all this does. Nothing here knows what the operator above
requires, so checking
+ * that the rebuilt node still meets it is the caller's job.
*/
def withKeyPositionsFor(newChild: SparkPlan): Option[GroupPartitionsExec] = {
// An aligned node's slot order comes from the parent distribution rather
than the child, so the
// factory cannot re-decide it for `newChild`.
if (expectedKeyCount.isDefined) {
return None
}
+ // The frozen merge ordering was checked against this node's child, not
against `newChild`.
+ if (enableSortedMerge) {
Review Comment:
The scaladoc says no merging node reaches this path, since
`CombineAdjacentAggregation` only rebuilds over a subtree that holds an
aggregate. So this guard and its test cover an unreachable state. It also
changes behavior: before, the rebuild carried `enableSortedMerge` over, and now
it is refused.
If a future caller does reach it, the combine optimization is skipped
silently. Would it be better to rebuild without the frozen ordering and let
`tryEnableSortedMerge` decide again over `newChild`? Or, if refusing is the
intent, an assertion would at least make a newly reachable path visible.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -253,23 +256,33 @@ case class GroupPartitionsExec(
case _ => true
}
+ /** Whether the k-way merge was planned for this node. See
`sortedMergeOrdering`. */
+ def enableSortedMerge: Boolean = sortedMergeOrdering.isDefined
+
+ /**
+ * The ordering the k-way merge merges by. Before the merge is planned it is
the child's. After
+ * that it is the one the merge was planned with, see `sortedMergeOrdering`.
+ */
+ private def mergeOrdering: Seq[SortOrder] =
sortedMergeOrdering.getOrElse(child.outputOrdering)
Review Comment:
`mergeOrdering` means the live child ordering before the merge is planned
and the frozen one after it. That makes the `mergeOrdering.nonEmpty` term in
`kWayMergeIsFeasible` vacuous for a planned node, or else a silent off switch.
Also, nothing prevents `sortedMergeOrdering = Some(Seq.empty)`. The factory,
`withJoinKeyPositions` and the tests all accept it, and the result claims
`enableSortedMerge` and `SortedMerge: true` while it concatenates.
Could we add a `require(sortedMergeOrdering.forall(_.nonEmpty))` to the
constructor, have `tryEnableSortedMerge` read `child.outputOrdering` directly,
and have the merge path use the frozen ordering only? Then the two meanings
would not share one accessor.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala:
##########
@@ -493,7 +505,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
withSQLConf(
SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key ->
configEnabled.toString) {
- val flagged = GroupPartitionsExec(child, enableSortedMerge = true)
+ val flagged = GroupPartitionsExec(child, sortedMergeOrdering =
Some(child.outputOrdering))
Review Comment:
These tests now hand `Some(child.outputOrdering)` in directly, and
`child.outputOrdering` does not change during the test, so they would pass with
or without the freeze. The `tryEnableSortedMerge` test (SPARK-56549 at L519)
checks only `enableSortedMerge` and `outputOrdering`, which would also pass if
the capture were `Some(Seq())` or if `withNewChildInternal` dropped the field,
since `outputOrdering` falls back to the same live value.
Could we add a unit test with a child whose reported ordering can be
changed? It would check that `tryEnableSortedMerge` captures the ordering, and
that `copy`, `withNewChildren` and `outputOrdering` keep it after the child's
ordering becomes empty. Today only the end-to-end test in
`KeyGroupedPartitioningSuite` covers this, and only for one table layout.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -75,14 +75,17 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
* [[REPLICATED_FOR_JOIN]] groups the splits first and
repeats each group
* into its expected slots, and `None` groups them,
one partition per key.
* See [[UngroupingOrigin]].
- * @param enableSortedMerge When true, uses [[SortedMergeCoalescedRDD]] to
perform a k-way merge
- * of the coalesced partitions, preserving the
child's output ordering
- * end-to-end. Set by [[EnsureRequirements]] when a
parent operator
- * requires the ordering that this node can satisfy
via sorted merge.
- * This flag is the decision, not a hint. Nothing
below re-reads the
- * config that produced it. A node planned with the
merge would otherwise
- * concatenate instead, and a sort-merge join above
it would lose rows
- * (SPARK-59279).
+ * @param sortedMergeOrdering When set, uses [[SortedMergeCoalescedRDD]] to
perform a k-way merge
+ * of the coalesced partitions by this ordering,
preserving it
+ * end-to-end. Set by [[EnsureRequirements]] when a
parent operator
+ * requires the ordering that this node can satisfy
via sorted merge.
+ * This is the decision, not a hint. Nothing below
re-reads the config
+ * that produced it. A node planned with the merge
would otherwise
+ * concatenate instead, and a sort-merge join above
it would lose rows
+ * (SPARK-59279). The ordering is captured from the
child when the merge
+ * is decided. That is the ordering the parent's
requirement was checked
+ * against. The child's ordering can follow a
config too, so a later
Review Comment:
The text here names the cause: the child's ordering follows a config. That
config is read live in `DataSourceV2ScanExecBase.outputOrdering` (`case (_, k:
KeyedPartitioning) if conf.v2BucketingPartitionKeyOrderingEnabled`), while
`reportedKeyedPartitioning` on the same scan already freezes
`conf.v2BucketingEnabled` per instance.
This PR fixes the one consumer whose execution depends on that ordering.
Other readers still see the ordering disappear after planning. For example,
under AQE `ValidateRequirements.validate` re-checks
`SortOrder.orderingSatisfies(child.outputOrdering, ordering)` for every node
after an `AQEShuffleReadRule` rewrite. A parent whose sort was elided because
of the key-derived ordering then fails validation, and a valid
`CoalesceShufflePartitions` or local-read rewrite gets reverted.
Not blocking for this PR. But would capturing the config once in the scan,
as `reportedKeyedPartitioning` does, be the more general fix? It may be worth a
separate JIRA.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala:
##########
@@ -450,7 +460,8 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
withSQLConf(SQLConf.V2_BUCKETING_PRESERVE_KEY_ORDERING_ON_COALESCE_ENABLED.key
-> "true") {
// Even though enableSortedMerge = true, the child is not safe for k-way
merge,
Review Comment:
nit. This comment, the one at L450, the test names at L447 and L470
("enableSortedMerge with ..." and "coalescing with enableSortedMerge = true"),
and the SPARK-59279 comment at L495 ("Once the flag is set") still describe a
boolean constructor flag. Could they say that the tests pass the merge
ordering, `sortedMergeOrdering`?
--
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]
