sunchao commented on code in PR #5972:
URL: https://github.com/apache/datafusion-comet/pull/5972#discussion_r4178987788
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1046,6 +1055,225 @@ 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
+ case shuffle: CometShuffleExchangeExec => shuffle.originalPlan
+ case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan
+ case _ => node
+ }
+
+ def evaluatedExpressions(node: SparkPlan): Seq[Expression] = node match {
+ // 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
+ }
+
+ // Match Spark's whole-stage eligibility before relying on deferred
generated expressions.
+ // In interpreted execution, or across an input adapter, a Project
materializes every output.
+ def supportsWholeStage(node: SparkPlan): Boolean = node match {
+ case codegen: CodegenSupport
+ if conf.wholeStageEnabled &&
+ !QueryPlanSerde.usesInterpretedProjection(conf) &&
+ codegen.supportCodegen =>
+ !node.expressions.exists(_.exists {
+ case _: LeafExpression => false
+ case _: CodegenFallback => true
+ case _ => false
+ }) && !WholeStageCodegenExec.isTooManyFields(conf, node.schema) &&
+ !node.children.exists(child =>
WholeStageCodegenExec.isTooManyFields(conf, child.schema))
+ case _ => false
+ }
+
+ def eagerReferences(expr: Expression): AttributeSet = expr match {
+ case attribute: AttributeReference => AttributeSet(Seq(attribute))
+ case _ =>
+
AttributeSet(QueryPlanSerde.eagerlyEvaluatedChildren(expr).flatMap(eagerReferences))
+ }
+
+ def protect(
+ node: SparkPlan,
+ belowLimit: Boolean,
+ hasLimitAncestor: Boolean,
+ deferredInputs: AttributeSet): (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
+ case topK: TakeOrderedAndProjectExec =>
+ SortOrder.orderingSatisfies(node.children.head.outputOrdering,
topK.sortOrder)
+ case windowLimit
+ if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(
+ _.isInstance(windowLimit)) =>
+ // Partitioned limits drain each group, but an outer LIMIT can still
stop them.
+ // Keep the conservative behavior if a future rank function cannot
be extracted.
+
ShimCometWindowGroupLimit.extract(windowLimit).forall(_.partitionSpec.isEmpty)
&&
+ 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
+ }
+ // A top-K is still a LIMIT even when its current input requires sorting.
+ val limitAncestor = hasLimitAncestor || startsLimit ||
+ original.isInstanceOf[TakeOrderedAndProjectExec]
+ val childDeferredInputs = original match {
+ // AQE can present a previously converted subtree. Batch bridges must
not hide
+ // the original Project from its Spark consumer's deferred-input
contract.
+ case _: RowToColumnarExec | _: ColumnarToRowExec | _:
CometColumnarToRowExec |
+ _: CometNativeColumnarToRowExec | _: CometSparkToColumnarExec =>
+ deferredInputs
+ // Spark gives both children of these joins separate generated stages.
Stage roots
+ // and input adapters likewise materialize rows before handing them to
a consumer.
+ case _: SortMergeJoinExec | _: ShuffledHashJoinExec | _:
WholeStageCodegenExec |
+ _: InputAdapter =>
+ AttributeSet.empty
+ case project: ProjectExec if supportsWholeStage(project) =>
+ val eagerOutputs = project.projectList.filterNot(expr =>
+ expr.deterministic && deferredInputs.contains(expr.toAttribute))
+ project.inputSet -- project.usedInputs --
+ AttributeSet(eagerOutputs.flatMap(eagerReferences))
+ case codegen: CodegenSupport if !materializesInput &&
supportsWholeStage(original) =>
+ original.inputSet -- codegen.usedInputs
Review Comment:
[P2] Preserve native execution for filter inputs that the predicate
necessarily consumes. `FilterExec.usedInputs` is deliberately empty because
Spark evaluates predicates inside `doConsume`, not because every input can be
skipped. For `SELECT x, r FROM (SELECT next_day(d,dow) AS x, rand(1L) AS r FROM
t) WHERE x >= DATE '2020-01-01'`, Spark must evaluate `x` on every input row.
This branch nevertheless marks it deferred, protects the lower projection, and
replaces both `CometProject` and `CometFilter` with Spark operators. This
introduces unnecessary fallback under default settings and a measured slowdown.
Exclude guaranteed predicate inputs from the deferred set while retaining
protection for genuinely short-circuited inputs.
Evidence: Reproduced with exact-head JVM sources and its verified CI native
library on Spark 4.1.3, ANSI and whole-stage codegen enabled, AQE disabled.
Store `(DATE '2024-01-01', 'Monday')` in Parquet table `t(d,dow)` and execute
the query above. Head uses Spark Project/Filter; base planner/serde controls
use CometProject/CometFilter. Setting
`spark.comet.exec.preserveEvaluationMasks.enabled=false` also restores both
native operators. An invalid-weekday control raises in Spark even when the
predicate rejects the row, confirming evaluation is required. A bounded
one-million-row probe used `d=date_add(DATE '2024-01-01', cast(id % 14 AS
int))`, `dow='Monday'`, one Parquet file, and an RDD sum over the date output.
Eight alternating runs per setting, discarding the first two, measured median
action times of 185.8 ms with protection versus 138.9 ms without it, with
identical checksums. These are local measurements, not a general throughput
claim.
--
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]
