improved codegen

apache · Jun 24, 2015 · 638596f · 638596f
1 parent 4d8049b
commit 638596f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 25 deletions.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetime.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetime.scala
@@ -79,18 +79,6 @@ case class DateFormatClass(left: Expression, right: Expression) extends BinaryEx
     val eval1 = left.gen(ctx)
     val eval2 = right.gen(ctx)
 
-    val calc = left.dataType match {
-      case TimestampType =>
-        s""""$utf8.fromString(sdf.format(new java.sql.Date(${eval1.primitive} / 10000)));"""
-      case DateType =>
-        s"""$utf8.fromString(
-          sdf.format($dtUtils.toJavaDate(${eval1.primitive})));"""
-      case StringType =>
-        s"""
-           $utf8.fromString(sdf.format(new java.sql.Date($dtUtils.stringToTime(${eval1.primitive}.toString()).getTime())));
-         """
-    }
-
     s"""
       ${eval1.code}
       boolean ${ev.isNull} = ${eval1.isNull};
@@ -99,7 +87,14 @@ case class DateFormatClass(left: Expression, right: Expression) extends BinaryEx
         ${eval2.code}
         if (!${eval2.isNull}) {
           $sdf sdf = new $sdf(${eval2.primitive}.toString());
-          ${ev.primitive} = $calc
+          Object o = ${eval1.primitive};
+          if (o instanceof ${ctx.boxedType(TimestampType)}) {
+            ${ev.primitive} = $utf8.fromString(sdf.format(new java.sql.Date(Long.parseLong(o.toString()) / 10000)));
+          } else if (o instanceof Integer) {
+            ${ev.primitive} = $utf8.fromString(sdf.format($dtUtils.toJavaDate(Integer.parseInt(o.toString()))));
+          } else {
+            ${ev.primitive} = $utf8.fromString(sdf.format(new java.sql.Date($dtUtils.stringToTime(o.toString()).getTime())));
+          }
         } else {
           ${ev.isNull} = true;
         }
@@ -108,12 +103,10 @@ case class DateFormatClass(left: Expression, right: Expression) extends BinaryEx
   }
 }
 
-case class Year(child: Expression) extends UnaryExpression with ExpectsInputTypes {
+case class Year(child: Expression) extends UnaryExpression {
 
   override def dataType: DataType = IntegerType
 
-  override def expectedChildTypes: Seq[DataType] = Seq(DateType, StringType, TimestampType)
-
   override def foldable: Boolean = child.foldable
 
   override def nullable: Boolean = true
@@ -138,12 +131,10 @@ case class Year(child: Expression) extends UnaryExpression with ExpectsInputType
 
 }
 
-case class Month(child: Expression) extends UnaryExpression with ExpectsInputTypes {
+case class Month(child: Expression) extends UnaryExpression {
 
   override def dataType: DataType = IntegerType
 
-  override def expectedChildTypes: Seq[DataType] = Seq(DateType, StringType, TimestampType)
-
   override def foldable: Boolean = child.foldable
 
   override def nullable: Boolean = true
@@ -167,12 +158,10 @@ case class Month(child: Expression) extends UnaryExpression with ExpectsInputTyp
     }
 }
 
-case class Day(child: Expression) extends UnaryExpression with ExpectsInputTypes {
+case class Day(child: Expression) extends UnaryExpression {
 
   override def dataType: DataType = IntegerType
 
-  override def expectedChildTypes: Seq[DataType] = Seq(DateType, StringType, TimestampType)
-
   override def foldable: Boolean = child.foldable
 
   override def nullable: Boolean = true

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
@@ -172,15 +172,17 @@ class DataFrameFunctionsSuite extends QueryTest {
     val d = new Date(sdf.parse("2015/04/08 13:10:15").getTime)
     val ts = new Timestamp(sdf.parse("2013/04/08 13:10:15").getTime)
 
-    val df = Seq((d, d.toString, ts)).toDF("a", "b", "c")
+    val df = Seq((d, d.toString, ts), (ts, d, d.toString), (d.toString, ts, d)).toDF("a", "b", "c")
 
     checkAnswer(
       df.select(dateFormat("a", "y"), dateFormat("b", "y"), dateFormat("c", "y")),
-      Row("2015", "2015", "2013"))
+      Row("2015", "2015", "2013") :: Row("2013", "2015", "2015") ::
+        Row("2015", "2013", "2015") :: Nil)
 
     checkAnswer(
       df.selectExpr("dateFormat(a, 'y')", "dateFormat(b, 'y')", "dateFormat(c, 'y')"),
-      Row("2015", "2015", "2013"))
+      Row("2015", "2015", "2013") :: Row("2013", "2015", "2015") ::
+        Row("2015", "2013", "2015") :: Nil)
   }
 
   test("year") {