sunchao commented on code in PR #5972:
URL: https://github.com/apache/datafusion-comet/pull/5972#discussion_r4179088418
##########
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:
Fixed in 8ec0da7af711442b8c540ed0eb21f95dcc409b9c. The filter analysis now
follows the first predicate Spark actually emits, including reordered/inferred
null checks. Its guaranteed inputs stay native; inputs behind an earlier
rejecting predicate remain protected.
The Parquet regression reproduces the previous Project/Filter fallback, with
preserveEvaluationMasks=false as the native control. It now keeps both
operators native, verifies required invalid weekdays still throw, and verifies
an earlier rejecting conjunct still skips the invalid value. Eleven plan cases
cover conjunct ordering, eager variable loading and direct/compound null guards.
All 35 affected Spark 4.1.3 tests pass. Both new regressions also pass on
Spark 3.5.9 with strict main/test compilation, including version-specific
exception class and SQLSTATE parity. The strengthened error assertions were
rerun successfully on Spark 4.1.3. Native inputs are unchanged and tests use
the matching verified baseline CI library.
--
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]
