[CARMEL-6752] Support more subexpression elimination cases (apache#266)

Author: wakun

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

@@ -21,7 +21,9 @@ import java.util.Objects
 
 import scala.collection.mutable
 
+import org.apache.spark.sql.catalyst.expressions.EquivalentExpressions.supportedExpression
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.expressions.codegen.ExprValue
 import org.apache.spark.sql.catalyst.expressions.objects.LambdaVariable
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.util.Utils
@@ -151,32 +153,9 @@ class EquivalentExpressions(
   //   2. ConditionalExpression: use its children that will always be evaluated.
   private def childrenToRecurse(expr: Expression): Seq[Expression] = expr match {
     case _: CodegenFallback => Nil
-    case c: ConditionalExpression => c.alwaysEvaluatedInputs.map(skipForShortcut)
     case other => skipForShortcut(other).children
   }
 
-  // For some special expressions we cannot just recurse into all of its children, but we can
-  // recursively add the common expressions shared between all of its children.
-  private def commonChildrenToRecurse(expr: Expression): Seq[Seq[Expression]] = expr match {
-    case _: CodegenFallback => Nil
-    case c: ConditionalExpression => c.branchGroups
-    case _ => Nil
-  }
-
-  private def supportedExpression(e: Expression) = {
-    !e.exists {
-      // `LambdaVariable` is usually used as a loop variable, which can't be evaluated ahead of the
-      // loop. So we can't evaluate sub-expressions containing `LambdaVariable` at the beginning.
-      case _: LambdaVariable => true
-
-      // `PlanExpression` wraps query plan. To compare query plans of `PlanExpression` on executor,
-      // can cause error like NPE.
-      case _: PlanExpression[_] => Utils.isInRunningSparkTask
-
-      case _ => false
-    }
-  }
-
   /**
    * Adds the expression to this data structure recursively. Stops if a matching expression
    * is found. That is, if `expr` has already been added, its children are not added.
@@ -198,7 +177,6 @@ class EquivalentExpressions(
     if (!skip && !updateExprInMap(expr, map, useCount)) {
       val uc = useCount.signum
       childrenToRecurse(expr).foreach(updateExprTree(_, map, uc))
-      commonChildrenToRecurse(expr).filter(_.nonEmpty).foreach(updateCommonExprs(_, map, uc))
     }
   }
 
@@ -240,6 +218,23 @@ class EquivalentExpressions(
   }
 }
 
+object EquivalentExpressions {
+  def supportedExpression(e: Expression): Boolean = {
+    !e.exists {
+      // `LambdaVariable` is usually used as a loop variable and `NamedLambdaVariable` is used in
+      // higher-order functions, which can't be evaluated ahead of the execution.
+      case _: LambdaVariable => true
+      case _: NamedLambdaVariable => true
+
+      // `PlanExpression` wraps query plan. To compare query plans of `PlanExpression` on executor,
+      // can cause error like NPE.
+      case _: PlanExpression[_] => Utils.isInRunningSparkTask
+
+      case _ => false
+    }
+  }
+}
+
 /**
  * Wrapper around an Expression that provides semantic equality.
  */
@@ -267,4 +262,11 @@ case class ExpressionEquals(e: Expression) {
  * Instead of appending to a mutable list/buffer of Expressions, just update the "flattened"
  * useCount in this wrapper in-place.
  */
-case class ExpressionStats(expr: Expression)(var useCount: Int)
+case class ExpressionStats(expr: Expression)(
+  var useCount: Int,
+  var initialized: Option[String] = None,
+  var isNull: Option[ExprValue] = None,
+  var value: Option[ExprValue] = None,
+  var funcName: Option[String] = None,
+  var params: Option[Seq[Class[_]]] = None,
+  var addedFunction: Boolean = false)

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

@@ -196,9 +196,40 @@ abstract class Expression extends TreeNode[Expression] {
     }.getOrElse {
       val isNull = ctx.freshName("isNull")
       val value = ctx.freshName("value")
-      val eval = doGenCode(ctx, ExprCode(
-        JavaCode.isNullVariable(isNull),
-        JavaCode.variable(value, dataType)))
+      val eval =
+        if (EquivalentExpressions.supportedExpression(this)) {
+          ctx.commonExpressions.get(ExpressionEquals(this)) match {
+            case Some(stats) =>
+              // We should reuse the currentVar references which code is not empty
+              val nonEmptyRefs = this.exists {
+                case BoundReference(ordinal, _, _) =>
+                  ctx.currentVars != null && ctx.currentVars(ordinal) != null &&
+                    ctx.currentVars(ordinal).code != EmptyBlock
+                case _ => false
+              }
+              val eval = doGenCode(ctx, ExprCode(
+                JavaCode.isNullVariable(isNull),
+                JavaCode.variable(value, dataType)))
+              if (eval.code != EmptyBlock && !nonEmptyRefs) {
+                ctx.genReusedCode(stats, eval)
+              } else {
+                eval
+              }
+
+            case None =>
+              doGenCode(ctx, ExprCode(
+                JavaCode.isNullVariable(isNull),
+                JavaCode.variable(value, dataType)))
+          }
+        } else {
+          doGenCode(ctx, ExprCode(
+            JavaCode.isNullVariable(isNull),
+            JavaCode.variable(value, dataType)))
+        }
+
+
+
+
       reduceCodeSize(ctx, eval)
       if (eval.code.toString.nonEmpty) {
         // Add `this` in the comment.