address comments.

lianhuiwang · lianhuiwang · commit 120662e4e283 · 2017-06-11T16:46:07.000+08:00
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PruneFileSourcePartitions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PruneFileSourcePartitions.scala
@@ -60,6 +60,7 @@ private[sql] object PruneFileSourcePartitions extends Rule[LogicalPlan] {
         val prunedFileIndex = catalogFileIndex.filterPartitions(partitionKeyFilters.toSeq)
         val prunedFsRelation =
           fsRelation.copy(location = prunedFileIndex)(sparkSession)
+        // Change table stats based on the sizeInBytes of pruned files
         val withStats = logicalRelation.catalogTable.map(_.copy(
           stats = Some(CatalogStatistics(sizeInBytes = BigInt(prunedFileIndex.sizeInBytes)))))
         val prunedLogicalRelation = logicalRelation.copy(
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/PruneFileSourcePartitionsSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/PruneFileSourcePartitionsSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.hive.execution
 
 import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.catalyst.TableIdentifier
 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}
@@ -81,19 +82,26 @@ class PruneFileSourcePartitionsSuite extends QueryTest with SQLTestUtils with Te
             """.stripMargin)
         }
 
+        val tableName = "partTbl"
+        sql(s"analyze table partTbl compute STATISTICS")
+
+        val tableStats =
+          spark.sessionState.catalog.getTableMetadata(TableIdentifier(tableName)).stats
+        assert(tableStats.isDefined && tableStats.get.sizeInBytes > 0, "tableStats is lost")
+
         withSQLConf(SQLConf.ENABLE_FALL_BACK_TO_HDFS_FOR_STATS.key -> "true") {
           val df = sql("SELECT * FROM partTbl where part = 1")
           val query = df.queryExecution.analyzed.analyze
           val sizes1 = query.collect {
             case relation: LogicalRelation => relation.computeStats(conf).sizeInBytes
           }
           assert(sizes1.size === 1, s"Size wrong for:\n ${df.queryExecution}")
-          assert(sizes1(0) > 5000, s"expected > 5000 for test table 'src', got: ${sizes1(0)}")
+          assert(sizes1(0) == tableStats.get.sizeInBytes)
           val sizes2 = Optimize.execute(query).collect {
             case relation: LogicalRelation => relation.computeStats(conf).sizeInBytes
           }
           assert(sizes2.size === 1, s"Size wrong for:\n ${df.queryExecution}")
-          assert(sizes2(0) < 5000, s"expected < 5000 for test table 'src', got: ${sizes2(0)}")
+          assert(sizes2(0) < tableStats.get.sizeInBytes)
         }
       }
     }

Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@`
`18`	`18`	`package org.apache.spark.sql.hive.execution`
`19`	`19`
`20`	`20`	`import org.apache.spark.sql.QueryTest`
	`21`	`+import org.apache.spark.sql.catalyst.TableIdentifier`
`21`	`22`	`import org.apache.spark.sql.catalyst.dsl.expressions._`
`22`	`23`	`import org.apache.spark.sql.catalyst.dsl.plans._`
`23`	`24`	`import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, Project}`
`@@ -81,19 +82,26 @@ class PruneFileSourcePartitionsSuite extends QueryTest with SQLTestUtils with Te`
`81`	`82`	`""".stripMargin)`
`82`	`83`	`}`
`83`	`84`
	`85`	`+ val tableName = "partTbl"`
	`86`	`+ sql(s"analyze table partTbl compute STATISTICS")`
	`87`	`+`
	`88`	`+ val tableStats =`
	`89`	`+ spark.sessionState.catalog.getTableMetadata(TableIdentifier(tableName)).stats`
	`90`	`+ assert(tableStats.isDefined && tableStats.get.sizeInBytes > 0, "tableStats is lost")`
	`91`	`+`
`84`	`92`	`withSQLConf(SQLConf.ENABLE_FALL_BACK_TO_HDFS_FOR_STATS.key -> "true") {`
`85`	`93`	`val df = sql("SELECT * FROM partTbl where part = 1")`
`86`	`94`	`val query = df.queryExecution.analyzed.analyze`
`87`	`95`	`val sizes1 = query.collect {`
`88`	`96`	`case relation: LogicalRelation => relation.computeStats(conf).sizeInBytes`
`89`	`97`	`}`
`90`	`98`	`assert(sizes1.size === 1, s"Size wrong for:\n ${df.queryExecution}")`
`91`		`- assert(sizes1(0) > 5000, s"expected > 5000 for test table 'src', got: ${sizes1(0)}")`
	`99`	`+ assert(sizes1(0) == tableStats.get.sizeInBytes)`
`92`	`100`	`val sizes2 = Optimize.execute(query).collect {`
`93`	`101`	`case relation: LogicalRelation => relation.computeStats(conf).sizeInBytes`
`94`	`102`	`}`
`95`	`103`	`assert(sizes2.size === 1, s"Size wrong for:\n ${df.queryExecution}")`
`96`		`- assert(sizes2(0) < 5000, s"expected < 5000 for test table 'src', got: ${sizes2(0)}")`
	`104`	`+ assert(sizes2(0) < tableStats.get.sizeInBytes)`
`97`	`105`	`}`
`98`	`106`	`}`
`99`	`107`	`}`