feat(sparksql): support decimal format_number

bolt/functions/sparksql/FormatNumber.cpp

@@ -18,9 +18,13 @@
 
 #include <fmt/format.h>
 #include <algorithm>
+#include <array>
 #include <cmath>
 #include <cstdint>
 #include <locale>
+#include "bolt/expression/DecodedArgs.h"
+#include "bolt/expression/VectorFunction.h"
+#include "bolt/type/DecimalUtil.h"
 
 namespace bytedance::bolt::functions::sparksql::detail {
 
@@ -49,6 +53,112 @@ const std::locale& usLocale() {
   return loc;
 }
 
+std::string_view zeroPadding(int32_t zeroCount) {
+  static const std::string kZeros(kMaxFractionDigits, '0');
+  return std::string_view(kZeros.data(), zeroCount);
+}
+
+void appendGroupedIntegerDigits(
+    std::string_view digits,
+    exec::StringWriter<false>& out) {
+  if (digits.empty()) {
+    out.append(std::string_view("0"));
+    return;
+  }
+
+  const auto firstGroupSize =
+      digits.size() <= 3 ? digits.size() : ((digits.size() - 1) % 3) + 1;
+  out.append(digits.substr(0, firstGroupSize));
+
+  for (size_t offset = firstGroupSize; offset < digits.size(); offset += 3) {
+    out.append(std::string_view(","));
+    out.append(digits.substr(offset, 3));
+  }
+}
+
+template <typename T>
+T roundDecimalHalfEven(T value, int32_t fromScale, int32_t toScale) {
+  if (toScale >= fromScale) {
+    return value;
+  }
+
+  const auto scaleFactor =
+      DecimalUtil::getPowersOfTen(static_cast<uint8_t>(fromScale - toScale));
+  int128_t roundedValue = value;
+  int128_t remainder = roundedValue % scaleFactor;
+  roundedValue /= scaleFactor;
+
+  const int128_t absRemainder = remainder < 0 ? -remainder : remainder;
+  const int128_t twiceRemainder = absRemainder * 2;
+  if (twiceRemainder > scaleFactor) {
+    roundedValue += value >= 0 ? 1 : -1;
+  } else if (
+      twiceRemainder == scaleFactor &&
+      ((roundedValue < 0 ? -roundedValue : roundedValue) % 2 == 1)) {
+    roundedValue += value >= 0 ? 1 : -1;
+  }
+
+  return static_cast<T>(roundedValue);
+}
+
+template <typename T>
+size_t estimateFormattedDecimalSize(
+    T value,
+    int32_t inputPrecision,
+    int32_t inputScale,
+    int32_t decimalPlaces) {
+  if (decimalPlaces < 0) {
+    return 0;
+  }
+
+  const int32_t cappedPlaces = std::min(decimalPlaces, kMaxFractionDigits);
+  const int32_t outputScale = std::min(cappedPlaces, inputScale);
+  const int32_t integerDigits = std::max(inputPrecision - outputScale, 1);
+  const int32_t groupingSeparators = (integerDigits - 1) / 3;
+  return (value < 0 ? 1 : 0) + integerDigits + groupingSeparators +
+      (cappedPlaces > 0 ? 1 + cappedPlaces : 0);
+}
+
+template <typename T>
+void formatDecimal(
+    T value,
+    int32_t inputScale,
+    int32_t decimalPlaces,
+    exec::StringWriter<false>& out) {
+  const int32_t cappedPlaces = std::min(decimalPlaces, kMaxFractionDigits);
+  const int32_t outputScale = std::min(cappedPlaces, inputScale);
+  const T roundedValue = roundDecimalHalfEven(value, inputScale, outputScale);
+
+  std::array<char, 64> plainBuffer;
+  const auto plainSize = DecimalUtil::convertToPlainString<T>(
+      roundedValue, outputScale, plainBuffer.size(), plainBuffer.data());
+  const std::string_view plainNumber(plainBuffer.data(), plainSize);
+  size_t start = 0;
+  if (!plainNumber.empty() && plainNumber.front() == '-') {
+    out.append(std::string_view("-"));
+    start = 1;
+  }
+
+  const auto decimalPoint = plainNumber.find('.', start);
+  appendGroupedIntegerDigits(
+      decimalPoint == std::string_view::npos
+          ? plainNumber.substr(start)
+          : plainNumber.substr(start, decimalPoint - start),
+      out);
+
+  if (decimalPoint != std::string_view::npos) {
+    out.append(plainNumber.substr(decimalPoint));
+  }
+
+  const auto trailingZeros = cappedPlaces - outputScale;
+  if (trailingZeros > 0) {
+    if (decimalPoint == std::string_view::npos) {
+      out.append(std::string_view("."));
+    }
+    out.append(zeroPadding(trailingZeros));
+  }
+}
+
 } // namespace
 
 void formatInteger(
@@ -100,4 +210,95 @@ void formatFloatingPoint(
   out.append(std::string_view(buf.data(), buf.size()));
 }
 
+namespace {
+
+template <typename T>
+class DecimalFormatNumberFunction final : public exec::VectorFunction {
+ public:
+  DecimalFormatNumberFunction(int32_t inputPrecision, int32_t inputScale)
+      : inputPrecision_(inputPrecision), inputScale_(inputScale) {}
+
+  void apply(
+      const SelectivityVector& rows,
+      std::vector<VectorPtr>& args,
+      const TypePtr& outputType,
+      exec::EvalCtx& context,
+      VectorPtr& result) const override {
+    exec::DecodedArgs decodedArgs(rows, args, context);
+    auto* decimalValues = decodedArgs.at(0);
+    auto* decimalPlaces = decodedArgs.at(1);
+
+    context.ensureWritable(rows, outputType, result);
+    result->clearNulls(rows);
+    auto* flatResult = result->asFlatVector<StringView>();
+
+    size_t estimatedBufferSize = 0;
+    rows.applyToSelected([&](vector_size_t row) {
+      estimatedBufferSize += estimateFormattedDecimalSize(
+          decimalValues->valueAt<T>(row),
+          inputPrecision_,
+          inputScale_,
+          decimalPlaces->valueAt<int32_t>(row));
+    });
+    if (estimatedBufferSize > 0) {
+      flatResult->getBufferWithSpace(estimatedBufferSize, true);
+    }
+
+    context.applyToSelectedNoThrow(rows, [&](vector_size_t row) {
+      const auto places = decimalPlaces->valueAt<int32_t>(row);
+      if (places < 0) {
+        result->setNull(row, true);
+        return;
+      }
+      exec::StringWriter<false> writer(flatResult, row);
+      formatDecimal(
+          decimalValues->valueAt<T>(row), inputScale_, places, writer);
+      writer.finalize();
+    });
+  }
+
+ private:
+  const int32_t inputPrecision_;
+  const int32_t inputScale_;
+};
+
+std::vector<std::shared_ptr<exec::FunctionSignature>>
+formatNumberDecimalSignatures() {
+  return {exec::FunctionSignatureBuilder()
+              .integerVariable("a_precision")
+              .integerVariable("a_scale")
+              .returnType("varchar")
+              .argumentType("DECIMAL(a_precision, a_scale)")
+              .argumentType("integer")
+              .build()};
+}
+
+std::shared_ptr<exec::VectorFunction> createDecimalFormatNumberFunction(
+    const std::string& /*name*/,
+    const std::vector<exec::VectorFunctionArg>& inputArgs,
+    const core::QueryConfig& /*config*/) {
+  BOLT_CHECK_EQ(inputArgs.size(), 2);
+  const auto& decimalType = inputArgs[0].type;
+  if (decimalType->isShortDecimal()) {
+    return std::make_shared<DecimalFormatNumberFunction<int64_t>>(
+        decimalType->asShortDecimal().precision(),
+        decimalType->asShortDecimal().scale());
+  }
+  BOLT_USER_CHECK(decimalType->isLongDecimal(), "Expect decimal input type.");
+  return std::make_shared<DecimalFormatNumberFunction<int128_t>>(
+      decimalType->asLongDecimal().precision(),
+      decimalType->asLongDecimal().scale());
+}
+
+} // namespace
+
 } // namespace bytedance::bolt::functions::sparksql::detail
