sunchao commented on code in PR #5533:
URL: https://github.com/apache/datafusion-comet/pull/5533#discussion_r4041853820
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -714,6 +722,132 @@ case class CometExecRule(session: SparkSession)
}
}
+ /** Keep operators that can skip malformed unbase64 inputs in Spark's row
pipeline. */
+ private def preserveUnbase64EvaluationMasks(plan: SparkPlan): SparkPlan = {
+ val limitReason = "unbase64 requires Spark evaluation below LIMIT"
+ val joinReason = "unbase64 requires Spark evaluation in first-match join
conditions"
+
+ def containsUnbase64(expr: Expression): Boolean =
+ expr.exists(_.isInstanceOf[UnBase64])
+
+ def firstMatch(joinType: JoinType): Boolean = joinType match {
+ case LeftSemi | LeftAnti => true
+ case _ => false
+ }
+
+ // A reused native Final may need to fall back along with its incompatible
Partial buffer.
+ // Restore only that buffer-producing chain so tagUnsafePartialAggregates
can protect it.
+ // Do not cross a materialized query stage or descend below the input of a
pure Partial.
+ def restoreNativeAggregateBuffers(node: SparkPlan): Option[SparkPlan] = {
+ val original = node match {
+ case agg: CometHashAggregateExec => agg.originalPlan
+ case shuffle: CometShuffleExchangeExec => shuffle.originalPlan
+ case _ => node
+ }
+ def restore(children: Seq[SparkPlan]): SparkPlan = {
+ val restored = original.withNewChildren(children)
+
node.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).foreach(restored.setLogicalLink)
+ restored
+ }
+ original match {
+ case agg: BaseAggregateExec
+ if agg.aggregateExpressions.nonEmpty &&
+ agg.aggregateExpressions.forall(_.mode == Partial) =>
+ if (original ne node) Some(restore(node.children)) else None
+ case agg: BaseAggregateExec
+ if agg.aggregateExpressions.forall(e =>
+ e.mode == Partial || e.mode == PartialMerge) =>
+ restoreNativeAggregateBuffers(node.children.head).map(child =>
restore(Seq(child)))
+ case _: ShuffleExchangeLike =>
+ restoreNativeAggregateBuffers(node.children.head).map(child =>
restore(Seq(child)))
+ case _ => None
+ }
+ }
+
+ def protect(node: SparkPlan, belowLimit: Boolean): (SparkPlan,
Option[String]) = {
+ val original = node match {
+ case scan: CometScanExec =>
+ scan.wrapped
+ .copy(partitionFilters = scan.partitionFilters, dataFilters =
scan.dataFilters)
+ case comet: CometExec => comet.originalPlan
+ case shuffle: CometShuffleExchangeExec => shuffle.originalPlan
+ case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan
+ case _ => node
+ }
+ val startsLimit = original match {
+ case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
true
+ case topK: TakeOrderedAndProjectExec =>
+ SortOrder.orderingSatisfies(node.children.head.outputOrdering,
topK.sortOrder)
+ case windowLimit
+ if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(
+ _.isInstance(windowLimit)) =>
+ SortOrder.orderingSatisfies(
+ node.children.head.outputOrdering,
+ windowLimit.requiredChildOrdering.head)
+ case _ => false
+ }
+ // These operators consume their input before yielding rows. Still visit
their children:
+ // an inner LocalLimit below an exchange must establish its own
evaluation boundary.
+ val materializesInput = original match {
+ case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec |
+ _: ShuffleExchangeLike | _: BroadcastExchangeLike | _:
QueryStageExec |
+ _: ReusedExchangeExec =>
+ true
+ case _ => false
+ }
+ val protectedChildren = node.children.map { child =>
+ protect(child, startsLimit || (belowLimit && !materializesInput))
+ }
+ val childReason = protectedChildren.flatMap(_._2).headOption
+ val condition = original match {
+ case join: HashJoin if firstMatch(join.joinType) => join.condition
+ case join: SortMergeJoinExec if firstMatch(join.joinType) =>
join.condition
+ case join: BroadcastNestedLoopJoinExec if firstMatch(join.joinType) =>
join.condition
+ case _ => None
+ }
+ val ownReason = if (belowLimit &&
original.expressions.exists(containsUnbase64)) {
+ Some(limitReason)
+ } else if (condition.exists(containsUnbase64)) {
+ Some(joinReason)
+ } else {
+ node.getTagValue(CometExecRule.UNSAFE_UNBASE64_EVALUATION)
+ }
+ // Exchanges can restart native execution after consuming a Spark row
pipeline.
+ val restartsNative = original.isInstanceOf[ShuffleExchangeLike] ||
+ original.isInstanceOf[BroadcastExchangeLike]
+ val reason = ownReason.orElse(if (restartsNative) None else childReason)
+ val children = protectedChildren.map(_._1).map { child =>
+ original match {
+ case agg: BaseAggregateExec
+ if reason.isDefined &&
agg.aggregateExpressions.map(_.mode).distinct == Seq(
+ Final) &&
+
!QueryPlanSerde.allAggsSupportMixedExecution(agg.aggregateExpressions) =>
+ restoreNativeAggregateBuffers(child).getOrElse(child)
+ case _ => child
+ }
+ }
+ val prepared = node match {
+ // Do not refill a batch between the row decoder and its
short-circuiting consumer.
+ case _: RowToColumnarExec | _: CometSparkToColumnarExec if
childReason.isDefined =>
+ children.head
+ case _: ColumnarToRowExec | _: CometColumnarToRowExec | _:
CometNativeColumnarToRowExec
+ if childReason.isDefined && !children.head.supportsColumnar =>
+ children.head
+ case _ if (original ne node) && (reason.isDefined || children !=
node.children) =>
+ // AQE can reuse an existing native subtree. Rebuild affected
ancestors as well so
+ // their serialized native plans do not retain the decoder that just
fell back.
+ val restored = original.withNewChildren(children)
+
node.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).foreach(restored.setLogicalLink)
+ restored
+ case _ => node.withNewChildren(children)
+ }
+
reason.foreach(prepared.setTagValue(CometExecRule.UNSAFE_UNBASE64_EVALUATION,
_))
+ (prepared, reason)
+ }
+
+ protect(plan, belowLimit = false)._1
Review Comment:
Addressed in `7df125174`.
Added the early bailout before the recursive protection pass. It returns the
input plan when there is neither an enrolled expression nor a sticky protection
tag. The check inspects native operators' original Spark expressions too, so
reused AQE/native subtrees still receive protection. The disabled configuration
returns immediately before that scan. Existing native-reuse and
repeated-planning regressions pass.
--
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]
