Address comments

HyukjinKwon · HyukjinKwon · commit b4888367fe22 · 2018-07-19T18:57:25.000+08:00
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala
@@ -166,6 +166,8 @@ case class FileSourceScanExec(
     override val tableIdentifier: Option[TableIdentifier])
   extends DataSourceScanExec with ColumnarBatchScan  {
 
+  // Note that some vals referring the file-based relation are lazy intentionally
+  // so that this plan can be canonicalized on executor side too. See SPARK-23731.
   override lazy val supportsBatch: Boolean = relation.fileFormat.supportBatch(
     relation.sparkSession, StructType.fromAttributes(output))
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/FileSourceScanExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/FileSourceScanExecSuite.scala
@@ -20,12 +20,12 @@ package org.apache.spark.sql.execution
 import org.apache.spark.sql.test.SharedSQLContext
 
 class FileSourceScanExecSuite extends SharedSQLContext {
-  test("FileSourceScanExec should be canonicalizable in executor side") {
+  test("FileSourceScanExec should be canonicalizable on executor side") {
     withTempPath { path =>
-      spark.range(1).toDF().write.parquet(path.getAbsolutePath)
+      spark.range(1).write.parquet(path.getAbsolutePath)
       val df = spark.read.parquet(path.getAbsolutePath)
       val fileSourceScanExec =
-        df.queryExecution.sparkPlan.find(_.isInstanceOf[FileSourceScanExec]).get
+        df.queryExecution.sparkPlan.collectFirst { case p: FileSourceScanExec => p }.get
       try {
         spark.range(1).foreach(_ => fileSourceScanExec.canonicalized)
       } catch {
