[SPARK-36665][SQL][FOLLOWUP] Avoid Optimizing Not(InSubquery)

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

@@ -466,9 +466,18 @@ object NotPropagation extends Rule[LogicalPlan] {
     case _ => false
   }
 
+  private def hasInSubquery(x: Expression): Boolean = x match {
+    case _: InSubquery => true
+    case Not(e) => hasInSubquery(e)
+    case _ => false
+  }
+
   def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning(
     _.containsPattern(NOT), ruleId) {
     case q: LogicalPlan => q.transformExpressionsDownWithPruning(_.containsPattern(NOT), ruleId) {
+      // [SPARK-36665][SPARK-38085] Do not simplify Not(InSubquery)
+      case e @ Equality(a, b) if hasInSubquery(a) || hasInSubquery(b) => e
+
       // Move `Not` from one side of `EqualTo`/`EqualNullSafe` to the other side if it's beneficial.
       // E.g. `EqualTo(Not(a), b)` where `b = Not(c)`, it will become
       // `EqualTo(a, Not(b))` => `EqualTo(a, Not(Not(c)))` => `EqualTo(a, c)`
@@ -483,12 +492,12 @@ object NotPropagation extends Rule[LogicalPlan] {
 
       // Push `Not` to one side of `EqualTo`/`EqualNullSafe` if it's beneficial.
       // E.g. Not(EqualTo(x, false)) => EqualTo(x, true)
-      case Not(EqualTo(a, b)) if canSimplifyNot(b) => EqualTo(a, Not(b))
-      case Not(EqualTo(a, b)) if canSimplifyNot(a) => EqualTo(Not(a), b)
-      case Not(EqualNullSafe(a, b)) if !a.nullable && !b.nullable && canSimplifyNot(b) =>
-        EqualNullSafe(a, Not(b))
-      case Not(EqualNullSafe(a, b)) if !a.nullable && !b.nullable && canSimplifyNot(a) =>
-        EqualNullSafe(Not(a), b)
+      case Not(EqualTo(a, b)) if canSimplifyNot(b) && !hasInSubquery(b) => EqualTo(a, Not(b))
+      case Not(EqualTo(a, b)) if canSimplifyNot(a) && !hasInSubquery(a) => EqualTo(Not(a), b)
+      case Not(EqualNullSafe(a, b)) if !a.nullable && !b.nullable && canSimplifyNot(b) &&
+        !hasInSubquery(b) => EqualNullSafe(a, Not(b))
+      case Not(EqualNullSafe(a, b)) if !a.nullable && !b.nullable && canSimplifyNot(a) &&
+        !hasInSubquery(a) => EqualNullSafe(Not(a), b)
     }
   }
 }

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

@@ -173,4 +173,53 @@ class NotPropagationSuite extends PlanTest with ExpressionEvalHelper {
     checkCondition(('a === 'b) =!= ('a === 'c), ('a === 'b) =!= ('a === 'c))
     checkCondition(('a === 'b) =!= ('c in(1, 2, 3)), ('a === 'b) =!= ('c in(1, 2, 3)))
   }
