Support TRUNC (number)

Author: wangyum

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -355,7 +355,6 @@ object FunctionRegistry {
     expression[ParseToDate]("to_date"),
     expression[ToUnixTimestamp]("to_unix_timestamp"),
     expression[ToUTCTimestamp]("to_utc_timestamp"),
-    expression[TruncDate]("trunc"),
     expression[UnixTimestamp]("unix_timestamp"),
     expression[WeekOfYear]("weekofyear"),
     expression[Year]("year"),
@@ -388,6 +387,7 @@ object FunctionRegistry {
     expression[CurrentDatabase]("current_database"),
     expression[CallMethodViaReflection]("reflect"),
     expression[CallMethodViaReflection]("java_method"),
+    expression[Trunc]("trunc"),
 
     // grouping sets
     expression[Cube]("cube"),

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

@@ -1227,86 +1227,6 @@ case class ParseToTimestamp(left: Expression, format: Option[Expression], child:
   override def dataType: DataType = TimestampType
 }
 
-/**
- * Returns date truncated to the unit specified by the format.
- */
-// scalastyle:off line.size.limit
-@ExpressionDescription(
-  usage = "_FUNC_(date, fmt) - Returns `date` with the time portion of the day truncated to the unit specified by the format model `fmt`.",
-  extended = """
-    Examples:
-      > SELECT _FUNC_('2009-02-12', 'MM');
-       2009-02-01
-      > SELECT _FUNC_('2015-10-27', 'YEAR');
-       2015-01-01
-  """)
-// scalastyle:on line.size.limit
-case class TruncDate(date: Expression, format: Expression)
-  extends BinaryExpression with ImplicitCastInputTypes {
-  override def left: Expression = date
-  override def right: Expression = format
-
-  override def inputTypes: Seq[AbstractDataType] = Seq(DateType, StringType)
-  override def dataType: DataType = DateType
-  override def nullable: Boolean = true
-  override def prettyName: String = "trunc"
-
-  private lazy val truncLevel: Int =
-    DateTimeUtils.parseTruncLevel(format.eval().asInstanceOf[UTF8String])
-
-  override def eval(input: InternalRow): Any = {
-    val level = if (format.foldable) {
-      truncLevel
-    } else {
-      DateTimeUtils.parseTruncLevel(format.eval().asInstanceOf[UTF8String])
-    }
-    if (level == -1) {
-      // unknown format
-      null
-    } else {
-      val d = date.eval(input)
-      if (d == null) {
-        null
-      } else {
-        DateTimeUtils.truncDate(d.asInstanceOf[Int], level)
-      }
-    }
-  }
-
-  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
-
-    if (format.foldable) {
-      if (truncLevel == -1) {
-        ev.copy(code = s"""
-          boolean ${ev.isNull} = true;
-          ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};""")
-      } else {
-        val d = date.genCode(ctx)
-        ev.copy(code = s"""
-          ${d.code}
-          boolean ${ev.isNull} = ${d.isNull};
-          ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
-          if (!${ev.isNull}) {
-            ${ev.value} = $dtu.truncDate(${d.value}, $truncLevel);
-          }""")
-      }
-    } else {
-      nullSafeCodeGen(ctx, ev, (dateVal, fmt) => {
-        val form = ctx.freshName("form")
-        s"""
-          int $form = $dtu.parseTruncLevel($fmt);
-          if ($form == -1) {
-            ${ev.isNull} = true;
-          } else {
-            ${ev.value} = $dtu.truncDate($dateVal, $form);
-          }
-        """
-      })
-    }
-  }
-}
-
 /**
  * Returns the number of days from startDate to endDate.
  */

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

@@ -19,7 +19,9 @@ package org.apache.spark.sql.catalyst.expressions
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.codegen._
+import org.apache.spark.sql.catalyst.util.{BigDecimalUtils, DateTimeUtils}
 import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
 
 /**
  * Print the result of an expression to stderr (used for debugging codegen).
@@ -104,3 +106,138 @@ case class CurrentDatabase() extends LeafExpression with Unevaluable {
   override def nullable: Boolean = false
   override def prettyName: String = "current_database"
 }
+
+/**
+ * Returns date truncated to the unit specified by the format or
+ * numeric truncated to scale decimal places.
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+  usage = """
+      _FUNC_(data, fmt) - Returns `data` truncated by the format model `fmt`.
+        If `data` is DateType, returns `data` with the time portion of the day truncated to the unit specified by the format model `fmt`.
+        If `data` is DoubleType, returns `data` truncated to `fmt` decimal places.
+  """,
+  extended = """
+    Examples:
+      > SELECT _FUNC_('2009-02-12', 'MM');
+       2009-02-01
+      > SELECT _FUNC_('2015-10-27', 'YEAR');
+       2015-01-01
+      > SELECT _FUNC_(1234567891.1234567891, 4);
+       1234567891.1234
+      > SELECT _FUNC_(1234567891.1234567891, -4);
+       1234560000
+             """)
+// scalastyle:on line.size.limit
+case class Trunc(data: Expression, format: Expression = Literal(0))
+  extends BinaryExpression with ImplicitCastInputTypes {
+
+  def this(numeric: Expression) = {
+    this(numeric, Literal(0))
+  }
+
+  override def left: Expression = data
+  override def right: Expression = format
+
+  override def dataType: DataType = data.dataType
+
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(TypeCollection(DoubleType, DateType), TypeCollection(StringType, IntegerType))
+
+  override def nullable: Boolean = true
+  override def prettyName: String = "trunc"
+
+  private lazy val truncFormat: Int = dataType match {
+    case doubleType: DoubleType =>
+      format.eval().asInstanceOf[Int]
+    case dateType: DateType =>
+      DateTimeUtils.parseTruncLevel(format.eval().asInstanceOf[UTF8String])
+  }
+
+  override def eval(input: InternalRow): Any = {
+    val d = data.eval(input)
+    val form = format.eval()
+    if (null == d || null == form) {
+      null
+    } else {
+      dataType match {
+        case doubleType: DoubleType =>
+          val scale = if (format.foldable) {
+            truncFormat
+          } else {
+            format.eval().asInstanceOf[Int]
+          }
+          BigDecimalUtils.trunc(d.asInstanceOf[Double], scale).doubleValue()
+        case dateType: DateType =>
+          val level = if (format.foldable) {
+            truncFormat
+          } else {
+            DateTimeUtils.parseTruncLevel(format.eval().asInstanceOf[UTF8String])
+          }
+          if (level == -1) {
+            // unknown format
+            null
+          } else {
+            DateTimeUtils.truncDate(d.asInstanceOf[Int], level)
+          }
+      }
+    }
+
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+
+    dataType match {
+      case doubleType: DoubleType =>
+        val bdu = BigDecimalUtils.getClass.getName.stripSuffix("$")
+
+        if (format.foldable) {
+          val d = data.genCode(ctx)
+          ev.copy(code = s"""
+            ${d.code}
+            boolean ${ev.isNull} = ${d.isNull};
+            ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+            if (!${ev.isNull}) {
+              ${ev.value} = $bdu.trunc(${d.value}, $truncFormat).doubleValue();
+            }""")
+        } else {
+          nullSafeCodeGen(ctx, ev, (doubleVal, fmt) => {
+            s"${ev.value} = $bdu.trunc($doubleVal, $fmt).doubleValue();"
+          })
+        }
+      case dateType: DateType =>
+        val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
+
+        if (format.foldable) {
+          if (truncFormat == -1) {
+            ev.copy(code = s"""
+              boolean ${ev.isNull} = true;
+              ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};""")
+          } else {
+            val d = data.genCode(ctx)
+            ev.copy(code = s"""
+              ${d.code}
+              boolean ${ev.isNull} = ${d.isNull};
+              ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+              if (!${ev.isNull}) {
+                ${ev.value} = $dtu.truncDate(${d.value}, $truncFormat);
+              }""")
+          }
+        } else {
+          nullSafeCodeGen(ctx, ev, (dateVal, fmt) => {
+            val form = ctx.freshName("form")
+            s"""
+              int $form = $dtu.parseTruncLevel($fmt);
+              if ($form == -1) {
+                ${ev.isNull} = true;
+              } else {
+                ${ev.value} = $dtu.truncDate($dateVal, $form);
+              }
+            """
+          })
+        }
+    }
+
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/BigDecimalUtils.scala

@@ -0,0 +1,56 @@
+/*
+ * 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.util
+
+import java.math.{BigDecimal => JBigDecimal}
+
+/**
+ * Helper functions for BigDecimal.
+ */
+object BigDecimalUtils {
+
+  /**
+   * Returns double type input truncated to scale decimal places.
+   */
+  def trunc(input: Double, scale: Int): JBigDecimal = {
+    trunc(JBigDecimal.valueOf(input), scale)
+  }
+
+  /**
+   * Returns BigDecimal type input truncated to scale decimal places.
+   */
+  def trunc(input: JBigDecimal, scale: Int): JBigDecimal = {
+
+    val pow = if (scale >= 0) {
+      JBigDecimal.valueOf(Math.pow(10, scale))
+    } else {
+      JBigDecimal.valueOf(Math.pow(10, Math.abs(scale)))
+    }
+
+    val out = if (scale > 0) {
+      val longValue = input.multiply(pow).longValue()
+      JBigDecimal.valueOf(longValue).divide(pow)
+    } else if (scale == 0) {
+      JBigDecimal.valueOf(input.longValue())
+    } else {
+      val longValue = input.divide(pow).longValue()
+      JBigDecimal.valueOf(longValue).multiply(pow)
+    }
+    out
+  }
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala

@@ -513,27 +513,6 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
       NextDay(Literal(Date.valueOf("2015-07-23")), Literal.create(null, StringType)), null)
   }
 
