init

Author: ulysses-you

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

@@ -235,26 +235,39 @@ object ReorderAssociativeOperator extends Rule[LogicalPlan] {
  *    [[InSet (value, HashSet[Literal])]] which is much faster.
  */
 object OptimizeIn extends Rule[LogicalPlan] {
+  def optimizeIn(expr: In, v: Expression, list: Seq[Expression]): Expression = {
+    val newList = ExpressionSet(list).toSeq
+    if (newList.length == 1
+      // TODO: `EqualTo` for structural types are not working. Until SPARK-24443 is addressed,
+      // TODO: we exclude them in this rule.
+      && !v.isInstanceOf[CreateNamedStruct]
+      && !newList.head.isInstanceOf[CreateNamedStruct]) {
+      EqualTo(v, newList.head)
+    } else if (newList.length > SQLConf.get.optimizerInSetConversionThreshold) {
+      val hSet = newList.map(e => e.eval(EmptyRow))
+      InSet(v, HashSet() ++ hSet)
+    } else if (newList.length < list.length) {
+      expr.copy(list = newList)
+    } else { // newList.length == list.length && newList.length > 1
+      expr
+    }
+  }
+
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case q: LogicalPlan => q transformExpressionsDown {
       case In(v, list) if list.isEmpty =>
         // When v is not nullable, the following expression will be optimized
         // to FalseLiteral which is tested in OptimizeInSuite.scala
         If(IsNotNull(v), FalseLiteral, Literal(null, BooleanType))
-      case expr @ In(v, list) if expr.inSetConvertible =>
-        val newList = ExpressionSet(list).toSeq
-        if (newList.length == 1
-          // TODO: `EqualTo` for structural types are not working. Until SPARK-24443 is addressed,
-          // TODO: we exclude them in this rule.
-          && !v.isInstanceOf[CreateNamedStruct]
-          && !newList.head.isInstanceOf[CreateNamedStruct]) {
-          EqualTo(v, newList.head)
-        } else if (newList.length > SQLConf.get.optimizerInSetConversionThreshold) {
-          val hSet = newList.map(e => e.eval(EmptyRow))
-          InSet(v, HashSet() ++ hSet)
-        } else if (newList.length < list.length) {
-          expr.copy(list = newList)
-        } else { // newList.length == list.length && newList.length > 1
+      case expr @ In(v, list) =>
+        // split list to 2 parts so that we can push down convertible part
+        val (convertible, nonConvertible) = list.partition(_.isInstanceOf[Literal])
+        if (convertible.nonEmpty && nonConvertible.isEmpty) {
+          optimizeIn(expr, v, list)
+        } else if (convertible.nonEmpty && nonConvertible.nonEmpty) {
+          val optimizedIn = optimizeIn(In(v, convertible), v, convertible)
+          And(optimizedIn, In(v, nonConvertible))
+        } else {
           expr
         }
     }

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

@@ -238,4 +238,31 @@ class OptimizeInSuite extends PlanTest {
 
     comparePlans(optimized, correctAnswer)
   }
+
+  test("SPARK-32196: Extract convertible part if In 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(
+          And(EqualTo(UnresolvedAttribute("a"), Literal(1)),
+            In(UnresolvedAttribute("a"), Seq(UnresolvedAttribute("b"))))
+        )
+        .analyze
+    comparePlans(optimized1, correctAnswer1)
+
+    val originalQuery2 =
+      testRelation
+        .where(In(UnresolvedAttribute("a"), Seq(UnresolvedAttribute("b"))))
+        .analyze
+    val optimized2 = Optimize.execute(originalQuery2)
+    val correctAnswer2 =
+      testRelation
+        .where(In(UnresolvedAttribute("a"), Seq(UnresolvedAttribute("b"))))
+        .analyze
+    comparePlans(optimized2, correctAnswer2)
+  }
 }