follow comment

Author: AngersZhuuuu

sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala

@@ -205,7 +205,7 @@ case class FileSourceScanExec(
   private def isDynamicPruningFilter(e: Expression): Boolean =
     e.find(_.isInstanceOf[PlanExpression[_]]).isDefined
 
-  @transient private lazy val selectedPartitions: Array[PartitionDirectory] = {
+  @transient lazy val selectedPartitions: Array[PartitionDirectory] = {
     val optimizerMetadataTimeNs = relation.location.metadataOpsTimeNs.getOrElse(0L)
     val startTime = System.nanoTime()
     val ret =

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/PruneFileSourcePartitionsSuite.scala

@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
-import org.apache.spark.sql.execution.FileSourceScanExec
+import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan}
 import org.apache.spark.sql.execution.datasources.{CatalogFileIndex, HadoopFsRelation, LogicalRelation, PruneFileSourcePartitions}
 import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
 import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec
@@ -109,9 +109,9 @@ class PruneFileSourcePartitionsSuite extends PrunePartitionSuiteBase {
     }
   }
 
-  override def getScanExecPartitionSize(query: String): Long = {
-    sql(query).queryExecution.sparkPlan.collectFirst {
+  override def getScanExecPartitionSize(plan: SparkPlan): Long = {
+    plan.collectFirst {
       case p: FileSourceScanExec => p
-    }.get.relation.location.inputFiles.length
+    }.get.selectedPartitions.length
   }
 }

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/PruneHiveTablePartitionsSuite.scala

@@ -20,6 +20,7 @@ package org.apache.spark.sql.hive.execution
 import org.apache.spark.sql.catalyst.analysis.EliminateSubqueryAliases
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.execution.SparkPlan
 
 class PruneHiveTablePartitionsSuite extends PrunePartitionSuiteBase {
 
@@ -54,8 +55,8 @@ class PruneHiveTablePartitionsSuite extends PrunePartitionSuiteBase {
     }
   }
 
-  override def getScanExecPartitionSize(query: String): Long = {
-    sql(query).queryExecution.sparkPlan.collectFirst {
+  override def getScanExecPartitionSize(plan: SparkPlan): Long = {
+    plan.collectFirst {
       case p: HiveTableScanExec => p
     }.get.prunedPartitions.size
   }

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/PrunePartitionSuiteBase.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.hive.execution
 
 import org.apache.spark.sql.QueryTest
-import org.apache.spark.sql.hive.HiveUtils
+import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.hive.test.TestHiveSingleton
 import org.apache.spark.sql.test.SQLTestUtils
 
@@ -68,8 +68,9 @@ abstract class PrunePartitionSuiteBase extends QueryTest with SQLTestUtils with
   }
 
   protected def assertPrunedPartitions(query: String, expected: Long): Unit = {
-    assert(getScanExecPartitionSize(query) == expected)
+    val plan = sql(query).queryExecution.sparkPlan
+    assert(getScanExecPartitionSize(plan) == expected)
   }
 
-  protected def getScanExecPartitionSize(query: String): Long
+  protected def getScanExecPartitionSize(plan: SparkPlan): Long
 }