-  test("function trunc") {
-    def testTrunc(input: Date, fmt: String, expected: Date): Unit = {
-      checkEvaluation(TruncDate(Literal.create(input, DateType), Literal.create(fmt, StringType)),
-        expected)
-      checkEvaluation(
-        TruncDate(Literal.create(input, DateType), NonFoldableLiteral.create(fmt, StringType)),
-        expected)
-    }
-    val date = Date.valueOf("2015-07-22")
-    Seq("yyyy", "YYYY", "year", "YEAR", "yy", "YY").foreach { fmt =>
-      testTrunc(date, fmt, Date.valueOf("2015-01-01"))
-    }
-    Seq("month", "MONTH", "mon", "MON", "mm", "MM").foreach { fmt =>
-      testTrunc(date, fmt, Date.valueOf("2015-07-01"))
-    }
-    testTrunc(date, "DD", null)
-    testTrunc(date, null, null)
-    testTrunc(null, "MON", null)
-    testTrunc(null, null, null)
-  }
-
   test("from_unixtime") {
     val sdf1 = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss", Locale.US)
     val fmt2 = "yyyy-MM-dd HH:mm:ss.SSS"

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscExpressionsSuite.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
+import java.sql.Date
+
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.types._
 
@@ -39,4 +41,50 @@ class MiscExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(AssertTrue(Cast(Literal(1), BooleanType)), null)
   }
 
