check thread safety of expression

Author: Davies Liu

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

@@ -460,7 +460,7 @@ case class Cast(child: Expression, dataType: DataType) extends UnaryExpression w
       case (BooleanType, dt: NumericType) =>
         defineCodeGen(ctx, ev, c => s"(${ctx.javaType(dt)})($c ? 1 : 0)")
       case (dt: DecimalType, BooleanType) =>
-        defineCodeGen(ctx, ev, c => s"$c.isZero()")
+        defineCodeGen(ctx, ev, c => s"!($c).equals(0)")
       case (dt: NumericType, BooleanType) =>
         defineCodeGen(ctx, ev, c => s"$c != 0")
 

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

@@ -52,6 +52,14 @@ abstract class Expression extends TreeNode[Expression] {
   /** Returns the result of evaluating this expression on a given input Row */
   def eval(input: Row = null): Any
 
+  /**
+   * Return true if this expression is thread-safe, which means it could be used by multiple
+   * threads in the same time.
+   *
+   * An expression that is not thread-safe can not be cached and re-used, especially for codegen.
+   */
+  def isThreadSafe: Boolean = true
+
   /**
    * Returns an [[GeneratedExpressionCode]], which contains Java source code that
    * can be used to generate the result of evaluating the expression on an input row.
@@ -60,6 +68,9 @@ abstract class Expression extends TreeNode[Expression] {
    * @return [[GeneratedExpressionCode]]
    */
   def gen(ctx: CodeGenContext): GeneratedExpressionCode = {
+    if (!isThreadSafe) {
+      throw new Exception(s"$this is not thread-safe, can not be used in codegen")
+    }
     val isNull = ctx.freshName("isNull")
     val primitive = ctx.freshName("primitive")
     val ve = GeneratedExpressionCode("", isNull, primitive)
@@ -156,6 +167,7 @@ abstract class BinaryExpression extends Expression with trees.BinaryNode[Express
 
   override def toString: String = s"($left $symbol $right)"
 
+  override def isThreadSafe: Boolean = left.isThreadSafe && right.isThreadSafe
   /**
    * Short hand for generating binary evaluation code, which depends on two sub-evaluations of
    * the same type.  If either of the sub-expressions is null, the result of this computation
@@ -203,6 +215,8 @@ abstract class LeafExpression extends Expression with trees.LeafNode[Expression]
 abstract class UnaryExpression extends Expression with trees.UnaryNode[Expression] {
   self: Product =>
 
+  override def isThreadSafe: Boolean = child.isThreadSafe
+
   /**
    * Called by unary expressions to generate a code block that returns null if its parent returns
    * null, and if not not null, use `f` to generate the expression.
@@ -240,6 +254,7 @@ case class GroupExpression(children: Seq[Expression]) extends Expression {
   override def nullable: Boolean = false
   override def foldable: Boolean = false
   override def dataType: DataType = throw new UnsupportedOperationException
+  override def isThreadSafe: Boolean = children.forall(_.isThreadSafe)
 }
 
 /**

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

@@ -347,31 +347,29 @@ case class MaxOf(left: Expression, right: Expression) extends BinaryArithmetic {
   }
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    if (ctx.isNativeType(left.dataType)) {
-      val eval1 = left.gen(ctx)
-      val eval2 = right.gen(ctx)
-      eval1.code + eval2.code + s"""
-        boolean ${ev.isNull} = false;
-        ${ctx.javaType(left.dataType)} ${ev.primitive} =
-          ${ctx.defaultValue(left.dataType)};
-
-        if (${eval1.isNull}) {
-          ${ev.isNull} = ${eval2.isNull};
-          ${ev.primitive} = ${eval2.primitive};
-        } else if (${eval2.isNull}) {
-          ${ev.isNull} = ${eval1.isNull};
+    val eval1 = left.gen(ctx)
+    val eval2 = right.gen(ctx)
+    val compCode = ctx.compFunc(dataType)(eval1.primitive, eval2.primitive)
+
+    eval1.code + eval2.code + s"""
+      boolean ${ev.isNull} = false;
+      ${ctx.javaType(left.dataType)} ${ev.primitive} =
+        ${ctx.defaultValue(left.dataType)};
+
+      if (${eval1.isNull}) {
+        ${ev.isNull} = ${eval2.isNull};
+        ${ev.primitive} = ${eval2.primitive};
+      } else if (${eval2.isNull}) {
+        ${ev.isNull} = ${eval1.isNull};
+        ${ev.primitive} = ${eval1.primitive};
+      } else {
+        if ($compCode > 0) {
           ${ev.primitive} = ${eval1.primitive};
         } else {
-          if (${eval1.primitive} > ${eval2.primitive}) {
-            ${ev.primitive} = ${eval1.primitive};
-          } else {
-            ${ev.primitive} = ${eval2.primitive};
-          }
+          ${ev.primitive} = ${eval2.primitive};
         }
-      """
-    } else {
-      super.genCode(ctx, ev)
-    }
+      }
+    """
   }
   override def toString: String = s"MaxOf($left, $right)"
 }
@@ -401,33 +399,29 @@ case class MinOf(left: Expression, right: Expression) extends BinaryArithmetic {
   }
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    if (ctx.isNativeType(left.dataType)) {
-
-      val eval1 = left.gen(ctx)
-      val eval2 = right.gen(ctx)
-
-      eval1.code + eval2.code + s"""
-        boolean ${ev.isNull} = false;
-        ${ctx.javaType(left.dataType)} ${ev.primitive} =
-          ${ctx.defaultValue(left.dataType)};
+    val eval1 = left.gen(ctx)
+    val eval2 = right.gen(ctx)
+    val compCode = ctx.compFunc(dataType)(eval1.primitive, eval2.primitive)
 
-        if (${eval1.isNull}) {
-          ${ev.isNull} = ${eval2.isNull};
-          ${ev.primitive} = ${eval2.primitive};
-        } else if (${eval2.isNull}) {
-          ${ev.isNull} = ${eval1.isNull};
+    eval1.code + eval2.code + s"""
+      boolean ${ev.isNull} = false;
+      ${ctx.javaType(left.dataType)} ${ev.primitive} =
+        ${ctx.defaultValue(left.dataType)};
+
+      if (${eval1.isNull}) {
+        ${ev.isNull} = ${eval2.isNull};
+        ${ev.primitive} = ${eval2.primitive};
+      } else if (${eval2.isNull}) {
+        ${ev.isNull} = ${eval1.isNull};
+        ${ev.primitive} = ${eval1.primitive};
+      } else {
+        if ($compCode < 0) {
           ${ev.primitive} = ${eval1.primitive};
         } else {
-          if (${eval1.primitive} < ${eval2.primitive}) {
-            ${ev.primitive} = ${eval1.primitive};
-          } else {
-            ${ev.primitive} = ${eval2.primitive};
-          }
+          ${ev.primitive} = ${eval2.primitive};
         }
-      """
-    } else {
-      super.genCode(ctx, ev)
-    }
+      }
+    """
   }
 
   override def toString: String = s"MinOf($left, $right)"

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala

@@ -164,12 +164,27 @@ class CodeGenContext {
   def equalFunc(dataType: DataType): ((String, String) => String) = dataType match {
     case BinaryType => { case (eval1, eval2) =>
       s"java.util.Arrays.equals($eval1, $eval2)" }
-    case IntegerType | BooleanType | LongType | DoubleType | FloatType | ShortType | ByteType =>
+    case IntegerType | BooleanType | LongType | DoubleType | FloatType | ShortType | ByteType
+         | DateType =>
       { case (eval1, eval2) => s"$eval1 == $eval2" }
     case other =>
       { case (eval1, eval2) => s"$eval1.equals($eval2)" }
   }
 
+  /**
+   * Return a function to generate compare expression in Java
+   */
+  def compFunc(dataType: DataType): (String, String) => String = dataType match {
+    case BinaryType => {
+      case (c1, c2) =>
+        s"org.apache.spark.sql.catalyst.util.TypeUtils.compareBinary($c1, $c2)"
+    }
+    case IntegerType | LongType | DoubleType | FloatType | ShortType | ByteType | DateType => {
+      case (c1, c2) => s"$c1 - $c2"
+    }
+    case other => { case (c1, c2) => s"$c1.compare($c2)" }
+  }
+
   /**
    * List of data types that have special accessors and setters in [[Row]].
    */

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateOrdering.scala

@@ -21,7 +21,7 @@ import org.apache.spark.Logging
 import org.apache.spark.annotation.Private
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.types.{DateType, BinaryType, NumericType}
+import org.apache.spark.sql.types.{DateType, BinaryType, DecimalType, NumericType}
 
 /**
  * Inherits some default implementation for Java from `Ordering[Row]`
@@ -66,7 +66,16 @@ object GenerateOrdering extends CodeGenerator[Seq[SortOrder], Ordering[Row]] wit
                 return d;
               }
             }"""
-        case _: NumericType | DateType =>
+        case dt: NumericType if !dt.isInstanceOf[DecimalType] =>
+          s"""
+            if (${evalA.primitive} != ${evalB.primitive}) {
+              if (${evalA.primitive} > ${evalB.primitive}) {
+                return ${if (asc) "1" else "-1"};
+              } else {
+                return ${if (asc) "-1" else "1"};
+              }
+            }"""
+        case DateType =>
           s"""
             if (${evalA.primitive} != ${evalB.primitive}) {
               if (${evalA.primitive} > ${evalB.primitive}) {

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

@@ -116,6 +116,8 @@ case class Alias(child: Expression, name: String)(
 
   override def eval(input: Row): Any = child.eval(input)
 
+  override def isThreadSafe: Boolean = child.isThreadSafe
+
   override def gen(ctx: CodeGenContext): GeneratedExpressionCode = child.gen(ctx)
 
   override def dataType: DataType = child.dataType

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

@@ -17,9 +17,8 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
-import org.apache.spark.sql.catalyst.expressions.codegen.{GeneratedExpressionCode, CodeGenContext}
-import org.apache.spark.sql.catalyst.trees
 import org.apache.spark.sql.catalyst.analysis.UnresolvedException
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodeGenContext, GeneratedExpressionCode}
 import org.apache.spark.sql.types.DataType
 
 case class Coalesce(children: Seq[Expression]) extends Expression {
@@ -53,6 +52,8 @@ case class Coalesce(children: Seq[Expression]) extends Expression {
     result
   }
 
+  override def isThreadSafe: Boolean = children.forall(_.isThreadSafe)
+
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     s"""
       boolean ${ev.isNull} = true;
@@ -73,7 +74,7 @@ case class Coalesce(children: Seq[Expression]) extends Expression {
   }
 }
 
-case class IsNull(child: Expression) extends Predicate with trees.UnaryNode[Expression] {
+case class IsNull(child: Expression) extends UnaryExpression with Predicate {
   override def foldable: Boolean = child.foldable
   override def nullable: Boolean = false
 
@@ -91,7 +92,7 @@ case class IsNull(child: Expression) extends Predicate with trees.UnaryNode[Expr
   override def toString: String = s"IS NULL $child"
 }
 
-case class IsNotNull(child: Expression) extends Predicate with trees.UnaryNode[Expression] {
+case class IsNotNull(child: Expression) extends UnaryExpression with Predicate {
   override def foldable: Boolean = child.foldable
   override def nullable: Boolean = false
   override def toString: String = s"IS NOT NULL $child"

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

@@ -19,8 +19,7 @@ package org.apache.spark.sql.catalyst.expressions
 
 import org.apache.spark.sql.catalyst.analysis.UnresolvedException
 import org.apache.spark.sql.catalyst.errors.TreeNodeException
-import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.types.{NumericType, DataType}
+import org.apache.spark.sql.types.{DataType, NumericType}
 
 /**
  * The trait of the Window Specification (specified in the OVER clause or WINDOW clause) for

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala

@@ -154,7 +154,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
       expressions: Seq[Expression], inputSchema: Seq[Attribute]): Projection = {
     log.debug(
       s"Creating Projection: $expressions, inputSchema: $inputSchema, codegen:$codegenEnabled")
-    if (codegenEnabled) {
+    if (codegenEnabled && expressions.forall(_.isThreadSafe)) {
       GenerateProjection.generate(expressions, inputSchema)
     } else {
       new InterpretedProjection(expressions, inputSchema)
@@ -166,7 +166,8 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
       inputSchema: Seq[Attribute]): () => MutableProjection = {
     log.debug(
       s"Creating MutableProj: $expressions, inputSchema: $inputSchema, codegen:$codegenEnabled")
-    if(codegenEnabled) {
+    if(codegenEnabled && expressions.forall(_.isThreadSafe)) {
+
       GenerateMutableProjection.generate(expressions, inputSchema)
     } else {
       () => new InterpretedMutableProjection(expressions, inputSchema)
@@ -176,7 +177,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
 
   protected def newPredicate(
       expression: Expression, inputSchema: Seq[Attribute]): (Row) => Boolean = {
-    if (codegenEnabled) {
+    if (codegenEnabled && expression.isThreadSafe) {
       GeneratePredicate.generate(expression, inputSchema)
     } else {
       InterpretedPredicate.create(expression, inputSchema)

sql/core/src/main/scala/org/apache/spark/sql/execution/expressions/MonotonicallyIncreasingID.scala

@@ -49,11 +49,5 @@ private[sql] case class MonotonicallyIncreasingID() extends LeafExpression {
     (TaskContext.get().partitionId().toLong << 33) + currentCount
   }
 
-  /**
-   * This expression is stateful and not thread safe, it can not be reused. Change equals() to
-   * always return `false`, make the generated expressions to not be cached.
-   */
-  override def equals(other: Any): Boolean = {
-    false
-  }
+  override def isThreadSafe: Boolean = false
 }