address comments

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

@@ -33,7 +33,7 @@ import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRela
  * 1. aggregate expression is partition columns.
  *  e.g. SELECT col FROM tbl GROUP BY col.
  * 2. aggregate function on partition columns with DISTINCT.
- *  e.g. SELECT count(DISTINCT col) FROM tbl GROUP BY col.
+ *  e.g. SELECT count(DISTINCT col1) FROM tbl GROUP BY col2.
  * 3. aggregate function on partition columns which have same result w or w/o DISTINCT keyword.
  *  e.g. SELECT Max(col2) FROM tbl GROUP BY col1.
  */

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -2867,13 +2867,19 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
   }
 
   test("spark-15752 metadata only optimizer for datasource table") {
-    val data = (1 to 10).map(i => (i, s"data-$i", i % 2, if ((i % 2) == 0) "even" else "odd"))
-      .toDF("id", "data", "partId", "part")
-    data.write.partitionBy("partId", "part").mode("append").saveAsTable("srcpart_15752")
-    checkAnswer(sql("select max(part) from srcpart_15752"), Row("odd"))
-    checkAnswer(sql("select max(part) from srcpart_15752 where partId = 0"), Row("even"))
-    checkAnswer(sql("select max(part) from (select part from srcpart_15752) t"), Row("odd"))
-    checkAnswer(sql("select distinct part from srcpart_15752"), Row("even") :: Row("odd") :: Nil)
-    checkAnswer(sql("select distinct part from srcpart_15752 where partId = 0"), Row("even"))
+    withSQLConf(SQLConf.OPTIMIZER_METADATA_ONLY.key -> "true") {
+      withTable("srcpart_15752") {
+        val data = (1 to 10).map(i => (i, s"data-$i", i % 2, if ((i % 2) == 0) "even" else "odd"))
+          .toDF("id", "data", "partId", "part")
+        data.write.partitionBy("partId", "part").mode("append").saveAsTable("srcpart_15752")
+        checkAnswer(sql("select max(part) from srcpart_15752"), Row("odd"))
+        checkAnswer(sql("select max(part) from srcpart_15752 where partId = 0"), Row("even"))
+        checkAnswer(sql("select max(part) from (select part from srcpart_15752) t"), Row("odd"))
+        checkAnswer(
+          sql("select distinct part from srcpart_15752"),
+          Row("even") :: Row("odd") :: Nil)
+        checkAnswer(sql("select distinct part from srcpart_15752 where partId = 0"), Row("even"))
+      }
+    }
   }
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/MetadataOnlyOptimizerSuite.scala

@@ -32,6 +32,14 @@ class MetadataOnlyOptimizerSuite extends QueryTest with SharedSQLContext {
     data.write.partitionBy("partId", "part").mode("append").saveAsTable("srcpart_15752")
   }
 
