Support decimal type for Spark unaryminus function (facebookincubator…

…#11454) Summary: Add support for decimal type to Spark unaryminus function. Pull Request resolved: facebookincubator#11454 Reviewed By: xiaoxmeng Differential Revision: D65894435 Pulled By: kevinwilfong fbshipit-source-id: 708045e9d31bc48e0c0c4a52b41902be1405dfd3
yuandagits · Nov 13, 2024 · f34035b · f34035b
1 parent a993e05
commit f34035b
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 6 deletions.
diff --git a/velox/docs/functions/spark/decimal.rst b/velox/docs/functions/spark/decimal.rst
@@ -115,6 +115,13 @@ Returns NULL when the actual result cannot be represented with the calculated de
 Decimal Functions
 -----------------
 
+.. spark:function:: unaryminus(x: decimal(p, s)) -> r: decimal(p, s)
+
+    Returns negated value of x (r = -x). Corresponds to Spark's operator ``-``.
+
+    ::
+        SELECT unaryminus(cast(-9999999999999999999.9999999999999999999 as DECIMAL(38, 19))); -- 9999999999999999999.9999999999999999999
+
 .. spark:function:: unscaled_value(x) -> bigint
 
     Return the unscaled bigint value of a short decimal ``x``.

diff --git a/velox/functions/sparksql/RegisterArithmetic.cpp b/velox/functions/sparksql/RegisterArithmetic.cpp
@@ -38,6 +38,14 @@ void registerArithmeticFunctions(const std::string& prefix) {
   registerFunction<DivideFunction, double, double, double>({prefix + "divide"});
   registerBinaryNumeric<RemainderFunction>({prefix + "remainder"});
   registerUnaryNumeric<UnaryMinusFunction>({prefix + "unaryminus"});
+  registerFunction<
+      UnaryMinusFunction,
+      LongDecimal<P1, S1>,
+      LongDecimal<P1, S1>>({prefix + "unaryminus"});
+  registerFunction<
+      UnaryMinusFunction,
+      ShortDecimal<P1, S1>,
+      ShortDecimal<P1, S1>>({prefix + "unaryminus"});
   // Math functions.
   registerUnaryNumeric<AbsFunction>({prefix + "abs"});
   registerFunction<

diff --git a/velox/functions/sparksql/tests/DecimalArithmeticTest.cpp b/velox/functions/sparksql/tests/DecimalArithmeticTest.cpp
@@ -37,13 +37,15 @@ class DecimalArithmeticTest : public SparkFunctionBaseTest {
       const std::string& functionName,
       const VectorPtr& expected,
       const std::vector<VectorPtr>& inputs) {
-    VELOX_USER_CHECK_EQ(
+    VELOX_USER_CHECK_GE(
         inputs.size(),
-        2,
-        "Two input vectors are needed for arithmetic function test.");
-    std::vector<core::TypedExprPtr> inputExprs = {
-        std::make_shared<core::FieldAccessTypedExpr>(inputs[0]->type(), "c0"),
-        std::make_shared<core::FieldAccessTypedExpr>(inputs[1]->type(), "c1")};
+        1,
+        "At least one input vector is needed for arithmetic function test.");
+    std::vector<core::TypedExprPtr> inputExprs;
+    for (int i = 0; i < inputs.size(); ++i) {
+      inputExprs.emplace_back(std::make_shared<core::FieldAccessTypedExpr>(
+          inputs[i]->type(), fmt::format("c{}", i)));
+    }
     auto expr = std::make_shared<const core::CallTypedExpr>(
         expected->type(), std::move(inputExprs), functionName);
     assertEqualVectors(expected, evaluate(expr, makeRowVector(inputs)));
@@ -576,5 +578,43 @@ TEST_F(DecimalArithmeticTest, denyPrecisionLoss) {
       {makeConstant<int128_t>(500, 1, DECIMAL(20, 2)),
        makeConstant<int64_t>(1000, 1, DECIMAL(7, 3))});
 }
+
+TEST_F(DecimalArithmeticTest, unaryMinus) {
+  testArithmeticFunction(
+      "unaryminus",
+      makeFlatVector<int64_t>(
+          {1111,
+           -1112,
+           9999,
+           0,
+           -DecimalUtil::kShortDecimalMin,
+           -DecimalUtil::kShortDecimalMax},
+          DECIMAL(18, 9)),
+      {makeFlatVector<int64_t>(
+          {-1111,
+           1112,
+           -9999,
+           0,
+           DecimalUtil::kShortDecimalMin,
+           DecimalUtil::kShortDecimalMax},
+          DECIMAL(18, 9))});
+
+  testArithmeticFunction(
+      "unaryminus",
+      makeFlatVector<int128_t>(
+          {11111111,
+           -11112112,
+           99999999,
+           -DecimalUtil::kLongDecimalMax,
+           -DecimalUtil::kLongDecimalMin},
+          DECIMAL(38, 19)),
+      {makeFlatVector<int128_t>(
+          {-11111111,
+           11112112,
+           -99999999,
+           DecimalUtil::kLongDecimalMax,
+           DecimalUtil::kLongDecimalMin},
+          DECIMAL(38, 19))});
+}
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test