[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,214 @@ 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;
+            int ${ev.value} = ${ctx.defaultValue(DateType)};""")
+        } else {
+          val d = data.genCode(ctx)
+          val dt = ctx.freshName("dt")
+          val pre = s"""
+            ${d.code}
+            boolean ${ev.isNull} = ${d.isNull};
+            int ${ev.value} = ${ctx.defaultValue(DateType)};"""
+          data.dataType match {
+            case DateType =>
+              ev.copy(code = pre + s"""
+                if (!${ev.isNull}) {
+                  ${ev.value} = $dtu.truncDate(${d.value}, $truncFormat);
+                }""")
+            case TimestampType =>
+              val ts = ctx.freshName("ts")
+              ev.copy(code = pre + s"""
+                String $ts = $dtu.timestampToString(${d.value});
+                scala.Option<SQLDate> $dt = $dtu.stringToDate(UTF8String.fromString($ts));
+                if (!${ev.isNull}) {
+                  ${ev.value} = $dtu.truncDate((Integer)dt.get(), $truncFormat);
+                }""")
+            case StringType =>
+              ev.copy(code = pre + s"""
+                scala.Option<SQLDate> $dt = $dtu.stringToDate(${d.value});
+                if (!${ev.isNull} && $dt.isDefined()) {
+                  ${ev.value} = $dtu.truncDate((Integer)$dt.get(), $truncFormat);
+                }""")
+          }
+        }
+      } else {
+        nullSafeCodeGen(ctx, ev, (dateVal, fmt) => {
+          val truncParam = ctx.freshName("truncParam")
+          val dt = ctx.freshName("dt")
+          val pre = s"int $truncParam = $dtu.parseTruncLevel($fmt);"
+          data.dataType match {
+            case DateType =>
+              pre + s"""
+                if ($truncParam == -1) {
+                  ${ev.isNull} = true;
+                } else {
+                  ${ev.value} = $dtu.truncDate($dateVal, $truncParam);
+                }"""
+            case TimestampType =>
+              val ts = ctx.freshName("ts")
+              pre + s"""
+                String $ts = $dtu.timestampToString($dateVal);
+                scala.Option<SQLDate> $dt = $dtu.stringToDate(UTF8String.fromString($ts));
+                if ($truncParam == -1 || $dt.isEmpty()) {
+                  ${ev.isNull} = true;
+                } else {
+                  ${ev.value} = $dtu.truncDate((Integer)$dt.get(), $truncParam);
+                }"""
+            case StringType =>
+              pre + s"""
+                scala.Option<SQLDate> $dt = $dtu.stringToDate($dateVal);
+                ${ev.value} = ${ctx.defaultValue(DateType)};
+                if ($truncParam == -1 || $dt.isEmpty()) {
+                  ${ev.isNull} = true;
+                } else {
+                  ${ev.value} = $dtu.truncDate((Integer)$dt.get(), $truncParam);
+                }"""
+          }
+        })
+      }
+    } else {
+      nullSafeCodeGen(ctx, ev, (dataVal, fmt) => s"${ev.isNull} = true;")
+    }
+  }
+}