diff --git a/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp b/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp index 58288df27bd..39601b12946 100644 --- a/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp +++ b/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp @@ -23,7 +23,6 @@ #include "velox/experimental/cudf/exec/ToCudf.h" #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" #include "velox/experimental/cudf/expression/ExpressionEvaluator.h" -#include "velox/experimental/cudf/expression/SubfieldFiltersToAst.h" #include "velox/experimental/cudf/vector/CudfVector.h" #include "velox/common/time/Timer.h" @@ -119,13 +118,8 @@ CudfHiveDataSource::CudfHiveDataSource( // readColumnNames_ } - // Build a combined AST for all subfield filters once. This is query-constant - // and doesn't depend on split-specific state. - if (!subfieldFilters_.empty()) { - auto const readerFilterType = getTableRowType(); - subfieldFilterExpr_ = &createAstFromSubfieldFilters( - subfieldFilters_, subfieldTree_, subfieldScalars_, readerFilterType); - } + // Subfield filter AST is built per split in CudfSplitReader after reading + // the Parquet schema so literal types match libcudf storage types. VELOX_CHECK_NOT_NULL(fileHandleFactory_, "No FileHandleFactory present"); @@ -152,7 +146,20 @@ std::unique_ptr CudfHiveDataSource::createCudfSplitReader() { ioStatistics_, ioStats_, useExperimentalCudfReader_, - subfieldFilterExpr_); + makeSubfieldFilterBuildState()); +} + +SubfieldFilterBuildState CudfHiveDataSource::makeSubfieldFilterBuildState() { + if (subfieldFilters_.empty()) { + return {}; + } + return SubfieldFilterBuildState{ + .filters = &subfieldFilters_, + .tree = &subfieldTree_, + .scalars = &subfieldScalars_, + .rowType = getTableRowType(), + .expr = &subfieldFilterExpr_, + }; } void CudfHiveDataSource::convertSplit(std::shared_ptr split) { diff --git a/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.h b/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.h index 0c6f3937813..0082bcc867c 100644 --- a/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.h +++ b/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.h @@ -111,6 +111,9 @@ class CudfHiveDataSource : public DataSource, public NvtxHelper { // Cached combined subfield filter expression owned by 'subfieldTree_'. cudf::ast::expression const* subfieldFilterExpr_{nullptr}; + // Build state passed to CudfSplitReader for per-split AST construction. + SubfieldFilterBuildState makeSubfieldFilterBuildState(); + private: // Construct and cache a RowTypePtr for the table column names and types. const RowTypePtr getTableRowType(); diff --git a/velox/experimental/cudf/connectors/hive/CudfSplitReader.cpp b/velox/experimental/cudf/connectors/hive/CudfSplitReader.cpp index 684b8555baa..20df016157e 100644 --- a/velox/experimental/cudf/connectors/hive/CudfSplitReader.cpp +++ b/velox/experimental/cudf/connectors/hive/CudfSplitReader.cpp @@ -18,6 +18,8 @@ #include "velox/experimental/cudf/connectors/hive/CudfSplitReader.h" #include "velox/experimental/cudf/connectors/hive/CudfSplitReaderHelpers.h" #include "velox/experimental/cudf/exec/GpuResources.h" +#include "velox/experimental/cudf/expression/ParquetSchemaUtils.h" +#include "velox/experimental/cudf/expression/SubfieldFiltersToAst.h" #include "velox/common/caching/CacheTTLController.h" #include "velox/common/time/Timer.h" @@ -74,7 +76,7 @@ CudfSplitReader::CudfSplitReader( const std::shared_ptr& ioStatistics, const std::shared_ptr& ioStats, bool useExperimentalCudfReader, - cudf::ast::expression const* subfieldFilterExpr) + SubfieldFilterBuildState subfieldFilterBuildState) : NvtxHelper( nvtx3::rgb{80, 171, 241}, std::nullopt, @@ -92,7 +94,7 @@ CudfSplitReader::CudfSplitReader( pool_(connectorQueryCtx->memoryPool()), useExperimentalCudfReader_(useExperimentalCudfReader), baseReaderOpts_(pool_), - subfieldFilterExpr_(subfieldFilterExpr) { + subfieldFilterBuildState_(std::move(subfieldFilterBuildState)) { baseReaderOpts_.setDataIoStats(ioStatistics_); baseReaderOpts_.setMetadataIoStats(ioStatistics_); } @@ -223,7 +225,44 @@ void CudfSplitReader::resetSplit() { } cudf::ast::expression const* CudfSplitReader::subfieldFilter() { - return subfieldFilterExpr_; + if (hasSubfieldFilters() && subfieldFilterBuildState_.expr != nullptr && + *subfieldFilterBuildState_.expr == nullptr && dataSource_ != nullptr) { + buildSubfieldFilterAst(); + } + if (subfieldFilterBuildState_.expr == nullptr || + *subfieldFilterBuildState_.expr == nullptr) { + return nullptr; + } + return *subfieldFilterBuildState_.expr; +} + +bool CudfSplitReader::hasSubfieldFilters() const { + return subfieldFilterBuildState_.filters != nullptr && + !subfieldFilterBuildState_.filters->empty(); +} + +void CudfSplitReader::buildSubfieldFilterAst() { + auto& state = subfieldFilterBuildState_; + if (state.filters == nullptr || state.filters->empty()) { + return; + } + VELOX_CHECK_NOT_NULL(state.tree); + VELOX_CHECK_NOT_NULL(state.scalars); + VELOX_CHECK_NOT_NULL(state.expr); + VELOX_CHECK_NOT_NULL(dataSource_); + + *state.tree = cudf::ast::tree{}; + state.scalars->clear(); + + auto sourceInfo = cudf::io::source_info{dataSource_.get()}; + auto metadata = cudf::io::read_parquet_metadata(sourceInfo); + auto parquetColumnTypes = parquetColumnTypesFromMetadata(metadata); + *state.expr = &createAstFromSubfieldFilters( + *state.filters, + *state.tree, + *state.scalars, + state.rowType, + &parquetColumnTypes); } void CudfSplitReader::setupCudfDataSource() { @@ -334,6 +373,7 @@ void CudfSplitReader::setupReaderOptions() { .allow_mismatched_pq_schemas( cudfHiveConfig_->isAllowMismatchedCudfHiveSchemas()) .timestamp_type(cudfHiveConfig_->timestampType()) + .case_sensitive_names(false) .build(); // Set skip_bytes and num_bytes if available @@ -344,14 +384,14 @@ void CudfSplitReader::setupReaderOptions() { readerOptions_.set_num_bytes(split_->size()); } - if (auto* filter = subfieldFilter(); filter != nullptr) { - readerOptions_.set_filter(*filter); - } - - // Set column projection if needed if (readColumnNames_.size()) { readerOptions_.set_column_names(readColumnNames_); } + + buildSubfieldFilterAst(); + if (auto* filter = subfieldFilter(); filter != nullptr) { + readerOptions_.set_filter(*filter); + } } rmm::device_async_resource_ref CudfSplitReader::determineCudfMemoryResource() { diff --git a/velox/experimental/cudf/connectors/hive/CudfSplitReader.h b/velox/experimental/cudf/connectors/hive/CudfSplitReader.h index d114c7b3d5d..8357182d424 100644 --- a/velox/experimental/cudf/connectors/hive/CudfSplitReader.h +++ b/velox/experimental/cudf/connectors/hive/CudfSplitReader.h @@ -27,6 +27,7 @@ #include "velox/connectors/hive/FileHandle.h" #include "velox/connectors/hive/TableHandle.h" #include "velox/dwio/common/Statistics.h" +#include "velox/type/Filter.h" #include "velox/type/Type.h" #include @@ -39,6 +40,16 @@ namespace facebook::velox::cudf_velox::connector::hive { using namespace facebook::velox::connector; +// Mutable state owned by CudfHiveDataSource; the split reader builds the AST +// once per split after reading the Parquet schema. +struct SubfieldFilterBuildState { + const common::SubfieldFilters* filters{nullptr}; + cudf::ast::tree* tree{nullptr}; + std::vector>* scalars{nullptr}; + RowTypePtr rowType; + cudf::ast::expression const** expr{nullptr}; +}; + using CudfParquetReader = cudf::io::chunked_parquet_reader; using CudfParquetReaderPtr = std::unique_ptr; @@ -61,7 +72,7 @@ class CudfSplitReader : public NvtxHelper { const std::shared_ptr& ioStatistics, const std::shared_ptr& ioStats, bool useExperimentalCudfReader, - cudf::ast::expression const* subfieldFilterExpr); + SubfieldFilterBuildState subfieldFilterBuildState = {}); virtual ~CudfSplitReader() = default; @@ -84,6 +95,9 @@ class CudfSplitReader : public NvtxHelper { // Return the subfield filter. virtual cudf::ast::expression const* subfieldFilter(); + // Whether subfield filters were provided for this scan. + bool hasSubfieldFilters() const; + // Determine the output memory resource for the cuDF reader. virtual rmm::device_async_resource_ref determineCudfMemoryResource(); @@ -127,6 +141,9 @@ class CudfSplitReader : public NvtxHelper { std::vector fileMetaData_; private: + // Build the subfield filter AST from the current split's Parquet schema. + void buildSubfieldFilterAst(); + // Setup the cuDF reader options void setupReaderOptions(); @@ -142,7 +159,7 @@ class CudfSplitReader : public NvtxHelper { bool useExperimentalCudfReader_; dwio::common::ReaderOptions baseReaderOpts_; - cudf::ast::expression const* subfieldFilterExpr_; + SubfieldFilterBuildState subfieldFilterBuildState_; struct TotalScanTimeCallbackData { uint64_t startTimeUs; diff --git a/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergDataSource.cpp b/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergDataSource.cpp index 8b39616dda4..f30cac8d92e 100644 --- a/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergDataSource.cpp +++ b/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergDataSource.cpp @@ -78,7 +78,7 @@ CudfIcebergDataSource::createCudfSplitReader() { ioStatistics_, ioStats_, useExperimentalCudfReader_, - subfieldFilterExpr_); + makeSubfieldFilterBuildState()); } } // namespace facebook::velox::cudf_velox::connector::hive::iceberg diff --git a/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergSplitReader.cpp b/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergSplitReader.cpp index 1a48a1f3e50..ae3485ecad8 100644 --- a/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergSplitReader.cpp +++ b/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergSplitReader.cpp @@ -93,7 +93,7 @@ CudfIcebergSplitReader::CudfIcebergSplitReader( const std::shared_ptr& ioStatistics, const std::shared_ptr& ioStats, bool useExperimentalCudfReader, - cudf::ast::expression const* subfieldFilterExpr) + SubfieldFilterBuildState subfieldFilterBuildState) : CudfSplitReader( std::move(split), std::move(tableHandle), @@ -106,7 +106,7 @@ CudfIcebergSplitReader::CudfIcebergSplitReader( ioStatistics, ioStats, useExperimentalCudfReader, - subfieldFilterExpr), + std::move(subfieldFilterBuildState)), icebergSplit_(std::move(icebergSplit)), hiveConfig_(hiveConfig) {} @@ -144,7 +144,7 @@ void CudfIcebergSplitReader::prepareSplit( // Defer subfield filter when it cannot evaluate on the physical parquet // table, or when positional deletes are present. - deferSubfieldFilter_ = CudfSplitReader::subfieldFilter() != nullptr and + deferSubfieldFilter_ = hasSubfieldFilters() and (noColumnsToRead_ or injectedColumns_.size() or // TODO(mh): Drop positional/DV deferral when cudf PR #23077 merges. deletionVectorReader_ or positionalDeleteFileReaders_.size()); diff --git a/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergSplitReader.h b/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergSplitReader.h index f0abe437a94..88f78970433 100644 --- a/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergSplitReader.h +++ b/velox/experimental/cudf/connectors/hive/iceberg/CudfIcebergSplitReader.h @@ -61,7 +61,7 @@ class CudfIcebergSplitReader : public CudfSplitReader { const std::shared_ptr& ioStatistics, const std::shared_ptr& ioStats, bool useExperimentalCudfReader, - cudf::ast::expression const* subfieldFilterExpr); + SubfieldFilterBuildState subfieldFilterBuildState); /// Override to setup delete file readers and column projection. /// @param runtimeStats Reference to the DataSource's runtime statistics, diff --git a/velox/experimental/cudf/exec/CudfGroupby.cpp b/velox/experimental/cudf/exec/CudfGroupby.cpp index e73c5ea2b0f..0f4a2ac89a5 100644 --- a/velox/experimental/cudf/exec/CudfGroupby.cpp +++ b/velox/experimental/cudf/exec/CudfGroupby.cpp @@ -266,8 +266,8 @@ struct GroupbyDecimalSumAggregator : GroupbyAggregator { uint32_t countIdx_{0}; std::unique_ptr decodedSum_; std::unique_ptr decodedCount_; - // Holds the DECIMAL64->DECIMAL128 cast of raw input (kPartial/kSingle), kept - // alive while the groupby request references its view. + // Holds DECIMAL32/DECIMAL64->DECIMAL128 cast of raw input (kPartial/kSingle), + // kept alive while the groupby request references its view. std::unique_ptr castedInput_; }; @@ -342,8 +342,8 @@ struct GroupbyDecimalAvgAggregator : GroupbyAggregator { uint32_t countIdx_{0}; std::unique_ptr decodedSum_; std::unique_ptr decodedCount_; - // Holds the DECIMAL64->DECIMAL128 cast of raw input (kPartial/kSingle), kept - // alive while the groupby request references its view. + // Holds DECIMAL32/DECIMAL64->DECIMAL128 cast of raw input (kPartial/kSingle), + // kept alive while the groupby request references its view. std::unique_ptr castedInput_; }; diff --git a/velox/experimental/cudf/exec/DecimalAggregationHostOps.cpp b/velox/experimental/cudf/exec/DecimalAggregationHostOps.cpp index 73bff493087..b367b2ef792 100644 --- a/velox/experimental/cudf/exec/DecimalAggregationHostOps.cpp +++ b/velox/experimental/cudf/exec/DecimalAggregationHostOps.cpp @@ -40,17 +40,54 @@ cudf::column_view castDecimal64InputToDecimal128( cudf::column_view inputCol, std::unique_ptr& holder, rmm::cuda_stream_view stream) { - if (inputCol.type().id() != cudf::type_id::DECIMAL64) { + auto inputType = inputCol.type().id(); + if (inputType == cudf::type_id::DECIMAL128) { return inputCol; } - holder = cudf::cast( - inputCol, - cudf::data_type{cudf::type_id::DECIMAL128, inputCol.type().scale()}, + if (inputType == cudf::type_id::DECIMAL32) { + holder = cudf::cast( + inputCol, + cudf::data_type{cudf::type_id::DECIMAL64, inputCol.type().scale()}, + stream, + get_temp_mr()); + inputCol = holder->view(); + inputType = inputCol.type().id(); + } + if (inputType == cudf::type_id::DECIMAL64) { + holder = cudf::cast( + inputCol, + cudf::data_type{cudf::type_id::DECIMAL128, inputCol.type().scale()}, + stream, + get_temp_mr()); + return holder->view(); + } + return inputCol; +} + +namespace { + +std::unique_ptr widenDecimalSumForSerialization( + std::unique_ptr sum, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + const auto sumType = sum->type().id(); + if (sumType == cudf::type_id::DECIMAL64 || + sumType == cudf::type_id::DECIMAL128) { + return sum; + } + VELOX_CHECK( + sumType == cudf::type_id::DECIMAL32, + "Unsupported decimal sum column type (type is {})", + cudf::type_to_name(sum->type())); + return cudf::cast( + sum->view(), + cudf::data_type{cudf::type_id::DECIMAL64, sum->type().scale()}, stream, - get_temp_mr()); - return holder->view(); + mr); } +} // namespace + std::unique_ptr castCountColumnToInt64( std::unique_ptr count, rmm::cuda_stream_view stream) { @@ -67,6 +104,7 @@ std::unique_ptr serializeDecimalPartialOrIntermediateState( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { count = castCountColumnToInt64(std::move(count), stream); + sum = widenDecimalSumForSerialization(std::move(sum), stream, mr); return serializeDecimalSumState(sum->view(), count->view(), stream, mr); } diff --git a/velox/experimental/cudf/exec/DecimalAggregationHostOps.h b/velox/experimental/cudf/exec/DecimalAggregationHostOps.h index 778f9120e78..fd57f9d1574 100644 --- a/velox/experimental/cudf/exec/DecimalAggregationHostOps.h +++ b/velox/experimental/cudf/exec/DecimalAggregationHostOps.h @@ -38,14 +38,14 @@ namespace facebook::velox::cudf_velox { void validateIntermediateColumnType(cudf::column_view const& column); /** - * Casts a DECIMAL64 column up to DECIMAL128 (scale preserved) so a subsequent - * SUM accumulates in 128 bits instead of wrapping. Allocates the casted column - * from the temporary memory resource into holder and returns its view. Lifetime - * stays valid only while holder is alive. + * Widens DECIMAL32/DECIMAL64 input to DECIMAL128 (scale preserved) so a + * subsequent SUM accumulates in 128 bits instead of wrapping. Allocates the + * casted column from the temporary memory resource into holder and returns its + * view. Lifetime stays valid only while holder is alive. * - * @param inputCol DECIMAL64 input column. - * @param holder receives ownership of the casted column when inputCol is - * DECIMAL64; unchanged otherwise. + * @param inputCol DECIMAL32, DECIMAL64, or DECIMAL128 input column. + * @param holder receives ownership of the casted column when widening is + * required; unchanged otherwise. * @param stream CUDA stream for device work. * @return view of inputCol or of the column stored in holder. */ diff --git a/velox/experimental/cudf/exec/DecimalAggregationState.cpp b/velox/experimental/cudf/exec/DecimalAggregationState.cpp index 523e23613d8..9abc890a512 100644 --- a/velox/experimental/cudf/exec/DecimalAggregationState.cpp +++ b/velox/experimental/cudf/exec/DecimalAggregationState.cpp @@ -26,9 +26,47 @@ #include #include +#include + #include namespace facebook::velox::cudf_velox { +namespace { + +int64_t readLastStringOffset( + cudf::column_view offsetsView, + cudf::size_type numRows, + rmm::cuda_stream_view stream) { + if (offsetsView.type().id() == cudf::type_id::INT32) { + int32_t last{}; + auto status = cudaMemcpyAsync( + &last, + offsetsView.data() + numRows, + sizeof(int32_t), + cudaMemcpyDeviceToHost, + stream.value()); + VELOX_CHECK_EQ(0, static_cast(status)); + stream.synchronize(); + return last; + } + if (offsetsView.type().id() == cudf::type_id::INT64) { + int64_t last{}; + auto status = cudaMemcpyAsync( + &last, + offsetsView.data() + numRows, + sizeof(int64_t), + cudaMemcpyDeviceToHost, + stream.value()); + VELOX_CHECK_EQ(0, static_cast(status)); + stream.synchronize(); + return last; + } + VELOX_FAIL( + "Decimal sum state requires INT32 or INT64 offsets (offset type is {})", + cudf::type_to_name(offsetsView.type())); +} + +} // namespace DecimalSumStateColumns deserializeDecimalSumState( const cudf::column_view& stateCol, @@ -80,16 +118,24 @@ DecimalSumStateColumns deserializeDecimalSumState( cudf::strings_column_view strings(stateCol); + auto offsetsView = strings.offsets(); auto const payloadSize = strings.chars_size(stream); auto const expectedPayloadSize = - static_cast(numRows) * detail::kDecimalSumStateSize; - VELOX_CHECK( - payloadSize == expectedPayloadSize, + readLastStringOffset(offsetsView, numRows, stream); + VELOX_CHECK_EQ( + payloadSize, + expectedPayloadSize, "Decimal sum state requires payload size {} (got {})", expectedPayloadSize, payloadSize); + VELOX_CHECK_LE( + expectedPayloadSize, + static_cast(numRows) * detail::kDecimalSumStateSize, + "Decimal sum state payload size {} exceeds max {} for {} rows", + expectedPayloadSize, + static_cast(numRows) * detail::kDecimalSumStateSize, + numRows); - auto offsetsView = strings.offsets(); auto charsPtr = reinterpret_cast(strings.chars_begin(stream)); auto sumCol = cudf::make_fixed_width_column( @@ -110,11 +156,6 @@ DecimalSumStateColumns deserializeDecimalSumState( // numRows is guaranteed positive here auto const offsetsType = offsetsView.type().id(); - VELOX_CHECK( - offsetsType == cudf::type_id::INT32 || - offsetsType == cudf::type_id::INT64, - "Decimal sum state requires INT32 or INT64 offsets (offset type is {})", - cudf::type_to_name(offsetsView.type())); detail::unpackDecimalSumState( offsetsType, offsetsView, charsPtr, sumView, countView, numRows, stream); diff --git a/velox/experimental/cudf/expression/AstUtils.h b/velox/experimental/cudf/expression/AstUtils.h index 489d473564d..b42ba088426 100644 --- a/velox/experimental/cudf/expression/AstUtils.h +++ b/velox/experimental/cudf/expression/AstUtils.h @@ -15,8 +15,10 @@ */ #pragma once +#include #include "velox/experimental/cudf/CudfConfig.h" #include "velox/experimental/cudf/CudfNoDefaults.h" +#include "velox/experimental/cudf/expression/ParquetSchemaUtils.h" #include "velox/expression/ConstantExpr.h" #include "velox/type/Timestamp.h" @@ -33,22 +35,96 @@ namespace facebook::velox::cudf_velox { +// libcudf Parquet I/O uses DECIMAL32 when precision <= 9 (see parquet_common.hpp). +constexpr uint8_t kCudfParquetDecimal32MaxPrecision = 9; + +inline bool useCudfDecimal32ForVeloxDecimal(const TypePtr& type) { + if (!type->isShortDecimal()) { + return false; + } + const auto [precision, _] = getDecimalPrecisionScale(*type); + return precision <= kCudfParquetDecimal32MaxPrecision; +} + +inline cudf::type_id veloxDecimalFallbackStorageType(const TypePtr& type) { + if (type->isShortDecimal()) { + return useCudfDecimal32ForVeloxDecimal(type) ? cudf::type_id::DECIMAL32 + : cudf::type_id::DECIMAL64; + } + if (type->isLongDecimal()) { + const auto [precision, _] = getDecimalPrecisionScale(*type); + return precision <= 18 ? cudf::type_id::DECIMAL64 : cudf::type_id::DECIMAL128; + } + VELOX_FAIL("Type is not decimal: {}", type->toString()); +} + +inline bool isParquetIntegerDecimalStorage(cudf::type_id typeId) { + return typeId == cudf::type_id::INT32 || typeId == cudf::type_id::INT64; +} + +inline std::pair decimalStorageTypeAndScale( + const TypePtr& type, + std::optional parquetColumnType = std::nullopt) { + if (parquetColumnType.has_value()) { + const auto parquetId = parquetColumnType->id(); + if (isParquetDecimalType(parquetId)) { + return {parquetId, + numeric::scale_type{parquetColumnType->scale()}}; + } + // Some Parquet files store scaled decimals as plain integers without a + // DECIMAL logical type. Bloom filters require literals to match that + // physical storage exactly. + if (type->isDecimal() && isParquetIntegerDecimalStorage(parquetId)) { + return {parquetId, numeric::scale_type{0}}; + } + } + + const auto [_, veloxScale] = getDecimalPrecisionScale(*type); + return {veloxDecimalFallbackStorageType(type), + numeric::scale_type{-veloxScale}}; +} + +inline cudf::type_id timestampUnitFromParquetColumnType( + std::optional parquetColumnType) { + if (!parquetColumnType.has_value() || + !cudf::is_timestamp(*parquetColumnType)) { + return CudfConfig::getInstance().timestampUnit; + } + return parquetColumnType->id(); +} + template cudf::ast::literal makeLiteralFromScalar( cudf::scalar& scalar, const TypePtr& type) { if constexpr (cudf::is_fixed_width()) { if (type->isDecimal()) { - if (type->kind() == TypeKind::BIGINT) { - using CudfScalarType = cudf::fixed_point_scalar; - return cudf::ast::literal{*static_cast(&scalar)}; + switch (scalar.type().id()) { + case cudf::type_id::DECIMAL32: { + using CudfScalarType = cudf::fixed_point_scalar; + return cudf::ast::literal{*static_cast(&scalar)}; + } + case cudf::type_id::DECIMAL64: { + using CudfScalarType = cudf::fixed_point_scalar; + return cudf::ast::literal{*static_cast(&scalar)}; + } + case cudf::type_id::DECIMAL128: { + using CudfScalarType = cudf::fixed_point_scalar; + return cudf::ast::literal{*static_cast(&scalar)}; + } + case cudf::type_id::INT32: { + using CudfScalarType = cudf::numeric_scalar; + return cudf::ast::literal{*static_cast(&scalar)}; + } + case cudf::type_id::INT64: { + using CudfScalarType = cudf::numeric_scalar; + return cudf::ast::literal{*static_cast(&scalar)}; + } + default: + VELOX_FAIL( + "Unsupported scalar type for decimal literal: {}", + static_cast(scalar.type().id())); } - if (type->kind() == TypeKind::HUGEINT) { - using CudfScalarType = cudf::fixed_point_scalar; - return cudf::ast::literal{*static_cast(&scalar)}; - } - VELOX_UNREACHABLE( - "Invalid Decimal Type (bad TypeKind: {})", type->kind()); } else if (type->isIntervalDayTime()) { using CudfDurationType = cudf::duration_ms; if constexpr (std::is_same_v) { @@ -92,7 +168,8 @@ std::unique_ptr makeScalarFromValue( const TypePtr& type, T value, bool isNull, - std::optional toType = std::nullopt) { + std::optional toType = std::nullopt, + std::optional parquetColumnType = std::nullopt) { auto stream = cudf::get_default_stream(cudf::allow_default_stream); auto mr = get_temp_mr(); @@ -104,7 +181,7 @@ std::unique_ptr makeScalarFromValue( // cheap (one-time cost per scalar) and guarantees the memory is // available on every stream. if constexpr (std::is_same_v) { - auto unit = CudfConfig::getInstance().timestampUnit; + auto unit = timestampUnitFromParquetColumnType(parquetColumnType); if (unit == cudf::type_id::TIMESTAMP_MICROSECONDS) { using CudfTimestampType = cudf::timestamp_us; auto micros = isNull ? 0 : value.toMicros(); @@ -140,30 +217,44 @@ std::unique_ptr makeScalarFromValue( if (type->isDecimal()) { // Velox DECIMAL scale is positive for fractional digits // cuDF scale is negative for fractional digits - // @TODO check the bigger picture here! - if (type->kind() == TypeKind::BIGINT) { - auto const decimalType = - std::dynamic_pointer_cast(type); - VELOX_CHECK(decimalType, "Invalid Decimal Type (failed dynamic_cast)"); - auto const cudfScale = numeric::scale_type{-decimalType->scale()}; + const auto [storageType, cudfScale] = + decimalStorageTypeAndScale(type, parquetColumnType); + if (storageType == cudf::type_id::DECIMAL32) { + using CudfDecimalType = cudf::fixed_point_scalar; + auto scalar = std::make_unique( + static_cast(value), cudfScale, !isNull, stream, mr); + stream.synchronize(); + return scalar; + } + if (storageType == cudf::type_id::DECIMAL64) { using CudfDecimalType = cudf::fixed_point_scalar; auto scalar = std::make_unique( - value, cudfScale, !isNull, stream, mr); + static_cast(value), cudfScale, !isNull, stream, mr); stream.synchronize(); return scalar; - } else if (type->kind() == TypeKind::HUGEINT) { - auto const decimalType = - std::dynamic_pointer_cast(type); - VELOX_CHECK(decimalType, "Invalid Decimal Type (failed dynamic_cast)"); - auto const cudfScale = numeric::scale_type{-decimalType->scale()}; + } + if (storageType == cudf::type_id::DECIMAL128) { using CudfDecimalType = cudf::fixed_point_scalar; auto scalar = std::make_unique( - value, cudfScale, !isNull, stream, mr); + static_cast(value), cudfScale, !isNull, stream, mr); stream.synchronize(); return scalar; } - VELOX_UNREACHABLE( - "Invalid Decimal Type (bad TypeKind: {})", type->kind()); + if (storageType == cudf::type_id::INT64) { + auto scalar = std::make_unique>( + static_cast(value), !isNull, stream, mr); + stream.synchronize(); + return scalar; + } + if (storageType == cudf::type_id::INT32) { + auto scalar = std::make_unique>( + static_cast(value), !isNull, stream, mr); + stream.synchronize(); + return scalar; + } + VELOX_FAIL( + "Unsupported decimal storage type {}", + static_cast(storageType)); } else if (type->isIntervalYearMonth()) { VELOX_FAIL("Interval year month not supported"); } else if (type->isIntervalDayTime()) { @@ -234,16 +325,19 @@ cudf::ast::literal makeScalarAndLiteral( const TypePtr& type, const variant& var, bool isNull, - std::vector>& scalars) { + std::vector>& scalars, + std::optional parquetColumnType = std::nullopt) { using T = typename TypeTraits::NativeType; if constexpr (cudf::is_fixed_width() || kind == TypeKind::VARCHAR) { if (isNull) { - auto scalar = makeScalarFromValue(type, T{}, true); + auto scalar = makeScalarFromValue( + type, T{}, true, std::nullopt, parquetColumnType); scalars.emplace_back(std::move(scalar)); return makeLiteralFromScalar(*(scalars.back()), type); } auto value = var.value(); - auto scalar = makeScalarFromValue(type, value, false); + auto scalar = makeScalarFromValue( + type, value, false, std::nullopt, parquetColumnType); scalars.emplace_back(std::move(scalar)); return makeLiteralFromScalar(*(scalars.back()), type); } diff --git a/velox/experimental/cudf/expression/DecimalExpressionKernels.cpp b/velox/experimental/cudf/expression/DecimalExpressionKernels.cpp index 9cb4c0077b4..f40069856f2 100644 --- a/velox/experimental/cudf/expression/DecimalExpressionKernels.cpp +++ b/velox/experimental/cudf/expression/DecimalExpressionKernels.cpp @@ -32,6 +32,11 @@ namespace { __int128_t getDecimalScalarValue( const cudf::scalar& s, rmm::cuda_stream_view stream) { + if (s.type().id() == cudf::type_id::DECIMAL32) { + auto const& dec = + static_cast const&>(s); + return static_cast<__int128_t>(static_cast(dec.value(stream))); + } if (s.type().id() == cudf::type_id::DECIMAL64) { auto const& dec = static_cast const&>(s); @@ -55,9 +60,12 @@ std::unique_ptr makeAllNullDecimalColumn( void checkDecimalDivideTypes(cudf::type_id inType, cudf::type_id outType) { VELOX_CHECK( - inType == cudf::type_id::DECIMAL64 || inType == cudf::type_id::DECIMAL128, + inType == cudf::type_id::DECIMAL32 || + inType == cudf::type_id::DECIMAL64 || + inType == cudf::type_id::DECIMAL128, "Unsupported input type for decimal divide"); - if (inType == cudf::type_id::DECIMAL64) { + if (inType == cudf::type_id::DECIMAL32 || + inType == cudf::type_id::DECIMAL64) { VELOX_CHECK( outType == cudf::type_id::DECIMAL64 || outType == cudf::type_id::DECIMAL128, diff --git a/velox/experimental/cudf/expression/ExpressionEvaluator.cpp b/velox/experimental/cudf/expression/ExpressionEvaluator.cpp index 9f45b4915c3..bc2341fb68f 100644 --- a/velox/experimental/cudf/expression/ExpressionEvaluator.cpp +++ b/velox/experimental/cudf/expression/ExpressionEvaluator.cpp @@ -95,6 +95,11 @@ bool decimalScalarIsZero( if (!scalar.is_valid(stream)) { return false; } + if (scalar.type().id() == cudf::type_id::DECIMAL32) { + auto const& dec = + static_cast const&>(scalar); + return dec.value(stream) == 0; + } if (scalar.type().id() == cudf::type_id::DECIMAL64) { auto const& dec = static_cast const&>( @@ -119,9 +124,15 @@ bool hasDecimalZero( } std::unique_ptr zero; auto scale = numeric::scale_type{col.type().scale()}; - if (col.type().id() == cudf::type_id::DECIMAL64) { - zero = - cudf::make_fixed_point_scalar(0, scale, stream, mr); + if (col.type().id() == cudf::type_id::DECIMAL32 || + col.type().id() == cudf::type_id::DECIMAL64) { + if (col.type().id() == cudf::type_id::DECIMAL32) { + zero = cudf::make_fixed_point_scalar( + 0, scale, stream, mr); + } else { + zero = + cudf::make_fixed_point_scalar(0, scale, stream, mr); + } } else if (col.type().id() == cudf::type_id::DECIMAL128) { zero = cudf::make_fixed_point_scalar( 0, scale, stream, mr); @@ -154,20 +165,24 @@ std::unique_ptr castDecimalScalar( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { if (!src.is_valid(stream)) { - VELOX_CHECK( - targetType.id() == cudf::type_id::DECIMAL64 || - targetType.id() == cudf::type_id::DECIMAL128, - "castDecimalScalar: target must be DECIMAL64 or DECIMAL128"); if (targetType.id() == cudf::type_id::DECIMAL128) { return std::make_unique>( 0, numeric::scale_type{targetType.scale()}, false, stream, mr); } + if (targetType.id() == cudf::type_id::DECIMAL32) { + return std::make_unique>( + 0, numeric::scale_type{targetType.scale()}, false, stream, mr); + } return std::make_unique>( 0, numeric::scale_type{targetType.scale()}, false, stream, mr); } __int128_t rep; - if (src.type().id() == cudf::type_id::DECIMAL64) { + if (src.type().id() == cudf::type_id::DECIMAL32) { + auto const& dec = + static_cast const&>(src); + rep = static_cast(dec.value(stream)); + } else if (src.type().id() == cudf::type_id::DECIMAL64) { auto const& dec = static_cast const&>(src); rep = static_cast(dec.value(stream)); @@ -192,6 +207,13 @@ std::unique_ptr castDecimalScalar( return cudf::make_fixed_point_scalar( rep, numeric::scale_type{targetType.scale()}, stream, mr); } + if (targetType.id() == cudf::type_id::DECIMAL32) { + return cudf::make_fixed_point_scalar( + static_cast(rep), + numeric::scale_type{targetType.scale()}, + stream, + mr); + } return cudf::make_fixed_point_scalar( static_cast(rep), numeric::scale_type{targetType.scale()}, @@ -199,6 +221,124 @@ std::unique_ptr castDecimalScalar( mr); } +int32_t decimalTypeRank(cudf::type_id id) { + switch (id) { + case cudf::type_id::DECIMAL32: + return 1; + case cudf::type_id::DECIMAL64: + return 2; + case cudf::type_id::DECIMAL128: + return 3; + default: + return 0; + } +} + +cudf::type_id widerDecimalTypeId(cudf::type_id lhs, cudf::type_id rhs) { + return decimalTypeRank(lhs) >= decimalTypeRank(rhs) ? lhs : rhs; +} + +std::optional widerDecimalTypeIfMismatch( + const cudf::data_type& lhs, + const cudf::data_type& rhs) { + if (!cudf::is_fixed_point(lhs) || !cudf::is_fixed_point(rhs)) { + return std::nullopt; + } + if (lhs == rhs) { + return std::nullopt; + } + VELOX_CHECK_EQ( + lhs.scale(), + rhs.scale(), + "Decimal operands must share scale ({} vs {})", + lhs.scale(), + rhs.scale()); + return cudf::data_type{widerDecimalTypeId(lhs.id(), rhs.id()), lhs.scale()}; +} + +cudf::column_view castDecimalColumnIfNeeded( + cudf::column_view view, + const cudf::data_type& targetType, + std::unique_ptr& holder, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (view.type() == targetType) { + return view; + } + holder = cudf::cast(view, targetType, stream, mr); + return holder->view(); +} + +const cudf::scalar* castDecimalScalarIfNeeded( + const cudf::scalar& scalar, + const cudf::data_type& targetType, + std::unique_ptr& holder, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (scalar.type() == targetType) { + return &scalar; + } + holder = castDecimalScalar(scalar, targetType, stream, mr); + return holder.get(); +} + +cudf::data_type commonDecimalOperandType( + const cudf::data_type& lhs, + const cudf::data_type& rhs, + const cudf::data_type& resultType) { + auto lhsScale = -lhs.scale(); + auto rhsScale = -rhs.scale(); + auto targetScale = lhsScale > rhsScale ? lhsScale : rhsScale; + auto targetTypeId = widerDecimalTypeId(lhs.id(), rhs.id()); + targetTypeId = widerDecimalTypeId(targetTypeId, resultType.id()); + return cudf::data_type{targetTypeId, numeric::scale_type{-targetScale}}; +} + +void alignDecimalColumnOperands( + cudf::column_view& lhsView, + cudf::column_view& rhsView, + const cudf::data_type& targetType, + std::unique_ptr& lhsCast, + std::unique_ptr& rhsCast, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + lhsView = castDecimalColumnIfNeeded(lhsView, targetType, lhsCast, stream, mr); + rhsView = castDecimalColumnIfNeeded(rhsView, targetType, rhsCast, stream, mr); +} + +std::unique_ptr decimalComparison( + cudf::column_view value, + cudf::column_view boundCol, + const cudf::scalar* boundScalar, + cudf::binary_operator op, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + static constexpr cudf::data_type kBoolType{cudf::type_id::BOOL8}; + std::unique_ptr valueCast; + std::unique_ptr boundCast; + std::unique_ptr boundScalarCast; + if (boundScalar != nullptr) { + if (cudf::is_fixed_point(value.type()) && + cudf::is_fixed_point(boundScalar->type())) { + auto targetType = commonDecimalOperandType( + value.type(), boundScalar->type(), value.type()); + value = castDecimalColumnIfNeeded(value, targetType, valueCast, stream, mr); + boundScalar = castDecimalScalarIfNeeded( + *boundScalar, targetType, boundScalarCast, stream, mr); + } + return cudf::binary_operation( + value, *boundScalar, op, kBoolType, stream, mr); + } + if (cudf::is_fixed_point(value.type()) && + cudf::is_fixed_point(boundCol.type())) { + auto targetType = + commonDecimalOperandType(value.type(), boundCol.type(), value.type()); + alignDecimalColumnOperands( + value, boundCol, targetType, valueCast, boundCast, stream, mr); + } + return cudf::binary_operation(value, boundCol, op, kBoolType, stream, mr); +} + struct CudfExpressionEvaluatorEntry { int priority; CudfExpressionEvaluatorCanEvaluate canEvaluate; @@ -561,19 +701,12 @@ class BinaryFunction : public CudfFunction { auto rhsView = asView(inputColumns[1]); std::unique_ptr lhsCast; std::unique_ptr rhsCast; - if (type_.id() == cudf::type_id::DECIMAL128) { - if (lhsView.type().id() == cudf::type_id::DECIMAL64) { - auto castType = cudf::data_type{ - cudf::type_id::DECIMAL128, lhsView.type().scale()}; - lhsCast = cudf::cast(lhsView, castType, stream, mr); - lhsView = lhsCast->view(); - } - if (rhsView.type().id() == cudf::type_id::DECIMAL64) { - auto castType = cudf::data_type{ - cudf::type_id::DECIMAL128, rhsView.type().scale()}; - rhsCast = cudf::cast(rhsView, castType, stream, mr); - rhsView = rhsCast->view(); - } + if (cudf::is_fixed_point(lhsView.type()) && + cudf::is_fixed_point(rhsView.type())) { + auto targetType = commonDecimalOperandType( + lhsView.type(), rhsView.type(), type_); + alignDecimalColumnOperands( + lhsView, rhsView, targetType, lhsCast, rhsCast, stream, mr); } auto lhsScale = -lhsView.type().scale(); auto rhsScale = -rhsView.type().scale(); @@ -588,10 +721,8 @@ class BinaryFunction : public CudfFunction { auto lhsScale = -lhsView.type().scale(); auto rhsScale = -rhsView.type().scale(); auto targetScale = lhsScale > rhsScale ? lhsScale : rhsScale; - auto targetTypeId = (lhsView.type().id() == cudf::type_id::DECIMAL128 || - rhsView.type().id() == cudf::type_id::DECIMAL128) - ? cudf::type_id::DECIMAL128 - : cudf::type_id::DECIMAL64; + auto targetTypeId = widerDecimalTypeId( + lhsView.type().id(), rhsView.type().id()); auto targetType = cudf::data_type{targetTypeId, numeric::scale_type{-targetScale}}; std::unique_ptr lhsCast; @@ -628,19 +759,12 @@ class BinaryFunction : public CudfFunction { if (op_ == cudf::binary_operator::MUL) { std::unique_ptr lhsCast; std::unique_ptr rhsCast; - if (type_.id() == cudf::type_id::DECIMAL128) { - if (lhsView.type().id() == cudf::type_id::DECIMAL64) { - auto castType = cudf::data_type{ - cudf::type_id::DECIMAL128, lhsView.type().scale()}; - lhsCast = cudf::cast(lhsView, castType, stream, mr); - lhsView = lhsCast->view(); - } - if (rhsView.type().id() == cudf::type_id::DECIMAL64) { - auto castType = cudf::data_type{ - cudf::type_id::DECIMAL128, rhsView.type().scale()}; - rhsCast = cudf::cast(rhsView, castType, stream, mr); - rhsView = rhsCast->view(); - } + if (cudf::is_fixed_point(lhsView.type()) && + cudf::is_fixed_point(rhsView.type())) { + auto targetType = commonDecimalOperandType( + lhsView.type(), rhsView.type(), type_); + alignDecimalColumnOperands( + lhsView, rhsView, targetType, lhsCast, rhsCast, stream, mr); } // @TODO Check for divide-by-zero as in the DECIMAL case above? return cudf::binary_operation( @@ -655,11 +779,23 @@ class BinaryFunction : public CudfFunction { VELOX_USER_FAIL("Division by zero"); } auto lhsView = asView(inputColumns[0]); + std::unique_ptr lhsCast; + std::unique_ptr rhsScalarCast; + const cudf::scalar* rhs = right_.get(); + if (cudf::is_fixed_point(lhsView.type()) && + cudf::is_fixed_point(rhs->type())) { + auto targetType = + commonDecimalOperandType(lhsView.type(), rhs->type(), type_); + lhsView = + castDecimalColumnIfNeeded(lhsView, targetType, lhsCast, stream, mr); + rhs = castDecimalScalarIfNeeded( + *rhs, targetType, rhsScalarCast, stream, mr); + } auto lhsScale = -lhsView.type().scale(); - auto rhsScale = -right_->type().scale(); + auto rhsScale = -rhs->type().scale(); auto outScale = -type_.scale(); auto aRescale = outScale - lhsScale + rhsScale; - return decimalDivide(lhsView, *right_, type_, aRescale, stream, mr); + return decimalDivide(lhsView, *rhs, type_, aRescale, stream, mr); } auto lhsView = asView(inputColumns[0]); if (isComparisonOp(op_) && cudf::is_fixed_point(lhsView.type()) && @@ -667,10 +803,8 @@ class BinaryFunction : public CudfFunction { auto lhsScale = -lhsView.type().scale(); auto rhsScale = -right_->type().scale(); auto targetScale = lhsScale > rhsScale ? lhsScale : rhsScale; - auto targetTypeId = (lhsView.type().id() == cudf::type_id::DECIMAL128 || - right_->type().id() == cudf::type_id::DECIMAL128) - ? cudf::type_id::DECIMAL128 - : cudf::type_id::DECIMAL64; + auto targetTypeId = widerDecimalTypeId( + lhsView.type().id(), right_->type().id()); auto targetType = cudf::data_type{targetTypeId, numeric::scale_type{-targetScale}}; std::unique_ptr lhsCast; @@ -704,27 +838,19 @@ class BinaryFunction : public CudfFunction { } if (op_ == cudf::binary_operator::MUL) { std::unique_ptr lhsCast; - std::unique_ptr rhsScalar; - if (type_.id() == cudf::type_id::DECIMAL128) { - if (lhsView.type().id() == cudf::type_id::DECIMAL64) { - auto castType = cudf::data_type{ - cudf::type_id::DECIMAL128, lhsView.type().scale()}; - lhsCast = cudf::cast(lhsView, castType, stream, mr); - lhsView = lhsCast->view(); - } - if (right_->type().id() == cudf::type_id::DECIMAL64) { - auto castType = cudf::data_type{ - cudf::type_id::DECIMAL128, right_->type().scale()}; - rhsScalar = castDecimalScalar(*right_, castType, stream, mr); - } + std::unique_ptr rhsScalarCast; + const cudf::scalar* rhs = right_.get(); + if (cudf::is_fixed_point(lhsView.type()) && + cudf::is_fixed_point(rhs->type())) { + auto targetType = + commonDecimalOperandType(lhsView.type(), rhs->type(), type_); + lhsView = castDecimalColumnIfNeeded( + lhsView, targetType, lhsCast, stream, mr); + rhs = castDecimalScalarIfNeeded( + *rhs, targetType, rhsScalarCast, stream, mr); } return cudf::binary_operation( - lhsView, - rhsScalar ? *rhsScalar : *right_, - op_, - type_, - stream, - mr); + lhsView, *rhs, op_, type_, stream, mr); } } return cudf::binary_operation( @@ -732,11 +858,23 @@ class BinaryFunction : public CudfFunction { } if (op_ == cudf::binary_operator::DIV && cudf::is_fixed_point(type_)) { auto rhsView = asView(inputColumns[0]); - auto lhsScale = -left_->type().scale(); + std::unique_ptr rhsCast; + std::unique_ptr lhsScalarCast; + const cudf::scalar* lhs = left_.get(); + if (cudf::is_fixed_point(lhs->type()) && + cudf::is_fixed_point(rhsView.type())) { + auto targetType = + commonDecimalOperandType(lhs->type(), rhsView.type(), type_); + lhs = castDecimalScalarIfNeeded( + *lhs, targetType, lhsScalarCast, stream, mr); + rhsView = + castDecimalColumnIfNeeded(rhsView, targetType, rhsCast, stream, mr); + } + auto lhsScale = -lhs->type().scale(); auto rhsScale = -rhsView.type().scale(); auto outScale = -type_.scale(); auto aRescale = outScale - lhsScale + rhsScale; - return decimalDivide(*left_, rhsView, type_, aRescale, stream, mr); + return decimalDivide(*lhs, rhsView, type_, aRescale, stream, mr); } auto rhsView = asView(inputColumns[0]); if (isComparisonOp(op_) && cudf::is_fixed_point(left_->type()) && @@ -744,10 +882,8 @@ class BinaryFunction : public CudfFunction { auto lhsScale = -left_->type().scale(); auto rhsScale = -rhsView.type().scale(); auto targetScale = lhsScale > rhsScale ? lhsScale : rhsScale; - auto targetTypeId = (left_->type().id() == cudf::type_id::DECIMAL128 || - rhsView.type().id() == cudf::type_id::DECIMAL128) - ? cudf::type_id::DECIMAL128 - : cudf::type_id::DECIMAL64; + auto targetTypeId = widerDecimalTypeId( + left_->type().id(), rhsView.type().id()); auto targetType = cudf::data_type{targetTypeId, numeric::scale_type{-targetScale}}; std::unique_ptr rhsCast; @@ -780,22 +916,18 @@ class BinaryFunction : public CudfFunction { } if (op_ == cudf::binary_operator::MUL) { std::unique_ptr rhsCast; - std::unique_ptr lhsScalar; - if (type_.id() == cudf::type_id::DECIMAL128) { - if (rhsView.type().id() == cudf::type_id::DECIMAL64) { - auto castType = cudf::data_type{ - cudf::type_id::DECIMAL128, rhsView.type().scale()}; - rhsCast = cudf::cast(rhsView, castType, stream, mr); - rhsView = rhsCast->view(); - } - if (left_->type().id() == cudf::type_id::DECIMAL64) { - auto castType = cudf::data_type{ - cudf::type_id::DECIMAL128, left_->type().scale()}; - lhsScalar = castDecimalScalar(*left_, castType, stream, mr); - } + std::unique_ptr lhsScalarCast; + const cudf::scalar* lhs = left_.get(); + if (cudf::is_fixed_point(lhs->type()) && + cudf::is_fixed_point(rhsView.type())) { + auto targetType = + commonDecimalOperandType(lhs->type(), rhsView.type(), type_); + lhs = castDecimalScalarIfNeeded( + *lhs, targetType, lhsScalarCast, stream, mr); + rhsView = + castDecimalColumnIfNeeded(rhsView, targetType, rhsCast, stream, mr); } - return cudf::binary_operation( - lhsScalar ? *lhsScalar : *left_, rhsView, op_, type_, stream, mr); + return cudf::binary_operation(*lhs, rhsView, op_, type_, stream, mr); } } return cudf::binary_operation(*left_, rhsView, op_, type_, stream, mr); @@ -1006,38 +1138,40 @@ class BetweenFunction : public CudfFunction { rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { // return (value >= min) && (value <= max) - std::unique_ptr geResultColumn, leResultColumn; + auto value = asView(inputColumns[0]); + std::unique_ptr geResultColumn; + std::unique_ptr leResultColumn; if (minLiteral_) { - geResultColumn = cudf::binary_operation( - asView(inputColumns[0]), - *minLiteral_, + geResultColumn = decimalComparison( + value, + cudf::column_view{}, + minLiteral_.get(), cudf::binary_operator::GREATER_EQUAL, - kBoolType, stream, mr); } else { - geResultColumn = cudf::binary_operation( - asView(inputColumns[0]), + geResultColumn = decimalComparison( + value, asView(inputColumns[1]), + nullptr, cudf::binary_operator::GREATER_EQUAL, - kBoolType, stream, mr); } if (maxLiteral_) { - leResultColumn = cudf::binary_operation( - asView(inputColumns[0]), - *maxLiteral_, + leResultColumn = decimalComparison( + value, + cudf::column_view{}, + maxLiteral_.get(), cudf::binary_operator::LESS_EQUAL, - kBoolType, stream, mr); } else { - leResultColumn = cudf::binary_operation( - asView(inputColumns[0]), + leResultColumn = decimalComparison( + value, asView(inputColumns[2]), + nullptr, cudf::binary_operator::LESS_EQUAL, - kBoolType, stream, mr); } @@ -1173,27 +1307,64 @@ class SwitchFunction : public CudfFunction { std::vector& inputColumns, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { + auto mask = asView(inputColumns[0]); + std::unique_ptr lhsCast; + std::unique_ptr rhsCast; + std::unique_ptr lhsScalarCast; + std::unique_ptr rhsScalarCast; + if (left_ == nullptr && right_ == nullptr) { + auto lhs = asView(inputColumns[1]); + auto rhs = asView(inputColumns[2]); + if (auto wider = widerDecimalTypeIfMismatch(lhs.type(), rhs.type())) { + lhs = castDecimalColumnIfNeeded(lhs, *wider, lhsCast, stream, mr); + rhs = castDecimalColumnIfNeeded(rhs, *wider, rhsCast, stream, mr); + } + return cudf::copy_if_else(lhs, rhs, mask, stream, mr); + } + if (left_ == nullptr) { + auto lhs = asView(inputColumns[1]); + const auto& rhs = *right_; + if (auto wider = widerDecimalTypeIfMismatch(lhs.type(), rhs.type())) { + lhs = castDecimalColumnIfNeeded(lhs, *wider, lhsCast, stream, mr); + return cudf::copy_if_else( + lhs, + *castDecimalScalarIfNeeded( + rhs, *wider, rhsScalarCast, stream, mr), + mask, + stream, + mr); + } + return cudf::copy_if_else(lhs, rhs, mask, stream, mr); + } + if (right_ == nullptr) { + const auto& lhs = *left_; + auto rhs = asView(inputColumns[1]); + if (auto wider = widerDecimalTypeIfMismatch(lhs.type(), rhs.type())) { + rhs = castDecimalColumnIfNeeded(rhs, *wider, rhsCast, stream, mr); + return cudf::copy_if_else( + *castDecimalScalarIfNeeded( + lhs, *wider, lhsScalarCast, stream, mr), + rhs, + mask, + stream, + mr); + } + return cudf::copy_if_else(lhs, rhs, mask, stream, mr); + } + const auto& lhs = *left_; + const auto& rhs = *right_; + if (auto wider = widerDecimalTypeIfMismatch(lhs.type(), rhs.type())) { return cudf::copy_if_else( - asView(inputColumns[1]), - asView(inputColumns[2]), - asView(inputColumns[0]), - stream, - mr); - } else if (left_ == nullptr) { - return cudf::copy_if_else( - asView(inputColumns[1]), - *right_, - asView(inputColumns[0]), + *castDecimalScalarIfNeeded( + lhs, *wider, lhsScalarCast, stream, mr), + *castDecimalScalarIfNeeded( + rhs, *wider, rhsScalarCast, stream, mr), + mask, stream, mr); - } else if (right_ == nullptr) { - return cudf::copy_if_else( - *left_, asView(inputColumns[1]), asView(inputColumns[0]), stream, mr); } - // right != null and left != null - return cudf::copy_if_else( - *left_, *right_, asView(inputColumns[0]), stream, mr); + return cudf::copy_if_else(lhs, rhs, mask, stream, mr); } private: @@ -1256,14 +1427,34 @@ class CoalesceFunction : public CudfFunction { !inputColumns.empty(), "coalesce requires at least one non-literal input"); ColumnOrView result = asView(inputColumns[0]); + std::unique_ptr lhsCast; + std::unique_ptr rhsCast; + std::unique_ptr literalCast; size_t stop = std::min(numColumnsBeforeLiteral_, inputColumns.size()); for (size_t i = 1; i < stop && asView(result).has_nulls(); ++i) { - result = cudf::replace_nulls( - asView(result), asView(inputColumns[i]), stream, mr); + auto lhs = asView(result); + auto rhs = asView(inputColumns[i]); + if (auto wider = widerDecimalTypeIfMismatch(lhs.type(), rhs.type())) { + lhs = castDecimalColumnIfNeeded(lhs, *wider, lhsCast, stream, mr); + rhs = castDecimalColumnIfNeeded(rhs, *wider, rhsCast, stream, mr); + } + result = cudf::replace_nulls(lhs, rhs, stream, mr); } if (literalScalar_ && asView(result).has_nulls()) { - result = cudf::replace_nulls(asView(result), *literalScalar_, stream, mr); + auto lhs = asView(result); + const auto& rhs = *literalScalar_; + if (auto wider = widerDecimalTypeIfMismatch(lhs.type(), rhs.type())) { + lhs = castDecimalColumnIfNeeded(lhs, *wider, lhsCast, stream, mr); + result = cudf::replace_nulls( + lhs, + *castDecimalScalarIfNeeded( + rhs, *wider, literalCast, stream, mr), + stream, + mr); + } else { + result = cudf::replace_nulls(lhs, rhs, stream, mr); + } } return result; diff --git a/velox/experimental/cudf/expression/ParquetSchemaUtils.h b/velox/experimental/cudf/expression/ParquetSchemaUtils.h new file mode 100644 index 00000000000..996965536c0 --- /dev/null +++ b/velox/experimental/cudf/expression/ParquetSchemaUtils.h @@ -0,0 +1,71 @@ +/* + * 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. + */ +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace facebook::velox::cudf_velox { + +using ParquetColumnTypeMap = std::unordered_map; + +inline bool isParquetDecimalType(cudf::type_id typeId) { + return typeId == cudf::type_id::DECIMAL32 || + typeId == cudf::type_id::DECIMAL64 || typeId == cudf::type_id::DECIMAL128; +} + +// Maps top-level Parquet column names to libcudf storage types. Uses the same +// schema resolution as cudf::io::read_parquet_metadata. +inline ParquetColumnTypeMap parquetColumnTypesFromMetadata( + const cudf::io::parquet_metadata& metadata) { + ParquetColumnTypeMap columnTypes; + const auto& root = metadata.schema().root(); + for (int i = 0; i < root.num_children(); ++i) { + const auto& child = root.child(i); + if (!child.name().empty()) { + columnTypes.emplace(child.name(), child.cudf_type()); + } + } + return columnTypes; +} + +inline std::optional lookupParquetColumnType( + const ParquetColumnTypeMap& columnTypes, + const std::string& fieldName) { + if (auto it = columnTypes.find(fieldName); it != columnTypes.end()) { + return it->second; + } + for (const auto& [name, dataType] : columnTypes) { + if (name.size() == fieldName.size() && + std::equal( + name.begin(), + name.end(), + fieldName.begin(), + [](char a, char b) { + return std::tolower(a) == std::tolower(b); + })) { + return dataType; + } + } + return std::nullopt; +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/SubfieldFiltersToAst.cpp b/velox/experimental/cudf/expression/SubfieldFiltersToAst.cpp index 1c2507a56d7..dc4dd201a0a 100644 --- a/velox/experimental/cudf/expression/SubfieldFiltersToAst.cpp +++ b/velox/experimental/cudf/expression/SubfieldFiltersToAst.cpp @@ -39,6 +39,20 @@ std::pair getInt128BoundsForType(const TypePtr& type) { std::numeric_limits::max()}; } +template +cudf::ast::literal makeSubfieldFilterLiteral( + const TypePtr& columnTypePtr, + const variant& veloxVariant, + std::vector>& scalars, + std::optional parquetColumnType = std::nullopt) { + return makeScalarAndLiteral( + columnTypePtr, + veloxVariant, + false, + scalars, + parquetColumnType); +} + template < typename RangeT, typename ScalarT, @@ -116,7 +130,8 @@ std::reference_wrapper buildIntegerRangeExpr( cudf::ast::tree& tree, std::vector>& scalars, const cudf::ast::expression& columnRef, - const TypePtr& columnTypePtr) { + const TypePtr& columnTypePtr, + std::optional parquetColumnType = std::nullopt) { using NativeT = typename TypeTraits::NativeType; if constexpr ( @@ -146,8 +161,8 @@ std::reference_wrapper buildIntegerRangeExpr( auto addLiteral = [&](ValueT value) -> const cudf::ast::expression& { variant veloxVariant = static_cast(value); - const auto& literal = - makeScalarAndLiteral(columnTypePtr, veloxVariant, scalars); + const auto& literal = makeSubfieldFilterLiteral( + columnTypePtr, veloxVariant, scalars, parquetColumnType); return tree.push(literal); }; @@ -196,9 +211,10 @@ std::reference_wrapper buildBigintRangeExpr( cudf::ast::tree& tree, std::vector>& scalars, const cudf::ast::expression& columnRef, - const TypePtr& columnTypePtr) { + const TypePtr& columnTypePtr, + std::optional parquetColumnType = std::nullopt) { return buildIntegerRangeExpr( - filter, tree, scalars, columnRef, columnTypePtr); + filter, tree, scalars, columnRef, columnTypePtr, parquetColumnType); } std::reference_wrapper buildHugeintRangeExpr( @@ -206,9 +222,10 @@ std::reference_wrapper buildHugeintRangeExpr( cudf::ast::tree& tree, std::vector>& scalars, const cudf::ast::expression& columnRef, - const TypePtr& columnTypePtr) { + const TypePtr& columnTypePtr, + std::optional parquetColumnType = std::nullopt) { return buildIntegerRangeExpr( - filter, tree, scalars, columnRef, columnTypePtr); + filter, tree, scalars, columnRef, columnTypePtr, parquetColumnType); } template @@ -218,7 +235,8 @@ const cudf::ast::expression& buildValuesListExpr( const cudf::ast::expression& columnRef, std::vector>& scalars, const TypePtr& columnTypePtr, - bool isNegated = false) { + bool isNegated = false, + std::optional parquetColumnType = std::nullopt) { using Op = cudf::ast::ast_operator; using Operation = cudf::ast::operation; @@ -230,8 +248,8 @@ const cudf::ast::expression& buildValuesListExpr( std::vector exprVec; for (const auto& value : values) { variant veloxVariant = static_cast(value); - auto const& literal = tree.push( - makeScalarAndLiteral(columnTypePtr, veloxVariant, scalars)); + auto const& literal = tree.push(makeSubfieldFilterLiteral( + columnTypePtr, veloxVariant, scalars, parquetColumnType)); auto const& equalExpr = tree.push( Operation{isNegated ? Op::NOT_EQUAL : Op::EQUAL, columnRef, literal}); exprVec.push_back(&equalExpr); @@ -286,7 +304,8 @@ std::reference_wrapper buildIntegerInListExpr( const cudf::ast::expression& columnRef, rmm::cuda_stream_view /*stream*/, rmm::device_async_resource_ref /*mr*/, - const TypePtr& columnTypePtr) { + const TypePtr& columnTypePtr, + std::optional parquetColumnType = std::nullopt) { using NativeT = typename TypeTraits::NativeType; if constexpr (std::is_integral_v) { @@ -308,8 +327,8 @@ std::reference_wrapper buildIntegerInListExpr( } variant veloxVariant = static_cast(value); - const auto& literal = - makeScalarAndLiteral(columnTypePtr, veloxVariant, scalars); + const auto& literal = makeSubfieldFilterLiteral( + columnTypePtr, veloxVariant, scalars, parquetColumnType); auto const& cudfLiteral = tree.push(literal); auto const& equalExpr = tree.push(Operation{Op::EQUAL, columnRef, cudfLiteral}); @@ -341,7 +360,8 @@ cudf::ast::expression const& createAstFromSubfieldFilter( const common::Filter& filter, cudf::ast::tree& tree, std::vector>& scalars, - const RowTypePtr& inputRowSchema) { + const RowTypePtr& inputRowSchema, + const ParquetColumnTypeMap* parquetColumnTypes) { // First, create column reference from subfield // For now, only support simple field references if (subfield.path().empty() || @@ -358,6 +378,12 @@ cudf::ast::expression const& createAstFromSubfieldFilter( VELOX_FAIL("Field '{}' not found in input schema", fieldName); } + std::optional parquetColumnType; + if (parquetColumnTypes != nullptr) { + parquetColumnType = + lookupParquetColumnType(*parquetColumnTypes, fieldName); + } + auto columnIndex = inputRowSchema->getChildIdx(fieldName); auto const& columnRef = tree.push(cudf::ast::column_reference(columnIndex)); @@ -377,14 +403,15 @@ cudf::ast::expression const& createAstFromSubfieldFilter( tree, scalars, columnRef, - columnType); + columnType, + parquetColumnType); return result.get(); } case common::FilterKind::kHugeintRange: { auto const& columnType = inputRowSchema->childAt(columnIndex); - auto const& expr = - buildHugeintRangeExpr(filter, tree, scalars, columnRef, columnType); + auto const& expr = buildHugeintRangeExpr( + filter, tree, scalars, columnRef, columnType, parquetColumnType); return expr.get(); } @@ -393,7 +420,14 @@ cudf::ast::expression const& createAstFromSubfieldFilter( return buildValuesListExpr< TypeKind::BIGINT, common::BigintValuesUsingHashTable, - int64_t>(filter, tree, columnRef, scalars, columnType); + int64_t>( + filter, + tree, + columnRef, + scalars, + columnType, + false, + parquetColumnType); } case common::FilterKind::kBigintValuesUsingBitmask: { @@ -408,7 +442,8 @@ cudf::ast::expression const& createAstFromSubfieldFilter( columnRef, stream, mr, - columnType); + columnType, + parquetColumnType); return result.get(); } @@ -417,7 +452,14 @@ cudf::ast::expression const& createAstFromSubfieldFilter( return buildValuesListExpr< TypeKind::HUGEINT, common::HugeintValuesUsingHashTable, - int128_t>(filter, tree, columnRef, scalars, columnType); + int128_t>( + filter, + tree, + columnRef, + scalars, + columnType, + false, + parquetColumnType); } case common::FilterKind::kBytesValues: { @@ -495,7 +537,12 @@ cudf::ast::expression const& createAstFromSubfieldFilter( exprRefs.reserve(subFilters.size()); for (const auto* subFilter : subFilters) { auto const& subExpr = createAstFromSubfieldFilter( - subfield, *subFilter, tree, scalars, inputRowSchema); + subfield, + *subFilter, + tree, + scalars, + inputRowSchema, + parquetColumnTypes); exprRefs.push_back(&subExpr); } @@ -520,7 +567,8 @@ cudf::ast::expression const& createAstFromSubfieldFilters( const common::SubfieldFilters& subfieldFilters, cudf::ast::tree& tree, std::vector>& scalars, - const RowTypePtr& inputRowSchema) { + const RowTypePtr& inputRowSchema, + const ParquetColumnTypeMap* parquetColumnTypes) { using Op = cudf::ast::ast_operator; using Operation = cudf::ast::operation; @@ -532,7 +580,12 @@ cudf::ast::expression const& createAstFromSubfieldFilters( continue; } auto const& expr = createAstFromSubfieldFilter( - subfield, *filterPtr, tree, scalars, inputRowSchema); + subfield, + *filterPtr, + tree, + scalars, + inputRowSchema, + parquetColumnTypes); exprRefs.push_back(&expr); } diff --git a/velox/experimental/cudf/expression/SubfieldFiltersToAst.h b/velox/experimental/cudf/expression/SubfieldFiltersToAst.h index e44e2bea15b..29ede589dbd 100644 --- a/velox/experimental/cudf/expression/SubfieldFiltersToAst.h +++ b/velox/experimental/cudf/expression/SubfieldFiltersToAst.h @@ -19,10 +19,13 @@ #include "velox/type/Subfield.h" #include "velox/type/Type.h" +#include "velox/experimental/cudf/expression/ParquetSchemaUtils.h" + #include #include #include +#include #include namespace cudf { @@ -39,7 +42,8 @@ cudf::ast::expression const& createAstFromSubfieldFilter( const common::Filter& filter, cudf::ast::tree& tree, std::vector>& scalars, - const RowTypePtr& inputRowSchema); + const RowTypePtr& inputRowSchema, + const ParquetColumnTypeMap* parquetColumnTypes = nullptr); // Build a single AST expression representing logical AND of all filters in // 'subfieldFilters'. The resulting expression reference is owned by the passed @@ -48,6 +52,7 @@ cudf::ast::expression const& createAstFromSubfieldFilters( const common::SubfieldFilters& subfieldFilters, cudf::ast::tree& tree, std::vector>& scalars, - const RowTypePtr& inputRowSchema); + const RowTypePtr& inputRowSchema, + const ParquetColumnTypeMap* parquetColumnTypes = nullptr); } // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/tests/DecimalAggregationTest.cpp b/velox/experimental/cudf/tests/DecimalAggregationTest.cpp index e7d499621ca..0fc6f7bc113 100644 --- a/velox/experimental/cudf/tests/DecimalAggregationTest.cpp +++ b/velox/experimental/cudf/tests/DecimalAggregationTest.cpp @@ -15,6 +15,8 @@ */ #include "velox/experimental/cudf/CudfConfig.h" +#include "velox/experimental/cudf/exec/DecimalAggregationHostOps.h" +#include "velox/experimental/cudf/exec/DecimalAggregationDevice.h" #include "velox/experimental/cudf/exec/DecimalAggregationState.h" #include "velox/experimental/cudf/exec/ToCudf.h" #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" @@ -157,6 +159,15 @@ std::unique_ptr makeDecimalColumn( return makeFixedWidthColumn(type, values, valid, stream); } +std::unique_ptr makeDecimal32Column( + const std::vector& values, + int32_t scale, + const std::vector* valid, + rmm::cuda_stream_view stream) { + cudf::data_type type{cudf::type_id::DECIMAL32, -scale}; + return makeFixedWidthColumn(type, values, valid, stream); +} + std::unique_ptr makeInt64Column( const std::vector& values, const std::vector* valid, @@ -1074,6 +1085,33 @@ TEST_F(CudfDecimalTest, decimalSumGlobalIntermediateVarbinaryAllNulls) { facebook::velox::test::assertEqualVectors(expected, result); } +TEST_F(CudfDecimalTest, decimalSerializeSumStateDecimal32) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector sums = {100, -200, 300}; + std::vector counts = {1, 2, 0}; + std::vector sumValid = {true, false, true}; + std::vector countValid = {true, true, true}; + + auto sumCol = makeDecimal32Column(sums, 2, &sumValid, stream); + auto countCol = makeInt64Column(counts, &countValid, stream); + auto stateCol = serializeDecimalPartialOrIntermediateState( + std::move(sumCol), std::move(countCol), stream, mr); + auto sumAndCount = deserializeDecimalSumState(stateCol->view(), 2, stream); + auto stateMask = copyNullMask(stateCol->view(), stream); + auto sumMask = copyNullMask(sumAndCount.sum->view(), stream); + EXPECT_EQ(stateMask, sumMask); + + auto outSum = copyColumnData<__int128_t>(sumAndCount.sum->view(), stream); + for (size_t i = 0; i < sums.size(); ++i) { + bool expectedValid = sumValid[i] && countValid[i] && counts[i] != 0; + EXPECT_EQ(isValidAt(sumMask, i), expectedValid); + if (expectedValid) { + EXPECT_EQ(outSum[i], static_cast<__int128_t>(sums[i])); + } + } +} + TEST_F(CudfDecimalTest, decimalDeserializeSumStateDecimal64) { auto stream = cudf::get_default_stream(); auto mr = cudf::get_current_device_resource_ref(); @@ -1101,6 +1139,46 @@ TEST_F(CudfDecimalTest, decimalDeserializeSumStateDecimal64) { } } +TEST_F(CudfDecimalTest, decimalDeserializeSumStateAfterVarbinaryRoundTrip) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector sums = {100, -200, 300, 400}; + std::vector counts = {1, 0, 2, 3}; + std::vector sumValid = {true, true, false, true}; + std::vector countValid = {true, true, true, false}; + + auto sumCol = makeDecimalColumn(sums, 2, &sumValid, stream); + auto countCol = makeInt64Column(counts, &countValid, stream); + auto stateCol = + serializeDecimalSumState(sumCol->view(), countCol->view(), stream, mr); + + auto rowType = ROW({{"state", VARBINARY()}}); + cudf::table_view stateTable{{stateCol->view()}}; + auto velox = with_arrow::toVeloxColumn( + stateTable, pool(), rowType, "rt_", stream, mr); + auto cudfAgain = with_arrow::toCudfTable(velox, pool(), stream, mr); + auto const stateView = cudfAgain->view().column(0); + + cudf::strings_column_view strings(stateView); + EXPECT_LT( + strings.chars_size(stream), + static_cast(sums.size()) * detail::kDecimalSumStateSize); + + auto sumAndCount = deserializeDecimalSumState(stateView, 2, stream); + auto stateMask = copyNullMask(stateView, stream); + auto sumMask = copyNullMask(sumAndCount.sum->view(), stream); + EXPECT_EQ(stateMask, sumMask); + + auto outSum = copyColumnData<__int128_t>(sumAndCount.sum->view(), stream); + for (size_t i = 0; i < sums.size(); ++i) { + bool expectedValid = sumValid[i] && countValid[i] && counts[i] != 0; + EXPECT_EQ(isValidAt(sumMask, i), expectedValid); + if (expectedValid) { + EXPECT_EQ(outSum[i], static_cast<__int128_t>(sums[i])); + } + } +} + TEST_F(CudfDecimalTest, decimalDeserializeSumStateDecimal128) { auto stream = cudf::get_default_stream(); auto mr = cudf::get_current_device_resource_ref(); diff --git a/velox/experimental/cudf/tests/SubfieldFilterAstTest.cpp b/velox/experimental/cudf/tests/SubfieldFilterAstTest.cpp index a2d1317f772..d0063cbc0a6 100644 --- a/velox/experimental/cudf/tests/SubfieldFilterAstTest.cpp +++ b/velox/experimental/cudf/tests/SubfieldFilterAstTest.cpp @@ -16,6 +16,7 @@ #include "velox/experimental/cudf/exec/ToCudf.h" #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" +#include "velox/experimental/cudf/expression/ParquetSchemaUtils.h" #include "velox/experimental/cudf/expression/SubfieldFiltersToAst.h" #include "velox/dwio/common/tests/utils/BatchMaker.h" @@ -27,6 +28,8 @@ #include #include +#include + #include using namespace facebook::velox; @@ -565,6 +568,87 @@ TEST_F(SubfieldFilterAstTest, DecimalRange) { testFilterExecution(rowType, columnName, *filter, vec, expr); } +TEST_F(SubfieldFilterAstTest, ShortDecimalFilterUsesDecimal32Literals) { + const std::string columnName = "c0"; + auto rowType = ROW({{columnName, DECIMAL(9, 2)}}); + auto filter = std::make_unique( + int64_t{100}, int64_t{500}, /*nullAllowed*/ false); + + common::Subfield subfield(columnName); + cudf::ast::tree tree; + std::vector> scalars; + const auto& expr = + createAstFromSubfieldFilter(subfield, *filter, tree, scalars, rowType); + + ASSERT_EQ(scalars.size(), 2UL); + EXPECT_EQ(scalars[0]->type().id(), cudf::type_id::DECIMAL32); + EXPECT_EQ(scalars[1]->type().id(), cudf::type_id::DECIMAL32); + EXPECT_GT(tree.size(), 0UL); +} + +TEST_F(SubfieldFilterAstTest, ShortDecimalFilterUsesDecimal64LiteralsFromParquetSchema) { + const std::string columnName = "c0"; + auto rowType = ROW({{columnName, DECIMAL(9, 2)}}); + auto filter = std::make_unique( + int64_t{100}, int64_t{500}, /*nullAllowed*/ false); + ParquetColumnTypeMap parquetTypes{ + {columnName, cudf::data_type{cudf::type_id::DECIMAL64, -2}}}; + + common::Subfield subfield(columnName); + cudf::ast::tree tree; + std::vector> scalars; + const auto& expr = createAstFromSubfieldFilter( + subfield, *filter, tree, scalars, rowType, &parquetTypes); + + ASSERT_EQ(scalars.size(), 2UL); + EXPECT_EQ(scalars[0]->type().id(), cudf::type_id::DECIMAL64); + EXPECT_EQ(scalars[0]->type().scale(), numeric::scale_type{-2}); + EXPECT_EQ(scalars[1]->type().id(), cudf::type_id::DECIMAL64); + EXPECT_EQ(scalars[1]->type().scale(), numeric::scale_type{-2}); + EXPECT_GT(tree.size(), 0UL); +} + +TEST_F(SubfieldFilterAstTest, LongDecimalFilterUsesDecimal64LiteralsFromParquetSchema) { + const std::string columnName = "c0"; + auto rowType = ROW({{columnName, DECIMAL(18, 2)}}); + auto filter = std::make_unique( + int128_t{100}, int128_t{500}, /*nullAllowed*/ false); + ParquetColumnTypeMap parquetTypes{ + {columnName, cudf::data_type{cudf::type_id::DECIMAL64, -2}}}; + + common::Subfield subfield(columnName); + cudf::ast::tree tree; + std::vector> scalars; + const auto& expr = createAstFromSubfieldFilter( + subfield, *filter, tree, scalars, rowType, &parquetTypes); + + ASSERT_EQ(scalars.size(), 2UL); + EXPECT_EQ(scalars[0]->type().id(), cudf::type_id::DECIMAL64); + EXPECT_EQ(scalars[0]->type().scale(), numeric::scale_type{-2}); + EXPECT_EQ(scalars[1]->type().id(), cudf::type_id::DECIMAL64); + EXPECT_EQ(scalars[1]->type().scale(), numeric::scale_type{-2}); + EXPECT_GT(tree.size(), 0UL); +} + +TEST_F(SubfieldFilterAstTest, ShortDecimalFilterUsesInt64LiteralsFromParquetSchema) { + const std::string columnName = "c0"; + auto rowType = ROW({{columnName, DECIMAL(9, 2)}}); + auto filter = std::make_unique( + int64_t{100}, int64_t{500}, /*nullAllowed*/ false); + ParquetColumnTypeMap parquetTypes{{columnName, cudf::data_type{cudf::type_id::INT64}}}; + + common::Subfield subfield(columnName); + cudf::ast::tree tree; + std::vector> scalars; + const auto& expr = createAstFromSubfieldFilter( + subfield, *filter, tree, scalars, rowType, &parquetTypes); + + ASSERT_EQ(scalars.size(), 2UL); + EXPECT_EQ(scalars[0]->type().id(), cudf::type_id::INT64); + EXPECT_EQ(scalars[1]->type().id(), cudf::type_id::INT64); + EXPECT_GT(tree.size(), 0UL); +} + TEST_F(SubfieldFilterAstTest, DecimalInList) { const std::string columnName = "c0"; auto rowType = ROW({{columnName, DECIMAL(20, 2)}}); diff --git a/velox/experimental/cudf/tests/TableScanTest.cpp b/velox/experimental/cudf/tests/TableScanTest.cpp index 6b3a366896c..e1f7aa79475 100644 --- a/velox/experimental/cudf/tests/TableScanTest.cpp +++ b/velox/experimental/cudf/tests/TableScanTest.cpp @@ -43,7 +43,11 @@ #include "velox/type/Type.h" #include "velox/type/tests/SubfieldFiltersBuilder.h" +#include #include +#include + +#include #include @@ -88,6 +92,126 @@ StatsFilterMetrics readParquetWithStatsFilter( result.metadata.num_row_groups_after_stats_filter, result.tbl->num_rows()}; } + +template +void copyHostToDeviceColumn( + cudf::mutable_column_view view, + const std::vector& values, + rmm::cuda_stream_view stream) { + if (values.empty()) { + return; + } + auto status = cudaMemcpyAsync( + view.data(), + values.data(), + values.size() * sizeof(T), + cudaMemcpyHostToDevice, + stream.value()); + VELOX_CHECK_EQ(0, static_cast(status)); + stream.synchronize(); +} + +std::unique_ptr makeDecimal32Column( + const std::vector& values, + int32_t scale, + rmm::cuda_stream_view stream) { + cudf::data_type type{cudf::type_id::DECIMAL32, -scale}; + auto col = cudf::make_fixed_width_column( + type, + static_cast(values.size()), + cudf::mask_state::UNALLOCATED, + stream); + copyHostToDeviceColumn(col->mutable_view(), values, stream); + return col; +} + +void writeDecimal32ParquetFile( + const std::string& filePath, + const RowTypePtr& rowType, + const std::vector& condValues, + const std::vector& decimalValues, + int32_t decimalScale, + rmm::cuda_stream_view stream) { + auto condCol = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::BOOL8}, + static_cast(condValues.size()), + cudf::mask_state::UNALLOCATED, + stream); + copyHostToDeviceColumn(condCol->mutable_view(), condValues, stream); + auto decimalCol = + makeDecimal32Column(decimalValues, decimalScale, stream); + std::vector> columns; + columns.push_back(std::move(condCol)); + columns.push_back(std::move(decimalCol)); + auto table = std::make_unique(std::move(columns)); + + auto sinkInfo = cudf::io::sink_info(filePath); + auto tableInputMetadata = cudf::io::table_input_metadata(table->view()); + for (int32_t i = 0; i < rowType->size(); ++i) { + tableInputMetadata.column_metadata[i].set_name(rowType->nameOf(i)); + } + auto options = + cudf::io::parquet_writer_options::builder(sinkInfo, table->view()) + .metadata(tableInputMetadata) + .build(); + cudf::io::write_parquet(options, stream); +} + +void writeDecimal32AndInt64ParquetFile( + const std::string& filePath, + const RowTypePtr& rowType, + const std::vector& decimalValues, + int32_t decimalScale, + const std::vector& intValues, + rmm::cuda_stream_view stream) { + auto decimalCol = + makeDecimal32Column(decimalValues, decimalScale, stream); + auto intCol = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::INT64}, + static_cast(intValues.size()), + cudf::mask_state::UNALLOCATED, + stream); + copyHostToDeviceColumn(intCol->mutable_view(), intValues, stream); + std::vector> columns; + columns.push_back(std::move(decimalCol)); + columns.push_back(std::move(intCol)); + auto table = std::make_unique(std::move(columns)); + + auto sinkInfo = cudf::io::sink_info(filePath); + auto tableInputMetadata = cudf::io::table_input_metadata(table->view()); + for (int32_t i = 0; i < rowType->size(); ++i) { + tableInputMetadata.column_metadata[i].set_name(rowType->nameOf(i)); + } + auto options = + cudf::io::parquet_writer_options::builder(sinkInfo, table->view()) + .metadata(tableInputMetadata) + .build(); + cudf::io::write_parquet(options, stream); +} + +void writeSingleDecimal32ParquetFile( + const std::string& filePath, + const RowTypePtr& rowType, + const std::vector& decimalValues, + int32_t decimalScale, + rmm::cuda_stream_view stream) { + auto decimalCol = + makeDecimal32Column(decimalValues, decimalScale, stream); + std::vector> columns; + columns.push_back(std::move(decimalCol)); + auto table = std::make_unique(std::move(columns)); + + auto sinkInfo = cudf::io::sink_info(filePath); + auto tableInputMetadata = cudf::io::table_input_metadata(table->view()); + for (int32_t i = 0; i < rowType->size(); ++i) { + tableInputMetadata.column_metadata[i].set_name(rowType->nameOf(i)); + } + auto options = + cudf::io::parquet_writer_options::builder(sinkInfo, table->view()) + .metadata(tableInputMetadata) + .build(); + cudf::io::write_parquet(options, stream); +} } // namespace class TableScanTest : public virtual CudfHiveConnectorTestBase { @@ -818,3 +942,132 @@ TEST_F(TableScanTest, decimalRemainingFilter) { {filePath}, "SELECT c0, c1 FROM tmp WHERE c0 = CAST('-5.00' AS DECIMAL(5, 2))"); } + +// Parquet I/O stores small decimals as DECIMAL32 while Velox literals become +// DECIMAL64 in cuDF. IF/CASE must widen operands before copy_if_else. +TEST_F(TableScanTest, decimal32IfWithLiteral) { + auto rowType = ROW({{"cond", BOOLEAN()}, {"d", DECIMAL(9, 2)}}); + auto filePath = TempFilePath::create(); + auto stream = cudf::get_default_stream(); + writeDecimal32ParquetFile( + filePath->getPath(), + rowType, + {1, 0, 1, 0}, + {1000, 2000, 3000, 4000}, + 2, + stream); + + auto vectors = {makeRowVector( + {"cond", "d"}, + { + makeFlatVector({true, false, true, false}), + makeFlatVector({1000, 2000, 3000, 4000}, DECIMAL(9, 2)), + })}; + createDuckDbTable(vectors); + + auto assignments = + facebook::velox::exec::test::HiveConnectorTestBase::allRegularColumns( + rowType); + + auto plan = PlanBuilder(pool_.get()) + .startTableScan() + .connectorId(kCudfHiveConnectorId) + .outputType(rowType) + .dataColumns(rowType) + .assignments(assignments) + .endTableScan() + .project( + {"if(cond, CAST('5.00' AS DECIMAL(9, 2)), d) AS result"}) + .planNode(); + + assertQuery( + plan, + {filePath}, + "SELECT IF(cond, CAST('5.00' AS DECIMAL(9, 2)), d) AS result FROM tmp"); +} + +// Parquet scan columns are DECIMAL32 for precision <= 9; subfield filter +// literals must match or cuDF AST parsing fails. +TEST_F(TableScanTest, decimal32SubfieldFilter) { + auto rowType = ROW({"c0", "c1"}, {DECIMAL(9, 2), BIGINT()}); + auto filePath = TempFilePath::create(); + auto stream = cudf::get_default_stream(); + writeDecimal32AndInt64ParquetFile( + filePath->getPath(), + rowType, + {10000, -50000, -70000, -50000}, + 2, + {1, 2, 3, 4}, + stream); + + auto vectors = {makeRowVector( + {"c0", "c1"}, + { + makeFlatVector( + {10000, -50000, -70000, -50000}, DECIMAL(9, 2)), + makeFlatVector({1, 2, 3, 4}), + })}; + createDuckDbTable(vectors); + + common::SubfieldFilters subfieldFilters = + common::test::SubfieldFiltersBuilder() + .add( + "c0", + std::make_unique( + int64_t{-50000}, int64_t{-50000}, /*nullAllowed*/ false)) + .build(); + + auto tableHandle = makeTableHandle( + "parquet_table", rowType, std::move(subfieldFilters), nullptr); + auto assignments = + facebook::velox::exec::test::HiveConnectorTestBase::allRegularColumns( + rowType); + + auto plan = PlanBuilder() + .startTableScan() + .outputType(rowType) + .tableHandle(tableHandle) + .assignments(assignments) + .endTableScan() + .planNode(); + + assertQuery( + plan, + {filePath}, + "SELECT c0, c1 FROM tmp WHERE c0 = CAST('-500.00' AS DECIMAL(9, 2))"); +} + +// DECIMAL32 Parquet columns with DECIMAL64 literals in multiply filters must +// widen operands before cuDF binary_operation. +TEST_F(TableScanTest, decimal32FilterWithMultiply) { + auto rowType = ROW({{"d", DECIMAL(9, 2)}}); + auto filePath = TempFilePath::create(); + auto stream = cudf::get_default_stream(); + writeSingleDecimal32ParquetFile( + filePath->getPath(), rowType, {1000, 200, 5000}, 2, stream); + + auto vectors = {makeRowVector( + {"d"}, + {makeFlatVector({1000, 200, 5000}, DECIMAL(9, 2))})}; + createDuckDbTable(vectors); + + auto assignments = + facebook::velox::exec::test::HiveConnectorTestBase::allRegularColumns( + rowType); + + auto plan = PlanBuilder(pool_.get()) + .startTableScan() + .connectorId(kCudfHiveConnectorId) + .outputType(rowType) + .dataColumns(rowType) + .assignments(assignments) + .endTableScan() + .filter( + "d * CAST(0.2 AS DECIMAL(9, 2)) > CAST(1.00 AS DECIMAL(18, 4))") + .planNode(); + + assertQuery( + plan, + {filePath}, + "SELECT d FROM tmp WHERE d * CAST(0.2 AS DECIMAL(9, 2)) > CAST(1.00 AS DECIMAL(18, 4))"); +} diff --git a/velox/vector/arrow/Bridge.cpp b/velox/vector/arrow/Bridge.cpp index 1ccbd8b6f12..09a0b32bdda 100644 --- a/velox/vector/arrow/Bridge.cpp +++ b/velox/vector/arrow/Bridge.cpp @@ -1396,7 +1396,7 @@ void exportToArrowImpl( // Parses the velox decimal format from the given arrow format. // The input format string should be in the form "d:precision,scale<,bitWidth>". -// bitWidth is optional and may be 64 or 128 if provided. +// bitWidth is optional and may be 32, 64, or 128 if provided. int32_t parseDecimalBitWidthOrDefault(const std::string_view format) { auto firstCommaIdx = format.find(',', 2); @@ -1431,16 +1431,16 @@ TypePtr parseDecimalFormat(const std::string_view format) { int precision = std::stoi(&format[2], &sz); int scale = std::stoi(&format[firstCommaIdx + 1], &sz); if (secondCommaIdx != std::string_view::npos) { - // BitWidth is provided. We only support 64 or 128. + // BitWidth is provided. We support 32, 64, and 128. int bitWidth = std::stoi(&format[secondCommaIdx + 1], &sz); // Return type depends on bitWidth. - if (bitWidth == 64) { + if (bitWidth == 32 || bitWidth == 64) { return std::make_shared(precision, scale); } else if (bitWidth == 128) { return std::make_shared(precision, scale); } VELOX_USER_FAIL( - "Conversion failed for '{}'. Only 64-bit and 128-bit decimal types are supported.", + "Conversion failed for '{}'. Only 32-bit, 64-bit, and 128-bit decimal types are supported.", format); } // Otherwise return type depends on precision. @@ -2258,6 +2258,23 @@ VectorPtr createShortDecimalVectorFromLongDecimals( pool, type, std::move(nulls), length, values, nullCount); } +VectorPtr createShortDecimalVectorFrom32BitDecimals( + memory::MemoryPool* pool, + const TypePtr& type, + BufferPtr nulls, + const int32_t* input, + vector_size_t length, + int64_t nullCount) { + auto values = AlignedBuffer::allocate(length, pool); + auto rawValues = values->asMutable(); + for (size_t i = 0; i < length; ++i) { + rawValues[i] = input[i]; + } + + return createFlatVector( + pool, type, std::move(nulls), length, values, nullCount); +} + // Arrow uses two uint64_t values to represent a 128-bit decimal value. The // memory allocated by Arrow might not be 16-byte aligned, so we need to copy // the values to a new buffer to ensure 16-byte alignment. @@ -2408,6 +2425,15 @@ VectorPtr importFromArrowImpl( } else if (type->isShortDecimal()) { // Validate the format bitWidth. const auto bitWidth = parseDecimalBitWidthOrDefault(arrowSchema.format); + if (bitWidth == 32) { + return createShortDecimalVectorFrom32BitDecimals( + pool, + type, + nulls, + static_cast(arrowArray.buffers[1]), + arrowArray.length, + arrowArray.null_count); + } if (bitWidth == 64) { return createShortDecimalVector( pool, @@ -2418,20 +2444,19 @@ VectorPtr importFromArrowImpl( arrowArray.null_count, wrapInBufferView); } - // Otherwise convert to 128. - VELOX_USER_CHECK_EQ( - bitWidth, - 128, + if (bitWidth == 128) { + return createShortDecimalVectorFromLongDecimals( + pool, + type, + nulls, + static_cast(arrowArray.buffers[1]), + arrowArray.length, + arrowArray.null_count); + } + VELOX_USER_FAIL( "Unsupported decimal bitWidth {} for '{}'", bitWidth, arrowSchema.format); - return createShortDecimalVectorFromLongDecimals( - pool, - type, - nulls, - static_cast(arrowArray.buffers[1]), - arrowArray.length, - arrowArray.null_count); } else if (type->isLongDecimal()) { // Validate that the format is actually 128. const int32_t bitWidth = parseDecimalBitWidthOrDefault(arrowSchema.format); diff --git a/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp b/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp index e681e2073c9..71ba3e4f0df 100644 --- a/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp +++ b/velox/vector/arrow/tests/ArrowBridgeArrayTest.cpp @@ -1272,6 +1272,25 @@ class ArrowBridgeArrayImportTest : public ArrowBridgeArrayExportTest { std::is_same_v && std::is_same_v) { assertShortDecimalVectorContent( inputValues, output, arrowArray.null_count); + } else if constexpr ( + std::is_same_v && std::is_same_v) { + if (format[0] == 'd') { + std::vector> widenedValues; + widenedValues.reserve(inputValues.size()); + for (const auto& value : inputValues) { + if (value.has_value()) { + widenedValues.emplace_back(static_cast(value.value())); + } else { + widenedValues.emplace_back(std::nullopt); + } + } + assertVectorContent(widenedValues, output, arrowArray.null_count); + } else if (format[0] == 't' && format[1] == 't') { + assertTimeVectorContent( + inputValues, output, arrowArray.null_count, format); + } else { + assertVectorContent(inputValues, output, arrowArray.null_count); + } } else if constexpr ( std::is_same_v && (std::is_same_v || std::is_same_v)) { @@ -1393,6 +1412,16 @@ class ArrowBridgeArrayImportTest : public ArrowBridgeArrayExportTest { testArrowImport( "d:5,2", {1, -1, 0, 12345, -12345, std::nullopt}); + testArrowImport( + "d:9,2,32", + {1, + -1, + 0, + 12345, + -12345, + std::numeric_limits::max(), + std::numeric_limits::min(), + std::nullopt}); testArrowImport( "d:36,2", {HugeInt::parse("20000000000000000"), diff --git a/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp b/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp index 7a581e9a87a..5f35d5d46f1 100644 --- a/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp +++ b/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp @@ -509,10 +509,11 @@ TEST_F(ArrowBridgeSchemaImportTest, scalar) { *testSchemaImport("d2,15"), "Unable to convert 'd2,15' ArrowSchema decimal format to Velox decimal"); EXPECT_EQ(*DECIMAL(10, 4), *testSchemaImport("d:10,4,64")); + EXPECT_EQ(*DECIMAL(9, 2), *testSchemaImport("d:9,2,32")); EXPECT_EQ(*DECIMAL(20, 15), *testSchemaImport("d:20,15,128")); VELOX_ASSERT_THROW( *testSchemaImport("d:10,4,256"), - "Conversion failed for 'd:10,4,256'. Only 64-bit and 128-bit decimal types are supported."); + "Conversion failed for 'd:10,4,256'. Only 32-bit, 64-bit, and 128-bit decimal types are supported."); VELOX_ASSERT_THROW( *testSchemaImport("d:10,4,"), "Unable to convert 'd:10,4,' ArrowSchema decimal format to Velox decimal");