sunchao commented on code in PR #5533:
URL: https://github.com/apache/datafusion-comet/pull/5533#discussion_r4178943216
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1046,6 +1054,159 @@ 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)
+ case comet: CometExec => comet.originalPlan
Review Comment:
[P2] Handle `CometLocalTopKExec` separately before unwrapping every
`CometExec`. Its `originalPlan` is the enclosing `TakeOrderedAndProjectExec`,
including that parent's projection and offset, although the local node only
selects candidates. With `spark.comet.exec.topK.fusion.enabled=true`, the outer
top-K sees ordered input on a later AQE pass. This walk then mistakes the local
node for a decoder below that limit and restores both nodes as full Spark top-K
operators. A valid-input broadcast-join query fails during stage creation.
Explicit replanning also exposes the duplicated projection as
`INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND`. Model the local node's actual evaluation
and materialization behavior, and ensure restoration applies the enclosing
projection and offset only once.
Evidence: Reproduced on Spark 4.1.3 using current-head JVM sources and its
verified CI native artifact. Write one Parquet file named `review5533` with
columns `_1`, `_2` and rows `(1,'YWJj'),(2,'YWJj'),(3,'YWJj'),(4,'YWJj')`. With
Comet native scanning/shuffle enabled, AQE=true, shuffle.partitions=1 and
topK.fusion.enabled=true, run: `SELECT /*+ BROADCAST(t) */ t._1, t.decoded FROM
(SELECT _1, hex(unbase64(_2)) AS decoded FROM (SELECT * FROM review5533 ORDER
BY _1 LIMIT 2 OFFSET 1) t) t JOIN range(1,5) r ON t._1=r.id`. Expected rows are
`(2,'616263'),(3,'616263')`. Actual result is `[INTERNAL_ERROR] Custom columnar
rules cannot transform broadcast node to something else`. Applying the rule
twice to the underlying top-K produces two nested
`TakeOrderedAndProject(limit=3, offset=1)` nodes and fails to bind `_2`. The
same probes pass with the base planner, fusion disabled, or evaluation-mask
protection disabled.
--
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]
