rely on cast mechanism of Spark. Simplified implementation

Author: tarekbecker

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

@@ -26,115 +26,48 @@ import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
-abstract class DateFormatExpression extends Expression { self: Product =>
+abstract class DateFormatExpression extends UnaryExpression with ExpectsInputTypes {
+  self: Product =>
 
   protected val format: String
 
-  protected val caller: String
-  
-  protected val date: Expression
-
-  override def foldable: Boolean = date.foldable
-
-  override def nullable: Boolean = true
-
-  override def children: Seq[Expression] = Seq(date)
+  override def expectedChildTypes: Seq[DataType] = Seq(TimestampType)
 
   override def eval(input: InternalRow): Any = {
-    val valueLeft = date.eval(input)
+    val valueLeft = child.eval(input)
     if (valueLeft == null) {
       null
     } else {
       if (format == null) {
         null
       } else {
         val sdf = new SimpleDateFormat(format)
-        date.dataType match {
-          case TimestampType =>
-            UTF8String.fromString(sdf.format(new Date(valueLeft.asInstanceOf[Long] / 10000)))
-          case DateType =>
-            UTF8String.fromString(sdf.format(DateTimeUtils.toJavaDate(valueLeft.asInstanceOf[Int])))
-          case StringType =>
-            UTF8String.fromString(
-              sdf.format(DateTimeUtils.stringToTime(valueLeft.toString)))
-        }
+        UTF8String.fromString(sdf.format(new Date(valueLeft.asInstanceOf[Long] / 10000)))
       }
     }
   }
 
-  override def checkInputDataTypes(): TypeCheckResult =
-    date.dataType match {
-      case null => TypeCheckResult.TypeCheckSuccess
-      case _: DateType => TypeCheckResult.TypeCheckSuccess
-      case _: TimestampType => TypeCheckResult.TypeCheckSuccess
-      case _: StringType => TypeCheckResult.TypeCheckSuccess
-      case _ =>
-        TypeCheckResult.TypeCheckFailure(s"$caller accepts date types as argument, " +
-          s" not ${date.dataType}")
-    }
-
-
-  /**
-   * Called by date format expressions to generate a code block that returns the result
-   *
-   * As an example, the following parse the result to int
-   * {{{
-   *   defineCodeGen(ctx, ev, c => s"Integer.parseInt($c.toString())")
-   * }}}
-   *
-   * @param f function that accepts a variable name and returns Java code to parse an
-   *          [[UTF8String]] to the expected output type
-   */
-
-  protected def defineCodeGen(
+  override protected def defineCodeGen(
       ctx: CodeGenContext,
       ev: GeneratedExpressionCode,
       f: String => String): String = {
 
     val sdf = classOf[SimpleDateFormat].getName
-    val dtUtils = "org.apache.spark.sql.catalyst.util.DateTimeUtils"
-
-    val eval1 = date.gen(ctx)
-
-    val parseInput = date.dataType match {
-      case StringType => s"new java.sql.Date($dtUtils.stringToTime(${eval1.primitive}.toString()).getTime())"
-      case TimestampType => s"new java.sql.Date(${eval1.primitive} / 10000)"
-      case DateType => s"$dtUtils.toJavaDate(${eval1.primitive})"
-    }
-
-    s"""
-      ${eval1.code}
-      boolean ${ev.isNull} = ${eval1.isNull};
-      ${ctx.javaType(dataType)} ${ev.primitive} = ${ctx.defaultValue(dataType)};
-      if (!${ev.isNull}) {
-        $sdf sdf = new $sdf("$format");
-        ${ctx.stringType} s = ${ctx.stringType}.fromString(sdf.format($parseInput));
-        ${ev.primitive} = ${f("s")};
-      } else {
-        ${ev.isNull} = true;
-      }
-    """
+    super.defineCodeGen(ctx, ev, (x) => {
+      f(s"""${ctx.stringType}.fromString((new $sdf("$format")).format(new java.sql.Date($x / 10000)))""")
+    })
   }
 
 }
 
-case class DateFormatClass(left: Expression, right: Expression) extends BinaryExpression {
+case class DateFormatClass(left: Expression, right: Expression) extends BinaryExpression
+    with ExpectsInputTypes {
 
   override def dataType: DataType = StringType
 
-  override def checkInputDataTypes(): TypeCheckResult =
-    (left.dataType, right.dataType) match {
-      case (null, _) => TypeCheckResult.TypeCheckSuccess
-      case (_, null) => TypeCheckResult.TypeCheckSuccess
-      case (_: DateType, _: StringType) => TypeCheckResult.TypeCheckSuccess
-      case (_: TimestampType, _: StringType) => TypeCheckResult.TypeCheckSuccess
-      case (_: StringType, _: StringType) => TypeCheckResult.TypeCheckSuccess
-      case _ =>
-        TypeCheckResult.TypeCheckFailure(s"DateFormat accepts date types as first argument, " +
-          s"and string types as second, not ${left.dataType} and ${right.dataType}")
-    }
-
   override def toString: String = s"DateFormat($left, $right)"
+  
+  override def expectedChildTypes: Seq[DataType] = Seq(TimestampType, StringType)
 
   override def eval(input: InternalRow): Any = {
     val valueLeft = left.eval(input)
@@ -146,53 +79,23 @@ case class DateFormatClass(left: Expression, right: Expression) extends BinaryEx
         null
       } else {
         val sdf = new SimpleDateFormat(valueRight.toString)
-        left.dataType match {
-          case TimestampType =>
-            UTF8String.fromString(sdf.format(new Date(valueLeft.asInstanceOf[Long] / 10000)))
-          case DateType =>
-            UTF8String.fromString(sdf.format(DateTimeUtils.toJavaDate(valueLeft.asInstanceOf[Int])))
-          case StringType =>
-            UTF8String.fromString(
-              sdf.format(DateTimeUtils.stringToTime(valueLeft.toString)))
-        }
+        UTF8String.fromString(sdf.format(new Date(valueLeft.asInstanceOf[Long] / 10000)))
       }
     }   
   }
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val sdf = classOf[SimpleDateFormat].getName
-    val dtUtils = "org.apache.spark.sql.catalyst.util.DateTimeUtils"
-
-    val eval1 = left.gen(ctx)
-    val eval2 = right.gen(ctx)
-
-    val parseInput = left.dataType match {
-      case StringType => s"new java.sql.Date($dtUtils.stringToTime(${eval1.primitive}.toString()).getTime())"
-      case TimestampType => s"new java.sql.Date(${eval1.primitive} / 10000)"
-      case DateType => s"$dtUtils.toJavaDate(${eval1.primitive})"
-    }
-
-    s"""
-      ${eval1.code}
-      ${eval2.code}
-      boolean ${ev.isNull} = ${eval1.isNull} || ${eval2.isNull};
-      ${ctx.javaType(dataType)} ${ev.primitive} = ${ctx.defaultValue(dataType)};
-      if (!${ev.isNull}) {
-        $sdf sdf = new $sdf(${eval2.primitive}.toString());
-        ${ev.primitive} = ${ctx.stringType}.fromString(sdf.format($parseInput));
-      } else {
-        ${ev.isNull} = true;
-      }
-    """
+    defineCodeGen(ctx, ev, (x, y) => {
+      s"""${ctx.stringType}.fromString((new $sdf($y.toString())).format(new java.sql.Date($x / 10000)))"""
+    })
   }
 }
 
