[SPARK-37324][SQL] Adds support for decimal rounding mode up, down, half_down

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

@@ -1338,12 +1338,26 @@ case class Logarithm(left: Expression, right: Expression)
  * @param mode rounding mode (e.g. HALF_UP, HALF_EVEN)
  * @param modeStr rounding mode string name (e.g. "ROUND_HALF_UP", "ROUND_HALF_EVEN")
  */
-abstract class RoundBase(child: Expression, scale: Expression,
-    mode: BigDecimal.RoundingMode.Value, modeStr: String)
-  extends BinaryExpression with Serializable with ImplicitCastInputTypes {
+abstract class RoundBase(child: Expression, scale: Expression, modeExpr: Expression)
+  extends TernaryExpression with Serializable with ImplicitCastInputTypes {
 
-  override def left: Expression = child
-  override def right: Expression = scale
+  override def first: Expression = child
+
+  override def second: Expression = scale
+
+  override def third: Expression = modeExpr
+
+  val (mode: BigDecimal.RoundingMode.Value, modeStr: String) = modeExpr.toString match {
+    case "up" => (BigDecimal.RoundingMode.UP, "ROUND_UP")
+    case "down" => (BigDecimal.RoundingMode.DOWN, "ROUND_DOWN")
+    case "half_up" => (BigDecimal.RoundingMode.HALF_UP, "ROUND_HALF_UP")
+    case "half_down" => (BigDecimal.RoundingMode.HALF_DOWN, "ROUND_HALF_DOWN")
+    case "half_even" => (BigDecimal.RoundingMode.HALF_EVEN, "ROUND_HALF_EVEN")
+    case _ =>
+      val supported = Seq("up", "down", "half_up", "half_down", "half_even")
+      throw new IllegalArgumentException(s"Unsupported rounding mode '${modeExpr.toString}'. " +
+        "Supported rounding modes include: " + supported.mkString("'", "', '", "'") + ".")
+  }
 
   // round of Decimal would eval to null if it fails to `changePrecision`
   override def nullable: Boolean = true
@@ -1357,7 +1371,7 @@ abstract class RoundBase(child: Expression, scale: Expression,
     case t => t
   }
 
-  override def inputTypes: Seq[AbstractDataType] = Seq(NumericType, IntegerType)
+  override def inputTypes: Seq[AbstractDataType] = Seq(NumericType, IntegerType, StringType)
 
   override def checkInputDataTypes(): TypeCheckResult = {
     super.checkInputDataTypes() match {
@@ -1500,26 +1514,47 @@ abstract class RoundBase(child: Expression, scale: Expression,
 }
 
 /**
- * Round an expression to d decimal places using HALF_UP rounding mode.
- * round(2.5) == 3.0, round(3.5) == 4.0.
+ * Round an expression to d decimal places using given rounding mode, default rounding mode HALF_UP.
+ * round(2.5) == 3.0, round(3.5) == 4.0,
+ * round(3.6, 0, 'half_down') == 4.0, round(3.6, 0, 'down') == 3.0.
  */
 // scalastyle:off line.size.limit
 @ExpressionDescription(
-  usage = "_FUNC_(expr, d) - Returns `expr` rounded to `d` decimal places using HALF_UP rounding mode.",
+  usage = "_FUNC_(expr, d) - Returns `expr` rounded to `d` decimal places using given rounding mode.",
   examples = """
     Examples:
       > SELECT _FUNC_(2.5, 0);
        3
+      > SELECT _FUNC_(2.5, 0, 'half_up');
+       3
+      > SELECT _FUNC_(2.5, 0, 'half_even');
+       2
+      > SELECT _FUNC_(2.6, 0, 'half_down');
+       3
+      > SELECT _FUNC_(2.1, 0, 'up');
+       3
+      > SELECT _FUNC_(2.9, 0, 'down');
+       2
+      > SELECT _FUNC_(-2.9, 0, 'down');
+       -2
   """,
   since = "1.5.0",
   group = "math_funcs")
 // scalastyle:on line.size.limit
-case class Round(child: Expression, scale: Expression)
-  extends RoundBase(child, scale, BigDecimal.RoundingMode.HALF_UP, "ROUND_HALF_UP")
-    with Serializable with ImplicitCastInputTypes {
+case class Round(
+    child: Expression,
+    scale: Expression,
+    modeExpr: Expression)
+  extends RoundBase(child, scale, modeExpr) with Serializable with ImplicitCastInputTypes {
+  def this(child: Expression, scale: Expression) = this(child, scale, Literal("half_up"))
   def this(child: Expression) = this(child, Literal(0))
-  override protected def withNewChildrenInternal(newLeft: Expression, newRight: Expression): Round =
-    copy(child = newLeft, scale = newRight)
+  override protected def withNewChildrenInternal(
+      newFirst: Expression, newSecond: Expression, newThird: Expression): Expression =
+    copy(child = newFirst, scale = newSecond, modeExpr = newThird)
+}
+
+object Round {
+  def apply(child: Expression, scale: Expression): Round = new Round(child, scale)
 }
 
 /**
@@ -1538,12 +1573,15 @@ case class Round(child: Expression, scale: Expression)
   since = "2.0.0",
   group = "math_funcs")
 // scalastyle:on line.size.limit
-case class BRound(child: Expression, scale: Expression)
-  extends RoundBase(child, scale, BigDecimal.RoundingMode.HALF_EVEN, "ROUND_HALF_EVEN")
-    with Serializable with ImplicitCastInputTypes {
+class BRound(child: Expression, scale: Expression, modeExpr: Expression)
+  extends Round(child, scale, modeExpr) {
+  def this(child: Expression, scale: Expression) = this(child, scale, Literal("half_even"))
   def this(child: Expression) = this(child, Literal(0))
-  override protected def withNewChildrenInternal(
-    newLeft: Expression, newRight: Expression): BRound = copy(child = newLeft, scale = newRight)
+  override def nodeName: String = "round"
+}
+
+object BRound {
+  def apply(child: Expression, scale: Expression): Round = new BRound(child, scale)
 }
 
 object WidthBucket {

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

@@ -668,7 +668,7 @@ object PushFoldableIntoBranches extends Rule[LogicalPlan] with PredicateHelper {
     case _: BinaryMathExpression => true
     case _: AddMonths | _: DateAdd | _: DateAddInterval | _: DateDiff | _: DateSub |
          _: DateAddYMInterval | _: TimestampAddYMInterval | _: TimeAdd => true
-    case _: FindInSet | _: RoundBase => true
+    case _: FindInSet => true
     case _ => false
   }
 
@@ -715,6 +715,21 @@ object PushFoldableIntoBranches extends Rule[LogicalPlan] with PredicateHelper {
         c.copy(
           branches.map(e => e.copy(_2 = b.withNewChildren(Array(left, e._2)))),
           Some(b.withNewChildren(Array(left, elseValue.getOrElse(Literal(null, c.dataType))))))
+
+      case r @ Round(i @ If(_, trueValue, falseValue), second, third)
+        if second.foldable && third.foldable &&
+          atMostOneUnfoldable(Seq(trueValue, falseValue)) =>
+        i.copy(
+          trueValue = r.withNewChildren(Array(trueValue, second, third)),
+          falseValue = r.withNewChildren(Array(falseValue, second, third)))
+
+      case r @ Round(c @ CaseWhen(branches, elseValue), second, third)
+        if second.foldable && third.foldable &&
+          atMostOneUnfoldable(branches.map(_._2) ++ elseValue) =>
+        c.copy(
+          branches.map(e => e.copy(_2 = r.withNewChildren(Array(e._2, second, third)))),
+          Some(r.withNewChildren(
+            Array(elseValue.getOrElse(Literal(null, c.dataType)), second, third))))
     }
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/types/Decimal.scala

@@ -531,6 +531,9 @@ final class Decimal extends Ordered[Decimal] with Serializable {
 
 @Unstable
 object Decimal {
+  val ROUND_UP = BigDecimal.RoundingMode.UP
+  val ROUND_DOWN = BigDecimal.RoundingMode.DOWN
+  val ROUND_HALF_DOWN = BigDecimal.RoundingMode.HALF_DOWN
   val ROUND_HALF_UP = BigDecimal.RoundingMode.HALF_UP
   val ROUND_HALF_EVEN = BigDecimal.RoundingMode.HALF_EVEN
   val ROUND_CEILING = BigDecimal.RoundingMode.CEILING

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/PhysicalAggregationSuite.scala

@@ -47,7 +47,7 @@ class PhysicalAggregationSuite extends PlanTest {
 
     // Verify that Round's scale parameter is a Literal.
     resultExpressions(1) match {
-      case Alias(Round(_, _: Literal), _) =>
+      case Alias(Round(_, _: Literal, _: Literal), _) =>
       case other => fail("unexpected result expression: " + other)
     }
   }

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

@@ -2197,13 +2197,23 @@ object functions {
   def round(e: Column): Column = round(e, 0)
 
   /**
-   * Round the value of `e` to `scale` decimal places with HALF_UP round mode
-   * if `scale` is greater than or equal to 0 or at integral part when `scale` is less than 0.
+   * Returns the value of the column `e` rounded to 0 decimal places with HALF_UP round mode.
    *
    * @group math_funcs
    * @since 1.5.0
    */
-  def round(e: Column, scale: Int): Column = withExpr { Round(e.expr, Literal(scale)) }
+  def round(e: Column, scale: Int): Column = round(e, lit(scale), lit("half_up"))
+
+  /**
+   * Round the value of `e` to `scale` decimal places with given round mode, default: HALF_UP
+   * if `scale` is greater than or equal to 0 or at integral part when `scale` is less than 0.
+   *
+   * @group math_funcs
+   * @since 3.3.0
+   */
+  def round(e: Column, scale: Column, mode: Column): Column = withExpr {
+    Round(e.expr, scale.expr, mode.expr)
+  }
 
   /**
    * Returns the value of the column `e` rounded to 0 decimal places with HALF_EVEN round mode.
@@ -2220,7 +2230,7 @@ object functions {
    * @group math_funcs
    * @since 2.0.0
    */
-  def bround(e: Column, scale: Int): Column = withExpr { BRound(e.expr, Literal(scale)) }
+  def bround(e: Column, scale: Int): Column = round(e, lit(scale), lit("half_even"))
 
   /**
    * @param e angle in radians

sql/core/src/test/resources/sql-functions/sql-expression-schema.md

@@ -40,7 +40,7 @@
 | org.apache.spark.sql.catalyst.expressions.Atan | atan | SELECT atan(0) | struct<ATAN(0):double> |
 | org.apache.spark.sql.catalyst.expressions.Atan2 | atan2 | SELECT atan2(0, 0) | struct<ATAN2(0, 0):double> |
 | org.apache.spark.sql.catalyst.expressions.Atanh | atanh | SELECT atanh(0) | struct<ATANH(0):double> |
-| org.apache.spark.sql.catalyst.expressions.BRound | bround | SELECT bround(2.5, 0) | struct<bround(2.5, 0):decimal(2,0)> |
+| org.apache.spark.sql.catalyst.expressions.BRound | bround | SELECT bround(2.5, 0) | struct<round(2.5, 0, half_even):decimal(2,0)> |
 | org.apache.spark.sql.catalyst.expressions.Base64 | base64 | SELECT base64('Spark SQL') | struct<base64(Spark SQL):string> |
 | org.apache.spark.sql.catalyst.expressions.Bin | bin | SELECT bin(13) | struct<bin(13):string> |
 | org.apache.spark.sql.catalyst.expressions.BitLength | bit_length | SELECT bit_length('Spark SQL') | struct<bit_length(Spark SQL):int> |
@@ -242,7 +242,7 @@
 | org.apache.spark.sql.catalyst.expressions.Reverse | reverse | SELECT reverse('Spark SQL') | struct<reverse(Spark SQL):string> |
 | org.apache.spark.sql.catalyst.expressions.Right | right | SELECT right('Spark SQL', 3) | struct<right(Spark SQL, 3):string> |
 | org.apache.spark.sql.catalyst.expressions.Rint | rint | SELECT rint(12.3456) | struct<rint(12.3456):double> |
-| org.apache.spark.sql.catalyst.expressions.Round | round | SELECT round(2.5, 0) | struct<round(2.5, 0):decimal(2,0)> |
+| org.apache.spark.sql.catalyst.expressions.Round | round | SELECT round(2.5, 0) | struct<round(2.5, 0, half_up):decimal(2,0)> |
 | org.apache.spark.sql.catalyst.expressions.RowNumber | row_number | SELECT a, b, row_number() OVER (PARTITION BY a ORDER BY b) FROM VALUES ('A1', 2), ('A1', 1), ('A2', 3), ('A1', 1) tab(a, b) | struct<a:string,b:int,row_number() OVER (PARTITION BY a ORDER BY b ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW):int> |
 | org.apache.spark.sql.catalyst.expressions.SchemaOfCsv | schema_of_csv | SELECT schema_of_csv('1,abc') | struct<schema_of_csv(1,abc):string> |
 | org.apache.spark.sql.catalyst.expressions.SchemaOfJson | schema_of_json | SELECT schema_of_json('[{"col":0}]') | struct<schema_of_json([{"col":0}]):string> |

sql/core/src/test/resources/sql-tests/results/postgreSQL/numeric.sql.out

@@ -3810,7 +3810,7 @@ SELECT t1.id1, t1.id2, t1.result, round(t2.expected, 40)
     WHERE t1.id1 = t2.id1 AND t1.id2 = t2.id2
     AND t1.result != round(t2.expected, 40)
 -- !query schema
-struct<id1:int,id2:int,result:decimal(38,10),round(expected, 40):decimal(38,10)>
+struct<id1:int,id2:int,result:decimal(38,10),round(expected, 40, half_up):decimal(38,10)>
 -- !query output
 
 
@@ -4404,7 +4404,7 @@ struct<>
 -- !query
 SELECT a, ceil(a), ceiling(a), floor(a), round(a) FROM ceil_floor_round
 -- !query schema
-struct<a:decimal(38,18),CEIL(a):decimal(21,0),ceiling(a):decimal(21,0),FLOOR(a):decimal(21,0),round(a, 0):decimal(38,0)>
+struct<a:decimal(38,18),CEIL(a):decimal(21,0),ceiling(a):decimal(21,0),FLOOR(a):decimal(21,0),round(a, 0, half_up):decimal(38,0)>
 -- !query output
 -0.000001000000000000	0	0	-1	0
 -5.499999000000000000	-5	-5	-6	-5

sql/core/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q2.sf100/explain.txt

@@ -195,14 +195,14 @@ Right keys [1]: [(d_week_seq2#63 - 53)]
 Join condition: None
 
 (35) Project [codegen id : 12]
-Output [8]: [d_week_seq1#45, round(CheckOverflow((promote_precision(sun_sales1#46) / promote_precision(sun_sales2#64)), DecimalType(37,20), true), 2) AS round((sun_sales1 / sun_sales2), 2)#72, round(CheckOverflow((promote_precision(mon_sales1#47) / promote_precision(mon_sales2#65)), DecimalType(37,20), true), 2) AS round((mon_sales1 / mon_sales2), 2)#73, round(CheckOverflow((promote_precision(tue_sales1#48) / promote_precision(tue_sales2#66)), DecimalType(37,20), true), 2) AS round((tue_sales1 / tue_sales2), 2)#74, round(CheckOverflow((promote_precision(wed_sales1#49) / promote_precision(wed_sales2#67)), DecimalType(37,20), true), 2) AS round((wed_sales1 / wed_sales2), 2)#75, round(CheckOverflow((promote_precision(thu_sales1#50) / promote_precision(thu_sales2#68)), DecimalType(37,20), true), 2) AS round((thu_sales1 / thu_sales2), 2)#76, round(CheckOverflow((promote_precision(fri_sales1#51) / promote_precision(fri_sales2#69)), DecimalType(37,20), true), 2) AS round((fri_sales1 / fri_sales2), 2)#77, round(CheckOverflow((promote_precision(sat_sales1#52) / promote_precision(sat_sales2#70)), DecimalType(37,20), true), 2) AS round((sat_sales1 / sat_sales2), 2)#78]
+Output [8]: [d_week_seq1#45, round(CheckOverflow((promote_precision(sun_sales1#46) / promote_precision(sun_sales2#64)), DecimalType(37,20), true), 2, half_up) AS round((sun_sales1 / sun_sales2), 2, half_up)#72, round(CheckOverflow((promote_precision(mon_sales1#47) / promote_precision(mon_sales2#65)), DecimalType(37,20), true), 2, half_up) AS round((mon_sales1 / mon_sales2), 2, half_up)#73, round(CheckOverflow((promote_precision(tue_sales1#48) / promote_precision(tue_sales2#66)), DecimalType(37,20), true), 2, half_up) AS round((tue_sales1 / tue_sales2), 2, half_up)#74, round(CheckOverflow((promote_precision(wed_sales1#49) / promote_precision(wed_sales2#67)), DecimalType(37,20), true), 2, half_up) AS round((wed_sales1 / wed_sales2), 2, half_up)#75, round(CheckOverflow((promote_precision(thu_sales1#50) / promote_precision(thu_sales2#68)), DecimalType(37,20), true), 2, half_up) AS round((thu_sales1 / thu_sales2), 2, half_up)#76, round(CheckOverflow((promote_precision(fri_sales1#51) / promote_precision(fri_sales2#69)), DecimalType(37,20), true), 2, half_up) AS round((fri_sales1 / fri_sales2), 2, half_up)#77, round(CheckOverflow((promote_precision(sat_sales1#52) / promote_precision(sat_sales2#70)), DecimalType(37,20), true), 2, half_up) AS round((sat_sales1 / sat_sales2), 2, half_up)#78]
 Input [16]: [d_week_seq1#45, sun_sales1#46, mon_sales1#47, tue_sales1#48, wed_sales1#49, thu_sales1#50, fri_sales1#51, sat_sales1#52, d_week_seq2#63, sun_sales2#64, mon_sales2#65, tue_sales2#66, wed_sales2#67, thu_sales2#68, fri_sales2#69, sat_sales2#70]
 
 (36) Exchange
-Input [8]: [d_week_seq1#45, round((sun_sales1 / sun_sales2), 2)#72, round((mon_sales1 / mon_sales2), 2)#73, round((tue_sales1 / tue_sales2), 2)#74, round((wed_sales1 / wed_sales2), 2)#75, round((thu_sales1 / thu_sales2), 2)#76, round((fri_sales1 / fri_sales2), 2)#77, round((sat_sales1 / sat_sales2), 2)#78]
+Input [8]: [d_week_seq1#45, round((sun_sales1 / sun_sales2), 2, half_up)#72, round((mon_sales1 / mon_sales2), 2, half_up)#73, round((tue_sales1 / tue_sales2), 2, half_up)#74, round((wed_sales1 / wed_sales2), 2, half_up)#75, round((thu_sales1 / thu_sales2), 2, half_up)#76, round((fri_sales1 / fri_sales2), 2, half_up)#77, round((sat_sales1 / sat_sales2), 2, half_up)#78]
 Arguments: rangepartitioning(d_week_seq1#45 ASC NULLS FIRST, 5), ENSURE_REQUIREMENTS, [id=#79]
 
 (37) Sort [codegen id : 13]
-Input [8]: [d_week_seq1#45, round((sun_sales1 / sun_sales2), 2)#72, round((mon_sales1 / mon_sales2), 2)#73, round((tue_sales1 / tue_sales2), 2)#74, round((wed_sales1 / wed_sales2), 2)#75, round((thu_sales1 / thu_sales2), 2)#76, round((fri_sales1 / fri_sales2), 2)#77, round((sat_sales1 / sat_sales2), 2)#78]
+Input [8]: [d_week_seq1#45, round((sun_sales1 / sun_sales2), 2, half_up)#72, round((mon_sales1 / mon_sales2), 2, half_up)#73, round((tue_sales1 / tue_sales2), 2, half_up)#74, round((wed_sales1 / wed_sales2), 2, half_up)#75, round((thu_sales1 / thu_sales2), 2, half_up)#76, round((fri_sales1 / fri_sales2), 2, half_up)#77, round((sat_sales1 / sat_sales2), 2, half_up)#78]
 Arguments: [d_week_seq1#45 ASC NULLS FIRST], true, 0