PullOutJoinCondition

Author: wangyum

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

@@ -151,6 +151,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
       ReplaceExpressions,
       RewriteNonCorrelatedExists,
       PullOutGroupingExpressions,
+      PullOutJoinCondition,
       ComputeCurrentTime,
       ReplaceCurrentLike(catalogManager)) ::
     //////////////////////////////////////////////////////////////////////////////////////////

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

@@ -0,0 +1,75 @@
+/*
+ * 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.expressions.{Alias, Expression, PredicateHelper}
+import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan, Project}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern.JOIN
+
+/**
+ * This rule ensures that [[Join]] condition doesn't contain complex expressions in the
+ * optimization phase.
+ *
+ * Complex condition expressions are pulled out to a [[Project]] node under [[Join]] and are
+ * referenced in join condition.
+ *
+ * {{{
+ *   SELECT * FROM t1 JOIN t2 ON t1.a + 10 = t2.x ==>
+ *   Project [a#0, b#1, x#2, y#3]
+ *   +- Join Inner, ((spark_catalog.default.t1.a + 10)#8 = x#2)
+ *      :- Project [a#0, b#1, (a#0 + 10) AS (spark_catalog.default.t1.a + 10)#8]
+ *      :  +- Filter isnotnull((a#0 + 10))
+ *      :     +- Relation default.t1[a#0,b#1] parquet
+ *      +- Filter isnotnull(x#2)
+ *         +- Relation default.t2[x#2,y#3] parquet
+ * }}}
+ */
+object PullOutJoinCondition extends Rule[LogicalPlan]
+  with JoinSelectionHelper with PredicateHelper {
+
+  def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning(_.containsPattern(JOIN)) {
+    case j @ Join(left, right, _, Some(condition), _) if j.resolved =>
+      val complexExpressions = splitConjunctivePredicates(condition).flatMap(_.children).flatMap {
+        case e: Expression if !e.foldable && e.children.nonEmpty => Seq(e)
+        case _ => Nil
+      }
+
+      val leftComplexExpressions = complexExpressions.filter(canEvaluate(_, left))
+      val rightComplexExpressions = complexExpressions.filter(canEvaluate(_, right))
+
+      val leftComplexExpressionMap =
+        leftComplexExpressions.map(e => e.canonicalized -> Alias(e, e.sql)()).toMap
+      val rightComplexExpressionMap =
+        rightComplexExpressions.map(e => e.canonicalized -> Alias(e, e.sql)()).toMap
+      val allComplexExpressionMap = leftComplexExpressionMap ++ rightComplexExpressionMap
+
+      if (allComplexExpressionMap.nonEmpty) {
+        val newCondition = condition.transformDown {
+          case e: Expression
+              if e.children.nonEmpty && allComplexExpressionMap.contains(e.canonicalized) =>
+            allComplexExpressionMap.get(e.canonicalized).map(_.toAttribute).getOrElse(e)
+        }
+        val newLeft = Project(left.output ++ leftComplexExpressionMap.values, left)
+        val newRight = Project(right.output ++ rightComplexExpressionMap.values, right)
+        j.copy(left = newLeft, right = newRight, condition = Some(newCondition))
+      } else {
+        j
+      }
+  }
+}

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

@@ -144,7 +144,6 @@ object RuleIdCollection {
       "org.apache.spark.sql.catalyst.optimizer.ReplaceIntersectWithSemiJoin" ::
       "org.apache.spark.sql.catalyst.optimizer.RewriteExceptAll" ::
       "org.apache.spark.sql.catalyst.optimizer.RewriteIntersectAll" ::
-      "org.apache.spark.sql.catalyst.optimizer.PushDownJoinConditionEvaluation" ::
       "org.apache.spark.sql.catalyst.optimizer.SimplifyBinaryComparison" ::
       "org.apache.spark.sql.catalyst.optimizer.SimplifyCaseConversionExpressions" ::
       "org.apache.spark.sql.catalyst.optimizer.SimplifyCasts" ::

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

@@ -0,0 +1,83 @@
+/*
+ * 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._
+import org.apache.spark.sql.catalyst.expressions.{Alias, Coalesce, Substring, Upper}
+import org.apache.spark.sql.catalyst.plans._
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules._
+
+class PullOutJoinConditionSuite extends PlanTest {
+
+  private object Optimize extends RuleExecutor[LogicalPlan] {
+    val batches =
+      Batch("Pull out join condition", Once,
+        PullOutJoinCondition,
+        CollapseProject) :: Nil
+  }
+
+  private val testRelation = LocalRelation('a.string, 'b.int, 'c.int)
+  private val testRelation1 = LocalRelation('d.string, 'e.int)
+  private val x = testRelation.subquery('x)
+  private val y = testRelation1.subquery('y)
+
+  test("Push down join condition evaluation(String expressions)") {
+    val joinType = Inner
+    Seq(Upper("y.d".attr), Substring("y.d".attr, 1, 5)).foreach { udf =>
+      val originalQuery = x.join(y, joinType, Option("x.a".attr === udf))
+        .select("x.a".attr, "y.e".attr)
+      val correctAnswer = x.select("x.a".attr, "x.b".attr, "x.c".attr)
+        .join(y.select("y.d".attr, "y.e".attr, Alias(udf, udf.sql)()),
+          joinType, Option("x.a".attr === s"`${udf.sql}`".attr)).select("x.a".attr, "y.e".attr)
+
+      comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze)
+    }
+  }
+
+  test("Push down join condition evaluation(null expressions)") {
+    val joinType = Inner
+    val udf = Coalesce(Seq("x.b".attr, "x.c".attr))
+    val originalQuery = x.join(y, joinType, Option(udf === "y.e".attr))
+      .select("x.a".attr, "y.e".attr)
+    val correctAnswer =
+      x.select("x.a".attr, "x.b".attr, "x.c".attr, Alias(udf, udf.sql)()).join(
+        y.select("y.d".attr, "y.e".attr),
+        joinType, Option(s"`${udf.sql}`".attr === "y.e".attr))
+        .select("x.a".attr, "y.e".attr)
+
+    comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze)
+  }
+
+  test("Negative case: all children are Attributes") {
+    val condition = Option("x.a".attr === "y.d".attr)
+    val originalQuery = x.join(y, Inner, condition)
+    val correctAnswer = x.join(y, Inner, condition)
+
+    comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze)
+  }
+
+  test("Negative case: contains Literal") {
+    val condition = Option("x.a".attr === "string")
+    val originalQuery = x.join(y, Inner, condition)
+    val correctAnswer = x.join(y, Inner, condition)
+
+    comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze)
+  }
+}

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

@@ -48,10 +48,7 @@ class SparkOptimizer(
       PushDownPredicates) :+
     Batch("Cleanup filters that cannot be pushed down", Once,
       CleanupDynamicPruningFilters,
-      PruneFilters) :+
-    Batch("Push down join condition evaluation", Once,
-      PushDownJoinConditionEvaluation,
-      CollapseProject)) ++
+      PruneFilters)) ++
     postHocOptimizationBatches :+
     Batch("Extract Python UDFs", Once,
       ExtractPythonUDFFromJoinCondition,