-case class Year(date: Expression) extends DateFormatExpression {
+case class Year(child: Expression) extends DateFormatExpression with ExpectsInputTypes {
 
   override protected val format: String = "y"
 
-  override protected val caller: String = "Year"
-
   override def dataType: DataType = IntegerType
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
@@ -205,14 +108,14 @@ case class Year(date: Expression) extends DateFormatExpression {
       case s: UTF8String => s.toString.toInt
     }
   }
+
+  override def toString: String = s"Year($child)"
 }
 
-case class Quarter(date: Expression) extends DateFormatExpression {
+case class Quarter(child: Expression) extends DateFormatExpression {
 
   override protected val format: String = "M"
 
-  override protected val caller: String = "Quarter"
-
   override def dataType: DataType = IntegerType
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
@@ -225,18 +128,16 @@ case class Quarter(date: Expression) extends DateFormatExpression {
       case s: UTF8String => (s.toString.toInt - 1) / 3 + 1
     }
   }
+
+  override def toString: String = s"Quarter($child)"
 }
 
-case class Month(date: Expression) extends DateFormatExpression {
+case class Month(child: Expression) extends DateFormatExpression with ExpectsInputTypes {
 
   override protected val format: String = "M"
 
-  override protected val caller: String = "Month"
-
   override def dataType: DataType = IntegerType
 
-  override def nullable: Boolean = true
-
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     defineCodeGen(ctx, ev, c => s"Integer.parseInt($c.toString())")
   }
@@ -247,14 +148,14 @@ case class Month(date: Expression) extends DateFormatExpression {
       case s: UTF8String => s.toString.toInt
     }
   }
