LinSimon-901101 commented on code in PR #7155:
URL: https://github.com/apache/datafusion-comet/pull/7155#discussion_r4511116969
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -600,7 +601,30 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
// A shuffle directly under the native write stays in the child stage
and is not wrapped
// again.
case op: IcebergWriteExec if
CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.get(op.conf) =>
- convertToComet(op, CometIcebergNativeWrite).getOrElse(op)
+ // Spark inserts AssertNotNull projections for nullable input to
required Iceberg fields.
+ // Keep their row-wise evaluation in Spark, including any sort or
limit above them, and
+ // convert the writer's input to Arrow only after Spark has evaluated
those operators.
+ val needsRowConversion = !op.child.isInstanceOf[CometNativeExec] &&
+ op.child.exists {
+ case project: ProjectExec =>
+
project.projectList.exists(_.exists(_.isInstanceOf[AssertNotNull]))
+ case _ => false
+ }
+ val rowConversionReasons = ListBuffer.empty[String]
+ val prepared =
+ if (needsRowConversion &&
+ CometSparkToColumnarExec.isRowConversionSupported(op.child,
rowConversionReasons) &&
+ CometIcebergNativeWrite.isRowInputEvaluationSupported(op,
rowConversionReasons)) {
+ val rows = ApplyColumnarRulesAndInsertTransitions(Seq.empty,
outputsColumnar = false)
+ .apply(op.child)
+ convertToComet(rows, CometSparkToColumnarExec)
+ .map(child => op.withNewChildren(Seq(child)))
+ .getOrElse(op)
+ } else {
+ rowConversionReasons.foreach(reason => withFallbackReason(op,
reason))
+ op
+ }
+ convertToComet(prepared, CometIcebergNativeWrite).getOrElse(op)
Review Comment:
Thanks for pointing this out. I’ve copied fallback reasons from`prepared`
back to `op`, with regression tests for AQE on and off.
--
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]
