[SPARK-20754][SQL] Support TRUNC (number)

python/pyspark/sql/functions.py

@@ -1028,20 +1028,29 @@ def to_timestamp(col, format=None):
 
 
 @since(1.5)
-def trunc(date, format):
+def trunc(data, truncParam):
     """
-    Returns date truncated to the unit specified by the format.
+    Returns date truncated to the unit specified by the truncParam or
+    numeric truncated by specified decimal places.
 
-    :param format: 'year', 'YYYY', 'yy' or 'month', 'mon', 'mm'
+    :param truncParam: 'year', 'YYYY', 'yy' or 'month', 'mon', 'mm' for date
+        and any int for numeric.
 
     >>> df = spark.createDataFrame([('1997-02-28',)], ['d'])
     >>> df.select(trunc(df.d, 'year').alias('year')).collect()
     [Row(year=datetime.date(1997, 1, 1))]
     >>> df.select(trunc(df.d, 'mon').alias('month')).collect()
     [Row(month=datetime.date(1997, 2, 1))]
+    >>> df = spark.createDataFrame([(1234567891.1234567891,)], ['d'])
+    >>> df.select(trunc(df.d, 4).alias('positive')).collect()
+    [Row(positive=1234567891.1234)]
+    >>> df.select(trunc(df.d, -4).alias('negative')).collect()
+    [Row(negative=1234560000.0)]
+    >>> df.select(trunc(df.d, 0).alias('zero')).collect()
+    [Row(zero=1234567891.0)]
     """
     sc = SparkContext._active_spark_context
-    return Column(sc._jvm.functions.trunc(_to_java_column(date), format))
+    return Column(sc._jvm.functions.trunc(_to_java_column(data), truncParam))
 
 
 @since(1.5)

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

@@ -389,7 +389,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[DayOfWeek]("dayofweek"),
     expression[WeekOfYear]("weekofyear"),
@@ -424,6 +423,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

@@ -1267,86 +1267,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

@@ -21,6 +21,7 @@ import java.util.UUID
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.codegen._
+import org.apache.spark.sql.catalyst.util.{DateTimeUtils, MathUtils}
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
@@ -132,3 +133,174 @@ case class Uuid() extends LeafExpression {
       s"UTF8String.fromString(java.util.UUID.randomUUID().toString());", isNull = "false")
   }
 }
