comphead commented on code in PR #5533:
URL: https://github.com/apache/datafusion-comet/pull/5533#discussion_r4277250793
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -931,9 +938,10 @@ case class CometExecRule(session: SparkSession)
// corresponding Final or PartialMerge cannot be converted and the
intermediate buffer
// formats are incompatible. This runs before transform() so the tags
are checked
// during the bottom-up conversion. Tags persist through AQE stage
creation.
- tagUnsafePartialAggregates(planWithJoinRewritten)
+ val planWithEvaluationMasks =
preserveEvaluationMasks(planWithJoinRewritten)
Review Comment:
The comment block above describes `tagUnsafePartialAggregates`, but it now
sits above `val planWithEvaluationMasks = preserveEvaluationMasks(...)`. Could
it move below that line, with a short comment on the new pass saying why it
must run first (it sets the early `Final` tag that `tagUnsafePartialAggregates`
and `canAggregateBeConverted` consume)?
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1046,6 +1054,167 @@ case class CometExecRule(session: SparkSession)
}
}
+ /** Keep opted-in expressions in Spark's row pipeline where an operator can
skip inputs. */
+ private def preserveEvaluationMasks(plan: SparkPlan): SparkPlan = {
+ if (!CometConf.COMET_EXEC_PRESERVE_EVALUATION_MASKS_ENABLED.get()) return
plan
+
+ def findEvaluationMaskName(expr: Expression): Option[String] = {
+ var name: Option[String] = None
+ expr.exists { child =>
+ name = QueryPlanSerde.evaluationMaskName(child)
+ name.isDefined
+ }
Review Comment:
`findEvaluationMaskName` descends into every child, including under
`TryEval`. By my reading `try_to_binary(x, 'base64')` is
`TryEval(ToBinary(...))` and holds `UnBase64(failOnError = true)`. `TryEval`
catches the failure and returns NULL, so evaluating a row that Spark would skip
cannot raise there, yet `SELECT try_to_binary(s, 'base64') FROM t LIMIT 10`
should now be forced to Spark. Could the traversal stop at `TryEval`? I read
`TryEval` and `TryToBinary` in Spark 4.1.3 and did not run the query. I could
not find a test in the suite or the SQL fixture for `try_to_binary`.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1046,6 +1054,167 @@ case class CometExecRule(session: SparkSession)
}
}
+ /** Keep opted-in expressions in Spark's row pipeline where an operator can
skip inputs. */
+ private def preserveEvaluationMasks(plan: SparkPlan): SparkPlan = {
+ if (!CometConf.COMET_EXEC_PRESERVE_EVALUATION_MASKS_ENABLED.get()) return
plan
+
+ def findEvaluationMaskName(expr: Expression): Option[String] = {
+ var name: Option[String] = None
+ expr.exists { child =>
+ name = QueryPlanSerde.evaluationMaskName(child)
+ name.isDefined
+ }
+ name
+ }
+
+ def originalPlan(node: SparkPlan): SparkPlan = node match {
+ case scan: CometScanExec =>
+ scan.wrapped
+ .copy(partitionFilters = scan.partitionFilters, dataFilters =
scan.dataFilters)
+ // The local selector saves its parent's original plan, but does not
evaluate that
+ // parent's projection or offset. Only the enclosing top-K owns those
operations.
+ case local: CometLocalTopKExec => local
+ case comet: CometExec => comet.originalPlan
+ case shuffle: CometShuffleExchangeExec => shuffle.originalPlan
+ case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan
+ case _ => node
+ }
+
+ def evaluatedExpressions(node: SparkPlan): Seq[Expression] = node match {
+ case local: CometLocalTopKExec => local.sortOrder
+ // V1 dataFilters include predicates Spark cannot push into Parquet.
Without native row
+ // filtering they only inform metadata pruning; the Spark Filter above
evaluates rows.
+ case _: FileSourceScanExec if
!CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get() =>
+ Seq.empty
+ case _ => node.expressions
+ }
+
+ // Most plans contain no opted-in expression. Include native originals and
sticky tags
+ // so repeated AQE passes still restore an already-protected subtree when
necessary.
+ if (!plan.exists(node =>
+ node.getTagValue(CometExecRule.UNSAFE_EXPRESSION_EVALUATION).isDefined
||
+
evaluatedExpressions(originalPlan(node)).exists(findEvaluationMaskName(_).isDefined)))
{
+ return plan
+ }
+
+ def firstMatch(joinType: JoinType): Boolean = joinType match {
+ case LeftSemi | LeftAnti => true
+ case _ => false
+ }
+
+ def protect(
+ node: SparkPlan,
+ belowLimit: Boolean,
+ hasLimitAncestor: Boolean): (SparkPlan, Option[String]) = {
+ val original = originalPlan(node)
+ val startsLimit = original match {
+ // Offset-only collection does not stop its input early.
+ case collect: CollectLimitExec => collect.limit >= 0
+ case _: LocalLimitExec | _: GlobalLimitExec => true
Review Comment:
The line above says offset-only collection does not stop its input early,
but `GlobalLimitExec` is counted as an early stop here without the same check.
In Spark 3.5.9 a non-root `OFFSET n` plans as `GlobalLimitExec` with the
default `limit = -1` and only does `iter.drop(offset)`, and the generated code
adds the count check only when `limit >= 0`. As written, an offset-only
subquery over `unbase64` falls back, and it also sets `hasLimitAncestor`. Would
`case g: GlobalLimitExec => g.limit >= 0` match the `CollectLimitExec` case?
This is read from `SparkStrategies` and `limit.scala`, not run.
--
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]
