[SPARK-23906][SQL] Add built-in UDF TRUNCATE(number)

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

@@ -262,6 +262,7 @@ object FunctionRegistry {
     expression[Tan]("tan"),
     expression[Cot]("cot"),
     expression[Tanh]("tanh"),
+    expression[Truncate]("truncate"),
 
     expression[Add]("+"),
     expression[Subtract]("-"),

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

@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{TypeCheckFailure, TypeCheckSuccess}
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
-import org.apache.spark.sql.catalyst.util.NumberConverter
+import org.apache.spark.sql.catalyst.util.{MathUtils, NumberConverter}
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
@@ -1245,3 +1245,65 @@ case class BRound(child: Expression, scale: Expression)
     with Serializable with ImplicitCastInputTypes {
   def this(child: Expression) = this(child, Literal(0))
 }
+
+/**
+ * The number truncated to scale decimal places.
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+  usage = "_FUNC_(number, scale) - Returns number truncated to scale decimal places. " +
+    "If scale is omitted, then number is truncated to 0 places. " +
+    "scale can be negative to truncate (make zero) scale digits left of the decimal point.",
+  examples = """
+    Examples:
+      > 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 Truncate(number: Expression, scale: Expression)
+  extends BinaryExpression with ImplicitCastInputTypes {
+
+  override def left: Expression = number
+  override def right: Expression = scale
+
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(TypeCollection(DoubleType, DecimalType), IntegerType)
+
+  override def dataType: DataType = left.dataType
+
+  private lazy val foldableTruncScale: Int = scale.eval().asInstanceOf[Int]
+
+  protected override def nullSafeEval(input1: Any, input2: Any): Any = {
+    val truncScale = if (scale.foldable) {
+      foldableTruncScale
+    } else {
+      scale.eval().asInstanceOf[Int]
+    }
+    number.dataType match {
+      case DoubleType => MathUtils.trunc(input1.asInstanceOf[Double], truncScale)
+      case DecimalType.Fixed(_, _) =>
+        MathUtils.trunc(input1.asInstanceOf[Decimal].toJavaBigDecimal, truncScale)
+    }
+  }
+
+  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val mu = MathUtils.getClass.getName.stripSuffix("$")
+    if (scale.foldable) {
+      val d = number.genCode(ctx)
+      ev.copy(code = code"""
+        ${d.code}
+        boolean ${ev.isNull} = ${d.isNull};
+        ${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+        if (!${ev.isNull}) {
+          ${ev.value} = $mu.trunc(${d.value}, $foldableTruncScale);
+        }""")
+    } else {
+      nullSafeCodeGen(ctx, ev, (doubleVal, truncParam) =>
+        s"${ev.value} = $mu.trunc($doubleVal, $truncParam);")
+    }
+  }
+}

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

@@ -0,0 +1,52 @@
+/*
+ * 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}
+
+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)))
+    }
+
+    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)
+    }
+  }
+}

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

@@ -644,4 +644,31 @@ class MathExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(BRound(-0.35, 1), -0.4)
     checkEvaluation(BRound(-35, -1), -40)
   }
+
+  test("Truncate number") {
+    def testTruncate(input: Double, fmt: Int, expected: Double): Unit = {
+      checkEvaluation(Truncate(Literal.create(input, DoubleType),
+        Literal.create(fmt, IntegerType)),
+        expected)
+      checkEvaluation(Truncate(Literal.create(input, DoubleType),
+        NonFoldableLiteral.create(fmt, IntegerType)),
+        expected)
+    }
+
+    testTruncate(1234567891.1234567891, 4, 1234567891.1234)
+    testTruncate(1234567891.1234567891, -4, 1234560000)
+    testTruncate(1234567891.1234567891, 0, 1234567891)
+    testTruncate(0.123, -1, 0)
+    testTruncate(0.123, 0, 0)
+
+    checkEvaluation(Truncate(Literal.create(1D, DoubleType),
+      NonFoldableLiteral.create(null, IntegerType)),
+      null)
+    checkEvaluation(Truncate(Literal.create(null, DoubleType),
+      NonFoldableLiteral.create(1, IntegerType)),
+      null)
+    checkEvaluation(Truncate(Literal.create(null, DoubleType),
+      NonFoldableLiteral.create(null, IntegerType)),
+      null)
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/functions.scala

@@ -2214,6 +2214,24 @@ object functions {
    */
   def radians(columnName: String): Column = radians(Column(columnName))
 
+  /**
+   * Returns number truncated to the unit specified by the scale.
+   *
+   * For example, `truncate(1234567891.1234567891, 4)` returns 1234567891.1234
+   *
+   * @param number The number to be truncated
+   * @param scale: A scale used to truncate number
+   *
+   * @return The number truncated to scale decimal places.
+   *         If scale is omitted, then number is truncated to 0 places.
+   *         scale can be negative to truncate (make zero) scale digits left of the decimal point.
+   * @group math_funcs
+   * @since 2.4.0
+   */
+  def truncate(number: Column, scale: Int): Column = withExpr {
+    Truncate(number.expr, Literal(scale))
+  }
+
   //////////////////////////////////////////////////////////////////////////////////////////////
   // Misc functions
   //////////////////////////////////////////////////////////////////////////////////////////////

sql/core/src/test/resources/sql-tests/inputs/operators.sql

@@ -96,3 +96,9 @@ select positive('-1.11'), positive(-1.11), negative('-1.11'), negative(-1.11);
 -- pmod
 select pmod(-7, 2), pmod(0, 2), pmod(7, 0), pmod(7, null), pmod(null, 2), pmod(null, null);
 select pmod(cast(3.13 as decimal), cast(0 as decimal)), pmod(cast(2 as smallint), cast(0 as smallint));
+
+-- truncate
+select truncate(1234567891.1234567891, -4), truncate(1234567891.1234567891, 0), truncate(1234567891.1234567891, 4);
+select truncate(cast(1234567891.1234567891 as decimal), -4), truncate(cast(1234567891.1234567891 as decimal), 0), truncate(cast(1234567891.1234567891 as decimal), 4);
+select truncate(cast(1234567891.1234567891 as long), -4), truncate(cast(1234567891.1234567891 as long), 0), truncate(cast(1234567891.1234567891 as long), 4);
+select truncate(cast(1234567891.1234567891 as long), 9.03)

sql/core/src/test/resources/sql-tests/results/operators.sql.out

@@ -1,5 +1,5 @@
 -- Automatically generated by SQLQueryTestSuite
--- Number of queries: 59
+-- Number of queries: 63
 
 
 -- !query 0
@@ -484,3 +484,35 @@ select pmod(cast(3.13 as decimal), cast(0 as decimal)), pmod(cast(2 as smallint)
 struct<pmod(CAST(3.13 AS DECIMAL(10,0)), CAST(0 AS DECIMAL(10,0))):decimal(10,0),pmod(CAST(2 AS SMALLINT), CAST(0 AS SMALLINT)):smallint>
 -- !query 58 output
 NULL	NULL
+
+
+-- !query 59
+select truncate(1234567891.1234567891, -4), truncate(1234567891.1234567891, 0), truncate(1234567891.1234567891, 4)
+-- !query 59 schema
+struct<truncate(1234567891.1234567891, -4):decimal(20,10),truncate(1234567891.1234567891, 0):decimal(20,10),truncate(1234567891.1234567891, 4):decimal(20,10)>
+-- !query 59 output
+1234560000	1234567891	1234567891.1234
+
+
+-- !query 60
+select truncate(cast(1234567891.1234567891 as decimal), -4), truncate(cast(1234567891.1234567891 as decimal), 0), truncate(cast(1234567891.1234567891 as decimal), 4)
+-- !query 60 schema
+struct<truncate(CAST(1234567891.1234567891 AS DECIMAL(10,0)), -4):decimal(10,0),truncate(CAST(1234567891.1234567891 AS DECIMAL(10,0)), 0):decimal(10,0),truncate(CAST(1234567891.1234567891 AS DECIMAL(10,0)), 4):decimal(10,0)>
+-- !query 60 output
+1234560000	1234567891	1234567891
+
+
+-- !query 61
+select truncate(cast(1234567891.1234567891 as long), -4), truncate(cast(1234567891.1234567891 as long), 0), truncate(cast(1234567891.1234567891 as long), 4)
+-- !query 61 schema
+struct<truncate(CAST(CAST(1234567891.1234567891 AS BIGINT) AS DOUBLE), -4):double,truncate(CAST(CAST(1234567891.1234567891 AS BIGINT) AS DOUBLE), 0):double,truncate(CAST(CAST(1234567891.1234567891 AS BIGINT) AS DOUBLE), 4):double>
+-- !query 61 output
+1.23456E9	1.234567891E9	1.234567891E9
+
+
+-- !query 62
+select truncate(cast(1234567891.1234567891 as long), 9.03)
+-- !query 62 schema
+struct<truncate(CAST(CAST(1234567891.1234567891 AS BIGINT) AS DOUBLE), CAST(9.03 AS INT)):double>
+-- !query 62 output
+1.234567891E9