sunchao commented on code in PR #7155:
URL: https://github.com/apache/datafusion-comet/pull/7155#discussion_r4347757342
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1325,37 +1411,69 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf)))
/**
- * Tags each `SerializeFromObjectExec` whose output an operator above it can
stop reading early
- * with that operator, so `transform` does not convert it. Spark computes
the rows of a typed
- * Dataset operation one at a time, as the operator above reads them, while
the conversion fills
- * a whole Arrow batch first. Below a limit, a `mapPartitions` function such
as `_.take(1)`, or
- * code reading `Dataset.rdd`, the conversion would run the user function on
rows Spark never
- * reaches, and a function that throws on one of them would fail a query
that succeeds in Spark.
- * An operator that reads all of its input before it returns a row ends the
search, since Spark
- * computes every row below it anyway: an exchange, a sort, a hash
aggregate, or a top-k over
- * input that is not already sorted.
+ * Tags operators whose output a consumer can skip evaluating. Spark
computes rows as the
+ * consumer reads them, while native expressions and typed Dataset
conversion fill a whole Arrow
+ * batch first. Below a limit, a `mapPartitions` function such as
`_.take(1)`, or code reading
+ * `Dataset.rdd`, an assertion or user function could throw on rows Spark
never reaches. Joins
+ * can also skip an input, and generated filters, samples and joins can
defer projected fields
+ * until a row passes the condition. Conversely, native joins can skip input
Spark would
+ * evaluate. An operator that reads all of its input before it returns a row
ends the search,
+ * since Spark computes every row below it anyway: an exchange, a sort, a
hash aggregate, or a
+ * top-k over input that is not already sorted.
*
* Conversion is bottom-up, so this runs first. TreeNode tags survive the
child copies made
- * during transformUp, while an identity set would not.
+ * during transformUp, while an identity set would not. Retain established
tags when AQE visits
+ * a fragment without its limiting ancestor.
*/
- private def tagPartiallyReadTypedDatasetOutputs(plan: SparkPlan): Unit = {
- def visit(op: SparkPlan, partialReader: Option[String]): Unit = {
- val childReader = op match {
- case serialize: SerializeFromObjectExec =>
- partialReader.foreach(
- serialize.setTagValue(CometExecRule.TYPED_DATASET_PARTIAL_READER,
_))
- partialReader
- case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
Some("a limit")
+ private def tagPartiallyReadOutputs(plan: SparkPlan): Unit = {
+ import CometExecRule.PartialReader
+
+ def visit(op: SparkPlan, partialReader: Option[PartialReader]): Unit = {
+ // AQE can revisit a fragment without the ancestor that could skip whole
rows. A new
+ // field-only reader must not weaken that established context or its
child propagation.
+ val inheritedReader = op
+ .getTagValue(CometExecRule.PARTIAL_READER)
+ .filter(_.canSkipInput)
+ .orElse(partialReader)
+ // A repeated transform can already contain Comet limits and consumption
barriers.
+ val original = op match {
+ case comet: CometExec => comet.originalPlan
+ case other => other
+ }
+ val childReader = original match {
+ case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
+ Some(PartialReader("a limit"))
// A top-k reads only its first rows when its input is already sorted.
case topK: TakeOrderedAndProjectExec
if SortOrder.orderingSatisfies(topK.child.outputOrdering,
topK.sortOrder) =>
- Some("a limit")
- case _: MapPartitionsExec => Some("a mapPartitions function")
+ Some(PartialReader("a limit"))
+ case _: MapPartitionsExec => Some(PartialReader("a mapPartitions
function"))
+ // These consumers defer projected fields, but still read every input
object. Keep
+ // typed Dataset conversion eligible unless an ancestor can skip whole
input rows.
+ case _: FilterExec =>
+ inheritedReader.orElse(Some(PartialReader("a filter", canSkipInput =
false)))
+ case _: SampleExec =>
+ inheritedReader.orElse(Some(PartialReader("a sample", canSkipInput =
false)))
+ case _: BroadcastHashJoinExec | _: ShuffledHashJoinExec | _:
SortMergeJoinExec |
+ _: BroadcastNestedLoopJoinExec =>
+ Some(PartialReader("a join"))
case _: Exchange | _: SortExec | _: HashAggregateExec | _:
ObjectHashAggregateExec |
_: TakeOrderedAndProjectExec =>
None
- case _ => partialReader
+ case _ => inheritedReader
Review Comment:
[P2] Preserve deferred assertions across adjacent projections. `ProjectExec`
can conditionally consume a child field, but this catch-all leaves its child
without a partial-reader tag. With whole-stage codegen enabled, put
`AssertNotNull(IF(id = 0, NULL, id)) AS checked` and `rand(0) AS nonce` in a
lower projection, then select `IF(id = 0, 0L, checked), nonce` above it for `id
= 0`. The nondeterministic sibling prevents projection collapse. Spark returns
the row without evaluating `checked`, whereas the native lower projection
throws `NOT_NULL_ASSERT_VIOLATION`. Before this PR, the unsupported assertion
kept these projections in Spark. Please treat such projection consumers as
field-only partial readers, with corresponding own-reader handling, or
otherwise preserve lazy child evaluation, and add a two-projection regression.
Evidence: Registered `ann` as Catalyst `AssertNotNull` in Spark 4.1.3 and
executed: `SELECT IF(id=0,CAST(0 AS BIGINT),checked) AS result, nonce FROM
(SELECT id, ann(IF(id=0,CAST(NULL AS BIGINT),id)) AS checked, rand(0) AS nonce
FROM range(1))`. With AQE disabled and whole-stage codegen enabled, optimized
and executed plans retain both projections and return `[0,0.7604953758285915]`.
Disabling whole-stage codegen raises the assertion. A bounded native
reproduction using exact-head `AssertNotNull`, `IfExpr`, `RandExpr`, and two
DataFusion 55.2.0 `ProjectionExec`s raises `NOT_NULL_ASSERT_VIOLATION` for the
same single-row expressions. `CometProjectExec.convert` preserves the two
native projections, and `jni_api.rs` executes the constructed plan without
merging them. Enable `spark.comet.exec.range.enabled=true` to provide native
input for this SQL example. Reproduction sources and logs are under
`/tmp/comet7155-head-review/`. A complete Comet JVM query was 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]