+
+  override def toString: String = s"Month($child)"
 }
 
-case class Day(date: Expression) extends DateFormatExpression {
+case class Day(child: Expression) extends DateFormatExpression with ExpectsInputTypes {
 
   override protected val format: String = "d"
 
-  override protected val caller: String = "Day"
-
   override def dataType: DataType = IntegerType
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
@@ -267,14 +168,14 @@ case class Day(date: Expression) extends DateFormatExpression {
       case s: UTF8String => s.toString.toInt
     }
   }
+
+  override def toString: String = s"Day($child)"
 }
 
-case class Hour(date: Expression) extends DateFormatExpression {
+case class Hour(child: Expression) extends DateFormatExpression with ExpectsInputTypes {
 
   override protected val format: String = "H"
 
-  override protected val caller: String = "Hour"
-
   override def dataType: DataType = IntegerType
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
@@ -287,14 +188,14 @@ case class Hour(date: Expression) extends DateFormatExpression {
       case s: UTF8String => s.toString.toInt
     }
   }
+
+  override def toString: String = s"Hour($child)"
 }
 
-case class Minute(date: Expression) extends DateFormatExpression {
+case class Minute(child: Expression) extends DateFormatExpression with ExpectsInputTypes {
 
   override protected val format: String = "m"
 
-  override protected val caller: String = "Minute"
-
   override def dataType: DataType = IntegerType
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
@@ -307,14 +208,14 @@ case class Minute(date: Expression) extends DateFormatExpression {
       case s: UTF8String => s.toString.toInt
     }
   }
+
+  override def toString: String = s"Minute($child)"
 }
 
-case class Second(date: Expression) extends DateFormatExpression {
+case class Second(child: Expression) extends DateFormatExpression with ExpectsInputTypes {
 
   override protected val format: String = "s"
 
-  override protected val caller: String = "Second"
-
   override def dataType: DataType = IntegerType
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
@@ -327,14 +228,14 @@ case class Second(date: Expression) extends DateFormatExpression {
       case s: UTF8String => s.toString.toInt
     }
   }
+
+  override def toString: String = s"Second($child)"
 }
 
-case class WeekOfYear(date: Expression) extends DateFormatExpression {
+case class WeekOfYear(child: Expression) extends DateFormatExpression with ExpectsInputTypes {
 
   override protected val format: String = "w"
 
-  override protected val caller: String = "WeekOfYear"
-
   override def dataType: DataType = IntegerType
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
@@ -347,4 +248,6 @@ case class WeekOfYear(date: Expression) extends DateFormatExpression {
       case s: UTF8String => s.toString.toInt
     }
   }
+
+  override def toString: String = s"WeekOfYear($child)"
 }