+
+  test("SPARK-36665: Do not simplify Not(InSubquery)") {
+    object Optimize extends RuleExecutor[LogicalPlan] {
+      val batches = Batch("AnalysisNodes", Once, EliminateSubqueryAliases) ::
+        Batch("Not Propagation", FixedPoint(50), NotPropagation, BooleanSimplification) :: Nil
+    }
+
+    // As checkCondition() changes the exprId of sub-queries, defining a stable checker here
+    def shouldNotOptimize(input: Expression): Unit = {
+      val plan = testRelation.where(input).analyze
+      val actual = Optimize.execute(plan)
+      assert(actual == plan)
+    }
+
+    // Check whether the input becomes a BinaryComparison and its one side matches the expected
+    def checkOneSide(input: Expression, expected: Expression): Unit = {
+      val plan = testRelation.where(input).analyze
+      val actual = Optimize.execute(plan)
+      val comp = actual.asInstanceOf[Filter].condition.asInstanceOf[BinaryComparison]
+      val side = testRelation.where(expected).analyze.asInstanceOf[Filter].condition
+      assert(comp.left == side || comp.right == side)
+    }
+
+    val inSubquery = InSubquery(Seq('a), ListQuery(testRelation.select("a")))
+    shouldNotOptimize(Not(inSubquery) === Literal(true))
+    shouldNotOptimize(Literal(true) === Not(inSubquery))
+    shouldNotOptimize(Not(inSubquery) <=> Literal(true))
+    shouldNotOptimize(Literal(true) <=> Not(inSubquery))
+
+    shouldNotOptimize(Not(inSubquery) === Not('a === 'b))
+    shouldNotOptimize(Not('a === 'b) === Not(inSubquery))
+    shouldNotOptimize(Not(inSubquery) <=> Not('a === 'b))
+    shouldNotOptimize(Not('a === 'b) <=> Not(inSubquery))
+
+    shouldNotOptimize(Not(Not(inSubquery) === ('a === 'b)))
+    shouldNotOptimize(Not(('a === 'b) === Not(inSubquery)))
+    shouldNotOptimize(Not(Not(inSubquery) <=> ('a === 'b)))
+    shouldNotOptimize(Not(('a === 'b) <=> Not(inSubquery)))
+
+    checkOneSide(Not(Not(inSubquery) === Literal(true)), Literal(false))
+    checkOneSide(Not(Literal(true) === Not(inSubquery)), Literal(false))
+    shouldNotOptimize(Not(Not(inSubquery) <=> Literal(true)))
+    shouldNotOptimize(Not(Literal(true) <=> Not(inSubquery)))
+
+    checkOneSide(Not(Not(inSubquery)) === Not('a === 'b), Not('a === 'b))
+    checkOneSide(Not('a === 'b) === Not(Not(inSubquery)), Not('a === 'b))
+    checkOneSide(Not(Not(inSubquery)) <=> Not('a === 'b), Not('a === 'b))
+    checkOneSide(Not('a === 'b) <=> Not(Not(inSubquery)), Not('a === 'b))
+  }
 }

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

@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql
+
+import org.apache.spark.sql.test.SharedSparkSession
+
+class NotPropagationEndToEndSuite extends QueryTest with SharedSparkSession {
+  import testImplicits._
+
+  val t = "test_table"
+
+  test("SPARK-36665: Do not simplify Not(InSubquery)") {
+    withTable(t) {
+      Seq[(Integer, Integer)](
+        (1, 1),
+        (2, 2),
+        (3, 3),
+        (4, null),
+        (null, 0))
+        .toDF("c1", "c2").write.saveAsTable(t)
+      val df = spark.table(t)
+
+      checkAnswer(df.where(s"(c1 NOT IN (SELECT c2 FROM $t)) = true"), Seq.empty)
+      checkAnswer(df.where(s"(c1 NOT IN (SELECT c2 FROM $t WHERE c2 IS NOT NULL)) = true"),
+        Row(4, null) :: Nil)
+      checkAnswer(df.where(s"(c1 NOT IN (SELECT c2 FROM $t)) <=> true"), Seq.empty)
+      checkAnswer(df.where(s"(c1 NOT IN (SELECT c2 FROM $t WHERE c2 IS NOT NULL)) <=> true"),
+        Row(4, null) :: Nil)
+      checkAnswer(df.where(s"(c1 NOT IN (SELECT c2 FROM $t)) != false"), Seq.empty)
+      checkAnswer(df.where(s"(c1 NOT IN (SELECT c2 FROM $t WHERE c2 IS NOT NULL)) != false"),
+        Row(4, null) :: Nil)
+      checkAnswer(df.where(s"NOT((c1 NOT IN (SELECT c2 FROM $t)) <=> false)"), Seq.empty)
+      checkAnswer(df.where(s"NOT((c1 NOT IN (SELECT c2 FROM $t WHERE c2 IS NOT NULL)) <=> false)"),
+        Row(4, null) :: Nil)
+    }
+  }
+}