[SPARK-58211][SQL] Peel all leading AND/OR operands in subexpression elimination shortcut

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

@@ -133,9 +133,15 @@ class EquivalentExpressions(
     if (skipForShortcutEnable) {
       // The subexpression may not need to eval even if it appears more than once.
       // e.g., `if(or(a, and(b, b)))`, the expression `b` would be skipped if `a` is true.
+      // `And`/`Or` short-circuit, so only the left operand is guaranteed to be evaluated.
+      // A chained predicate `a AND b AND c` is left-deep, i.e. `And(And(a, b), c)`, so we must
+      // keep peeling left operands until we reach the single always-evaluated leaf. Peeling only
+      // one level would leave operands past the first (which are conditionally evaluated) to be
+      // treated as always-evaluated, so a subexpression they share could be hoisted and eagerly
+      // evaluated, breaking short-circuit semantics (e.g. a spurious NPE or divide-by-zero).
       expr match {
-        case and: And => and.left
-        case or: Or => or.left
+        case and: And => skipForShortcut(and.left)
+        case or: Or => skipForShortcut(or.left)
         case other => other
       }
     } else {

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala

@@ -495,6 +495,44 @@ class SubexpressionEliminationSuite extends SparkFunSuite with ExpressionEvalHel
     checkShortcut(Not(And(equal, Literal(false))), 1)
   }
 
+  test("SPARK-58211: shortcut eliminate operand past the first in a chained AND/OR") {
+    val add = Add(Literal(1), Literal(0))
+    val equal = EqualTo(add, add)
+
+    def checkShortcut(expr: Expression, numCommonExpr: Int): Unit = {
+      val e1 = If(expr, Literal(1), Literal(2))
+      val ee1 = new EquivalentExpressions(true)
+      ee1.addExprTree(e1)
+      assert(ee1.getCommonSubexpressions.size == numCommonExpr)
+
+      val e2 = expr
+      val ee2 = new EquivalentExpressions(true)
+      ee2.addExprTree(e2)
+      assert(ee2.getCommonSubexpressions.size == numCommonExpr)
+    }
+
+    // A left-deep chain `a AND b AND c` is `And(And(a, b), c)`. Only `a` is always evaluated;
+    // `b` and `c` are short-circuited, so their subexpressions must not be eliminated. Peeling
+    // a single operand would wrongly recurse into `b` and eliminate its inner subexpression.
+    checkShortcut(And(And(Literal(false), equal), Literal(true)), 0)
+    checkShortcut(Or(Or(Literal(true), equal), Literal(false)), 0)
+    checkShortcut(And(And(Literal(false), Literal(true)), equal), 0)
+    checkShortcut(Or(Or(Literal(true), Literal(false)), equal), 0)
+
+    // The always-evaluated leftmost operand is still eligible for elimination.
+    checkShortcut(And(And(equal, Literal(false)), Literal(true)), 1)
+    checkShortcut(Or(Or(equal, Literal(true)), Literal(false)), 1)
+
+    // Deeper chain `a AND b AND c AND d` = `And(And(And(a, b), c), d)`. The subexpression lives
+    // in a conditional operand (`c`), so it must not be eliminated no matter the nesting depth.
+    checkShortcut(And(And(And(Literal(false), Literal(true)), equal), Literal(true)), 0)
+
+    // Mixed nesting: the always-evaluated leaf `equal` is reached through a left spine of both
+    // `And` and `Or`. Its internal subexpression (`add`, which appears twice inside `equal`) is
+    // still eliminated, so the recursive peel is not over-conservative on the leaf.
+    checkShortcut(And(Or(equal, Literal(true)), equal), 1)
+  }
+
   test("Equivalent ternary expressions have different children") {
     val add1 = Add(Add(Literal(1), Literal(2)), Literal(3))
     val add2 = Add(Add(Literal(3), Literal(1)), Literal(2))

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

@@ -3333,6 +3333,27 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper
     }
   }
 
+  test("SPARK-58211: subexpression elimination respects AND/OR short-circuit with chained ANDs") {
+    // A subexpression (1 / id) repeated inside a short-circuited operand of a chained AND/OR
+    // must not be hoisted and eagerly evaluated. For id = 0, `id != 0` is false, so the second
+    // operand should never be evaluated and no divide-by-zero should be raised. The subexpression
+    // is duplicated *inside* operand 2 (not split across operands 2 and 3) so the failure shape
+    // does not depend on how many operands the one-level peel happened to drop, and operand 3 is
+    // non-foldable to keep it in the plan. skipForShortcutExpr must peel every leading AND/OR
+    // operand, not just one, to make this hold for three or more operands.
+    withSQLConf(
+      SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true",
+      SQLConf.SUBEXPRESSION_ELIMINATION_SKIP_FOR_SHORTCUT_EXPR.key -> "true") {
+      val andQuery =
+        "select id != 0 and (1 / id + 1 / id) > 0 and id >= 0 from range(0, 1, 1, 1)"
+      checkAnswer(sql(andQuery), Row(false))
+
+      val orQuery =
+        "select id == 0 or (1 / id + 1 / id) > 0 or id < 0 from range(0, 1, 1, 1)"
+      checkAnswer(sql(orQuery), Row(true))
+    }
+  }
+
   test("SPARK-29213: FilterExec should not throw NPE") {
     // Under ANSI mode, casting string '' as numeric will cause runtime error
     if (!conf.ansiEnabled) {