Add decimal support to Round

avulanov · Jul 14, 2015 · 56db4bb · 56db4bb
1 parent 7e163ae
commit 56db4bb
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 18 deletions.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/math.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/math.scala
@@ -528,9 +528,13 @@ case class Round(children: Seq[Expression]) extends Expression {
 
   def nullable: Boolean = true
 
-  def dataType: DataType = {
+  private lazy val evalE2 = if (children.size == 2) children(1).eval(EmptyRow) else null
+  private lazy val _scale = if (evalE2 != null) evalE2.asInstanceOf[Int] else 0
+
+  override lazy val dataType: DataType = {
     children(0).dataType match {
       case StringType | BinaryType => DoubleType
+      case DecimalType.Fixed(p, s) => DecimalType(p, _scale)
       case t => t
     }
   }
@@ -564,23 +568,14 @@ case class Round(children: Seq[Expression]) extends Expression {
 
   def eval(input: InternalRow): Any = {
     val evalE1 = children(0).eval(input)
-    if (evalE1 == null) {
-      return null
-    }
 
-    var _scale: Int = 0
-    if (children.size == 2) {
-      val evalE2 = children(1).eval(input)
-      if (evalE2 == null) {
-        return null
-      } else {
-        _scale = evalE2.asInstanceOf[Int]
-      }
-    }
+    if (evalE1 == null) return null
+    if (children.size == 2 && evalE2 == null) return null
 
     children(0).dataType match {
       case decimalType: DecimalType =>
-        // TODO: Support Decimal Round
+        val decimal = evalE1.asInstanceOf[Decimal]
+        if (decimal.changePrecision(decimal.precision, _scale)) decimal else null
       case ByteType =>
         round(evalE1.asInstanceOf[Byte], _scale)
       case ShortType =>

diff --git a/...atibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala b/...atibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala
@@ -221,9 +221,6 @@ class HiveCompatibilitySuite extends HiveQueryFileTest with BeforeAndAfter {
     "udf_when",
     "udf_case",
 
-    // Needs constant object inspectors
-    "udf_round",
-
     // the table src(key INT, value STRING) is not the same as HIVE unittest. In Hive
     // is src(key STRING, value STRING), and in the reflect.q, it failed in
     // Integer.valueOf, which expect the first argument passed as STRING type not INT.
@@ -918,7 +915,7 @@ class HiveCompatibilitySuite extends HiveQueryFileTest with BeforeAndAfter {
     "udf_regexp_replace",
     "udf_repeat",
     "udf_rlike",
-    "udf_round",
+    // "udf_round",  turn this on after we figure out null vs nan vs infinity
     "udf_round_3",
     "udf_rpad",
     "udf_rtrim",