For review

apache · Mar 18, 2024 · 15eb659 · 15eb659
1 parent 4b78e4f
commit 15eb659
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 9 deletions.
diff --git a/core/src/execution/datafusion/planner.rs b/core/src/execution/datafusion/planner.rs
@@ -1131,8 +1131,9 @@ impl From<ExpressionError> for DataFusionError {
     }
 }
 
-/// Returns true if given operator probably returns input array as output array without
-/// modification.
+/// Returns true if given operator can return input array as output array without
+/// modification. This is used to determine if we need to copy the input batch to avoid
+/// data corruption from reusing the input batch.
 fn op_reuse_array(op: &Arc<dyn ExecutionPlan>) -> bool {
     op.as_any().downcast_ref::<ScanExec>().is_some()
         || op.as_any().downcast_ref::<LocalLimitExec>().is_some()

diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -224,13 +224,9 @@ class CometSparkSessionExtensions
     // spotless:on
     private def transform(plan: SparkPlan): SparkPlan = {
       def transform1(op: SparkPlan): Option[Operator] = {
-        val allNativeExec = op.children.map {
-          case childNativeOp: CometNativeExec => Some(childNativeOp.nativeOp)
-          case _ => None
-        }
-
-        if (allNativeExec.forall(_.isDefined)) {
-          QueryPlanSerde.operator2Proto(op, allNativeExec.map(_.get): _*)
+        if (op.children.forall(_.isInstanceOf[CometNativeExec])) {
+          QueryPlanSerde.operator2Proto(op,
+            op.children.map(_.asInstanceOf[CometNativeExec].nativeOp): _*)
         } else {
           None
         }