[SPARK-40177][SQL] Simplify condition of form (a==b) || (a==null&&b==null) to a<=>b

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

@@ -127,6 +127,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
         EliminateSerialization,
         RemoveRedundantAliases,
         RemoveRedundantAggregates,
+        SimplifyJoinCondition,
         UnwrapCastInBinaryComparison,
         RemoveNoopOperators,
         OptimizeUpdateFields,

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

@@ -306,6 +306,47 @@ object OptimizeIn extends Rule[LogicalPlan] {
   }
 }
 
+/**
+ * SimplifyJoinCondition:
+ * Changes the join condition of type (key1=key2) || (isnull(key1)&&isnull(key2))
+ * to key1 <=> key2 to prevent join from converting to BNLJ in physical planning
+ * which is costly implementation of join.
+ */
+object SimplifyJoinCondition extends Rule[LogicalPlan] with PredicateHelper {
+  override def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning(
+    _.containsPattern(JOIN), ruleId) {
+    case join@Join(_, _, _, Some(joinCondition), _) =>
+      val predicates = splitConjunctivePredicates(joinCondition)
+      val newPreds = predicates.map {
+        _.transformUp {
+          case or@Or(left, right) =>
+            val pattern1 = getNewPattern(left, right)
+            val pattern2 = getNewPattern(right, left)
+            if (pattern1.isDefined) {
+              pattern1.get
+            } else if (pattern2.isDefined) {
+              pattern2.get
+            } else {
+              or
+            }
+        }
+      }
+      join.copy(condition = Some(newPreds.reduce(And)))
+    }
+
+  private def getNewPattern(left: Expression, right: Expression): Option[Expression] = {
+    left match {
+      case And(IsNull(attrO1), IsNull(attrO2)) =>
+        right match {
+          case EqualTo(attrI1, attrI2) if (attrO1 == attrI1 && attrO2 == attrI2) ||
+            (attrO1 == attrI2 && attrO2 == attrI1) =>
+            Some(EqualNullSafe(attrI1, attrI2))
+          case _ => None
+        }
+      case _ => None
+    }
+  }
+}
 
 /**
  * Simplifies boolean expressions:

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala

@@ -96,6 +96,7 @@ object RuleIdCollection {
       "org.apache.spark.sql.catalyst.analysis.UpdateOuterReferences" ::
       "org.apache.spark.sql.catalyst.analysis.UpdateAttributeNullability" ::
       // Catalyst Optimizer rules
+      "org.apache.spark.sql.catalyst.optimizer.SimplifyJoinCondition" ::
       "org.apache.spark.sql.catalyst.optimizer.BooleanSimplification" ::
       "org.apache.spark.sql.catalyst.optimizer.CollapseProject" ::
       "org.apache.spark.sql.catalyst.optimizer.CollapseRepartition" ::

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

@@ -0,0 +1,126 @@
+/*
+ * 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.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans.DslLogicalPlan
+import org.apache.spark.sql.catalyst.expressions.{IsNotNull, IsNull}
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+
+class SimplifyJoinConditionSuite extends PlanTest {
+  object Optimize extends RuleExecutor[LogicalPlan] {
+    val batches =
+      Batch("Simplify Join Condition", Once,
+        SimplifyJoinCondition)  :: Nil
+  }
+
+  val testRelation = LocalRelation('a.int, 'b.int, 'c.int)
+  val testRelation1 = LocalRelation('d.int, 'e.int)
+
+  test("Simple condition with null check on right side of or") {
+    val originalQuery = testRelation
+      .join(testRelation1, condition = Some(('b === 'd)||(IsNull('b) && IsNull('d))))
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation
+      .join(testRelation1, condition = Some('b <=> 'd))
+      .analyze
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("Simple condition with null check on left side of or") {
+    val originalQuery = testRelation
+      .join(testRelation1, condition = Some((IsNull('b) && IsNull('d)) || ('b === 'd)))
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation
+      .join(testRelation1, condition = Some('b <=> 'd))
+      .analyze
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("Simple condition with is not null check on one column") {
+    val originalQuery = testRelation
+      .join(testRelation1, condition = Some((IsNull('b) && IsNotNull('d)) || ('b === 'd)))
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation
+      .join(testRelation1, condition = Some((IsNull('b) && IsNotNull('d)) || ('b === 'd)))
+      .analyze
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("multiple equal null safe conditions separated by and") {
+    val originalQuery = testRelation.join(testRelation1,
+      condition = Some(((IsNull('b) && IsNull('d)) || ('b === 'd)) &&
+        ((IsNull('a) && IsNull('e)) || ('a === 'e))))
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation
+      .join(testRelation1, condition = Some('b <=> 'd && 'a <=> 'e))
+      .analyze
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("multiple equal null safe conditions separated by or") {
+    val originalQuery = testRelation.join(testRelation1,
+      condition = Some(((IsNull('b) && IsNull('d)) || ('b === 'd)) ||
+        ((IsNull('a) && IsNull('e)) || ('a === 'e))))
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation
+      .join(testRelation1, condition = Some('b <=> 'd || 'a <=> 'e))
+      .analyze
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("Condition with another or in expression") {
+    val originalQuery = testRelation.join(testRelation1,
+      condition = Some((IsNull('b) && IsNull('d)) || ('b === 'd) || ('a === 'e)))
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation.join(testRelation1,
+      condition = Some('b <=> 'd || ('a === 'e)))
+      .analyze
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("Condition with another and in expression so that and gets calculated first") {
+    val originalQuery = testRelation.join(testRelation1,
+      condition = Some((IsNull('b) && IsNull('d)) || ('b === 'd) && ('a === 'e)))
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation.join(testRelation1,
+      condition = Some((IsNull('b) && IsNull('d)) || ('b === 'd) && ('a === 'e)))
+      .analyze
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("Condition with another and in expression") {
+    val originalQuery = testRelation.join(testRelation1,
+      condition = Some(((IsNull('b) && IsNull('d)) || ('b === 'd)) && ('a === 'e)))
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation.join(testRelation1,
+      condition = Some('b <=> 'd && ('a === 'e)))
+      .analyze
+    comparePlans(optimized, correctAnswer)
+  }
+}