LinSimon-901101 commented on code in PR #7155:
URL: https://github.com/apache/datafusion-comet/pull/7155#discussion_r4373553512
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1178,10 +1178,96 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
serde.convert(op, builder)
}
}
- return nativeOp.map { nativeOp =>
- val exec = serde.createExec(nativeOp, op)
- rollUpInfoMessages(op, exec)
- exec
+ return nativeOp
+ .filter { proto =>
+ // A serde may emit several native nodes for one Spark operator.
Inspect those nodes,
+ // but stop at real Spark children, whose consumption context was
tagged independently.
+ def containsOwnAssertion(node: Operator): Boolean =
+
CometAssertNotNull.containsNativeAssertion(node.toBuilder.clearChildren().build())
||
+ node.getChildrenList.asScala.exists { child =>
+ child.getPlanId == proto.getPlanId &&
containsOwnAssertion(child)
+ }
+
+ def hasNativeAssertion: Boolean = op match {
+ // TopK's sink proto omits its runtime projection over the
selected rows.
+ case topK: TakeOrderedAndProjectExec =>
+ QueryPlanSerde.withSharedSubexpressions(topK.projectList) {
+ topK.projectList.exists { expression =>
+ QueryPlanSerde
+ .exprToProto(expression, topK.child.output)
+ .exists(CometAssertNotNull.containsNativeAssertion)
+ }
+ }
+ case _: HashAggregateExec | _: ObjectHashAggregateExec =>
+ // Input/grouping expressions consume all rows; only the result
projection is lazy.
+ proto.hasProjection &&
CometAssertNotNull.containsNativeAssertion(
+ proto.getProjection)
+ case _ => containsOwnAssertion(proto)
+ }
+ val fallback = nativeAssertionJoinFallback(proto).orElse {
+ op.getTagValue(CometExecRule.PARTIAL_READER)
+ .filter(_ => hasNativeAssertion)
+ .map(reader =>
+ s"Native AssertNotNull below ${reader.description} could
change which values " +
+ "Spark checks")
+ }
+ fallback.foreach(reason => withFallbackReason(op, reason))
+ fallback.isEmpty
+ }
+ .map { nativeOp =>
+ val exec = serde.createExec(nativeOp, op)
+ rollUpInfoMessages(op, exec)
+ exec
+ }
+ }
+ None
+ }
+
+ private def nativeAssertionJoinFallback(proto: Operator): Option[String] = {
+ import OperatorOuterClass.JoinType
+
+ // These consumers can skip candidates even without a LIMIT above the
join. Inspect only
+ // the serialized condition. Child projections use the partial-reader tags
independently;
+ // assertions captured inside opaque JVM dispatch keep their existing
execution path.
+ val condition = if (proto.hasHashJoin) {
+ val join = proto.getHashJoin
+ Some(join.getJoinType -> join.getCondition)
+ } else if (proto.hasSortMergeJoin) {
+ val join = proto.getSortMergeJoin
+ Some(join.getJoinType -> join.getCondition)
+ } else if (proto.hasBroadcastNestedLoopJoin) {
+ val join = proto.getBroadcastNestedLoopJoin
+ Some(join.getJoinType -> join.getCondition)
+ } else None
+
+ if (condition.exists {
+ case (JoinType.LeftSemi | JoinType.LeftAnti, expression) =>
+ CometAssertNotNull.containsNativeAssertion(expression)
+ case _ => false
+ }) {
+ return Some(
+ "Native AssertNotNull in an early-exit join condition could check
candidates Spark skips")
+ }
+
+ if (proto.hasHashJoin) {
+ val join = proto.getHashJoin
+ val leftKeys = join.getLeftJoinKeysList.asScala
+ val rightKeys = join.getRightJoinKeysList.asScala
+ // Empty and all-NULL build relations can skip different streamed keys
in Spark and
+ // DataFusion. Null-aware anti joins also stop building at the first
NULL key in Spark,
+ // while native execution builds on the opposite side and evaluates keys
in batches.
+ val keys = if (join.getNullAwareAntiJoin) {
+ leftKeys ++ rightKeys
+ } else {
+ join.getJoinType match {
+ case JoinType.Inner | JoinType.LeftSemi | JoinType.LeftAnti =>
+ if (join.getBuildSide == OperatorOuterClass.BuildSide.BuildRight)
leftKeys
+ else rightKeys
+ case _ => Seq.empty
Review Comment:
Fixed in `4678792`, included in the current head `6e850d7`.
By using `HashJoin.rewriteKeyExpr`, `CometExecRule.scala:1413–1424` detects
packed keys and falls back to Spark when they contain native `AssertNotNull`,
in order to preserve the NULL short-circuit evaluation.
Regression tests cover codegen on/off in `CometJoinSuite.scala:2322–2339`.
--
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]