+  override protected def afterAll(): Unit = {
+    try {
+      sql("DROP TABLE IF EXISTS srcpart_15752")
+    } finally {
+      super.afterAll()
+    }
+  }
+
   private def checkWithMetadataOnly(df: DataFrame): Unit = {
     val localRelations = df.queryExecution.optimizedPlan.collect {
       case l @ LocalRelation(_, _) => l

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

@@ -1692,66 +1692,68 @@ class SQLQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
 
   test("spark-15752 metadata only optimizer for hive table") {
     withSQLConf(SQLConf.OPTIMIZER_METADATA_ONLY.key -> "true") {
-      val df = Seq((1, "2"), (3, "4")).toDF("key", "value")
-      df.createOrReplaceTempView("data_15752")
-      sql(
-        """
-          |CREATE TABLE srcpart_15752 (key INT, value STRING)
-          |PARTITIONED BY (ds STRING, hr INT) STORED AS parquet
-        """.stripMargin)
-      for (ds <- Seq("2008-04-08", "2008-04-09"); hr <- Seq(11, 12)) {
+      withTable("data_15752", "srcpart_15752", "srctext_15752") {
+        val df = Seq((1, "2"), (3, "4")).toDF("key", "value")
+        df.createOrReplaceTempView("data_15752")
         sql(
-          s"""
-             |INSERT OVERWRITE TABLE srcpart_15752 PARTITION (ds='$ds',hr='$hr')
-             |select key, value from data_15752
-        """.stripMargin)
-      }
-      checkAnswer(sql("select hr from srcpart_15752 where hr = 11 group by hr"), Row(11))
-      checkAnswer(sql("select max(hr) from srcpart_15752"), Row(12))
-      checkAnswer(sql("select max(hr) from srcpart_15752 where hr = 11"), Row(11))
-      checkAnswer(sql("select max(hr) from (select hr from srcpart_15752) t"), Row(12))
-      checkAnswer(
-        sql("select max(x) from (select hr + 1 as x from srcpart_15752 where hr = 12) t"),
-        Row(13))
-      checkAnswer(sql("select distinct hr from srcpart_15752"), Row(11) :: Row(12) :: Nil)
-      checkAnswer(sql("select distinct hr from srcpart_15752 where hr = 11"), Row(11))
-      checkAnswer(
-        sql("select distinct x from (select hr + 1 as x from srcpart_15752 where hr = 12) t"),
-        Row(13))
+          """
+            |CREATE TABLE srcpart_15752 (key INT, value STRING)
+            |PARTITIONED BY (ds STRING, hr INT) STORED AS parquet
+          """.stripMargin)
+        for (ds <- Seq("2008-04-08", "2008-04-09"); hr <- Seq(11, 12)) {
+          sql(
+            s"""
+              |INSERT OVERWRITE TABLE srcpart_15752 PARTITION (ds='$ds',hr='$hr')
+              |select key, value from data_15752
+            """.stripMargin)
+        }
+        checkAnswer(sql("select hr from srcpart_15752 where hr = 11 group by hr"), Row(11))
+        checkAnswer(sql("select max(hr) from srcpart_15752"), Row(12))
+        checkAnswer(sql("select max(hr) from srcpart_15752 where hr = 11"), Row(11))
+        checkAnswer(sql("select max(hr) from (select hr from srcpart_15752) t"), Row(12))
+        checkAnswer(
+          sql("select max(x) from (select hr + 1 as x from srcpart_15752 where hr = 12) t"),
+          Row(13))
+        checkAnswer(sql("select distinct hr from srcpart_15752"), Row(11) :: Row(12) :: Nil)
+        checkAnswer(sql("select distinct hr from srcpart_15752 where hr = 11"), Row(11))
+        checkAnswer(
+          sql("select distinct x from (select hr + 1 as x from srcpart_15752 where hr = 12) t"),
+          Row(13))
 
-      // Now donot support metadata only optimizer
-      checkAnswer(
-        sql("select hr from srcpart_15752 where hr = 12 group by rollup(hr)"),
-        Row(null) :: Row(12) :: Nil)
-      checkAnswer(
-        sql("select hr from (select hr from srcpart_15752 where hr = 11 union all " +
-          "select hr from srcpart_15752 where hr= 12)t group by hr"),
-        Row(11) :: Row(12) :: Nil)
+        // Now donot support metadata only optimizer
+        checkAnswer(
+          sql("select hr from srcpart_15752 where hr = 12 group by rollup(hr)"),
+          Row(null) :: Row(12) :: Nil)
+        checkAnswer(
+          sql("select hr from (select hr from srcpart_15752 where hr = 11 union all " +
+            "select hr from srcpart_15752 where hr= 12)t group by hr"),
+          Row(11) :: Row(12) :: Nil)
 
-      sql(
-        """
-          |CREATE TABLE srctext_15752 (key INT, value STRING)
-          |PARTITIONED BY (ds STRING, hr INT) STORED AS textfile
-        """.stripMargin)
-      for (ds <- Seq("2008-04-08", "2008-04-09"); hr <- Seq(11, 12)) {
         sql(
-          s"""
-             |INSERT OVERWRITE TABLE srctext_15752 PARTITION (ds='$ds',hr='$hr')
-             |select key, value from data_15752
-        """.stripMargin)
+          """
+            |CREATE TABLE srctext_15752 (key INT, value STRING)
+            |PARTITIONED BY (ds STRING, hr INT) STORED AS textfile
+          """.stripMargin)
+        for (ds <- Seq("2008-04-08", "2008-04-09"); hr <- Seq(11, 12)) {
+          sql(
+            s"""
+              |INSERT OVERWRITE TABLE srctext_15752 PARTITION (ds='$ds',hr='$hr')
+              |select key, value from data_15752
+            """.stripMargin)
+        }
+        checkAnswer(sql("select hr from srctext_15752 where hr = 11 group by hr"), Row(11))
+        checkAnswer(sql("select max(hr) from srctext_15752"), Row(12))
+        checkAnswer(sql("select max(hr) from srctext_15752 where hr = 11"), Row(11))
+        checkAnswer(sql("select max(hr) from (select hr from srctext_15752) t"), Row(12))
+        checkAnswer(
+          sql("select max(x) from (select hr + 1 as x from srctext_15752 where hr = 12) t"),
+          Row(13))
+        checkAnswer(sql("select distinct hr from srctext_15752"), Row(11) :: Row(12) :: Nil)
+        checkAnswer(sql("select distinct hr from srctext_15752 where hr = 11"), Row(11))
+        checkAnswer(
+          sql("select distinct x from (select hr + 1 as x from srctext_15752 where hr = 12) t"),
+          Row(13))
       }
-      checkAnswer(sql("select hr from srctext_15752 where hr = 11 group by hr"), Row(11))
-      checkAnswer(sql("select max(hr) from srctext_15752"), Row(12))
-      checkAnswer(sql("select max(hr) from srctext_15752 where hr = 11"), Row(11))
-      checkAnswer(sql("select max(hr) from (select hr from srctext_15752) t"), Row(12))
-      checkAnswer(
-        sql("select max(x) from (select hr + 1 as x from srctext_15752 where hr = 12) t"),
-        Row(13))
-      checkAnswer(sql("select distinct hr from srctext_15752"), Row(11) :: Row(12) :: Nil)
-      checkAnswer(sql("select distinct hr from srctext_15752 where hr = 11"), Row(11))
-      checkAnswer(
-        sql("select distinct x from (select hr + 1 as x from srctext_15752 where hr = 12) t"),
-        Row(13))
     }
   }
 }