add conf

Author: ulysses-you

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

@@ -265,7 +265,8 @@ object OptimizeIn extends Rule[LogicalPlan] {
         val (convertible, nonConvertible) = list.partition(_.isInstanceOf[Literal])
         if (convertible.nonEmpty && nonConvertible.isEmpty) {
           optimizeIn(expr, v, list)
-        } else if (convertible.nonEmpty && nonConvertible.nonEmpty) {
+        } else if (convertible.nonEmpty && nonConvertible.nonEmpty &&
+          SQLConf.get.optimizerInExtractLiteralPart) {
           val optimizedIn = optimizeIn(In(v, convertible), v, convertible)
           Or(optimizedIn, In(v, nonConvertible))
         } else {

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -197,6 +197,14 @@ object SQLConf {
     .intConf
     .createWithDefault(100)
 
+  val OPTIMIZER_IN_EXTRACT_LITERAL_PART =
+    buildConf("spark.sql.optimizer.inExtractLiteralPart")
+      .internal()
+      .doc("When true, we will extract and optimize the literal part of in if not all are literal.")
+      .version("3.1.0")
+      .booleanConf
+      .createWithDefault(true)
+
   val OPTIMIZER_INSET_CONVERSION_THRESHOLD =
     buildConf("spark.sql.optimizer.inSetConversionThreshold")
       .internal()
@@ -2761,6 +2769,8 @@ class SQLConf extends Serializable with Logging {
 
   def optimizerMaxIterations: Int = getConf(OPTIMIZER_MAX_ITERATIONS)
 
+  def optimizerInExtractLiteralPart: Boolean = getConf(OPTIMIZER_IN_EXTRACT_LITERAL_PART)
+
   def optimizerInSetConversionThreshold: Int = getConf(OPTIMIZER_INSET_CONVERSION_THRESHOLD)
 
   def optimizerInSetSwitchThreshold: Int = getConf(OPTIMIZER_INSET_SWITCH_THRESHOLD)

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeInSuite.scala

@@ -24,6 +24,7 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{Filter, LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.internal.SQLConf.OPTIMIZER_IN_EXTRACT_LITERAL_PART
 import org.apache.spark.sql.internal.SQLConf.OPTIMIZER_INSET_CONVERSION_THRESHOLD
 import org.apache.spark.sql.types._
 
@@ -225,19 +226,32 @@ class OptimizeInSuite extends PlanTest {
   }
 
   test("SPARK-32196: Extract In convertible part if it is not convertible") {
-    val originalQuery1 =
-      testRelation
-        .where(In(UnresolvedAttribute("a"), Seq(Literal(1), UnresolvedAttribute("b"))))
-        .analyze
-    val optimized1 = Optimize.execute(originalQuery1)
-    val correctAnswer1 =
-      testRelation
-        .where(
-          Or(EqualTo(UnresolvedAttribute("a"), Literal(1)),
-            In(UnresolvedAttribute("a"), Seq(UnresolvedAttribute("b"))))
-        )
-        .analyze
-    comparePlans(optimized1, correctAnswer1)
+    Seq("true", "false").foreach { enable =>
+      withSQLConf(OPTIMIZER_IN_EXTRACT_LITERAL_PART.key -> enable) {
+        val originalQuery1 =
+          testRelation
+            .where(In(UnresolvedAttribute("a"), Seq(Literal(1), UnresolvedAttribute("b"))))
+            .analyze
+        val optimized1 = Optimize.execute(originalQuery1)
+
+        if (enable.toBoolean) {
+          val correctAnswer1 =
+            testRelation
+              .where(
+                Or(EqualTo(UnresolvedAttribute("a"), Literal(1)),
+                  In(UnresolvedAttribute("a"), Seq(UnresolvedAttribute("b"))))
+              )
+              .analyze
+          comparePlans(optimized1, correctAnswer1)
+        } else {
+          val correctAnswer1 =
+            testRelation
+              .where(In(UnresolvedAttribute("a"), Seq(Literal(1), UnresolvedAttribute("b"))))
+              .analyze
+          comparePlans(optimized1, correctAnswer1)
+        }
+      }
+    }
 
     val originalQuery2 =
       testRelation