+  test("trunc") {
+
+    // numeric
+    def testTruncNumber(input: Double, fmt: Int, expected: Double): Unit = {
+      checkEvaluation(Trunc(Literal.create(input, DoubleType),
+        Literal.create(fmt, IntegerType)),
+        expected)
+      checkEvaluation(Trunc(Literal.create(input, DoubleType),
+        NonFoldableLiteral.create(fmt, IntegerType)),
+        expected)
+    }
+
+    testTruncNumber(1234567891.1234567891, 4, 1234567891.1234)
+    testTruncNumber(1234567891.1234567891, -4, 1234560000)
+    testTruncNumber(1234567891.1234567891, 0, 1234567891)
+
+    checkEvaluation(Trunc(Literal.create(1D, DoubleType),
+      NonFoldableLiteral.create(null, IntegerType)),
+      null)
+    checkEvaluation(Trunc(Literal.create(null, DoubleType),
+      NonFoldableLiteral.create(1, IntegerType)),
+      null)
+    checkEvaluation(Trunc(Literal.create(null, DoubleType),
+      NonFoldableLiteral.create(null, IntegerType)),
+      null)
+
+    // date
+    def testTruncDate(input: Date, fmt: String, expected: Date): Unit = {
+      checkEvaluation(Trunc(Literal.create(input, DateType), Literal.create(fmt, StringType)),
+        expected)
+      checkEvaluation(
+        Trunc(Literal.create(input, DateType), NonFoldableLiteral.create(fmt, StringType)),
+        expected)
+    }
+    val date = Date.valueOf("2015-07-22")
+    Seq("yyyy", "YYYY", "year", "YEAR", "yy", "YY").foreach { fmt =>
+      testTruncDate(date, fmt, Date.valueOf("2015-01-01"))
+    }
+    Seq("month", "MONTH", "mon", "MON", "mm", "MM").foreach { fmt =>
+      testTruncDate(date, fmt, Date.valueOf("2015-07-01"))
+    }
+    testTruncDate(date, "DD", null)
+    testTruncDate(date, null, null)
+    testTruncDate(null, "MON", null)
+    testTruncDate(null, null, null)
+  }
 }