+
+namespace bytedance::bolt::functions::sparksql {
+
+BOLT_DECLARE_STATEFUL_VECTOR_FUNCTION(
+    udf_decimal_format_number,
+    detail::formatNumberDecimalSignatures(),
+    detail::createDecimalFormatNumberFunction);
+
+} // namespace bytedance::bolt::functions::sparksql

bolt/functions/sparksql/registration/RegisterString.cpp

@@ -70,6 +70,8 @@ void registerStringFunctions(const std::string& prefix) {
   registerFunction<ToTitleFunction, Varchar, Varchar>(
       {prefix + "to_title", prefix + "initcap"});
 
+  BOLT_REGISTER_VECTOR_FUNCTION(
+      udf_decimal_format_number, prefix + "format_number");
   registerFunction<FormatNumberFunction, Varchar, int8_t, int32_t>(
       {prefix + "format_number"});
   registerFunction<FormatNumberFunction, Varchar, int16_t, int32_t>(

bolt/functions/sparksql/tests/FormatNumberTest.cpp

@@ -17,6 +17,7 @@
 #include <limits>
 
 #include "bolt/functions/sparksql/tests/SparkFunctionBaseTest.h"
+#include "bolt/type/HugeInt.h"
 
 namespace bytedance::bolt::functions::sparksql::test {
 
@@ -45,6 +46,28 @@ class FormatNumberTest : public SparkFunctionBaseTest {
   std::optional<std::string> formatFloat(float value, int32_t d) {
     return formatNumber<float>(value, d);
   }
+
+  std::optional<std::string> formatShortDecimal(
+      int64_t unscaledValue,
+      const TypePtr& type,
+      int32_t decimalPlaces) {
+    return evaluateOnce<std::string>(
+        "format_number(c0, c1)",
+        makeRowVector(
+            {makeFlatVector<int64_t>({unscaledValue}, type),
+             makeFlatVector<int32_t>({decimalPlaces})}));
+  }
+
+  std::optional<std::string> formatLongDecimal(
+      int128_t unscaledValue,
+      const TypePtr& type,
+      int32_t decimalPlaces) {
+    return evaluateOnce<std::string>(
+        "format_number(c0, c1)",
+        makeRowVector(
+            {makeFlatVector<int128_t>({unscaledValue}, type),
+             makeFlatVector<int32_t>({decimalPlaces})}));
+  }
 };
 
 TEST_F(FormatNumberTest, integers) {
@@ -117,6 +140,38 @@ TEST_F(FormatNumberTest, binaryTieRounding) {
   EXPECT_EQ(formatDouble(-2.675, 2), "-2.67");
 }
 
+TEST_F(FormatNumberTest, decimals) {
+  EXPECT_EQ(formatShortDecimal(1234567, DECIMAL(10, 2), 2), "12,345.67");
+  EXPECT_EQ(formatShortDecimal(1234567, DECIMAL(10, 2), 1), "12,345.7");
+  EXPECT_EQ(formatShortDecimal(1234567, DECIMAL(10, 2), 0), "12,346");
+  EXPECT_EQ(formatShortDecimal(-1234567, DECIMAL(10, 2), 3), "-12,345.670");
+  EXPECT_EQ(
+      formatShortDecimal(123456789012345678LL, DECIMAL(18, 4), 4),
+      "12,345,678,901,234.5678");
+
+  EXPECT_EQ(
+      formatLongDecimal(
+          HugeInt::parse("12345678901234567890123456789012345678"),
+          DECIMAL(38, 4),
+          4),
+      "1,234,567,890,123,456,789,012,345,678,901,234.5678");
+  EXPECT_EQ(
+      formatLongDecimal(
+          HugeInt::parse("12345678901234567890123456789012345678"),
+          DECIMAL(38, 20),
+          2),
+      "123,456,789,012,345,678.90");
+}
+
+TEST_F(FormatNumberTest, decimalHalfEvenRounding) {
+  EXPECT_EQ(formatShortDecimal(125, DECIMAL(4, 2), 1), "1.2");
+  EXPECT_EQ(formatShortDecimal(135, DECIMAL(4, 2), 1), "1.4");
+  EXPECT_EQ(formatShortDecimal(250, DECIMAL(4, 2), 0), "2");
+  EXPECT_EQ(formatShortDecimal(350, DECIMAL(4, 2), 0), "4");
+  EXPECT_EQ(formatShortDecimal(-250, DECIMAL(4, 2), 0), "-2");
+  EXPECT_EQ(formatShortDecimal(-350, DECIMAL(4, 2), 0), "-4");
+}
+
 TEST_F(FormatNumberTest, negativeDecimalPlaces) {
   // d < 0 returns null per Spark semantics.
   EXPECT_EQ(
@@ -127,6 +182,7 @@ TEST_F(FormatNumberTest, negativeDecimalPlaces) {
       (evaluateOnce<std::string, double, int32_t>(
           "format_number(c0, c1)", 12345.6, -2)),
       std::nullopt);
+  EXPECT_EQ(formatShortDecimal(12345, DECIMAL(10, 2), -1), std::nullopt);
 }
 
 TEST_F(FormatNumberTest, boundaryValues) {