From a6eeee59981d28f8a1084d5db6df221e302f3722 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Fri, 31 Jul 2026 08:33:12 +0000 Subject: [PATCH] feat(cudf): support Spark string-to-decimal try cast --- .../cudf/expression/CMakeLists.txt | 1 + .../cudf/expression/ExpressionEvaluator.cpp | 28 +- .../cudf/expression/StringToDecimal.cu | 431 ++++++++++++++++++ .../cudf/expression/StringToDecimal.h | 37 ++ .../cudf/tests/sparksql/CMakeLists.txt | 12 + .../tests/sparksql/StringToDecimalTest.cpp | 146 ++++++ 6 files changed, 654 insertions(+), 1 deletion(-) create mode 100644 velox/experimental/cudf/expression/StringToDecimal.cu create mode 100644 velox/experimental/cudf/expression/StringToDecimal.h create mode 100644 velox/experimental/cudf/tests/sparksql/StringToDecimalTest.cpp diff --git a/velox/experimental/cudf/expression/CMakeLists.txt b/velox/experimental/cudf/expression/CMakeLists.txt index 856e855c48c..4c98b55734b 100644 --- a/velox/experimental/cudf/expression/CMakeLists.txt +++ b/velox/experimental/cudf/expression/CMakeLists.txt @@ -25,6 +25,7 @@ add_library( PrestoFunctions.cpp prestosql/DatePlusIntervalFunction.cpp SparkFunctions.cpp + StringToDecimal.cu sparksql/DateAddFunction.cpp sparksql/HashFunction.cpp sparksql/SubStringFunction.cpp diff --git a/velox/experimental/cudf/expression/ExpressionEvaluator.cpp b/velox/experimental/cudf/expression/ExpressionEvaluator.cpp index 9a8f5ff66d7..e5628f50ad6 100644 --- a/velox/experimental/cudf/expression/ExpressionEvaluator.cpp +++ b/velox/experimental/cudf/expression/ExpressionEvaluator.cpp @@ -19,6 +19,7 @@ #include "velox/experimental/cudf/expression/DecimalExpressionKernels.h" #include "velox/experimental/cudf/expression/ExpressionEvaluator.h" #include "velox/experimental/cudf/expression/NullMask.h" +#include "velox/experimental/cudf/expression/StringToDecimal.h" #include "velox/common/base/Exceptions.h" #include "velox/common/memory/Memory.h" @@ -536,7 +537,7 @@ bool hasSupportedConstantDecodeCharset( } bool isIntegralNonDecimalVeloxType(const TypePtr& type) { - if (type == nullptr || type->isDate()) { + if (type == nullptr || type->isDate() || type->isDecimal()) { return false; } switch (type->kind()) { @@ -770,6 +771,12 @@ bool isPlainCudfCastSupported( return cudf::is_supported_cast(src, dst); } +bool isStringToDecimalTryCast(const std::shared_ptr& expr) { + return expr->name() == "try_cast" && expr->inputs().size() == 1 && + expr->inputs()[0]->type()->kind() == TypeKind::VARCHAR && + expr->type()->isDecimal(); +} + std::unique_ptr makeAllNullStringColumn( cudf::size_type size, rmm::cuda_stream_view stream, @@ -1371,6 +1378,7 @@ class CastFunction : public CudfFunction { kFloatToString, kStringToInt, kStringToFloat, + kStringToDecimalTryCast, kStringToBool, kStringToDate, kStringToTimestamp, @@ -1384,6 +1392,10 @@ class CastFunction : public CudfFunction { const auto& sourceVeloxType = expr->inputs()[0]->type(); const auto& targetVeloxType = expr->type(); + if (targetVeloxType->isDecimal()) { + targetDecimalPrecision_ = + getDecimalPrecisionScale(*targetVeloxType).first; + } targetCudfType_ = cudf_velox::veloxToCudfDataType(expr->type()); auto sourceType = cudf_velox::veloxToCudfDataType(expr->inputs()[0]->type()); @@ -1414,6 +1426,8 @@ class CastFunction : public CudfFunction { sourceVeloxType->kind() == TypeKind::VARCHAR && isFloatingPointVeloxType(targetVeloxType)) { castMode_ = CastMode::kStringToFloat; + } else if (isStringToDecimalTryCast(expr)) { + castMode_ = CastMode::kStringToDecimalTryCast; } else if (isStringToBooleanVeloxCast(sourceVeloxType, targetVeloxType)) { castMode_ = CastMode::kStringToBool; } else if (isStringToDateVeloxCast(sourceVeloxType, targetVeloxType)) { @@ -1482,6 +1496,14 @@ class CastFunction : public CudfFunction { case CastMode::kStringToFloat: return cudf::strings::to_floats( cudf::strings_column_view(inputCol), targetCudfType_, stream, mr); + case CastMode::kStringToDecimalTryCast: + return tryCastStringToDecimal( + cudf::strings_column_view(inputCol), + targetCudfType_, + targetDecimalPrecision_, + /*stripWhitespace=*/true, + stream, + mr); case CastMode::kStringToBool: { cudf::string_scalar stripChars("", true, stream, mr); auto trimmed = cudf::strings::strip( @@ -1581,6 +1603,7 @@ class CastFunction : public CudfFunction { CastMode castMode_{CastMode::kCudfCast}; bool numericToBoolSourceIsFloating_{false}; bool numericToTimestampSourceIsFloating_{false}; + uint8_t targetDecimalPrecision_{0}; }; class CardinalityFunction : public CudfFunction { @@ -5512,6 +5535,9 @@ bool FunctionExpression::canEvaluate(std::shared_ptr expr) { } if (opName == "cast" || opName == "try_cast") { + if (isStringToDecimalTryCast(expr)) { + return true; + } const auto& srcType = expr->inputs().empty() ? nullptr : expr->inputs()[0]->type(); const auto& dstType = expr->type(); diff --git a/velox/experimental/cudf/expression/StringToDecimal.cu b/velox/experimental/cudf/expression/StringToDecimal.cu new file mode 100644 index 00000000000..3bdf2e3bac0 --- /dev/null +++ b/velox/experimental/cudf/expression/StringToDecimal.cu @@ -0,0 +1,431 @@ +/* + * Copyright (c) 2022-2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Adapted from NVIDIA/spark-rapids-jni cast_string.cu at commit +// 80d402b9e2ac23aedd07cb49b2ea3d8ab929ad54. +#include "velox/experimental/cudf/expression/StringToDecimal.h" + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace facebook::velox::cudf_velox { +namespace { + +constexpr int32_t kThreadsPerBlock = 256; + +__host__ __device__ constexpr bool isWhitespace(char value) { + const auto c = static_cast(value); + return c <= 0x1f || c == ' '; +} + +template +__device__ bool willOverflow(T value, bool adding) { + if constexpr (cuda::std::is_signed_v) { + if (!adding) { + return value < cuda::std::numeric_limits::min() / 10; + } + } + return value > cuda::std::numeric_limits::max() / 10; +} + +template +__device__ bool willOverflow(T lhs, T rhs, bool adding) { + if constexpr (cuda::std::is_signed_v) { + if (!adding) { + return lhs < cuda::std::numeric_limits::min() + rhs; + } + } + return lhs > cuda::std::numeric_limits::max() - rhs; +} + +template +__device__ cuda::std::pair +appendDigit(bool first, T value, T digit, bool positive) { + if (!first) { + if (willOverflow(value, positive)) { + return {false, value}; + } + value *= 10; + } + if (willOverflow(value, digit, positive)) { + return {false, value}; + } + return {true, positive ? value + digit : value - digit}; +} + +template +__device__ cuda::std::optional> +validateAndLocateDecimal(const char* chars, int32_t length, bool strip) { + enum class State { + kDigits, + kExponent, + kDecimalPoint, + kExponentOrSign, + kExponentSign, + kTrailingWhitespace, + kInvalid, + }; + + T exponent = 0; + int32_t index = 0; + bool positive = true; + bool exponentPositive = true; + int32_t decimalLocation = -1; + + if (length == 0) { + return cuda::std::nullopt; + } + if (strip) { + while (index < length && isWhitespace(chars[index])) { + ++index; + } + } + if (index == length) { + return cuda::std::nullopt; + } + if (chars[index] == '-') { + positive = false; + ++index; + } else if (chars[index] == '+') { + ++index; + } + if (index == length) { + return cuda::std::nullopt; + } + + const auto firstDigit = index; + int32_t lastDigit = length; + auto state = State::kDigits; + bool sawMantissaDigit = false; + bool sawExponentDigit = false; + + for (int32_t i = index; i < length; ++i) { + const auto value = chars[i]; + const auto relativeIndex = i - index; + const auto previous = state; + + switch (state) { + case State::kTrailingWhitespace: + if (!isWhitespace(value)) { + state = State::kInvalid; + } + break; + case State::kDecimalPoint: + case State::kDigits: + if (value >= '0' && value <= '9') { + state = State::kDigits; + sawMantissaDigit = true; + } else if (value == '.' && decimalLocation == -1) { + decimalLocation = relativeIndex; + state = State::kDecimalPoint; + } else if ((value == 'e' || value == 'E') && sawMantissaDigit) { + state = State::kExponentOrSign; + } else if (strip && isWhitespace(value) && relativeIndex != 0) { + state = State::kTrailingWhitespace; + } else { + state = State::kInvalid; + } + break; + case State::kExponentOrSign: + if (value == '+') { + state = State::kExponentSign; + } else if (value == '-') { + exponentPositive = false; + state = State::kExponentSign; + } else if (value >= '0' && value <= '9') { + state = State::kExponent; + sawExponentDigit = true; + } else { + state = State::kInvalid; + } + break; + case State::kExponentSign: + case State::kExponent: + if (value >= '0' && value <= '9') { + state = State::kExponent; + sawExponentDigit = true; + } else if (strip && isWhitespace(value) && sawExponentDigit) { + state = State::kTrailingWhitespace; + } else { + state = State::kInvalid; + } + break; + case State::kInvalid: + break; + } + + if (state == State::kInvalid) { + return cuda::std::nullopt; + } + if (previous == State::kDigits && state != State::kDigits && + state != State::kDecimalPoint) { + lastDigit = i; + } + if (state == State::kExponent) { + const T digit = value - '0'; + auto [success, newValue] = + appendDigit(exponent == 0, exponent, digit, exponentPositive); + if (!success) { + return cuda::std::nullopt; + } + exponent = newValue; + } + } + + if (!sawMantissaDigit || state == State::kExponentOrSign || + state == State::kExponentSign) { + return cuda::std::nullopt; + } + if (decimalLocation < 0) { + decimalLocation = lastDigit - firstDigit; + } + decimalLocation += exponent; + return cuda::std::tuple{positive, decimalLocation, firstDigit}; +} + +template +CUDF_KERNEL void stringToDecimalKernel( + T* output, + cudf::bitmask_type* validity, + const char* chars, + const cudf::size_type* offsets, + const cudf::bitmask_type* inputNullMask, + cudf::size_type size, + int32_t scale, + int32_t precision, + bool stripWhitespace) { + auto block = cooperative_groups::this_thread_block(); + auto warp = + cooperative_groups::tiled_partition(block); + const auto row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= size) { + return; + } + + const auto rowStart = offsets[row]; + const auto length = offsets[row + 1] - rowStart; + bool valid = + (inputNullMask == nullptr || cudf::bit_is_set(inputNullMask, row)) && + length > 0; + + auto countSignificantDigits = + [](const char* input, int32_t inputLength, int32_t numDigits) { + int32_t count = 0; + int32_t digitsFound = 0; + for (int32_t i = 0; i < inputLength && digitsFound < numDigits; ++i) { + if (input[i] == 'e' || input[i] == 'E') { + break; + } + if (input[i] != '.') { + ++digitsFound; + if (count != 0 || input[i] != '0') { + ++count; + } + } + } + return count; + }; + + const auto validated = valid + ? validateAndLocateDecimal(chars + rowStart, length, stripWhitespace) + : cuda::std::nullopt; + valid = validated.has_value(); + + if (valid) { + auto [positive, decimalLocation, firstDigit] = *validated; + const auto maxDigitsBeforeDecimal = precision + scale; + const auto significantDigitsBeforeDecimalInString = countSignificantDigits( + chars + rowStart + firstDigit, length - firstDigit, decimalLocation); + const auto lastDigit = decimalLocation - scale; + + int32_t preciseDigits = 0; + int32_t totalDigits = 0; + T value = 0; + bool foundSignificantDigit = false; + int32_t roundingDigits = 0; + + if (lastDigit >= 0) { + for (int32_t i = firstDigit; i < length && valid; ++i) { + const auto c = chars[rowStart + i]; + if (c == '.') { + continue; + } + if (c < '0' || c > '9') { + break; + } + + const T digit = c - '0'; + if (preciseDigits + 1 > precision || totalDigits + 1 > lastDigit) { + if (digit >= 5) { + const auto previousValue = value; + if (willOverflow(value, static_cast(1), positive)) { + valid = false; + break; + } + value += positive ? 1 : -1; + + auto countDigits = [](T number) { + int32_t count = 0; + while (number != 0) { + ++count; + number /= 10; + } + return count; + }; + if (previousValue != 0 && + countDigits(value) > countDigits(previousValue)) { + ++totalDigits; + ++preciseDigits; + ++decimalLocation; + ++roundingDigits; + } + } + break; + } + + ++totalDigits; + if (foundSignificantDigit || totalDigits > decimalLocation || + digit != 0) { + foundSignificantDigit = true; + ++preciseDigits; + } + auto [success, newValue] = + appendDigit(i == firstDigit, value, digit, positive); + if (!success) { + valid = false; + break; + } + value = newValue; + } + } + + const auto significantPrecedingZeros = + decimalLocation < 0 ? -decimalLocation : 0; + const auto zerosToDecimal = cuda::std::max( + 0, + scale > 0 ? decimalLocation - totalDigits - scale + : decimalLocation - totalDigits); + const auto significantDigitsBeforeDecimal = + significantDigitsBeforeDecimalInString + zerosToDecimal + + roundingDigits; + const auto leadingZeros = totalDigits - preciseDigits; + if (maxDigitsBeforeDecimal < decimalLocation - leadingZeros) { + valid = false; + } + + for (int32_t i = 0; i < zerosToDecimal && valid; ++i) { + if (willOverflow(value, positive)) { + valid = false; + break; + } + value *= 10; + ++preciseDigits; + } + + const auto digitsAfterDecimal = preciseDigits - + significantDigitsBeforeDecimal + significantPrecedingZeros; + const auto digitsNeededAfterDecimal = + cuda::std::min(precision - significantDigitsBeforeDecimal, -scale); + for (int32_t i = digitsAfterDecimal; i < digitsNeededAfterDecimal && valid; + ++i) { + if (willOverflow(value, positive)) { + valid = false; + break; + } + value *= 10; + } + if (valid) { + output[row] = value; + } + } + + const auto validityWord = warp.ballot(static_cast(valid)); + if (warp.thread_rank() == 0) { + validity[warp.meta_group_rank() + blockIdx.x * warp.meta_group_size()] = + validityWord; + } +} + +template +std::unique_ptr launchStringToDecimal( + const cudf::strings_column_view& input, + cudf::data_type outputType, + int32_t precision, + bool stripWhitespace, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + using Storage = cudf::device_storage_type_t; + rmm::device_uvector data(input.size(), stream, mr); + const auto words = cudf::bitmask_allocation_size_bytes(input.size()) / + sizeof(cudf::bitmask_type); + rmm::device_uvector nullMask(words, stream, mr); + const dim3 blocks((input.size() + kThreadsPerBlock - 1) / kThreadsPerBlock); + const dim3 threads(kThreadsPerBlock); + stringToDecimalKernel<<>>( + data.data(), + nullMask.data(), + input.chars_begin(stream), + input.offsets().data(), + input.null_mask(), + input.size(), + outputType.scale(), + precision, + stripWhitespace); + CUDF_CUDA_TRY(cudaGetLastError()); + const auto nullCount = + cudf::null_count(nullMask.data(), 0, input.size(), stream); + return std::make_unique( + outputType, input.size(), data.release(), nullMask.release(), nullCount); +} + +} // namespace + +std::unique_ptr tryCastStringToDecimal( + const cudf::strings_column_view& input, + cudf::data_type outputType, + int32_t precision, + bool stripWhitespace, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (input.is_empty()) { + return std::make_unique( + outputType, 0, rmm::device_buffer{}, rmm::device_buffer{}, 0); + } + switch (outputType.id()) { + case cudf::type_id::DECIMAL64: + return launchStringToDecimal( + input, outputType, precision, stripWhitespace, stream, mr); + case cudf::type_id::DECIMAL128: + return launchStringToDecimal( + input, outputType, precision, stripWhitespace, stream, mr); + default: + CUDF_FAIL("String-to-decimal cast requires DECIMAL64 or DECIMAL128"); + } +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/StringToDecimal.h b/velox/experimental/cudf/expression/StringToDecimal.h new file mode 100644 index 00000000000..a2e632fa39e --- /dev/null +++ b/velox/experimental/cudf/expression/StringToDecimal.h @@ -0,0 +1,37 @@ +/* + * Copyright (c) 2022-2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include +#include + +#include + +namespace facebook::velox::cudf_velox { + +/// Spark-compatible, non-ANSI string-to-decimal cast. Invalid rows are null. +std::unique_ptr tryCastStringToDecimal( + const cudf::strings_column_view& input, + cudf::data_type outputType, + int32_t precision, + bool stripWhitespace, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/tests/sparksql/CMakeLists.txt b/velox/experimental/cudf/tests/sparksql/CMakeLists.txt index 3fbf8b4027b..ceea2da2064 100644 --- a/velox/experimental/cudf/tests/sparksql/CMakeLists.txt +++ b/velox/experimental/cudf/tests/sparksql/CMakeLists.txt @@ -40,3 +40,15 @@ velox_add_cudf_test( gflags::gflags TIMEOUT 3000 ) + +velox_add_cudf_test( + NAME velox_cudf_spark_string_to_decimal_test + SOURCES StringToDecimalTest.cpp Main.cpp + LIBS + velox_cudf_exec + velox_exec_test_lib + velox_functions_spark + velox_functions_test_lib + velox_vector_test_lib + gflags::gflags +) diff --git a/velox/experimental/cudf/tests/sparksql/StringToDecimalTest.cpp b/velox/experimental/cudf/tests/sparksql/StringToDecimalTest.cpp new file mode 100644 index 00000000000..5ebd015398e --- /dev/null +++ b/velox/experimental/cudf/tests/sparksql/StringToDecimalTest.cpp @@ -0,0 +1,146 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/expression/SparkFunctions.h" +#include "velox/experimental/cudf/tests/CudfFunctionBaseTest.h" + +#include "velox/functions/sparksql/registration/Register.h" +#include "velox/parse/TypeResolver.h" + +namespace facebook::velox::cudf_velox { +namespace { + +class StringToDecimalTest : public CudfFunctionBaseTest { + protected: + static void SetUpTestCase() { + parse::registerTypeResolver(); + functions::sparksql::registerFunctions(""); + memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{}); + registerCudf(); + registerSparkFunctions(""); + } + + static void TearDownTestCase() { + unregisterFunctions(); + unregisterCudf(); + } + + void assertMatchesCpu( + const std::string& expression, + std::vector> values) { + auto input = + makeRowVector({makeNullableFlatVector(std::move(values))}); + assertExpressionMatchesCpu(expression, input, input->rowType()); + } +}; + +TEST_F(StringToDecimalTest, shortDecimal) { + assertMatchesCpu( + "try_cast(c0 as decimal(12, 2))", + {"9999999999.99", + "15", + "1.5", + "-1.5", + "1.556", + "1.554", + "0000.123", + ".123", + "9.", + "3E2", + "-3E+2", + "3E-2", + "3.5E-2", + "31.523e-2", + " -3E+2 ", + "not-a-number", + "", + std::nullopt}); +} + +TEST_F(StringToDecimalTest, longDecimal) { + assertMatchesCpu( + "try_cast(c0 as decimal(38, 0))", + {"99999999999999999999999999999999999999", + "-99999999999999999999999999999999999999", + "100000000000000000000000000000000000000", + "1.5", + "-1.5", + "1.23e37", + "1.23e67", + " 1.23 ", + "1. 23", + std::nullopt}); +} + +TEST_F(StringToDecimalTest, longDecimalWithScale) { + assertMatchesCpu( + "try_cast(c0 as decimal(20, 4))", + {"112345612.23e-6", + "112345662.23e-6", + "1.23e-6", + "1.26e-3", + "1.23456781e3", + "1.23456789e3", + "1.23456789123451789123456789e9", + "1.23456789123456789123456789e9", + std::nullopt}); + + assertMatchesCpu( + "try_cast(c0 as decimal(38, 38))", + {"0.999999999999999999999999999999999999992", + "0.999999999999999999999999999999999999996", + "111111111111111111.23", + std::nullopt}); +} + +TEST_F(StringToDecimalTest, invalidValuesAndPrecisionOverflow) { + assertMatchesCpu( + "try_cast(c0 as decimal(5, 0))", + {"0", + "80", + "81", + "99999", + "100000", + "-100000", + "+", + ".", + "9e", + "-3E+", + "-3E+2.1", + " ", + std::nullopt}); +} + +TEST_F(StringToDecimalTest, validationPredicate) { + assertMatchesCpu( + "not(rlike(c0, '^[0-9]+$')) OR " + "isnull(try_cast(c0 as decimal(38, 0))) OR " + "decimal_greaterthan(" + "try_cast(c0 as decimal(38, 0)), " + "cast(80 as decimal(38, 0)))", + {"0", + "80", + "81", + "99999999999999999999999999999999999999", + "100000000000000000000000000000000000000", + "-1", + "not-a-number", + std::nullopt}); +} + +} // namespace +} // namespace facebook::velox::cudf_velox