dwsmith1983 commented on code in PR #6547:
URL: https://github.com/apache/datafusion-comet/pull/6547#discussion_r4171270398
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -608,6 +610,89 @@ case class CometExecRule(session: SparkSession)
}
}
+ /**
+ * AQE re-plans around a materialized stage by reusing the physical node
linked to it, so a
+ * native operator that shares its logical node with a shuffle stage (the
final aggregate of a
+ * two-phase aggregate) keeps the native plan it got while that input was a
bare exchange, read
+ * through a plain `Scan`. Once the input is a sink that reads the shuffle
directly, its
+ * `ShuffleScan` takes the place of the stale leaf. The leaf is patched in
place because
+ * converting the node again from `originalPlan` would drop the stage's
logical link that AQE
+ * relies on and re-run serde on a node that is already planned.
+ */
+ private def refreshStaleShuffleScans(op: SparkPlan): SparkPlan = op match {
+ // These build their own `Scan` over the child rather than embedding the
child's plan.
+ case _: CometNativeWriteExec | _: CometIcebergWriteExec | _:
CometWriteFilesExec => op
+ case native: CometNativeExec if native.children.nonEmpty =>
+ refreshedNativeOp(native) match {
+ case Some(newOp) =>
+ val refreshed = native.withRefreshedNativeOp(newOp)
+ // An operator that does not hold its native plan as a field cannot
take a new one.
+ if (refreshed.nativeOp eq newOp) refreshed else op
+ case None => op
+ }
+ case _ => op
+ }
+
+ /**
+ * The native plan of `native` with each `Scan` leaf whose input is now a
`ShuffleScan` of the
+ * same field types replaced by that `ShuffleScan`, or None if there is no
such leaf or the plan
+ * children cannot be matched to the leaves.
+ */
+ private def refreshedNativeOp(native: CometNativeExec): Option[Operator] = {
+ val children = native.children.collect { case child: CometNativeExec =>
child }
+ // Only a sink that reads a shuffle directly, or a native child that may
hold one, can feed
+ // a `ShuffleScan`.
+ val mayFeedShuffleScan = children.exists {
+ case sink: CometSinkPlaceHolder => sink.nativeOp.hasShuffleScan
+ case _ => true
+ }
+ if (children.length != native.children.length || !mayFeedShuffleScan)
return None
+ val leaves = nativeLeaves(native.nativeOp)
+ if (!leaves.exists(_.hasScan)) return None
+
+ // Each plan child feeds a run of leaves, in order: a sink feeds one, and
a native child
+ // feeds the leaves of its own native plan.
+ val current = children.flatMap {
+ case sink: CometSinkPlaceHolder => Seq(sink.nativeOp)
+ case child => nativeLeaves(child.nativeOp)
+ }
+ def isStale(leaf: Operator, input: Operator): Boolean = leaf.hasScan &&
input.hasShuffleScan
+ val stale = leaves.zip(current).filter { case (leaf, input) =>
isStale(leaf, input) }
+ val isRefreshable = current.length == leaves.length && stale.nonEmpty &&
+ stale.forall { case (leaf, input) =>
+ leaf.getScan.getFieldsList == input.getShuffleScan.getFieldsList
+ }
+ if (isRefreshable) {
+ val newLeaves = leaves.zip(current).map { case (leaf, input) =>
+ if (isStale(leaf, input)) input else leaf
+ }
+ Some(withLeaves(native.nativeOp, newLeaves))
+ } else {
+ None
Review Comment:
Added an info tag. When stale leaves are found but the counts or field types
don't line up, the node keeps its plan and gets `Shuffle direct read not
applied: the native plan's leaves do not match its inputs`, which shows in
verbose extended explain without a fallback reason, so coverage counts don't
change. A test covers both the field-type and the leaf-count mismatch and
checks the message.
--
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]
