fix: Use makeCopy to change relation in FileSourceScanExec

apache · Mar 13, 2024 · bc12dc3 · bc12dc3
1 parent 9aa42cd
commit bc12dc3
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala
@@ -439,8 +439,18 @@ case class CometScanExec(
 
 object CometScanExec {
   def apply(scanExec: FileSourceScanExec, session: SparkSession): CometScanExec = {
-    val wrapped = scanExec.copy(relation =
-      scanExec.relation.copy(fileFormat = new CometParquetFileFormat)(session))
+    // Replacing the relation in FileSourceScanExec by `copy` seems causing some issues
+    // on other Spark distributions if FileSourceScanExec constructor is changed.
+    // Using `makeCopy` to avoid the issue.
+    // https://github.com/apache/arrow-datafusion-comet/issues/190
+    def transform(arg: Any): AnyRef = arg match {
+      case _: HadoopFsRelation =>
+        scanExec.relation.copy(fileFormat = new CometParquetFileFormat)(session)
+      case other: AnyRef => other
+      case null => null
+    }
+    val newArgs = scanExec.mapProductIterator(transform)
+    val wrapped = scanExec.makeCopy(newArgs).asInstanceOf[FileSourceScanExec]
     val batchScanExec = CometScanExec(
       wrapped.relation,
       wrapped.output,