address comments

Author: gatorsmile

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -937,8 +937,7 @@ object SimplifyConditionals extends Rule[LogicalPlan] with PredicateHelper {
  */
 case class OptimizeCodegen(conf: CatalystConf) extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
-    case e: CaseWhen if canCodeGen(e) =>
-      e.toCodegen()
+    case e: CaseWhen if canCodeGen(e) => e.toCodegen()
   }
 
   private def canCodeGen(e: CaseWhen): Boolean = {

sql/core/src/test/scala/org/apache/spark/sql/internal/SQLConfSuite.scala

@@ -17,36 +17,14 @@
 
 package org.apache.spark.sql.internal
 
-import org.scalatest.BeforeAndAfterAll
-
 import org.apache.spark.sql.{QueryTest, Row, SparkSession, SQLContext}
 import org.apache.spark.sql.execution.WholeStageCodegenExec
 import org.apache.spark.sql.test.{SharedSQLContext, TestSQLContext}
 
-class SQLConfSuite extends QueryTest with SharedSQLContext with BeforeAndAfterAll {
-  import testImplicits._
-
+class SQLConfSuite extends QueryTest with SharedSQLContext {
   private val testKey = "test.key.0"
   private val testVal = "test.val.0"
 
-  override def beforeAll() {
-    super.beforeAll()
-    sql("DROP TABLE IF EXISTS testData")
-    spark
-      .range(10)
-      .select('id as 'a, 'id as 'b, 'id as 'c, 'id as 'd)
-      .write
-      .saveAsTable("testData")
-  }
-
-  override def afterAll(): Unit = {
-    try {
-      sql("DROP TABLE IF EXISTS testData")
-    } finally {
-      super.afterAll()
-    }
-  }
-
   test("propagate from spark conf") {
     // We create a new context here to avoid order dependence with other tests that might call
     // clear().
@@ -243,31 +221,40 @@ class SQLConfSuite extends QueryTest with SharedSQLContext with BeforeAndAfterAl
   }
 
   test("MAX_CASES_BRANCHES") {
+    import testImplicits._
+
     val original = spark.conf.get(SQLConf.MAX_CASES_BRANCHES)
     try {
-      val sql_one_branch_caseWhen = "SELECT CASE WHEN a = 1 THEN 1 END FROM testData"
-      val sql_two_branch_caseWhen = "SELECT CASE WHEN a = 1 THEN 1 ELSE 0 END FROM testData"
-
-      spark.conf.set(SQLConf.MAX_CASES_BRANCHES.key, "0")
-      assert(!sql(sql_one_branch_caseWhen)
-        .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
-      assert(!sql(sql_two_branch_caseWhen)
-        .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
-
-      spark.conf.set(SQLConf.MAX_CASES_BRANCHES.key, "1")
-      assert(sql(sql_one_branch_caseWhen)
-        .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
-      assert(!sql(sql_two_branch_caseWhen)
-        .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
-
-      spark.conf.set(SQLConf.MAX_CASES_BRANCHES.key, "2")
-      assert(sql(sql_one_branch_caseWhen)
-        .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
-      assert(sql(sql_two_branch_caseWhen)
-        .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
+      withTable("tab1") {
+        spark
+          .range(10)
+          .select('id as 'a, 'id as 'b, 'id as 'c, 'id as 'd)
+          .write
+          .saveAsTable("tab1")
+
+        val sql_one_branch_caseWhen = "SELECT CASE WHEN a = 1 THEN 1 END FROM tab1"
+        val sql_two_branch_caseWhen = "SELECT CASE WHEN a = 1 THEN 1 ELSE 0 END FROM tab1"
+
+        spark.conf.set(SQLConf.MAX_CASES_BRANCHES.key, "0")
+        assert(!sql(sql_one_branch_caseWhen)
+          .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
+        assert(!sql(sql_two_branch_caseWhen)
+          .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
+
+        spark.conf.set(SQLConf.MAX_CASES_BRANCHES.key, "1")
+        assert(sql(sql_one_branch_caseWhen)
+          .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
+        assert(!sql(sql_two_branch_caseWhen)
+          .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
+
+        spark.conf.set(SQLConf.MAX_CASES_BRANCHES.key, "2")
+        assert(sql(sql_one_branch_caseWhen)
+          .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
+        assert(sql(sql_two_branch_caseWhen)
+          .queryExecution.executedPlan.isInstanceOf[WholeStageCodegenExec])
+      }
     } finally {
       spark.conf.set(SQLConf.MAX_CASES_BRANCHES.key, s"$original")
     }
   }
-
 }