+
+/**
+ * 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[, trunc_param]) - Returns `data` truncated by the format model `trunc_param`.
+        If `data` is date/timestamp/string type, returns `data` with the time portion of the day truncated to the unit specified by the format model `trunc_param`. If `trunc_param` is omitted, then the default `trunc_param` is 'MM'.
+        If `data` is decimal/double type, returns `data` truncated to `trunc_param` decimal places. If `trunc_param` is omitted, then the default `trunc_param` is 0.
+  """,
+  extended = """
+    Examples:
+      > SELECT _FUNC_('2009-02-12', 'MM');
+       2009-02-01.
+      > SELECT _FUNC_('2015-10-27', 'YEAR');
+       2015-01-01
+      > SELECT _FUNC_('1989-03-13');
+       1989-03-01
+      > SELECT _FUNC_(1234567891.1234567891, 4);
+       1234567891.1234
+      > SELECT _FUNC_(1234567891.1234567891, -4);
+       1234560000
+      > SELECT _FUNC_(1234567891.1234567891);
+       1234567891
+  """)
+// scalastyle:on line.size.limit
+case class Trunc(data: Expression, truncExpr: Expression)
+  extends BinaryExpression with ExpectsInputTypes {
+
+  def this(data: Expression) = {
+    this(data, Literal(
+      if (data.dataType.isInstanceOf[DateType] ||
+        data.dataType.isInstanceOf[TimestampType] ||
+        data.dataType.isInstanceOf[StringType]) {
+      "MM"
+     } else {
+        0
+      })
+    )
+  }
+
+  override def left: Expression = data
+  override def right: Expression = truncExpr
+
+  private val isTruncNumber = truncExpr.dataType.isInstanceOf[IntegerType]
+  private val isTruncDate = truncExpr.dataType.isInstanceOf[StringType]
+
+  override def dataType: DataType = if (isTruncDate) DateType else data.dataType
+
+  override def inputTypes: Seq[AbstractDataType] = data.dataType match {
+    case NullType =>
+      Seq(dataType, TypeCollection(StringType, IntegerType))
+    case DateType | TimestampType | StringType =>
+      Seq(TypeCollection(DateType, TimestampType, StringType), StringType)
+    case DoubleType | DecimalType.Fixed(_, _) =>
+      Seq(TypeCollection(DoubleType, DecimalType), IntegerType)
+    case _ =>
+      Seq(TypeCollection(DateType, StringType, TimestampType, DoubleType, DecimalType),
+        TypeCollection(StringType, IntegerType))
+  }
+
+  override def nullable: Boolean = true
+
+  override def prettyName: String = "trunc"
+
+
+  private lazy val truncFormat: Int = if (isTruncNumber) {
+    truncExpr.eval().asInstanceOf[Int]
+  } else if (isTruncDate) {
+    DateTimeUtils.parseTruncLevel(truncExpr.eval().asInstanceOf[UTF8String])
+  } else {
+    0
+  }
+
+  override def eval(input: InternalRow): Any = {
+    val d = data.eval(input)
+    val truncParam = truncExpr.eval()
+    if (null == d || null == truncParam) {
+      null
+    } else {
+      if (isTruncNumber) {
+        val scale = if (truncExpr.foldable) truncFormat else truncExpr.eval().asInstanceOf[Int]
+        data.dataType match {
+          case DoubleType => MathUtils.trunc(d.asInstanceOf[Double], scale)
+          case DecimalType.Fixed(_, _) =>
+            MathUtils.trunc(d.asInstanceOf[Decimal].toJavaBigDecimal, scale)
+        }
+      } else if (isTruncDate) {
+        val level = if (truncExpr.foldable) {
+          truncFormat
+        } else {
+          DateTimeUtils.parseTruncLevel(truncExpr.eval().asInstanceOf[UTF8String])
+        }
+        if (level == -1) {
+          // unknown format
+          null
+        } else {
+          data.dataType match {
+            case DateType => DateTimeUtils.truncDate(d.asInstanceOf[Int], level)
+            case TimestampType =>
+              val ts = DateTimeUtils.timestampToString(d.asInstanceOf[Long])
+              val dt = DateTimeUtils.stringToDate(UTF8String.fromString(ts))
+              if (dt.isDefined) DateTimeUtils.truncDate(dt.get, level) else null
+            case StringType =>
+              val dt = DateTimeUtils.stringToDate(d.asInstanceOf[UTF8String])
+              if (dt.isDefined) DateTimeUtils.truncDate(dt.get, level) else null
+          }
+        }
+      } else {
+        null
+      }
+    }
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+
+    if (isTruncNumber) {
+      val bdu = MathUtils.getClass.getName.stripSuffix("$")
+
+      if (truncExpr.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);
+          }""")
+      } else {
+        nullSafeCodeGen(ctx, ev, (doubleVal, truncParam) =>
+          s"${ev.value} = $bdu.trunc($doubleVal, $truncParam);")
+      }
+    } else if (isTruncDate) {
+      val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
+
+      if (truncExpr.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 truncParam = ctx.freshName("truncParam")
+          s"""
+            int $truncParam = $dtu.parseTruncLevel($fmt);
+            if ($truncParam == -1) {
+              ${ev.isNull} = true;
+            } else {
+              ${ev.value} = $dtu.truncDate($dateVal, $truncParam);
+            }
+          """
+        })
+      }
+    } else {
+      nullSafeCodeGen(ctx, ev, (dataVal, fmt) => s"${ev.isNull} = true;")
+    }
+  }
+}

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

@@ -0,0 +1,57 @@
+/*
+ * 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 MathUtils {
+
+  /**
+   * Returns double type input truncated to scale decimal places.
+   */
+  def trunc(input: Double, scale: Int): Double = {
+    trunc(JBigDecimal.valueOf(input), scale).doubleValue()
+  }
+
+  /**
+   * Returns BigDecimal type input truncated to scale decimal places.
+   */
+  def trunc(input: JBigDecimal, scale: Int): JBigDecimal = {
+    // Copy from (https://github.com/apache/hive/blob/release-2.3.0-rc0
+    // /ql/src/java/org/apache/hadoop/hive/ql/udf/generic/GenericUDFTrunc.java#L471-L487)
+    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
+  }
+}