diff --git a/bolt/functions/sparksql/FormatNumber.cpp b/bolt/functions/sparksql/FormatNumber.cpp index aab2789da..c5bfee404 100644 --- a/bolt/functions/sparksql/FormatNumber.cpp +++ b/bolt/functions/sparksql/FormatNumber.cpp @@ -18,9 +18,13 @@ #include #include +#include #include #include #include +#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& 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 +T roundDecimalHalfEven(T value, int32_t fromScale, int32_t toScale) { + if (toScale >= fromScale) { + return value; + } + + const auto scaleFactor = + DecimalUtil::getPowersOfTen(static_cast(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(roundedValue); +} + +template +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 +void formatDecimal( + T value, + int32_t inputScale, + int32_t decimalPlaces, + exec::StringWriter& 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 plainBuffer; + const auto plainSize = DecimalUtil::convertToPlainString( + 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 +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& 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(); + + size_t estimatedBufferSize = 0; + rows.applyToSelected([&](vector_size_t row) { + estimatedBufferSize += estimateFormattedDecimalSize( + decimalValues->valueAt(row), + inputPrecision_, + inputScale_, + decimalPlaces->valueAt(row)); + }); + if (estimatedBufferSize > 0) { + flatResult->getBufferWithSpace(estimatedBufferSize, true); + } + + context.applyToSelectedNoThrow(rows, [&](vector_size_t row) { + const auto places = decimalPlaces->valueAt(row); + if (places < 0) { + result->setNull(row, true); + return; + } + exec::StringWriter writer(flatResult, row); + formatDecimal( + decimalValues->valueAt(row), inputScale_, places, writer); + writer.finalize(); + }); + } + + private: + const int32_t inputPrecision_; + const int32_t inputScale_; +}; + +std::vector> +formatNumberDecimalSignatures() { + return {exec::FunctionSignatureBuilder() + .integerVariable("a_precision") + .integerVariable("a_scale") + .returnType("varchar") + .argumentType("DECIMAL(a_precision, a_scale)") + .argumentType("integer") + .build()}; +} + +std::shared_ptr createDecimalFormatNumberFunction( + const std::string& /*name*/, + const std::vector& inputArgs, + const core::QueryConfig& /*config*/) { + BOLT_CHECK_EQ(inputArgs.size(), 2); + const auto& decimalType = inputArgs[0].type; + if (decimalType->isShortDecimal()) { + return std::make_shared>( + decimalType->asShortDecimal().precision(), + decimalType->asShortDecimal().scale()); + } + BOLT_USER_CHECK(decimalType->isLongDecimal(), "Expect decimal input type."); + return std::make_shared>( + 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 diff --git a/bolt/functions/sparksql/registration/RegisterString.cpp b/bolt/functions/sparksql/registration/RegisterString.cpp index f4a9e9cff..c8a7e4a2e 100644 --- a/bolt/functions/sparksql/registration/RegisterString.cpp +++ b/bolt/functions/sparksql/registration/RegisterString.cpp @@ -70,6 +70,8 @@ void registerStringFunctions(const std::string& prefix) { registerFunction( {prefix + "to_title", prefix + "initcap"}); + BOLT_REGISTER_VECTOR_FUNCTION( + udf_decimal_format_number, prefix + "format_number"); registerFunction( {prefix + "format_number"}); registerFunction( diff --git a/bolt/functions/sparksql/tests/FormatNumberTest.cpp b/bolt/functions/sparksql/tests/FormatNumberTest.cpp index 87dbe95b2..8b1fee3a5 100644 --- a/bolt/functions/sparksql/tests/FormatNumberTest.cpp +++ b/bolt/functions/sparksql/tests/FormatNumberTest.cpp @@ -17,6 +17,7 @@ #include #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 formatFloat(float value, int32_t d) { return formatNumber(value, d); } + + std::optional formatShortDecimal( + int64_t unscaledValue, + const TypePtr& type, + int32_t decimalPlaces) { + return evaluateOnce( + "format_number(c0, c1)", + makeRowVector( + {makeFlatVector({unscaledValue}, type), + makeFlatVector({decimalPlaces})})); + } + + std::optional formatLongDecimal( + int128_t unscaledValue, + const TypePtr& type, + int32_t decimalPlaces) { + return evaluateOnce( + "format_number(c0, c1)", + makeRowVector( + {makeFlatVector({unscaledValue}, type), + makeFlatVector({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( "format_number(c0, c1)", 12345.6, -2)), std::nullopt); + EXPECT_EQ(formatShortDecimal(12345, DECIMAL(10, 2), -1), std::nullopt); } TEST_F(FormatNumberTest, boundaryValues) {