diff --git a/velox/experimental/cudf/exec/AggregationRegistry.cpp b/velox/experimental/cudf/exec/AggregationRegistry.cpp index 490877f7b72..657bf669249 100644 --- a/velox/experimental/cudf/exec/AggregationRegistry.cpp +++ b/velox/experimental/cudf/exec/AggregationRegistry.cpp @@ -82,6 +82,37 @@ void registerCommonAggregationFunctions( .argumentType("double") .build()}; + auto decimalSumSingle = std::vector{ + FunctionSignatureBuilder() + .integerVariable("a_precision") + .integerVariable("a_scale") + .returnType("decimal(38, a_scale)") + .argumentType("decimal(a_precision, a_scale)") + .build()}; + auto decimalSumPartial = std::vector{ + FunctionSignatureBuilder() + .integerVariable("a_precision") + .integerVariable("a_scale") + .returnType("varbinary") + .argumentType("decimal(a_precision, a_scale)") + .build()}; + auto decimalSumFinal = std::vector{ + FunctionSignatureBuilder() + .integerVariable("a_scale") + .returnType("decimal(38, a_scale)") + .argumentType("varbinary") + .build()}; + auto decimalSumIntermediate = + std::vector{FunctionSignatureBuilder() + .returnType("varbinary") + .argumentType("varbinary") + .build()}; + + sumSingleSignatures.insert( + sumSingleSignatures.end(), + decimalSumSingle.begin(), + decimalSumSingle.end()); + registerAggregationFunctionForStep( registry, prefix + "sum", @@ -109,6 +140,12 @@ void registerCommonAggregationFunctions( .returnType("double") .argumentType("double") .build()}; + + sumPartialSignatures.insert( + sumPartialSignatures.end(), + decimalSumPartial.begin(), + decimalSumPartial.end()); + registerAggregationFunctionForStep( registry, prefix + "sum", @@ -125,16 +162,27 @@ void registerCommonAggregationFunctions( .argumentType("double") .build()}; + auto sumFinalSignatures = sumFinalIntermediateSignatures; + sumFinalSignatures.insert( + sumFinalSignatures.end(), decimalSumFinal.begin(), decimalSumFinal.end()); + registerAggregationFunctionForStep( registry, prefix + "sum", core::AggregationNode::Step::kFinal, - sumFinalIntermediateSignatures); + sumFinalSignatures); + + auto sumIntermediateSignatures = sumFinalIntermediateSignatures; + sumIntermediateSignatures.insert( + sumIntermediateSignatures.end(), + decimalSumIntermediate.begin(), + decimalSumIntermediate.end()); + registerAggregationFunctionForStep( registry, prefix + "sum", core::AggregationNode::Step::kIntermediate, - sumFinalIntermediateSignatures); + sumIntermediateSignatures); auto countSinglePartialSignatures = std::vector{ FunctionSignatureBuilder() @@ -234,6 +282,12 @@ void registerCommonAggregationFunctions( FunctionSignatureBuilder() .returnType("varchar") .argumentType("varchar") + .build(), + FunctionSignatureBuilder() + .integerVariable("p") + .integerVariable("s") + .returnType("decimal(p,s)") + .argumentType("decimal(p,s)") .build()}; registerAggregationFunctionForStep( @@ -296,6 +350,38 @@ void registerCommonAggregationFunctions( .argumentType("double") .build()}; + auto decimalAvgSingle = std::vector{ + FunctionSignatureBuilder() + .integerVariable("a_precision") + .integerVariable("a_scale") + .returnType("decimal(a_precision, a_scale)") + .argumentType("decimal(a_precision, a_scale)") + .build()}; + auto decimalAvgPartial = std::vector{ + FunctionSignatureBuilder() + .integerVariable("a_precision") + .integerVariable("a_scale") + .returnType("varbinary") + .argumentType("decimal(a_precision, a_scale)") + .build()}; + auto decimalAvgFinal = std::vector{ + FunctionSignatureBuilder() + .integerVariable("a_precision") + .integerVariable("a_scale") + .returnType("decimal(a_precision, a_scale)") + .argumentType("varbinary") + .build()}; + auto decimalAvgIntermediate = + std::vector{FunctionSignatureBuilder() + .returnType("varbinary") + .argumentType("varbinary") + .build()}; + + avgSingleSignatures.insert( + avgSingleSignatures.end(), + decimalAvgSingle.begin(), + decimalAvgSingle.end()); + registerAggregationFunctionForStep( registry, prefix + "avg", @@ -323,17 +409,28 @@ void registerCommonAggregationFunctions( .returnType("row(double,bigint)") .argumentType("double") .build()}; + + avgPartialSignatures.insert( + avgPartialSignatures.end(), + decimalAvgPartial.begin(), + decimalAvgPartial.end()); + registerAggregationFunctionForStep( registry, prefix + "avg", core::AggregationNode::Step::kPartial, avgPartialSignatures); - auto avgFinalSignatures = std::vector{ + auto avgFinalIntermediateSignatures = std::vector{ FunctionSignatureBuilder() .returnType("double") .argumentType("row(double,bigint)") .build()}; + + auto avgFinalSignatures = avgFinalIntermediateSignatures; + avgFinalSignatures.insert( + avgFinalSignatures.end(), decimalAvgFinal.begin(), decimalAvgFinal.end()); + registerAggregationFunctionForStep( registry, prefix + "avg", @@ -345,6 +442,12 @@ void registerCommonAggregationFunctions( .returnType("row(double,bigint)") .argumentType("row(double,bigint)") .build()}; + + avgIntermediateSignatures.insert( + avgIntermediateSignatures.end(), + decimalAvgIntermediate.begin(), + decimalAvgIntermediate.end()); + registerAggregationFunctionForStep( registry, prefix + "avg", diff --git a/velox/experimental/cudf/exec/CMakeLists.txt b/velox/experimental/cudf/exec/CMakeLists.txt index ec17982119c..e96daead1ed 100644 --- a/velox/experimental/cudf/exec/CMakeLists.txt +++ b/velox/experimental/cudf/exec/CMakeLists.txt @@ -33,6 +33,9 @@ add_library( CudfReduce.cpp CudfTopN.cpp DebugUtil.cpp + DecimalAggregationDevice.cu + DecimalAggregationHostOps.cpp + DecimalAggregationState.cpp GpuResources.cpp OperatorAdapters.cpp PrestoAggregateFunctions.cpp diff --git a/velox/experimental/cudf/exec/CudfAggregation.cpp b/velox/experimental/cudf/exec/CudfAggregation.cpp index d99513aaded..3feb3fe4df1 100644 --- a/velox/experimental/cudf/exec/CudfAggregation.cpp +++ b/velox/experimental/cudf/exec/CudfAggregation.cpp @@ -136,6 +136,8 @@ std::vector resolveAggregateInfos( const auto resultType = exec::isPartialOutput(companionStep) ? exec::resolveIntermediateType(originalName, aggregate.rawInputTypes) : outputType->childAt(numKeys + i); + const auto isDecimalAggregate = aggregate.rawInputTypes.size() == 1 && + aggregate.rawInputTypes[0]->isDecimal(); params.emplace_back( companionStep, @@ -145,7 +147,8 @@ std::vector resolveAggregateInfos( resultType, isCountFunctionName(aggregate.call->name()) ? std::make_optional(getCountInputKind(aggregate, constants[i])) - : std::nullopt); + : std::nullopt, + isDecimalAggregate); } return params; } diff --git a/velox/experimental/cudf/exec/CudfAggregation.h b/velox/experimental/cudf/exec/CudfAggregation.h index 48a79df9705..f00b52a4d60 100644 --- a/velox/experimental/cudf/exec/CudfAggregation.h +++ b/velox/experimental/cudf/exec/CudfAggregation.h @@ -62,6 +62,10 @@ struct ResolvedAggregateInfo { VectorPtr constant; TypePtr resultType; std::optional countInputKind; + // True if the aggregate was declared on a decimal raw input in the plan. + // Routing keys off the function family, not the physical batch type (which is + // VARBINARY/STRING on intermediate and final steps). + bool isDecimalAggregate; }; // Parse aggregate inputs from the aggregation node and resolve companion steps, diff --git a/velox/experimental/cudf/exec/CudfGroupby.cpp b/velox/experimental/cudf/exec/CudfGroupby.cpp index a6ebeb33571..172bf0e8456 100644 --- a/velox/experimental/cudf/exec/CudfGroupby.cpp +++ b/velox/experimental/cudf/exec/CudfGroupby.cpp @@ -18,6 +18,8 @@ #include "velox/experimental/cudf/CudfNoDefaults.h" #include "velox/experimental/cudf/exec/CudfFilterProject.h" #include "velox/experimental/cudf/exec/CudfGroupby.h" +#include "velox/experimental/cudf/exec/DecimalAggregationHostOps.h" +#include "velox/experimental/cudf/exec/DecimalAggregationState.h" #include "velox/experimental/cudf/exec/GpuResources.h" #include "velox/experimental/cudf/exec/Utilities.h" #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" @@ -33,59 +35,318 @@ #include #include #include +#include #include namespace { using namespace facebook::velox; +using cudf_velox::castDecimal64InputToDecimal128; using cudf_velox::CountInputKind; +using cudf_velox::finalizeDecimalAverage; using cudf_velox::get_output_mr; using cudf_velox::get_temp_mr; using cudf_velox::GroupbyAggregator; using cudf_velox::ResolvedAggregateInfo; - -#define DEFINE_SIMPLE_GROUPBY_AGGREGATOR(Name, name, KIND) \ - struct Groupby##Name##Aggregator : GroupbyAggregator { \ - Groupby##Name##Aggregator( \ - core::AggregationNode::Step step, \ - uint32_t inputIndex, \ - VectorPtr constant, \ - const TypePtr& resultType) \ - : GroupbyAggregator(step, inputIndex, constant, resultType) {} \ - \ - void addGroupbyRequest( \ - cudf::table_view const& tbl, \ - std::vector& requests) override { \ - VELOX_CHECK( \ - constant == nullptr, \ - #Name "Aggregator does not yet support constant input"); \ - auto& request = requests.emplace_back(); \ - output_idx = requests.size() - 1; \ - request.values = tbl.column(inputIndex); \ - request.aggregations.push_back( \ - cudf::make_##name##_aggregation()); \ - } \ - \ - std::unique_ptr makeOutputColumn( \ - std::vector& results, \ - rmm::cuda_stream_view stream) override { \ - auto col = std::move(results[output_idx].results[0]); \ - const auto cudfType = \ - cudf::data_type(cudf_velox::veloxToCudfTypeId(resultType)); \ - if (col->type() != cudfType) { \ - col = cudf::cast(*col, cudfType, stream, get_output_mr()); \ - } \ - return col; \ - } \ - \ - private: \ - uint32_t output_idx; \ +using cudf_velox::serializeDecimalPartialOrIntermediateState; +using cudf_velox::validateIntermediateColumnType; + +#define DEFINE_SIMPLE_GROUPBY_AGGREGATOR(Name, name, KIND) \ + struct Groupby##Name##Aggregator : GroupbyAggregator { \ + Groupby##Name##Aggregator( \ + core::AggregationNode::Step step, \ + uint32_t inputIndex, \ + VectorPtr constant, \ + const TypePtr& resultType) \ + : GroupbyAggregator(step, inputIndex, constant, resultType) {} \ + \ + void addGroupbyRequest( \ + cudf::table_view const& tbl, \ + std::vector& requests, \ + rmm::cuda_stream_view stream) override { \ + VELOX_CHECK( \ + constant == nullptr, \ + #Name "Aggregator does not yet support constant input"); \ + auto& request = requests.emplace_back(); \ + output_idx = requests.size() - 1; \ + request.values = tbl.column(inputIndex); \ + request.aggregations.push_back( \ + cudf::make_##name##_aggregation()); \ + } \ + \ + std::unique_ptr makeOutputColumn( \ + std::vector& results, \ + rmm::cuda_stream_view stream, \ + rmm::device_async_resource_ref mr) override { \ + auto col = std::move(results[output_idx].results[0]); \ + const auto cudfType = cudf_velox::veloxToCudfDataType(resultType); \ + if (col->type() != cudfType) { \ + col = cudf::cast(*col, cudfType, stream, mr); \ + } \ + return col; \ + } \ + \ + private: \ + uint32_t output_idx; \ }; DEFINE_SIMPLE_GROUPBY_AGGREGATOR(Sum, sum, SUM) DEFINE_SIMPLE_GROUPBY_AGGREGATOR(Min, min, MIN) DEFINE_SIMPLE_GROUPBY_AGGREGATOR(Max, max, MAX) +// Decimal SUM and AVG aggregators are separate implementations, as they need to +// handle the VARBINARY encoded intermediate state for streaming aggregation. +// Due to the packing and unpacking of that intermediate state, and the special +// handling required for the decimal divide, we cannot just use the existing +// cudf::make_mean_aggregation class. Also, unlike other aggregators, these +// classes hold state (the decoded intermediate sum and count columns and +// associated indices) in order to guarantee a lifetime constraint between +// aggregation steps. + +void addDecimalSumCountRequestsAfterDecode( + cudf::column_view encodedColumn, + int32_t scale, + std::vector& requests, + rmm::cuda_stream_view stream, + uint32_t& sumIdx, + uint32_t& countIdx, + std::unique_ptr& decodedSum, + std::unique_ptr& decodedCount) { + auto sumAndCount = + cudf_velox::deserializeDecimalSumState(encodedColumn, scale, stream); + decodedSum.swap(sumAndCount.sum); + decodedCount.swap(sumAndCount.count); + + sumIdx = requests.size(); + auto& sumRequest = requests.emplace_back(); + sumRequest.values = decodedSum->view(); + sumRequest.aggregations.push_back( + cudf::make_sum_aggregation()); + + countIdx = requests.size(); + auto& countRequest = requests.emplace_back(); + countRequest.values = decodedCount->view(); + countRequest.aggregations.push_back( + cudf::make_sum_aggregation()); +} + +// Decodes serialized state and adds sum + count groupby requests, used by the +// intermediate step (both SUM and AVG) and the final AVG step. resultType is +// DECIMAL for final AVG and DECIMAL or VARBINARY for intermediate; VARBINARY +// carries no scale, so decode at scale 0. +void addDecimalDecodedSumCountRequests( + cudf::table_view const& tbl, + uint32_t inputIndex, + const TypePtr& resultType, + std::vector& requests, + rmm::cuda_stream_view stream, + uint32_t& sumIdx, + uint32_t& countIdx, + std::unique_ptr& decodedSum, + std::unique_ptr& decodedCount) { + validateIntermediateColumnType(tbl.column(inputIndex)); + auto scale = resultType->isDecimal() + ? getDecimalPrecisionScale(*resultType).second + : 0; + addDecimalSumCountRequestsAfterDecode( + tbl.column(inputIndex), + scale, + requests, + stream, + sumIdx, + countIdx, + decodedSum, + decodedCount); +} + +void addDecimalFinalSumOnlyRequest( + cudf::table_view const& tbl, + uint32_t inputIndex, + const TypePtr& resultType, + std::vector& requests, + rmm::cuda_stream_view stream, + uint32_t& sumIdx, + std::unique_ptr& decodedSum) { + validateIntermediateColumnType(tbl.column(inputIndex)); + auto scale = getDecimalPrecisionScale(*resultType).second; + auto& request = requests.emplace_back(); + sumIdx = requests.size() - 1; + auto sumAndCount = cudf_velox::deserializeDecimalSumState( + tbl.column(inputIndex), scale, stream); + decodedSum.swap(sumAndCount.sum); + request.values = decodedSum->view(); + request.aggregations.push_back( + cudf::make_sum_aggregation()); +} + +void addDecimalRawPartialSingleSumRequest( + cudf::table_view const& tbl, + uint32_t inputIndex, + std::vector& requests, + bool includeCountAggregation, + rmm::cuda_stream_view stream, + uint32_t& sumIdx, + std::unique_ptr& castedInput) { + auto inputView = castDecimal64InputToDecimal128( + tbl.column(inputIndex), castedInput, stream); + auto& request = requests.emplace_back(); + sumIdx = requests.size() - 1; + request.values = inputView; + request.aggregations.push_back( + cudf::make_sum_aggregation()); + if (includeCountAggregation) { + request.aggregations.push_back( + cudf::make_count_aggregation( + cudf::null_policy::EXCLUDE)); + } +} + +struct GroupbyDecimalSumAggregator : GroupbyAggregator { + GroupbyDecimalSumAggregator( + core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + const TypePtr& resultType) + : GroupbyAggregator(step, inputIndex, constant, resultType) {} + + void addGroupbyRequest( + cudf::table_view const& tbl, + std::vector& requests, + rmm::cuda_stream_view stream) override { + if (step == core::AggregationNode::Step::kIntermediate) { + addDecimalDecodedSumCountRequests( + tbl, + inputIndex, + resultType, + requests, + stream, + sumIdx_, + countIdx_, + decodedSum_, + decodedCount_); + } else if (step == core::AggregationNode::Step::kFinal) { + addDecimalFinalSumOnlyRequest( + tbl, inputIndex, resultType, requests, stream, sumIdx_, decodedSum_); + } else { + addDecimalRawPartialSingleSumRequest( + tbl, + inputIndex, + requests, + step == core::AggregationNode::Step::kPartial, + stream, + sumIdx_, + castedInput_); + } + } + + std::unique_ptr makeOutputColumn( + std::vector& results, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) override { + auto col = std::move(results[sumIdx_].results[0]); + if (step == core::AggregationNode::Step::kPartial) { + auto count = std::move(results[sumIdx_].results[1]); + return serializeDecimalPartialOrIntermediateState( + std::move(col), std::move(count), stream, mr); + } + if (step == core::AggregationNode::Step::kIntermediate) { + auto count = std::move(results[countIdx_].results[0]); + return serializeDecimalPartialOrIntermediateState( + std::move(col), std::move(count), stream, mr); + } + auto const cudfResType = cudf_velox::veloxToCudfDataType(resultType); + if (col->type() != cudfResType) { + col = cudf::cast(*col, cudfResType, stream, mr); + } + return col; + } + + private: + uint32_t sumIdx_{0}; + 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. + std::unique_ptr castedInput_; +}; + +struct GroupbyDecimalAvgAggregator : GroupbyAggregator { + GroupbyDecimalAvgAggregator( + core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + const TypePtr& resultType) + : GroupbyAggregator(step, inputIndex, constant, resultType) {} + + void addGroupbyRequest( + cudf::table_view const& tbl, + std::vector& requests, + rmm::cuda_stream_view stream) override { + if (step == core::AggregationNode::Step::kIntermediate || + step == core::AggregationNode::Step::kFinal) { + addDecimalDecodedSumCountRequests( + tbl, + inputIndex, + resultType, + requests, + stream, + sumIdx_, + countIdx_, + decodedSum_, + decodedCount_); + } else { + addDecimalRawPartialSingleSumRequest( + tbl, + inputIndex, + requests, + step == core::AggregationNode::Step::kPartial || + step == core::AggregationNode::Step::kSingle, + stream, + sumIdx_, + castedInput_); + } + } + + std::unique_ptr makeOutputColumn( + std::vector& results, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) override { + auto col = std::move(results[sumIdx_].results[0]); + if (step == core::AggregationNode::Step::kSingle) { + auto count = std::move(results[sumIdx_].results[1]); + return finalizeDecimalAverage( + std::move(col), std::move(count), resultType, stream, mr); + } + if (step == core::AggregationNode::Step::kPartial) { + auto count = std::move(results[sumIdx_].results[1]); + return serializeDecimalPartialOrIntermediateState( + std::move(col), std::move(count), stream, mr); + } + if (step == core::AggregationNode::Step::kIntermediate) { + auto count = std::move(results[countIdx_].results[0]); + return serializeDecimalPartialOrIntermediateState( + std::move(col), std::move(count), stream, mr); + } + if (step == core::AggregationNode::Step::kFinal) { + auto count = std::move(results[countIdx_].results[0]); + return finalizeDecimalAverage( + std::move(col), std::move(count), resultType, stream, mr); + } + // All four aggregation steps are handled above. + VELOX_UNREACHABLE(); + } + + private: + uint32_t sumIdx_{0}; + 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. + std::unique_ptr castedInput_; +}; + struct GroupbyCountAggregator : GroupbyAggregator { GroupbyCountAggregator( core::AggregationNode::Step step, @@ -97,7 +358,8 @@ struct GroupbyCountAggregator : GroupbyAggregator { void addGroupbyRequest( cudf::table_view const& tbl, - std::vector& requests) override { + std::vector& requests, + rmm::cuda_stream_view stream) override { auto& request = requests.emplace_back(); outputIndex_ = requests.size() - 1; // kCountAll and kNullConstant both submit a count-all-rows request; @@ -120,18 +382,17 @@ struct GroupbyCountAggregator : GroupbyAggregator { std::unique_ptr makeOutputColumn( std::vector& results, - rmm::cuda_stream_view stream) override { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) override { auto col = std::move(results[outputIndex_].results[0]); if (inputKind_ == CountInputKind::kNullConstant) { auto zero = cudf::numeric_scalar(0, true, stream, get_temp_mr()); - col = cudf::make_column_from_scalar( - zero, col->size(), stream, get_output_mr()); + col = cudf::make_column_from_scalar(zero, col->size(), stream, mr); } // cudf produces int32 for count but velox expects int64. - const auto cudfOutputType = - cudf::data_type(cudf_velox::veloxToCudfTypeId(resultType)); + const auto cudfOutputType = cudf_velox::veloxToCudfDataType(resultType); if (col->type() != cudfOutputType) { - col = cudf::cast(*col, cudfOutputType, stream, get_output_mr()); + col = cudf::cast(*col, cudfOutputType, stream, mr); } return col; } @@ -151,7 +412,8 @@ struct GroupbyMeanAggregator : GroupbyAggregator { void addGroupbyRequest( cudf::table_view const& tbl, - std::vector& requests) override { + std::vector& requests, + rmm::cuda_stream_view stream) override { switch (step) { case core::AggregationNode::Step::kSingle: { auto& request = requests.emplace_back(); @@ -198,7 +460,8 @@ struct GroupbyMeanAggregator : GroupbyAggregator { std::unique_ptr makeOutputColumn( std::vector& results, - rmm::cuda_stream_view stream) override { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) override { const auto& outputType = asRowType(resultType); switch (step) { case core::AggregationNode::Step::kSingle: @@ -208,17 +471,16 @@ struct GroupbyMeanAggregator : GroupbyAggregator { auto count = std::move(results[sumIdx_].results[1]); auto const size = sum->size(); - auto const cudfSumType = cudf::data_type( - cudf_velox::veloxToCudfTypeId(outputType->childAt(0))); - auto const cudfCountType = cudf::data_type( - cudf_velox::veloxToCudfTypeId(outputType->childAt(1))); + auto const cudfSumType = + cudf_velox::veloxToCudfDataType(outputType->childAt(0)); + auto const cudfCountType = + cudf_velox::veloxToCudfDataType(outputType->childAt(1)); if (sum->type() != cudf::data_type(cudfSumType)) { - sum = cudf::cast( - *sum, cudf::data_type(cudfSumType), stream, get_output_mr()); + sum = cudf::cast(*sum, cudf::data_type(cudfSumType), stream, mr); } if (count->type() != cudf::data_type(cudfCountType)) { - count = cudf::cast( - *count, cudf::data_type(cudfCountType), stream, get_output_mr()); + count = + cudf::cast(*count, cudf::data_type(cudfCountType), stream, mr); } auto children = std::vector>(); @@ -246,17 +508,16 @@ struct GroupbyMeanAggregator : GroupbyAggregator { auto count = std::move(results[countIdx_].results[0]); auto size = sum->size(); - auto const cudfSumType = cudf::data_type( - cudf_velox::veloxToCudfTypeId(outputType->childAt(0))); - auto const cudfCountType = cudf::data_type( - cudf_velox::veloxToCudfTypeId(outputType->childAt(1))); + auto const cudfSumType = + cudf_velox::veloxToCudfDataType(outputType->childAt(0)); + auto const cudfCountType = + cudf_velox::veloxToCudfDataType(outputType->childAt(1)); if (sum->type() != cudf::data_type(cudfSumType)) { - sum = cudf::cast( - *sum, cudf::data_type(cudfSumType), stream, get_output_mr()); + sum = cudf::cast(*sum, cudf::data_type(cudfSumType), stream, mr); } if (count->type() != cudf::data_type(cudfCountType)) { - count = cudf::cast( - *count, cudf::data_type(cudfCountType), stream, get_output_mr()); + count = + cudf::cast(*count, cudf::data_type(cudfCountType), stream, mr); } auto children = std::vector>(); @@ -278,9 +539,9 @@ struct GroupbyMeanAggregator : GroupbyAggregator { *sum, *count, cudf::binary_operator::DIV, - cudf::data_type(cudf_velox::veloxToCudfTypeId(resultType)), + cudf_velox::veloxToCudfDataType(resultType), stream, - get_output_mr()); + mr); return avg; } default: @@ -306,7 +567,8 @@ struct GroupbyStddevSampAggregator : GroupbyAggregator { void addGroupbyRequest( cudf::table_view const& tbl, - std::vector& requests) override { + std::vector& requests, + rmm::cuda_stream_view stream) override { auto& request = requests.emplace_back(); outputIdx_ = requests.size() - 1; request.values = tbl.column(inputIndex); @@ -340,7 +602,8 @@ struct GroupbyStddevSampAggregator : GroupbyAggregator { std::unique_ptr makeOutputColumn( std::vector& results, - rmm::cuda_stream_view stream) override { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) override { switch (step) { case core::AggregationNode::Step::kSingle: return std::move(results[outputIdx_].results[0]); @@ -349,19 +612,19 @@ struct GroupbyStddevSampAggregator : GroupbyAggregator { auto mean = std::move(results[outputIdx_].results[1]); auto m2 = std::move(results[outputIdx_].results[2]); return makeM2StructColumn( - std::move(count), std::move(mean), std::move(m2), stream); + std::move(count), std::move(mean), std::move(m2), stream, mr); } case core::AggregationNode::Step::kIntermediate: { auto merged = std::move(results[outputIdx_].results[0]); // Check if types already match expected output - avoid copies if so const auto& outputType = asRowType(resultType); - auto const cudfCountType = cudf::data_type( - cudf_velox::veloxToCudfTypeId(outputType->childAt(0))); - auto const cudfMeanType = cudf::data_type( - cudf_velox::veloxToCudfTypeId(outputType->childAt(1))); - auto const cudfM2Type = cudf::data_type( - cudf_velox::veloxToCudfTypeId(outputType->childAt(2))); + auto const cudfCountType = + cudf_velox::veloxToCudfDataType(outputType->childAt(0)); + auto const cudfMeanType = + cudf_velox::veloxToCudfDataType(outputType->childAt(1)); + auto const cudfM2Type = + cudf_velox::veloxToCudfDataType(outputType->childAt(2)); auto mergedView = merged->view(); bool typesMatch = mergedView.child(0).type() == cudfCountType && @@ -375,14 +638,14 @@ struct GroupbyStddevSampAggregator : GroupbyAggregator { // Types don't match - need to copy and cast (use output_mr since // these become part of the output) - auto count = std::make_unique( - mergedView.child(0), stream, get_output_mr()); - auto mean = std::make_unique( - mergedView.child(1), stream, get_output_mr()); - auto m2 = std::make_unique( - mergedView.child(2), stream, get_output_mr()); + auto count = + std::make_unique(mergedView.child(0), stream, mr); + auto mean = + std::make_unique(mergedView.child(1), stream, mr); + auto m2 = + std::make_unique(mergedView.child(2), stream, mr); return makeM2StructColumn( - std::move(count), std::move(mean), std::move(m2), stream); + std::move(count), std::move(mean), std::move(m2), stream, mr); } case core::AggregationNode::Step::kFinal: { // MERGE_M2 returns struct(count, mean, m2) @@ -428,8 +691,7 @@ struct GroupbyStddevSampAggregator : GroupbyAggregator { // Apply mask: where count < 2, result is NULL cudf::numeric_scalar nullDouble( 0.0, false, stream, get_temp_mr()); - return cudf::copy_if_else( - *stddev, nullDouble, *validMask, stream, get_output_mr()); + return cudf::copy_if_else(*stddev, nullDouble, *validMask, stream, mr); } default: VELOX_NYI("Unsupported aggregation step for stddev_samp"); @@ -442,23 +704,23 @@ struct GroupbyStddevSampAggregator : GroupbyAggregator { std::unique_ptr count, std::unique_ptr mean, std::unique_ptr m2, - rmm::cuda_stream_view stream) { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { const auto& outputType = asRowType(resultType); auto const cudfCountType = - cudf::data_type(cudf_velox::veloxToCudfTypeId(outputType->childAt(0))); + cudf_velox::veloxToCudfDataType(outputType->childAt(0)); auto const cudfMeanType = - cudf::data_type(cudf_velox::veloxToCudfTypeId(outputType->childAt(1))); + cudf_velox::veloxToCudfDataType(outputType->childAt(1)); auto const cudfM2Type = - cudf::data_type(cudf_velox::veloxToCudfTypeId(outputType->childAt(2))); - + cudf_velox::veloxToCudfDataType(outputType->childAt(2)); if (count->type() != cudfCountType) { - count = cudf::cast(*count, cudfCountType, stream, get_output_mr()); + count = cudf::cast(*count, cudfCountType, stream, mr); } if (mean->type() != cudfMeanType) { - mean = cudf::cast(*mean, cudfMeanType, stream, get_output_mr()); + mean = cudf::cast(*mean, cudfMeanType, stream, mr); } if (m2->type() != cudfM2Type) { - m2 = cudf::cast(*m2, cudfM2Type, stream, get_output_mr()); + m2 = cudf::cast(*m2, cudfM2Type, stream, mr); } auto const size = count->size(); @@ -484,6 +746,10 @@ std::unique_ptr createGroupbyAggregator( auto const& kind = p.kind; auto prefix = cudf_velox::CudfConfig::getInstance().functionNamePrefix; if (kind.rfind(prefix + "sum", 0) == 0) { + if (p.isDecimalAggregate) { + return std::make_unique( + p.companionStep, p.inputIndex, p.constant, p.resultType); + } return std::make_unique( p.companionStep, p.inputIndex, p.constant, p.resultType); } else if (kind.rfind(prefix + "count", 0) == 0) { @@ -497,6 +763,10 @@ std::unique_ptr createGroupbyAggregator( return std::make_unique( p.companionStep, p.inputIndex, p.constant, p.resultType); } else if (kind.rfind(prefix + "avg", 0) == 0) { + if (p.isDecimalAggregate) { + return std::make_unique( + p.companionStep, p.inputIndex, p.constant, p.resultType); + } return std::make_unique( p.companionStep, p.inputIndex, p.constant, p.resultType); } else if (kind.rfind(prefix + "stddev_samp", 0) == 0) { @@ -702,7 +972,8 @@ void CudfGroupby::computePartialGroupbyStreaming(CudfVectorPtr tbl) { groupingKeyOutputChannels_, aggregators_, bufferedResultType_, - inputTableStream); + inputTableStream, + get_output_mr()); // If we already have partial output, concatenate the new results with it. if (bufferedResult_) { @@ -714,7 +985,7 @@ void CudfGroupby::computePartialGroupbyStreaming(CudfVectorPtr tbl) { std::move(tablesToConcat), bufferedResultType_, partialOutputStream, - get_output_mr()); + get_temp_mr()); // Now we have to groupby again but this time with intermediate aggregators. // Keep concatenatedTable alive while we use its view. @@ -723,7 +994,8 @@ void CudfGroupby::computePartialGroupbyStreaming(CudfVectorPtr tbl) { groupingKeyOutputChannels_, intermediateAggregators_, bufferedResultType_, - partialOutputStream); + partialOutputStream, + get_output_mr()); bufferedResult_ = compactedOutput; } else { // First time processing, just store the result of the input batch's groupby @@ -743,7 +1015,8 @@ void CudfGroupby::computeFinalGroupbyStreaming(CudfVectorPtr tbl) { groupingKeyOutputChannels_, intermediateAggregators_, bufferedResultType_, - inputTableStream); + inputTableStream, + get_output_mr()); if (!groupbyOnInput) { return; } @@ -768,7 +1041,8 @@ void CudfGroupby::computeFinalGroupbyStreaming(CudfVectorPtr tbl) { groupingKeyOutputChannels_, intermediateAggregators_, bufferedResultType_, - finalStream); + finalStream, + get_output_mr()); bufferedResult_ = compactedOutput; } @@ -781,7 +1055,8 @@ void CudfGroupby::computeSingleGroupbyStreaming(CudfVectorPtr tbl) { groupingKeyOutputChannels_, partialAggregators_, bufferedResultType_, - inputTableStream); + inputTableStream, + get_output_mr()); if (bufferedResult_) { auto partialOutputStream = bufferedResult_->stream(); @@ -799,7 +1074,8 @@ void CudfGroupby::computeSingleGroupbyStreaming(CudfVectorPtr tbl) { groupingKeyOutputChannels_, intermediateAggregators_, bufferedResultType_, - partialOutputStream); + partialOutputStream, + get_output_mr()); bufferedResult_ = compactedOutput; } else { bufferedResult_ = groupbyOnInput; @@ -837,7 +1113,8 @@ CudfVectorPtr CudfGroupby::doGroupByAggregation( std::vector const& groupByKeys, std::vector>& aggregators, TypePtr const& outputType, - rmm::cuda_stream_view stream) { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { auto groupbyKeyView = tableView.select(groupByKeys.begin(), groupByKeys.end()); @@ -850,11 +1127,10 @@ CudfVectorPtr CudfGroupby::doGroupByAggregation( std::vector requests; for (auto& aggregator : aggregators) { - aggregator->addGroupbyRequest(tableView, requests); + aggregator->addGroupbyRequest(tableView, requests, stream); } - auto [groupKeys, results] = - groupByOwner.aggregate(requests, stream, get_output_mr()); + auto [groupKeys, results] = groupByOwner.aggregate(requests, stream, mr); // flatten the results std::vector> resultColumns; @@ -867,7 +1143,7 @@ CudfVectorPtr CudfGroupby::doGroupByAggregation( // then fill the aggregation results for (auto& aggregator : aggregators) { - resultColumns.push_back(aggregator->makeOutputColumn(results, stream)); + resultColumns.push_back(aggregator->makeOutputColumn(results, stream, mr)); } // make a cudf table out of columns @@ -951,7 +1227,8 @@ RowVectorPtr CudfGroupby::doGetOutput() { groupingKeyOutputChannels_, aggs, outputType_, - stream); + stream, + get_output_mr()); stream.synchronize(); bufferedResult_.reset(); return result; @@ -964,7 +1241,7 @@ RowVectorPtr CudfGroupby::doGetOutput() { auto stream = cudfGlobalStreamPool().get_stream(); auto tbl = getConcatenatedTable( - std::exchange(inputs_, {}), inputType_, stream, get_output_mr()); + std::exchange(inputs_, {}), inputType_, stream, get_temp_mr()); // Release input data after synchronizing. stream.synchronize(); @@ -983,7 +1260,8 @@ RowVectorPtr CudfGroupby::doGetOutput() { groupingKeyOutputChannels_, aggregators_, outputType_, - stream); + stream, + get_output_mr()); } void CudfGroupby::doNoMoreInput() { diff --git a/velox/experimental/cudf/exec/CudfGroupby.h b/velox/experimental/cudf/exec/CudfGroupby.h index 27ca367fa50..0d946145b33 100644 --- a/velox/experimental/cudf/exec/CudfGroupby.h +++ b/velox/experimental/cudf/exec/CudfGroupby.h @@ -30,11 +30,13 @@ struct GroupbyAggregator { virtual void addGroupbyRequest( cudf::table_view const& tbl, - std::vector& requests) = 0; + std::vector& requests, + rmm::cuda_stream_view stream) = 0; virtual std::unique_ptr makeOutputColumn( std::vector& results, - rmm::cuda_stream_view stream) = 0; + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) = 0; virtual ~GroupbyAggregator() = default; @@ -100,7 +102,8 @@ class CudfGroupby : public CudfOperatorBase { std::vector const& groupByKeys, std::vector>& aggregators, TypePtr const& outputType, - rmm::cuda_stream_view stream); + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); CudfVectorPtr releaseAndResetBufferedResult(); diff --git a/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp b/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp index de49cf1eeae..82a88bd8950 100644 --- a/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp +++ b/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp @@ -731,9 +731,10 @@ std::unique_ptr CudfNestedLoopJoinProbe::emitProbeMismatchRows( for (size_t i = 0; i < buildColumnOutputIndices_.size(); ++i) { auto outIdx = buildColumnOutputIndices_[i]; auto buildChannel = buildColumnIndicesToGather_[i]; - auto buildCudfType = veloxToCudfTypeId(buildType_->childAt(buildChannel)); + auto buildCudfDataType = + veloxToCudfDataType(buildType_->childAt(buildChannel)); auto nullScalar = cudf::make_default_constructed_scalar( - cudf::data_type{buildCudfType}, stream, get_temp_mr()); + buildCudfDataType, stream, get_temp_mr()); outCols[outIdx] = cudf::make_column_from_scalar( *nullScalar, numUnmatched, stream, get_output_mr()); } @@ -776,9 +777,10 @@ RowVectorPtr CudfNestedLoopJoinProbe::emitBuildMismatchRows( for (size_t li = 0; li < probeColumnOutputIndices_.size(); ++li) { auto outIdx = probeColumnOutputIndices_[li]; auto probeChannel = probeColumnIndicesToGather_[li]; - auto probeCudfType = veloxToCudfTypeId(probeType_->childAt(probeChannel)); + auto probeCudfDataType = + veloxToCudfDataType(probeType_->childAt(probeChannel)); auto nullScalar = cudf::make_default_constructed_scalar( - cudf::data_type{probeCudfType}, stream, get_temp_mr()); + probeCudfDataType, stream, get_temp_mr()); outCols[outIdx] = cudf::make_column_from_scalar( *nullScalar, numUnmatched, stream, get_output_mr()); } diff --git a/velox/experimental/cudf/exec/CudfReduce.cpp b/velox/experimental/cudf/exec/CudfReduce.cpp index dbf71b5143f..ef2e2cdc701 100644 --- a/velox/experimental/cudf/exec/CudfReduce.cpp +++ b/velox/experimental/cudf/exec/CudfReduce.cpp @@ -19,6 +19,8 @@ #include "velox/experimental/cudf/exec/CudfAggregation.h" #include "velox/experimental/cudf/exec/CudfFilterProject.h" #include "velox/experimental/cudf/exec/CudfReduce.h" +#include "velox/experimental/cudf/exec/DecimalAggregationHostOps.h" +#include "velox/experimental/cudf/exec/DecimalAggregationState.h" #include "velox/experimental/cudf/exec/GpuResources.h" #include "velox/experimental/cudf/exec/Utilities.h" #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" @@ -27,6 +29,7 @@ #include "velox/exec/AggregateFunctionRegistry.h" #include "velox/exec/Task.h" #include "velox/expression/Expr.h" +#include "velox/type/Type.h" #include #include @@ -39,39 +42,42 @@ namespace { using namespace facebook::velox; +using facebook::velox::cudf_velox::castDecimal64InputToDecimal128; using facebook::velox::cudf_velox::CountInputKind; +using facebook::velox::cudf_velox::finalizeDecimalAverage; using facebook::velox::cudf_velox::get_output_mr; using facebook::velox::cudf_velox::get_temp_mr; using facebook::velox::cudf_velox::ReduceAggregator; using facebook::velox::cudf_velox::ResolvedAggregateInfo; - -#define DEFINE_SIMPLE_REDUCE_AGGREGATOR(Name, name) \ - struct Reduce##Name##Aggregator : ReduceAggregator { \ - Reduce##Name##Aggregator( \ - core::AggregationNode::Step step, \ - uint32_t inputIndex, \ - VectorPtr constant, \ - const TypePtr& resultType) \ - : ReduceAggregator(step, inputIndex, constant, resultType) {} \ - \ - std::unique_ptr doReduce( \ - cudf::table_view const& input, \ - TypePtr const& outputType, \ - rmm::cuda_stream_view stream, \ - vector_size_t /*inputRowCount*/) override { \ - auto const aggRequest = \ - cudf::make_##name##_aggregation(); \ - auto const cudfOutputType = \ - cudf::data_type(cudf_velox::veloxToCudfTypeId(outputType)); \ - auto const resultScalar = cudf::reduce( \ - input.column(inputIndex), \ - *aggRequest, \ - cudfOutputType, \ - stream, \ - get_temp_mr()); \ - return cudf::make_column_from_scalar( \ - *resultScalar, 1, stream, get_output_mr()); \ - } \ +using facebook::velox::cudf_velox::serializeDecimalPartialOrIntermediateState; +using facebook::velox::cudf_velox::validateIntermediateColumnType; + +#define DEFINE_SIMPLE_REDUCE_AGGREGATOR(Name, name) \ + struct Reduce##Name##Aggregator : ReduceAggregator { \ + Reduce##Name##Aggregator( \ + core::AggregationNode::Step step, \ + uint32_t inputIndex, \ + VectorPtr constant, \ + const TypePtr& resultType) \ + : ReduceAggregator(step, inputIndex, constant, resultType) {} \ + \ + std::unique_ptr doReduce( \ + cudf::table_view const& input, \ + TypePtr const& outputType, \ + vector_size_t /* inputRowCount */, \ + rmm::cuda_stream_view stream, \ + rmm::device_async_resource_ref mr) override { \ + auto const aggRequest = \ + cudf::make_##name##_aggregation(); \ + auto const cudfOutputType = cudf_velox::veloxToCudfDataType(outputType); \ + auto const resultScalar = cudf::reduce( \ + input.column(inputIndex), \ + *aggRequest, \ + cudfOutputType, \ + stream, \ + get_temp_mr()); \ + return cudf::make_column_from_scalar(*resultScalar, 1, stream, mr); \ + } \ }; DEFINE_SIMPLE_REDUCE_AGGREGATOR(Sum, sum) @@ -90,8 +96,9 @@ struct ReduceCountAggregator : ReduceAggregator { std::unique_ptr doReduce( cudf::table_view const& input, TypePtr const& outputType, + vector_size_t inputRowCount, rmm::cuda_stream_view stream, - vector_size_t inputRowCount) override { + rmm::device_async_resource_ref mr) override { if (exec::isRawInput(step)) { int64_t count; switch (inputKind_) { @@ -117,8 +124,7 @@ struct ReduceCountAggregator : ReduceAggregator { auto resultScalar = cudf::numeric_scalar(count, true, stream, get_temp_mr()); - return cudf::make_column_from_scalar( - resultScalar, 1, stream, get_output_mr()); + return cudf::make_column_from_scalar(resultScalar, 1, stream, mr); } else { // For non-raw input (intermediate/final), use sum aggregation auto const aggRequest = @@ -131,8 +137,7 @@ struct ReduceCountAggregator : ReduceAggregator { stream, get_temp_mr()); resultScalar->set_valid_async(true, stream); - return cudf::make_column_from_scalar( - *resultScalar, 1, stream, get_output_mr()); + return cudf::make_column_from_scalar(*resultScalar, 1, stream, mr); } } @@ -151,32 +156,29 @@ struct ReduceMeanAggregator : ReduceAggregator { std::unique_ptr doReduce( cudf::table_view const& input, TypePtr const& outputType, + vector_size_t /* inputRowCount */, rmm::cuda_stream_view stream, - vector_size_t /*inputRowCount*/) override { + rmm::device_async_resource_ref mr) override { switch (step) { case core::AggregationNode::Step::kSingle: { auto const aggRequest = cudf::make_mean_aggregation(); - auto const cudfOutputType = - cudf::data_type(cudf_velox::veloxToCudfTypeId(outputType)); + auto const cudfOutputType = cudf_velox::veloxToCudfDataType(outputType); auto const resultScalar = cudf::reduce( input.column(inputIndex), *aggRequest, cudfOutputType, stream, get_temp_mr()); - return cudf::make_column_from_scalar( - *resultScalar, 1, stream, get_output_mr()); + return cudf::make_column_from_scalar(*resultScalar, 1, stream, mr); } case core::AggregationNode::Step::kPartial: { VELOX_CHECK(outputType->isRow()); auto const& rowType = outputType->asRow(); auto const sumType = rowType.childAt(0); auto const countType = rowType.childAt(1); - auto const cudfSumType = - cudf::data_type(cudf_velox::veloxToCudfTypeId(sumType)); - auto const cudfCountType = - cudf::data_type(cudf_velox::veloxToCudfTypeId(countType)); + auto const cudfSumType = cudf_velox::veloxToCudfDataType(sumType); + auto const cudfCountType = cudf_velox::veloxToCudfDataType(countType); // sum auto const aggRequest = @@ -187,8 +189,8 @@ struct ReduceMeanAggregator : ReduceAggregator { cudfSumType, stream, get_temp_mr()); - auto sumCol = cudf::make_column_from_scalar( - *sumResultScalar, 1, stream, get_output_mr()); + auto sumCol = + cudf::make_column_from_scalar(*sumResultScalar, 1, stream, mr); // libcudf doesn't have a count agg for reduce. What we want is to // count the number of valid rows. @@ -201,7 +203,7 @@ struct ReduceMeanAggregator : ReduceAggregator { get_temp_mr()), 1, stream, - get_output_mr()); + mr); // Assemble into struct as expected by velox. auto children = std::vector>(); @@ -225,8 +227,8 @@ struct ReduceMeanAggregator : ReduceAggregator { cudf::make_sum_aggregation(); auto const sumResultScalar = cudf::reduce( sumCol, *sumAggRequest, sumCol.type(), stream, get_temp_mr()); - auto sumResultCol = cudf::make_column_from_scalar( - *sumResultScalar, 1, stream, get_output_mr()); + auto sumResultCol = + cudf::make_column_from_scalar(*sumResultScalar, 1, stream, mr); // sum the counts auto const countAggRequest = @@ -235,15 +237,14 @@ struct ReduceMeanAggregator : ReduceAggregator { countCol, *countAggRequest, countCol.type(), stream, get_temp_mr()); // divide the sums by the counts - auto const cudfOutputType = - cudf::data_type(cudf_velox::veloxToCudfTypeId(outputType)); + auto const cudfOutputType = cudf_velox::veloxToCudfDataType(outputType); return cudf::binary_operation( *sumResultCol, *countResultScalar, cudf::binary_operator::DIV, cudfOutputType, stream, - get_output_mr()); + mr); } default: VELOX_NYI("Unsupported aggregation step for mean"); @@ -251,6 +252,236 @@ struct ReduceMeanAggregator : ReduceAggregator { } }; +// Materializes reduced sum/count scalars into 1-row columns. +cudf_velox::DecimalSumStateColumns makeSumCountColumns( + cudf::scalar const& sumScalar, + cudf::scalar const& countScalar, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + cudf_velox::DecimalSumStateColumns cols; + cols.sum = cudf::make_column_from_scalar(sumScalar, 1, stream, mr); + cols.count = cudf::make_column_from_scalar(countScalar, 1, stream, mr); + return cols; +} + +std::unique_ptr partialDecimalSumCountToSerializedString( + cudf::column_view inputCol, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + std::unique_ptr castedInput; + inputCol = castDecimal64InputToDecimal128(inputCol, castedInput, stream); + auto const sumAgg = cudf::make_sum_aggregation(); + auto sumScalar = + cudf::reduce(inputCol, *sumAgg, inputCol.type(), stream, get_temp_mr()); + auto countAgg = cudf::make_count_aggregation( + cudf::null_policy::EXCLUDE); + auto countScalar = cudf::reduce( + inputCol, + *countAgg, + cudf::data_type{cudf::type_id::INT64}, + stream, + get_temp_mr()); + auto cols = makeSumCountColumns(*sumScalar, *countScalar, stream, mr); + return serializeDecimalPartialOrIntermediateState( + std::move(cols.sum), std::move(cols.count), stream, mr); +} + +// Decodes serialized decimal SUM state, sums the per-row partial sums and +// counts, and returns them as 1-row columns. Shared by the intermediate and +// final reduce steps before re-serializing or finalizing. The merged columns +// are consumed by the caller and never leave the operator, so they come from +// the temporary memory resource. +cudf_velox::DecimalSumStateColumns mergeSerializedDecimalSumState( + cudf::column_view inputCol, + int32_t scale, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto const sumAgg = cudf::make_sum_aggregation(); + auto sumAndCount = + cudf_velox::deserializeDecimalSumState(inputCol, scale, stream); + auto sumScalar = cudf::reduce( + sumAndCount.sum->view(), + *sumAgg, + sumAndCount.sum->view().type(), + stream, + mr); + auto countScalar = cudf::reduce( + sumAndCount.count->view(), + *sumAgg, + cudf::data_type{cudf::type_id::INT64}, + stream, + mr); + return makeSumCountColumns(*sumScalar, *countScalar, stream, mr); +} + +std::unique_ptr intermediateDecimalMergeSerializedString( + cudf::column_view inputCol, + int32_t scale, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto merged = mergeSerializedDecimalSumState(inputCol, scale, stream, mr); + return serializeDecimalPartialOrIntermediateState( + std::move(merged.sum), std::move(merged.count), stream, mr); +} + +std::unique_ptr finalDecimalAvgFromSerializedString( + cudf::column_view inputCol, + int32_t scale, + TypePtr const& resultType, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto merged = mergeSerializedDecimalSumState(inputCol, scale, stream, mr); + return finalizeDecimalAverage( + std::move(merged.sum), std::move(merged.count), resultType, stream, mr); +} + +std::unique_ptr singleDecimalAvgFromRawColumn( + cudf::column_view inputCol, + TypePtr const& resultType, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + std::unique_ptr castedInput; + inputCol = castDecimal64InputToDecimal128(inputCol, castedInput, stream); + auto const sumAgg = cudf::make_sum_aggregation(); + auto sumScalar = + cudf::reduce(inputCol, *sumAgg, inputCol.type(), stream, get_temp_mr()); + auto countAgg = cudf::make_count_aggregation( + cudf::null_policy::EXCLUDE); + auto countScalar = cudf::reduce( + inputCol, + *countAgg, + cudf::data_type{cudf::type_id::INT64}, + stream, + get_temp_mr()); + auto cols = makeSumCountColumns(*sumScalar, *countScalar, stream, mr); + return finalizeDecimalAverage( + std::move(cols.sum), std::move(cols.count), resultType, stream, mr); +} + +std::unique_ptr singleOrRawDecimalSumWithCast( + cudf::column_view inputCol, + TypePtr const& outputType, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto const sumAgg = cudf::make_sum_aggregation(); + auto const cudfOutType = cudf_velox::veloxToCudfDataType(outputType); + std::unique_ptr castedInput; + if (outputType->isDecimal() && inputCol.type() != cudfOutType) { + castedInput = cudf::cast(inputCol, cudfOutType, stream, get_temp_mr()); + inputCol = castedInput->view(); + } + auto const resultScalar = + cudf::reduce(inputCol, *sumAgg, cudfOutType, stream, get_temp_mr()); + return cudf::make_column_from_scalar(*resultScalar, 1, stream, mr); +} + +std::unique_ptr reduceIntermediateDecimalFromSerializedColumn( + cudf::column_view inputCol, + TypePtr const& outputType, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + validateIntermediateColumnType(inputCol); + // outputType here could be DECIMAL or VARBINARY + auto scale = outputType->isDecimal() + ? getDecimalPrecisionScale(*outputType).second + : 0; + return intermediateDecimalMergeSerializedString(inputCol, scale, stream, mr); +} + +std::unique_ptr reduceFinalDecimalSumFromSerializedColumn( + cudf::column_view inputCol, + TypePtr const& outputType, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + validateIntermediateColumnType(inputCol); + auto scale = getDecimalPrecisionScale(*outputType).second; + auto sumAndCount = + cudf_velox::deserializeDecimalSumState(inputCol, scale, stream); + return singleOrRawDecimalSumWithCast( + sumAndCount.sum->view(), outputType, stream, mr); +} + +std::unique_ptr reduceFinalDecimalAvgFromSerializedColumn( + cudf::column_view inputCol, + TypePtr const& outputType, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + validateIntermediateColumnType(inputCol); + auto scale = getDecimalPrecisionScale(*outputType).second; + return finalDecimalAvgFromSerializedString( + inputCol, scale, outputType, stream, mr); +} + +// Decimal SUM and AVG use dedicated aggregators rather than cudf::reduce's +// built-in sum/mean: partial/intermediate state is VARBINARY-encoded sum+count +// (see DecimalAggregationState), and the final divide needs decimal half-up +// rounding. +struct ReduceDecimalSumAggregator : ReduceAggregator { + ReduceDecimalSumAggregator( + core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + const TypePtr& resultType) + : ReduceAggregator(step, inputIndex, constant, resultType) {} + + std::unique_ptr doReduce( + cudf::table_view const& input, + TypePtr const& outputType, + vector_size_t /* inputRowCount */, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) override { + cudf::column_view inputCol = input.column(inputIndex); + switch (step) { + case core::AggregationNode::Step::kSingle: + return singleOrRawDecimalSumWithCast(inputCol, outputType, stream, mr); + case core::AggregationNode::Step::kPartial: + return partialDecimalSumCountToSerializedString(inputCol, stream, mr); + case core::AggregationNode::Step::kIntermediate: + return reduceIntermediateDecimalFromSerializedColumn( + inputCol, outputType, stream, mr); + case core::AggregationNode::Step::kFinal: + return reduceFinalDecimalSumFromSerializedColumn( + inputCol, outputType, stream, mr); + default: + VELOX_NYI("Unsupported aggregation step for decimal sum reduce"); + } + } +}; + +struct ReduceDecimalAvgAggregator : ReduceAggregator { + ReduceDecimalAvgAggregator( + core::AggregationNode::Step step, + uint32_t inputIndex, + VectorPtr constant, + const TypePtr& resultType) + : ReduceAggregator(step, inputIndex, constant, resultType) {} + + std::unique_ptr doReduce( + cudf::table_view const& input, + TypePtr const& outputType, + vector_size_t /* inputRowCount */, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) override { + cudf::column_view inputCol = input.column(inputIndex); + switch (step) { + case core::AggregationNode::Step::kSingle: + return singleDecimalAvgFromRawColumn(inputCol, resultType, stream, mr); + case core::AggregationNode::Step::kPartial: + return partialDecimalSumCountToSerializedString(inputCol, stream, mr); + case core::AggregationNode::Step::kIntermediate: + return reduceIntermediateDecimalFromSerializedColumn( + inputCol, outputType, stream, mr); + case core::AggregationNode::Step::kFinal: + VELOX_CHECK( + *outputType == *resultType, "outputType/resultType mismatch"); + return reduceFinalDecimalAvgFromSerializedColumn( + inputCol, outputType, stream, mr); + default: + VELOX_NYI("Unsupported aggregation step for decimal avg reduce"); + } + } +}; + struct ApproxDistinctAggregator : ReduceAggregator { static constexpr cudf::null_policy kNullPolicy = cudf::null_policy::EXCLUDE; static constexpr cudf::nan_policy kNanPolicy = cudf::nan_policy::NAN_IS_VALID; @@ -272,25 +503,27 @@ struct ApproxDistinctAggregator : ReduceAggregator { std::unique_ptr doReduce( cudf::table_view const& input, TypePtr const& outputType, + vector_size_t /* inputRowCount */, rmm::cuda_stream_view stream, - vector_size_t /*inputRowCount*/) override { + rmm::device_async_resource_ref mr) override { if (exec::isRawInput(step)) { - return doPartialReduce(input, stream); + return doPartialReduce(input, stream, mr); } else if (step == core::AggregationNode::Step::kIntermediate) { - return doIntermediateReduce(input, stream); + return doIntermediateReduce(input, stream, mr); } else { - return doFinalReduce(input, stream); + return doFinalReduce(input, stream, mr); } } private: std::unique_ptr makeSketchColumn( cuda::std::span sketch_bytes, - rmm::cuda_stream_view stream) { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { auto sketch_size = static_cast(sketch_bytes.size()); cudf::size_type offsets[2] = {0, sketch_size}; - rmm::device_buffer offsets_device{2 * sizeof(cudf::size_type), stream}; + rmm::device_buffer offsets_device{2 * sizeof(cudf::size_type), stream, mr}; CUDF_CUDA_TRY(cudaMemcpyAsync( offsets_device.data(), offsets, @@ -298,7 +531,7 @@ struct ApproxDistinctAggregator : ReduceAggregator { cudaMemcpyHostToDevice, stream.value())); - rmm::device_buffer chars_buffer{sketch_bytes.size(), stream}; + rmm::device_buffer chars_buffer{sketch_bytes.size(), stream, mr}; CUDF_CUDA_TRY(cudaMemcpyAsync( chars_buffer.data(), sketch_bytes.data(), @@ -390,35 +623,38 @@ struct ApproxDistinctAggregator : ReduceAggregator { std::unique_ptr doPartialReduce( cudf::table_view const& input, - rmm::cuda_stream_view stream) { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { auto inputTable = cudf::table_view({input.column(inputIndex)}); cudf::approx_distinct_count sketch{ inputTable, precision_, kNullPolicy, kNanPolicy, stream}; - return makeSketchColumn(sketch.sketch(), stream); + return makeSketchColumn(sketch.sketch(), stream, mr); } std::unique_ptr doIntermediateReduce( cudf::table_view const& input, - rmm::cuda_stream_view stream) { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { auto sketch_column = input.column(inputIndex); if (sketch_column.size() == 0) { - return makeSketchColumn({}, stream); + return makeSketchColumn({}, stream, mr); } return mergeSketchesAndApply( sketch_column, - [this, stream](cudf::approx_distinct_count& sketch) { - return makeSketchColumn(sketch.sketch(), stream); + [this, stream, mr](cudf::approx_distinct_count& sketch) { + return makeSketchColumn(sketch.sketch(), stream, mr); }, stream); } std::unique_ptr doFinalReduce( cudf::table_view const& input, - rmm::cuda_stream_view stream) { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { auto sketch_column = input.column(inputIndex); if (sketch_column.size() == 0) { @@ -426,19 +662,19 @@ struct ApproxDistinctAggregator : ReduceAggregator { cudf::numeric_scalar(0, true, stream, get_temp_mr()), 1, stream, - get_output_mr()); + mr); } return mergeSketchesAndApply( sketch_column, - [stream](cudf::approx_distinct_count& sketch) { + [stream, mr](cudf::approx_distinct_count& sketch) { std::size_t estimate = sketch.estimate(stream); return cudf::make_column_from_scalar( cudf::numeric_scalar( static_cast(estimate), true, stream, get_temp_mr()), 1, stream, - get_output_mr()); + mr); }, stream); } @@ -451,6 +687,10 @@ std::unique_ptr createReduceAggregator( auto const& kind = p.kind; auto prefix = cudf_velox::CudfConfig::getInstance().functionNamePrefix; if (kind.rfind(prefix + "sum", 0) == 0) { + if (p.isDecimalAggregate) { + return std::make_unique( + p.companionStep, p.inputIndex, p.constant, p.resultType); + } return std::make_unique( p.companionStep, p.inputIndex, p.constant, p.resultType); } else if (kind.rfind(prefix + "count", 0) == 0) { @@ -464,6 +704,10 @@ std::unique_ptr createReduceAggregator( return std::make_unique( p.companionStep, p.inputIndex, p.constant, p.resultType); } else if (kind.rfind(prefix + "avg", 0) == 0) { + if (p.isDecimalAggregate) { + return std::make_unique( + p.companionStep, p.inputIndex, p.constant, p.resultType); + } return std::make_unique( p.companionStep, p.inputIndex, p.constant, p.resultType); } else if (kind.rfind(prefix + "approx_distinct", 0) == 0) { @@ -605,13 +849,14 @@ void CudfReduce::doAddInput(RowVectorPtr input) { CudfVectorPtr CudfReduce::doGlobalAggregation( cudf::table_view tableView, - rmm::cuda_stream_view stream) { + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { std::vector> resultColumns; resultColumns.reserve(aggregators_.size()); for (auto i = 0; i < aggregators_.size(); i++) { resultColumns.push_back( aggregators_[i]->doReduce( - tableView, outputType_->childAt(i), stream, numInputRows_)); + tableView, outputType_->childAt(i), numInputRows_, stream, mr)); } return std::make_shared( @@ -640,7 +885,7 @@ RowVectorPtr CudfReduce::doGetOutput() { auto stream = cudfGlobalStreamPool().get_stream(); auto tbl = getConcatenatedTable( - std::move(inputs_), inputType_, stream, get_output_mr()); + std::move(inputs_), inputType_, stream, get_temp_mr()); // Release input data after synchronizing. stream.synchronize(); @@ -656,7 +901,7 @@ RowVectorPtr CudfReduce::doGetOutput() { ? tbl->view() : tbl->view().select( aggregationInputChannels_.begin(), aggregationInputChannels_.end()); - auto output = doGlobalAggregation(tableView, stream); + auto output = doGlobalAggregation(tableView, stream, get_output_mr()); if (isPartialOutput_ && !noMoreInput_) { numInputRows_ = 0; } diff --git a/velox/experimental/cudf/exec/CudfReduce.h b/velox/experimental/cudf/exec/CudfReduce.h index 117adc42f71..d983985afb6 100644 --- a/velox/experimental/cudf/exec/CudfReduce.h +++ b/velox/experimental/cudf/exec/CudfReduce.h @@ -29,8 +29,9 @@ struct ReduceAggregator { virtual std::unique_ptr doReduce( cudf::table_view const& input, TypePtr const& outputType, + vector_size_t inputRowCount, rmm::cuda_stream_view stream, - vector_size_t inputRowCount) = 0; + rmm::device_async_resource_ref mr) = 0; virtual ~ReduceAggregator() = default; @@ -91,7 +92,8 @@ class CudfReduce : public CudfOperatorBase { private: CudfVectorPtr doGlobalAggregation( cudf::table_view tableView, - rmm::cuda_stream_view stream); + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); std::shared_ptr aggregationNode_; std::vector> aggregators_; diff --git a/velox/experimental/cudf/exec/DecimalAggregationDevice.cu b/velox/experimental/cudf/exec/DecimalAggregationDevice.cu new file mode 100644 index 00000000000..16d26df831b --- /dev/null +++ b/velox/experimental/cudf/exec/DecimalAggregationDevice.cu @@ -0,0 +1,382 @@ +/* + * 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/DecimalAggregationDevice.h" + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace facebook::velox::cudf_velox { +namespace { + +// Mirrors the CPU LongDecimalWithOverflowState layout so serialized SUM state +// is interchangeable between CPU and GPU aggregation. +// TODO: Track int128 overflow as the CPU does (DecimalUtil::addWithOverflow); +// the `overflow` field is reserved for that and is always 0 until then. +struct DecimalSumState { + int64_t count; // count of non-null input rows aggregated + int64_t overflow; // net int128 carries (CPU parity); always 0 on GPU for now + uint64_t lower; // lower 64 bits of the decimal sum + int64_t upper; // upper 64 bits of the decimal sum (signed) +}; + +static_assert(sizeof(DecimalSumState) == detail::kDecimalSumStateSize); + +__device__ __forceinline__ void +splitToWords(int64_t value, int64_t& upper, uint64_t& lower) { + lower = static_cast(value); + upper = value < 0 ? -1 : 0; +} + +__device__ __forceinline__ void +splitToWords(__int128_t value, int64_t& upper, uint64_t& lower) { + lower = static_cast(value); + upper = static_cast(value >> 64); +} + +template +struct FillOffsetsFunctor { + cuda::std::span offsets; + + __device__ void operator()(cudf::size_type idx) const { + int64_t offset = static_cast(idx) * detail::kDecimalSumStateSize; + offsets[idx] = static_cast(offset); + } +}; + +template +struct PackStateFunctor { + cuda::std::span sums; + cuda::std::span counts; + cuda::std::span offsets; + uint8_t* chars; + + __device__ void operator()(cudf::size_type idx) const { + int64_t offset = static_cast(offsets[idx]); + auto* state = reinterpret_cast(chars + offset); + int64_t upper; + uint64_t lower; + splitToWords(sums[idx], upper, lower); + state->count = counts[idx]; + state->overflow = 0; + state->lower = lower; + state->upper = upper; + } +}; + +template +struct UnpackStateFunctor { + cuda::std::span offsets; + const uint8_t* chars; + cuda::std::span<__int128_t> sums; + cuda::std::span counts; + + __device__ void operator()(cudf::size_type idx) const { + int64_t offset = static_cast(offsets[idx]); + auto* state = reinterpret_cast(chars + offset); + counts[idx] = state->count; + sums[idx] = (static_cast<__int128_t>(state->upper) << 64) | state->lower; + } +}; + +// Half-up sum/count divide for AVG. +template +struct AvgRoundFunctor { + cuda::std::span sums; + cuda::std::span counts; + cuda::std::span out; + + __device__ void operator()(cudf::size_type idx) const { + auto count = counts[idx]; + if (count == 0) { + out[idx] = SumT{0}; + return; + } + auto sum = sums[idx]; + using U = cuda::std::make_unsigned_t; + U absSum = sum < 0 ? -static_cast(sum) : static_cast(sum); + U half = static_cast(count / 2); + U rounded = (absSum + half) / static_cast(count); + // Use `U{0} - rounded` below to avoid signed overflow + out[idx] = static_cast(sum < 0 ? U{0} - rounded : rounded); + } +}; + +template +void launchDeviceFor( + cudf::size_type size, + BuildOp buildOp, + rmm::cuda_stream_view stream) { + if (size == 0) { + return; + } + auto op = buildOp(); + cub::DeviceFor::ForEachN( + cuda::counting_iterator{0}, size, op, stream.value()); + CUDF_CUDA_TRY(cudaGetLastError()); +} + +struct StateValidPredicate { + cudf::column_device_view sum; + cudf::column_device_view count; + + __device__ bool operator()(cudf::size_type idx) const { + if (sum.is_null(idx) || count.is_null(idx)) { + return false; + } + return count.element(idx) != 0; + } +}; + +std::pair buildStateValidityMaskImpl( + const cudf::column_view& sumCol, + const cudf::column_view& countCol, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto numRows = sumCol.size(); + if (numRows == 0) { + return {rmm::device_buffer{}, 0}; + } + auto sumDeviceView = cudf::column_device_view::create(sumCol, stream); + auto countDeviceView = cudf::column_device_view::create(countCol, stream); + StateValidPredicate pred{*sumDeviceView, *countDeviceView}; + // Build a BOOL8 column of per-row validity, then convert via the public API. + auto bools = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::BOOL8}, + numRows, + cudf::mask_state::UNALLOCATED, + stream, + mr); + auto iter = cuda::counting_iterator{0}; + thrust::transform( + rmm::exec_policy(stream), + iter, + iter + numRows, + bools->mutable_view().begin(), + pred); + auto [mask, nullCount] = cudf::bools_to_mask(bools->view(), stream, mr); + return {std::move(*mask), nullCount}; +} + +} // namespace + +namespace detail { + +template +concept OffsetStorageType = + std::same_as || std::same_as; + +template +concept DecimalSumStorageType = + std::same_as || std::same_as; + +template +concept ValidDecimalPackStorageTypes = + DecimalSumStorageType && OffsetStorageType; + +struct fillOffsetsForDecimalSumStateKernel { + cudf::mutable_column_view offsetsView; + cudf::size_type numRows; + rmm::cuda_stream_view stream; + + template + requires OffsetStorageType + void operator()() const { + launchDeviceFor( + numRows + 1, + [&] { + return FillOffsetsFunctor{cuda::std::span{ + offsetsView.data(), static_cast(numRows) + 1}}; + }, + stream); + } + + template + requires(!OffsetStorageType) + void operator()() const { + CUDF_FAIL("Invalid offset type for decimal sum state"); + } +}; + +struct unpackDecimalSumStateKernel { + cudf::column_view offsetsView; + const uint8_t* chars; + cudf::mutable_column_view sumView; + cudf::mutable_column_view countView; + cudf::size_type numRows; + rmm::cuda_stream_view stream; + + template + requires OffsetStorageType + void operator()() const { + auto const n = static_cast(numRows); + launchDeviceFor( + numRows, + [&] { + return UnpackStateFunctor{ + cuda::std::span{offsetsView.data(), n}, + chars, + cuda::std::span<__int128_t>{sumView.data<__int128_t>(), n}, + cuda::std::span{countView.data(), n}}; + }, + stream); + } + + template + requires(!OffsetStorageType) + void operator()() const { + CUDF_FAIL("Invalid offset type for decimal sum state"); + } +}; + +struct averageRoundDecimalSumKernel { + cudf::column_view sumCol; + const int64_t* counts; + cudf::mutable_column_view outView; + cudf::size_type numRows; + rmm::cuda_stream_view stream; + + template + requires DecimalSumStorageType + void operator()() const { + auto const n = static_cast(numRows); + launchDeviceFor( + numRows, + [&] { + return AvgRoundFunctor{ + cuda::std::span{sumCol.data(), n}, + cuda::std::span{counts, n}, + cuda::std::span{outView.data(), n}}; + }, + stream); + } + + template + requires(!DecimalSumStorageType) + void operator()() const { + CUDF_FAIL("Invalid sum type for decimal average"); + } +}; + +struct packDecimalSumStateKernel { + cudf::column_view sumCol; + const int64_t* counts; + cudf::column_view offsetsView; + uint8_t* chars; + cudf::size_type numRows; + rmm::cuda_stream_view stream; + + template + requires ValidDecimalPackStorageTypes + void operator()() const { + auto const n = static_cast(numRows); + auto const sums = sumCol.data(); + launchDeviceFor( + numRows, + [&] { + return PackStateFunctor{ + cuda::std::span{sums, n}, + cuda::std::span{counts, n}, + cuda::std::span{offsetsView.data(), n}, + chars}; + }, + stream); + } + + template + requires(!ValidDecimalPackStorageTypes) + void operator()() const { + CUDF_FAIL("Invalid types for decimal sum state pack"); + } +}; + +void fillOffsetsForDecimalSumState( + cudf::type_id offsetType, + cudf::mutable_column_view offsetsView, + cudf::size_type numRows, + rmm::cuda_stream_view stream) { + cudf::type_dispatcher( + cudf::data_type{offsetType}, + fillOffsetsForDecimalSumStateKernel{offsetsView, numRows, stream}); +} + +void unpackDecimalSumState( + cudf::type_id offsetType, + cudf::column_view offsetsView, + const uint8_t* chars, + cudf::mutable_column_view sumView, + cudf::mutable_column_view countView, + cudf::size_type numRows, + rmm::cuda_stream_view stream) { + cudf::type_dispatcher( + cudf::data_type{offsetType}, + unpackDecimalSumStateKernel{ + offsetsView, chars, sumView, countView, numRows, stream}); +} + +void averageRoundDecimalSum( + cudf::type_id sumType, + cudf::column_view sumCol, + const int64_t* counts, + cudf::mutable_column_view outView, + cudf::size_type numRows, + rmm::cuda_stream_view stream) { + cudf::type_dispatcher( + cudf::data_type{sumType}, + averageRoundDecimalSumKernel{sumCol, counts, outView, numRows, stream}); +} + +void packDecimalSumState( + cudf::type_id sumType, + cudf::type_id offsetType, + cudf::column_view sumCol, + const int64_t* counts, + cudf::column_view offsetsView, + uint8_t* chars, + cudf::size_type numRows, + rmm::cuda_stream_view stream) { + cudf::double_type_dispatcher( + cudf::data_type{sumType}, + cudf::data_type{offsetType}, + packDecimalSumStateKernel{ + sumCol, counts, offsetsView, chars, numRows, stream}); +} + +std::pair buildStateValidityMask( + const cudf::column_view& sumCol, + const cudf::column_view& countCol, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return buildStateValidityMaskImpl(sumCol, countCol, stream, mr); +} + +} // namespace detail +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/DecimalAggregationDevice.h b/velox/experimental/cudf/exec/DecimalAggregationDevice.h new file mode 100644 index 00000000000..08cc37e4018 --- /dev/null +++ b/velox/experimental/cudf/exec/DecimalAggregationDevice.h @@ -0,0 +1,132 @@ +/* + * 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 +#include +#include + +namespace facebook::velox::cudf_velox::detail { + +// Size in bytes of each row's packed decimal SUM intermediate state in the +// strings payload (count, overflow placeholder, and 128-bit sum split into +// words). +constexpr size_t kDecimalSumStateSize = 32; + +/** + * Writes strings-style prefix offsets: offset[i] == i * kDecimalSumStateSize. + * + * @param offsetType INT32 or INT64; selects offset storage width via + * cudf::type_dispatcher. + * @param offsetsView output offsets column of numRows + 1 elements. + * @param numRows number of payload rows. + * @param stream CUDA stream for the launch. + */ +void fillOffsetsForDecimalSumState( + cudf::type_id offsetType, + cudf::mutable_column_view offsetsView, + cudf::size_type numRows, + rmm::cuda_stream_view stream); + +/** + * Encodes each row's partial sum and count into the fixed-width device layout + * used for VARBINARY interchange. + * + * @param sumType DECIMAL64 or DECIMAL128; selects sum storage width. + * @param offsetType INT32 or INT64; selects offset storage width. sumType and + * offsetType are dispatched via cudf::double_type_dispatcher. + * @param sumCol per-row sums. + * @param counts per-row int64 counts. + * @param offsetsView per-row byte offsets into chars. + * @param chars output payload buffer. + * @param numRows number of rows. + * @param stream CUDA stream for the launch. + */ +void packDecimalSumState( + cudf::type_id sumType, + cudf::type_id offsetType, + cudf::column_view sumCol, + const int64_t* counts, + cudf::column_view offsetsView, + uint8_t* chars, + cudf::size_type numRows, + rmm::cuda_stream_view stream); + +/** + * Inverse of packDecimalSumState. + * + * @param offsetType INT32 or INT64; selects offset storage width via + * cudf::type_dispatcher. + * @param offsetsView per-row byte offsets into chars. + * @param chars packed payload buffer. + * @param sumView output per-row DECIMAL128 sums. + * @param countView output per-row counts. + * @param numRows number of rows. + * @param stream CUDA stream for the launch. + */ +void unpackDecimalSumState( + cudf::type_id offsetType, + cudf::column_view offsetsView, + const uint8_t* chars, + cudf::mutable_column_view sumView, + cudf::mutable_column_view countView, + cudf::size_type numRows, + rmm::cuda_stream_view stream); + +/** + * Per-row half-up integer divide of sum by count; count == 0 writes zero + * (validity is applied separately). + * + * @param sumType DECIMAL64 or DECIMAL128; selects sum storage width via + * cudf::type_dispatcher. + * @param sumCol per-row sums. + * @param counts per-row counts. + * @param outView output per-row averages. + * @param numRows number of rows. + * @param stream CUDA stream for the launch. + */ +void averageRoundDecimalSum( + cudf::type_id sumType, + cudf::column_view sumCol, + const int64_t* counts, + cudf::mutable_column_view outView, + cudf::size_type numRows, + rmm::cuda_stream_view stream); + +/** + * Builds a null mask for rows where sum and count are both valid and count is + * non-zero, for serializing state or finalizing averages. + * + * @param sumCol decoded sum column. + * @param countCol decoded count column. + * @param stream CUDA stream for the launch. + * @param mr memory resource for the returned mask. + * @return {null mask buffer, null count}. + */ +std::pair buildStateValidityMask( + const cudf::column_view& sumCol, + const cudf::column_view& countCol, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace facebook::velox::cudf_velox::detail diff --git a/velox/experimental/cudf/exec/DecimalAggregationHostOps.cpp b/velox/experimental/cudf/exec/DecimalAggregationHostOps.cpp new file mode 100644 index 00000000000..73bff493087 --- /dev/null +++ b/velox/experimental/cudf/exec/DecimalAggregationHostOps.cpp @@ -0,0 +1,88 @@ +/* + * 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/CudfNoDefaults.h" +#include "velox/experimental/cudf/exec/DecimalAggregationHostOps.h" +#include "velox/experimental/cudf/exec/DecimalAggregationState.h" +#include "velox/experimental/cudf/exec/GpuResources.h" +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" + +#include "velox/common/base/Exceptions.h" + +#include + +namespace facebook::velox::cudf_velox { + +void validateIntermediateColumnType(cudf::column_view const& column) { + // fmt does not understand cudf::type_id enum class + auto const colType = static_cast(column.type().id()); + VELOX_CHECK_EQ( + colType, + static_cast(cudf::type_id::STRING), + "Expected serialized decimal aggregation state: Velox VARBINARY represented as cuDF STRING (got type {})", + colType); +} + +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) { + return inputCol; + } + holder = cudf::cast( + inputCol, + cudf::data_type{cudf::type_id::DECIMAL128, inputCol.type().scale()}, + stream, + get_temp_mr()); + return holder->view(); +} + +std::unique_ptr castCountColumnToInt64( + std::unique_ptr count, + rmm::cuda_stream_view stream) { + if (count->type().id() != cudf::type_id::INT64) { + count = cudf::cast( + *count, cudf::data_type{cudf::type_id::INT64}, stream, get_temp_mr()); + } + return count; +} + +std::unique_ptr serializeDecimalPartialOrIntermediateState( + std::unique_ptr sum, + std::unique_ptr count, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + count = castCountColumnToInt64(std::move(count), stream); + return serializeDecimalSumState(sum->view(), count->view(), stream, mr); +} + +std::unique_ptr finalizeDecimalAverage( + std::unique_ptr sum, + std::unique_ptr count, + const TypePtr& resultType, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + count = castCountColumnToInt64(std::move(count), stream); + auto avgCol = computeDecimalAverage(sum->view(), count->view(), stream, mr); + auto const cudfOutType = veloxToCudfDataType(resultType); + if (avgCol->type() != cudfOutType) { + avgCol = cudf::cast(avgCol->view(), cudfOutType, stream, mr); + } + return avgCol; +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/DecimalAggregationHostOps.h b/velox/experimental/cudf/exec/DecimalAggregationHostOps.h new file mode 100644 index 00000000000..778f9120e78 --- /dev/null +++ b/velox/experimental/cudf/exec/DecimalAggregationHostOps.h @@ -0,0 +1,108 @@ +/* + * 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 "velox/type/Type.h" + +#include +#include + +#include + +#include + +namespace facebook::velox::cudf_velox { + +/** + * Asserts that a column holds serialized decimal aggregate state in the form + * Velox uses for VARBINARY: a cuDF STRING column whose bytes are the packed + * sum/count payloads (see serializeDecimalSumState). The payload does not carry + * scale, so VARBINARY intermediate steps decode at scale 0; the real scale is + * applied at final cast time. + * + * @param column column to validate. + */ +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. + * + * @param inputCol DECIMAL64 input column. + * @param holder receives ownership of the casted column when inputCol is + * DECIMAL64; unchanged otherwise. + * @param stream CUDA stream for device work. + * @return view of inputCol or of the column stored in holder. + */ +cudf::column_view castDecimal64InputToDecimal128( + cudf::column_view inputCol, + std::unique_ptr& holder, + rmm::cuda_stream_view stream); + +/** + * Ensures the partial-row count column is INT64, casting with the temporary + * memory resource (the result is consumed internally, not part of operator + * output) when the incoming type differs. + * + * @param count partial-row count column. + * @param stream CUDA stream for device work. + * @return INT64 count column (moved through when already INT64). + */ +std::unique_ptr castCountColumnToInt64( + std::unique_ptr count, + rmm::cuda_stream_view stream); + +/** + * Normalizes the count column to INT64, then encodes sum and count into a + * single STRING column of fixed-width per-row payloads (delegates to + * serializeDecimalSumState). Used when emitting or persisting partial / + * intermediate decimal SUM state for the cuDF path. + * + * @param sum partial sum column (DECIMAL64 or DECIMAL128). + * @param count partial-row count column. + * @param stream CUDA stream for device work. + * @param mr memory resource for allocated columns. + * @return STRING column of serialized state. + */ +std::unique_ptr serializeDecimalPartialOrIntermediateState( + std::unique_ptr sum, + std::unique_ptr count, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/** + * Normalizes the count column to INT64, computes a per-row decimal average + * from intermediate sum/count (delegates to computeDecimalAverage), then casts + * the result to the Velox result type when its cuDF decimal encoding differs + * from the average column's type. + * + * @param sum intermediate sum column (DECIMAL64 or DECIMAL128). + * @param count intermediate count column. + * @param resultType Velox type of the finalized average. + * @param stream CUDA stream for device work. + * @param mr memory resource for allocated columns. + * @return finalized average column. + */ +std::unique_ptr finalizeDecimalAverage( + std::unique_ptr sum, + std::unique_ptr count, + const TypePtr& resultType, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/DecimalAggregationState.cpp b/velox/experimental/cudf/exec/DecimalAggregationState.cpp new file mode 100644 index 00000000000..523e23613d8 --- /dev/null +++ b/velox/experimental/cudf/exec/DecimalAggregationState.cpp @@ -0,0 +1,263 @@ +/* + * 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/CudfNoDefaults.h" +#include "velox/experimental/cudf/exec/DecimalAggregationDevice.h" +#include "velox/experimental/cudf/exec/DecimalAggregationState.h" +#include "velox/experimental/cudf/exec/GpuResources.h" + +#include "velox/common/base/Exceptions.h" + +#include +#include +#include +#include + +#include + +namespace facebook::velox::cudf_velox { + +DecimalSumStateColumns deserializeDecimalSumState( + const cudf::column_view& stateCol, + int32_t scale, + rmm::cuda_stream_view stream) { + VELOX_CHECK( + stateCol.type().id() == cudf::type_id::STRING, + "Decimal sum state requires STRING/VARBINARY column (type is {})", + cudf::type_to_name(stateCol.type())); + // The decoded sum/count columns are consumed by the next groupby/reduce and + // never leave the operator and should use the temporary memory resource. + auto const mr = get_temp_mr(); + auto numRows = stateCol.size(); + if (numRows == 0) { + DecimalSumStateColumns empty; + empty.sum = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::DECIMAL128, -scale}, + 0, + cudf::mask_state::UNALLOCATED, + stream, + mr); + empty.count = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::INT64}, + 0, + cudf::mask_state::UNALLOCATED, + stream, + mr); + return empty; + } + + // For fully-null state columns there is nothing to deserialize. Avoid + // launching unpack kernels over string payload buffers that may be empty. + if (stateCol.nullable() && stateCol.null_count() == numRows) { + DecimalSumStateColumns allNull; + allNull.sum = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::DECIMAL128, -scale}, + numRows, + cudf::mask_state::ALL_NULL, + stream, + mr); + allNull.count = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::INT64}, + numRows, + cudf::mask_state::ALL_NULL, + stream, + mr); + return allNull; + } + + cudf::strings_column_view strings(stateCol); + + auto const payloadSize = strings.chars_size(stream); + auto const expectedPayloadSize = + static_cast(numRows) * detail::kDecimalSumStateSize; + VELOX_CHECK( + payloadSize == expectedPayloadSize, + "Decimal sum state requires payload size {} (got {})", + expectedPayloadSize, + payloadSize); + + auto offsetsView = strings.offsets(); + auto charsPtr = reinterpret_cast(strings.chars_begin(stream)); + + auto sumCol = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::DECIMAL128, -scale}, + numRows, + cudf::mask_state::UNALLOCATED, + stream, + mr); + auto countCol = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::INT64}, + numRows, + cudf::mask_state::UNALLOCATED, + stream, + mr); + + auto sumView = sumCol->mutable_view(); + auto countView = countCol->mutable_view(); + + // 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); + + if (stateCol.nullable()) { + auto nullMask = cudf::copy_bitmask(stateCol, stream, mr); + auto nullCount = stateCol.null_count(); + sumCol->set_null_mask(std::move(nullMask), nullCount); + auto countMask = cudf::copy_bitmask(stateCol, stream, mr); + countCol->set_null_mask(std::move(countMask), nullCount); + } + + DecimalSumStateColumns result; + result.sum = std::move(sumCol); + result.count = std::move(countCol); + return result; +} + +std::unique_ptr serializeDecimalSumState( + const cudf::column_view& sumCol, + const cudf::column_view& countCol, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + VELOX_CHECK( + countCol.type().id() == cudf::type_id::INT64, + "Decimal sum state requires INT64 count column (type is {})", + cudf::type_to_name(countCol.type())); + auto numRows = sumCol.size(); + VELOX_CHECK_EQ( + numRows, + countCol.size(), + "Decimal sum state requires sum and count to be same size (sum size is {}, count size is {})", + sumCol.size(), + countCol.size()); + VELOX_CHECK_LE( + numRows, + static_cast(std::numeric_limits::max()), + "Too many rows to serialize decimal sum state (row count is {})", + numRows); + + if (numRows == 0) { + return cudf::make_empty_column(cudf::type_id::STRING); + } + + auto const rowCount = static_cast(numRows); + + auto const charsBytes = + static_cast(numRows) * detail::kDecimalSumStateSize; + auto const threshold = cudf::strings::get_offset64_threshold(); + auto const useLargeOffsets = charsBytes >= threshold; + // Previously this guard threw std::overflow_error; Velox uses + // VeloxRuntimeError for this guard. + VELOX_CHECK( + !useLargeOffsets || cudf::strings::is_large_strings_enabled(), + "Size of output ({}) exceeds the column size limit ({})", + charsBytes, + threshold); + + auto const offsetsType = + useLargeOffsets ? cudf::type_id::INT64 : cudf::type_id::INT32; + auto offsetsCol = cudf::make_fixed_width_column( + cudf::data_type{offsetsType}, + numRows + 1, + cudf::mask_state::UNALLOCATED, + stream, + mr); + auto offsetsView = offsetsCol->mutable_view(); + + rmm::device_buffer charsBuf( + static_cast(numRows) * detail::kDecimalSumStateSize, stream, mr); + + auto charsPtr = reinterpret_cast(charsBuf.data()); + detail::fillOffsetsForDecimalSumState( + offsetsType, offsetsView, rowCount, stream); + + const auto sumType = sumCol.type().id(); + VELOX_CHECK( + sumType == cudf::type_id::DECIMAL64 || + sumType == cudf::type_id::DECIMAL128, + "Unsupported decimal sum column type (type is {})", + cudf::type_to_name(sumCol.type())); + detail::packDecimalSumState( + sumType, + offsetsType, + sumCol, + countCol.data(), + offsetsView, + charsPtr, + rowCount, + stream); + + auto [nullMask, nullCount] = + detail::buildStateValidityMask(sumCol, countCol, stream, mr); + return cudf::make_strings_column( + static_cast(numRows), + std::move(offsetsCol), + std::move(charsBuf), + nullCount, + std::move(nullMask)); +} + +std::unique_ptr computeDecimalAverage( + const cudf::column_view& sumCol, + const cudf::column_view& countCol, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + VELOX_CHECK( + countCol.type().id() == cudf::type_id::INT64, + "Decimal average requires INT64 count column (type is {})", + cudf::type_to_name(countCol.type())); + VELOX_CHECK( + sumCol.type().id() == cudf::type_id::DECIMAL64 || + sumCol.type().id() == cudf::type_id::DECIMAL128, + "Decimal average requires DECIMAL64 or DECIMAL128 sum column (type is {})", + cudf::type_to_name(sumCol.type())); + VELOX_CHECK_EQ( + sumCol.size(), + countCol.size(), + "Decimal average requires sum and count to be same size (sum size is {}, count size is {})", + sumCol.size(), + countCol.size()); + + auto numRows = sumCol.size(); + auto out = cudf::make_fixed_width_column( + sumCol.type(), numRows, cudf::mask_state::UNALLOCATED, stream, mr); + + if (numRows > 0) { + auto const rowCount = static_cast(numRows); + const auto sumType = sumCol.type().id(); + detail::averageRoundDecimalSum( + sumType, + sumCol, + countCol.data(), + out->mutable_view(), + rowCount, + stream); + } + + auto [nullMask, nullCount] = + detail::buildStateValidityMask(sumCol, countCol, stream, mr); + if (nullCount > 0) { + out->set_null_mask(std::move(nullMask), nullCount); + } + return out; +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/DecimalAggregationState.h b/velox/experimental/cudf/exec/DecimalAggregationState.h new file mode 100644 index 00000000000..3813d58f45e --- /dev/null +++ b/velox/experimental/cudf/exec/DecimalAggregationState.h @@ -0,0 +1,90 @@ +/* + * 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 + +namespace facebook::velox::cudf_velox { + +struct DecimalSumStateColumns { + std::unique_ptr sum; + std::unique_ptr count; +}; + +/** + * Decodes intermediate decimal SUM aggregate state stored as a cuDF STRING + * column (fixed-size packed bytes per row, converted from Velox VARBINARY) into + * two device columns: a DECIMAL128 sum with scale -scale (matching Velox + * intermediate state) and an INT64 partial row count. Handles empty input, + * all-null state without touching payload buffers, and propagates the source + * null mask to both outputs when present. + * + * @param stateCol STRING column of packed sum/count payloads. + * @param scale decimal scale used to set the output sum column's scale to + * -scale. + * @param stream CUDA stream for device work. + * @return decoded sum and count columns. + */ +DecimalSumStateColumns deserializeDecimalSumState( + const cudf::column_view& stateCol, + int32_t scale, + rmm::cuda_stream_view stream); + +/** + * Encodes partial decimal SUM state (DECIMAL64 or DECIMAL128 sums plus INT64 + * counts) into a single STRING column (later converted to Velox VARBINARY): + * per-row fixed-width payloads and string offsets (INT32 or INT64 depending on + * total char size and cuDF large-strings settings). The output null mask + * matches buildStateValidityMask: a row is invalid if the sum or count is null, + * or the count is zero. + * + * @param sumCol per-row partial sums (DECIMAL64 or DECIMAL128). + * @param countCol per-row INT64 partial row counts. + * @param stream CUDA stream for device work. + * @param mr memory resource for allocated columns. + * @return STRING column of serialized state. + */ +std::unique_ptr serializeDecimalSumState( + const cudf::column_view& sumCol, + const cudf::column_view& countCol, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/** + * Finalizes AVG from intermediate SUM state: divides each sum by its count on + * device with decimal-specific rounding (see averageRoundDecimalSum), + * producing a column of the same decimal type as the sum. Rows are null where + * buildStateValidityMask marks them invalid (null sum/count or zero count), + * matching serializeDecimalSumState. + * + * @param sumCol per-row partial sums (DECIMAL64 or DECIMAL128). + * @param countCol per-row INT64 partial row counts. + * @param stream CUDA stream for device work. + * @param mr memory resource for allocated columns. + * @return per-row decimal average column. + */ +std::unique_ptr computeDecimalAverage( + const cudf::column_view& sumCol, + const cudf::column_view& countCol, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/exec/VeloxCudfInterop.cpp b/velox/experimental/cudf/exec/VeloxCudfInterop.cpp index 0a73bfcf34c..ab6418443d0 100644 --- a/velox/experimental/cudf/exec/VeloxCudfInterop.cpp +++ b/velox/experimental/cudf/exec/VeloxCudfInterop.cpp @@ -40,12 +40,6 @@ namespace facebook::velox::cudf_velox { -cudf::type_id veloxToCudfTypeId(const TypePtr& type) { - // Legacy helper retained for compatibility. Note: returning cudf::type_id - // discards decimal scale; prefer veloxToCudfDataType when scale matters. - return veloxToCudfDataType(type).id(); -} - cudf::data_type veloxToCudfDataType(const TypePtr& type) { switch (type->kind()) { case TypeKind::BOOLEAN: diff --git a/velox/experimental/cudf/exec/VeloxCudfInterop.h b/velox/experimental/cudf/exec/VeloxCudfInterop.h index 47f638a8b6a..03fe48e5494 100644 --- a/velox/experimental/cudf/exec/VeloxCudfInterop.h +++ b/velox/experimental/cudf/exec/VeloxCudfInterop.h @@ -26,8 +26,6 @@ namespace facebook::velox::cudf_velox { -cudf::type_id veloxToCudfTypeId(const TypePtr& type); - cudf::data_type veloxToCudfDataType(const TypePtr& type); namespace with_arrow { diff --git a/velox/experimental/cudf/tests/CMakeLists.txt b/velox/experimental/cudf/tests/CMakeLists.txt index 69522308374..e92ac454c2a 100644 --- a/velox/experimental/cudf/tests/CMakeLists.txt +++ b/velox/experimental/cudf/tests/CMakeLists.txt @@ -107,6 +107,13 @@ velox_add_cudf_test( LIBS velox_cudf_exec Folly::folly ) +velox_add_cudf_test( + NAME velox_cudf_decimal_aggregation_test + SOURCES Main.cpp DecimalAggregationTest.cpp + LIBS ${CUDF_TEST_DEFAULT_LIBS} velox_functions_test_lib + TIMEOUT 3000 +) + velox_add_cudf_test( NAME velox_cudf_vector_test SOURCES Main.cpp CudfVectorTest.cpp diff --git a/velox/experimental/cudf/tests/DecimalAggregationTest.cpp b/velox/experimental/cudf/tests/DecimalAggregationTest.cpp new file mode 100644 index 00000000000..e7d499621ca --- /dev/null +++ b/velox/experimental/cudf/tests/DecimalAggregationTest.cpp @@ -0,0 +1,1671 @@ +/* + * 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/CudfConfig.h" +#include "velox/experimental/cudf/exec/DecimalAggregationState.h" +#include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/exec/VeloxCudfInterop.h" + +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/common/file/FileSystems.h" +#include "velox/exec/tests/utils/AssertQueryBuilder.h" +#include "velox/exec/tests/utils/OperatorTestBase.h" +#include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h" +#include "velox/functions/prestosql/registration/RegistrationFunctions.h" +#include "velox/parse/TypeResolver.h" +#include "velox/type/DecimalUtil.h" + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace facebook::velox::cudf_velox { +namespace { + +int64_t computeAvgRaw(const std::vector& values) { + int128_t sum = 0; + for (auto value : values) { + sum += value; + } + int128_t avg = 0; + facebook::velox::DecimalUtil::computeAverage(avg, sum, values.size(), 0); + return static_cast(avg); +} + +constexpr int kBitsPerWord = 8 * sizeof(cudf::bitmask_type); + +std::pair makeNullMask( + const std::vector& valid, + rmm::cuda_stream_view stream) { + auto numBits = static_cast(valid.size()); + if (numBits == 0) { + return {rmm::device_buffer{}, 0}; + } + auto maskBytes = cudf::bitmask_allocation_size_bytes(numBits); + auto numWords = maskBytes / sizeof(cudf::bitmask_type); + std::vector host(numWords, 0); + cudf::size_type nullCount = 0; + for (cudf::size_type i = 0; i < numBits; ++i) { + if (valid[i]) { + auto word = i / kBitsPerWord; + auto bit = i % kBitsPerWord; + host[word] |= (cudf::bitmask_type{1} << bit); + } else { + ++nullCount; + } + } + rmm::device_buffer mask(maskBytes, stream); + if (!host.empty()) { + auto status = cudaMemcpyAsync( + mask.data(), + host.data(), + host.size() * sizeof(cudf::bitmask_type), + cudaMemcpyHostToDevice, + stream.value()); + VELOX_CHECK_EQ(0, static_cast(status)); + stream.synchronize(); + } + return {std::move(mask), nullCount}; +} + +class ScopedEnvVar { + public: + ScopedEnvVar(const char* key, const char* value) : key_(key) { + const char* existing = std::getenv(key); + if (existing) { + oldValue_ = std::string(existing); + } + if (value) { + setenv(key, value, 1); + } else { + unsetenv(key); + } + } + + ~ScopedEnvVar() { + if (oldValue_) { + setenv(key_.c_str(), oldValue_->c_str(), 1); + } else { + unsetenv(key_.c_str()); + } + } + + private: + std::string key_; + std::optional oldValue_; +}; + +template +std::unique_ptr makeFixedWidthColumn( + cudf::data_type type, + const std::vector& values, + const std::vector* valid, + rmm::cuda_stream_view stream) { + auto col = cudf::make_fixed_width_column( + type, + static_cast(values.size()), + cudf::mask_state::UNALLOCATED, + stream); + if (!values.empty()) { + auto status = cudaMemcpyAsync( + col->mutable_view().data(), + values.data(), + values.size() * sizeof(T), + cudaMemcpyHostToDevice, + stream.value()); + VELOX_CHECK_EQ(0, static_cast(status)); + stream.synchronize(); + } + if (valid) { + auto [mask, nullCount] = makeNullMask(*valid, stream); + col->set_null_mask(std::move(mask), nullCount); + } + return col; +} + +template +std::unique_ptr makeDecimalColumn( + const std::vector& values, + int32_t scale, + const std::vector* valid, + rmm::cuda_stream_view stream) { + cudf::type_id typeId = std::is_same_v ? cudf::type_id::DECIMAL64 + : cudf::type_id::DECIMAL128; + cudf::data_type type{typeId, -scale}; + return makeFixedWidthColumn(type, values, valid, stream); +} + +std::unique_ptr makeInt64Column( + const std::vector& values, + const std::vector* valid, + rmm::cuda_stream_view stream) { + return makeFixedWidthColumn( + cudf::data_type{cudf::type_id::INT64}, values, valid, stream); +} + +template +std::vector copyColumnData( + const cudf::column_view& view, + rmm::cuda_stream_view stream) { + std::vector host(view.size()); + if (view.size() == 0) { + return host; + } + auto status = cudaMemcpyAsync( + host.data(), + view.data(), + view.size() * sizeof(T), + cudaMemcpyDeviceToHost, + stream.value()); + VELOX_CHECK_EQ(0, static_cast(status)); + stream.synchronize(); + return host; +} + +std::vector copyNullMask( + const cudf::column_view& view, + rmm::cuda_stream_view stream) { + auto numWords = cudf::num_bitmask_words(view.size()); + std::vector host(numWords, 0); + if (!view.nullable() || numWords == 0) { + return host; + } + auto status = cudaMemcpyAsync( + host.data(), + view.null_mask(), + host.size() * sizeof(cudf::bitmask_type), + cudaMemcpyDeviceToHost, + stream.value()); + VELOX_CHECK_EQ(0, static_cast(status)); + stream.synchronize(); + return host; +} + +bool isValidAt(const std::vector& mask, size_t idx) { + if (mask.empty()) { + return true; + } + auto word = idx / kBitsPerWord; + auto bit = idx % kBitsPerWord; + return (mask[word] >> bit) & 1; +} + +class CudfDecimalTest : public exec::test::OperatorTestBase { + protected: + void SetUp() override { + exec::test::OperatorTestBase::SetUp(); + filesystems::registerLocalFileSystem(); + parse::registerTypeResolver(); + functions::prestosql::registerAllScalarFunctions(); + aggregate::prestosql::registerAllAggregateFunctions(); + CudfConfig::getInstance().allowCpuFallback = false; + // Ensure a CUDA device is selected and initialized (RMM asserts otherwise). + int deviceCount = 0; + auto status = cudaGetDeviceCount(&deviceCount); + if (status != cudaSuccess) { + GTEST_SKIP() << "cudaGetDeviceCount failed: " << static_cast(status) + << " (" << cudaGetErrorString(status) << ")"; + } + if (deviceCount == 0) { + GTEST_SKIP() << "No CUDA devices visible (check CUDA_VISIBLE_DEVICES)"; + } + VELOX_CHECK_EQ(0, static_cast(cudaSetDevice(0))); + VELOX_CHECK_EQ(0, static_cast(cudaFree(nullptr))); + registerCudf(); + } + + void TearDown() override { + unregisterCudf(); + exec::test::OperatorTestBase::TearDown(); + } +}; + +TEST_F(CudfDecimalTest, decimalAvgDecimalInput) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"d"}, + {makeFlatVector( + {100, 200, 300, 400}, // 1.00, 2.00, 3.00, 4.00 + DECIMAL(12, 2))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .singleAggregation({}, {"avg(d) AS avg_d"}) + .planNode(); + + auto expected = makeRowVector( + {"avg_d"}, {makeFlatVector({250}, DECIMAL(12, 2))}); // 2.50 + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgDecimalInputRounds) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + // Sum = 1.60, count = 7 => 0.22857..., rounds to 0.23 at scale 2. + std::vector rawValues = {100, 10, 10, 10, 10, 10, 10}; + auto input = makeRowVector( + {"d"}, {makeFlatVector(rawValues, DECIMAL(12, 2))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .singleAggregation({}, {"avg(d) AS avg_d"}) + .planNode(); + + auto expected = makeRowVector( + {"avg_d"}, + {makeFlatVector({computeAvgRaw(rawValues)}, DECIMAL(12, 2))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgPartialFinalVarbinaryRounds) { + auto rowType = ROW({ + {"k", INTEGER()}, + {"d", DECIMAL(12, 2)}, + }); + + std::vector keys = {1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 3}; + std::vector values = {100, 10, 10, 10, 10, 10, 10, 100, 1, -100, -1}; + + auto input = makeRowVector( + {"k", "d"}, + { + makeFlatVector(keys), + makeFlatVector(values, DECIMAL(12, 2)), + }); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({"k"}, {"avg(d) AS a"}) + .finalAggregation() + .orderBy({"k"}, false) + .planNode(); + + std::vector>> groups = { + {1, {100, 10, 10, 10, 10, 10, 10}}, + {2, {100, 1}}, + {3, {-100, -1}}, + }; + + auto expected = makeRowVector( + {"k", "a"}, + { + makeFlatVector({1, 2, 3}), + makeFlatVector( + {computeAvgRaw(groups[0].second), + computeAvgRaw(groups[1].second), + computeAvgRaw(groups[2].second)}, + DECIMAL(12, 2)), + }); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgIntermediateVarbinaryRounds) { + auto rowType = ROW({ + {"k", INTEGER()}, + {"d", DECIMAL(12, 2)}, + }); + + auto input1 = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 1, 2, 3}), + makeFlatVector({100, 10, 100, -100}, DECIMAL(12, 2)), + }); + auto input2 = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 1, 1, 1, 1, 2, 3}), + makeFlatVector({10, 10, 10, 10, 10, 1, -1}, DECIMAL(12, 2)), + }); + + std::vector vectors = {input1, input2}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({"k"}, {"avg(d) AS a"}) + .intermediateAggregation() + .finalAggregation() + .orderBy({"k"}, false) + .planNode(); + + std::vector>> groups = { + {1, {100, 10, 10, 10, 10, 10, 10}}, + {2, {100, 1}}, + {3, {-100, -1}}, + }; + + auto expected = makeRowVector( + {"k", "a"}, + { + makeFlatVector({1, 2, 3}), + makeFlatVector( + {computeAvgRaw(groups[0].second), + computeAvgRaw(groups[1].second), + computeAvgRaw(groups[2].second)}, + DECIMAL(12, 2)), + }); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgGlobalPartialFinalVarbinaryRounds) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input1 = makeRowVector( + {"d"}, {makeFlatVector({100, 10, 10}, DECIMAL(12, 2))}); + auto input2 = makeRowVector( + {"d"}, {makeFlatVector({10, 10, 10, 10}, DECIMAL(12, 2))}); + + std::vector vectors = {input1, input2}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"avg(d) AS a"}) + .finalAggregation() + .planNode(); + + std::vector allValues = {100, 10, 10, 10, 10, 10, 10}; + auto expected = makeRowVector( + {"a"}, + {makeFlatVector({computeAvgRaw(allValues)}, DECIMAL(12, 2))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgGlobalIntermediateVarbinaryRounds) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input1 = makeRowVector( + {"d"}, {makeFlatVector({100, 10, 10}, DECIMAL(12, 2))}); + auto input2 = makeRowVector( + {"d"}, {makeFlatVector({10, 10, 10, 10}, DECIMAL(12, 2))}); + + std::vector vectors = {input1, input2}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"avg(d) AS a"}) + .intermediateAggregation() + .finalAggregation() + .planNode(); + + std::vector allValues = {100, 10, 10, 10, 10, 10, 10}; + auto expected = makeRowVector( + {"a"}, + {makeFlatVector({computeAvgRaw(allValues)}, DECIMAL(12, 2))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgGlobalSingleRounds) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"d"}, + {makeFlatVector({100, 10, 10, 10, 10, 10, 10}, DECIMAL(12, 2))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .singleAggregation({}, {"avg(d) AS a"}) + .planNode(); + + std::vector allValues = {100, 10, 10, 10, 10, 10, 10}; + auto expected = makeRowVector( + {"a"}, + {makeFlatVector({computeAvgRaw(allValues)}, DECIMAL(12, 2))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgGlobalSingleDecimal64Overflow) { + // 12 values of 9e17 (DECIMAL(18,0)) sum to 1.08e19, past 2^63. The sum must + // accumulate in 128 bits or a DECIMAL64 accumulator wraps; avg is 9e17. + constexpr int64_t kBig = 900'000'000'000'000'000; + constexpr int kNumRows = 12; + std::vector values(kNumRows, kBig); + + auto input = + makeRowVector({"d"}, {makeFlatVector(values, DECIMAL(18, 0))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .singleAggregation({}, {"avg(d) AS a"}) + .planNode(); + + auto expected = + makeRowVector({"a"}, {makeFlatVector({kBig}, DECIMAL(18, 0))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgGroupbySingleDecimal64Overflow) { + // Same overflow within a single group, exercising the groupby raw sum path. + constexpr int64_t kBig = 900'000'000'000'000'000; + constexpr int kNumRows = 12; + std::vector keys(kNumRows, 1); + std::vector values(kNumRows, kBig); + + auto input = makeRowVector( + {"k", "d"}, + { + makeFlatVector(keys), + makeFlatVector(values, DECIMAL(18, 0)), + }); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .singleAggregation({"k"}, {"avg(d) AS a"}) + .planNode(); + + auto expected = makeRowVector( + {"k", "a"}, + { + makeFlatVector({1}), + makeFlatVector({kBig}, DECIMAL(18, 0)), + }); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgGlobalSingleAllNulls) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"d"}, + {makeNullableFlatVector( + {std::nullopt, std::nullopt, std::nullopt}, DECIMAL(12, 2))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .singleAggregation({}, {"avg(d) AS a"}) + .planNode(); + + auto expected = makeRowVector( + {"a"}, {makeNullableFlatVector({std::nullopt}, DECIMAL(12, 2))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgGlobalPartialFinalVarbinaryAllNulls) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"d"}, + {makeNullableFlatVector( + {std::nullopt, std::nullopt, std::nullopt}, DECIMAL(12, 2))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"avg(d) AS a"}) + .finalAggregation() + .planNode(); + + auto expected = makeRowVector( + {"a"}, {makeNullableFlatVector({std::nullopt}, DECIMAL(12, 2))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgGlobalIntermediateVarbinaryAllNulls) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"d"}, + {makeNullableFlatVector( + {std::nullopt, std::nullopt, std::nullopt}, DECIMAL(12, 2))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"avg(d) AS a"}) + .intermediateAggregation() + .finalAggregation() + .planNode(); + + auto expected = makeRowVector( + {"a"}, {makeNullableFlatVector({std::nullopt}, DECIMAL(12, 2))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgPartialFinalVarbinaryNullGroup) { + auto rowType = ROW({ + {"k", INTEGER()}, + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 1, 2, 2, 3, 3}), + makeNullableFlatVector( + {100, 200, std::nullopt, std::nullopt, 400, std::nullopt}, + DECIMAL(12, 2)), + }); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({"k"}, {"avg(d) AS a"}) + .finalAggregation() + .orderBy({"k"}, false) + .planNode(); + + auto expected = makeRowVector( + {"k", "a"}, + { + makeFlatVector({1, 2, 3}), + makeNullableFlatVector( + {150, std::nullopt, 400}, DECIMAL(12, 2)), + }); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalAvgIntermediateVarbinaryNullGroup) { + auto rowType = ROW({ + {"k", INTEGER()}, + {"d", DECIMAL(12, 2)}, + }); + + auto input1 = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 2, 3}), + makeNullableFlatVector( + {100, std::nullopt, 400}, DECIMAL(12, 2)), + }); + auto input2 = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 2, 3}), + makeNullableFlatVector( + {200, std::nullopt, std::nullopt}, DECIMAL(12, 2)), + }); + + std::vector vectors = {input1, input2}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({"k"}, {"avg(d) AS a"}) + .intermediateAggregation() + .finalAggregation() + .orderBy({"k"}, false) + .planNode(); + + auto expected = makeRowVector( + {"k", "a"}, + { + makeFlatVector({1, 2, 3}), + makeNullableFlatVector( + {150, std::nullopt, 400}, DECIMAL(12, 2)), + }); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalSumPartialFinalVarbinary) { + auto rowType = ROW({ + {"k", INTEGER()}, + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 1, 2, 2, 2}), + makeFlatVector( + {12345, -2500, 10000, 200, -300}, DECIMAL(12, 2)), + }); + + std::vector vectors = {input}; + createDuckDbTable(vectors); + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({"k"}, {"sum(d) AS s"}) + .finalAggregation() + .planNode(); + + facebook::velox::exec::test::AssertQueryBuilder(plan, duckDbQueryRunner_) + .assertResults("SELECT k, sum(d) AS s FROM tmp GROUP BY k"); +} + +TEST_F(CudfDecimalTest, decimalPartialSumVarbinaryToVeloxRoundTrip) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"d"}, {makeFlatVector({100, 200, 300}, DECIMAL(12, 2))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"sum(d) AS s"}) + .planNode(); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + VELOX_CHECK_NOT_NULL(result); + ASSERT_GT(result->size(), 0); + ASSERT_EQ(result->childAt(0)->type()->kind(), TypeKind::VARBINARY); +} + +TEST_F(CudfDecimalTest, decimalSumPartialFinalEmptyInput) { + auto rowType = ROW({ + {"k", INTEGER()}, + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 2, 3}), + makeFlatVector({100, 200, 300}, DECIMAL(12, 2)), + }); + + std::vector vectors = {input}; + createDuckDbTable(vectors); + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .filter("k < 0") + .partialAggregation({"k"}, {"sum(d) AS s"}) + .finalAggregation() + .planNode(); + + facebook::velox::exec::test::AssertQueryBuilder(plan, duckDbQueryRunner_) + .assertResults("SELECT k, sum(d) AS s FROM tmp WHERE k < 0 GROUP BY k"); +} + +TEST_F(CudfDecimalTest, decimalSumIntermediateVarbinary) { + auto rowType = ROW({ + {"k", INTEGER()}, + {"d", DECIMAL(12, 2)}, + }); + + auto input1 = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 1, 2}), + makeFlatVector({12345, -2500, 10000}, DECIMAL(12, 2)), + }); + auto input2 = makeRowVector( + {"k", "d"}, + { + makeFlatVector({2, 3}), + makeFlatVector({200, -300}, DECIMAL(12, 2)), + }); + + std::vector vectors = {input1, input2}; + createDuckDbTable(vectors); + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({"k"}, {"sum(d) AS s"}) + .intermediateAggregation() + .finalAggregation() + .planNode(); + + facebook::velox::exec::test::AssertQueryBuilder(plan, duckDbQueryRunner_) + .assertResults("SELECT k, sum(d) AS s FROM tmp GROUP BY k"); +} + +TEST_F(CudfDecimalTest, decimalSumGlobalPartialFinalVarbinary) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input1 = makeRowVector( + {"d"}, {makeFlatVector({12345, -2500, 10000}, DECIMAL(12, 2))}); + auto input2 = makeRowVector( + {"d"}, {makeFlatVector({200, -300}, DECIMAL(12, 2))}); + + std::vector vectors = {input1, input2}; + createDuckDbTable(vectors); + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"sum(d) AS s"}) + .finalAggregation() + .planNode(); + + facebook::velox::exec::test::AssertQueryBuilder(plan, duckDbQueryRunner_) + .assertResults("SELECT sum(d) AS s FROM tmp"); +} + +TEST_F(CudfDecimalTest, decimalSumGlobalIntermediateVarbinary) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input1 = makeRowVector( + {"d"}, {makeFlatVector({12345, -2500, 10000}, DECIMAL(12, 2))}); + auto input2 = makeRowVector( + {"d"}, {makeFlatVector({200, -300}, DECIMAL(12, 2))}); + + std::vector vectors = {input1, input2}; + createDuckDbTable(vectors); + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"sum(d) AS s"}) + .intermediateAggregation() + .finalAggregation() + .planNode(); + + facebook::velox::exec::test::AssertQueryBuilder(plan, duckDbQueryRunner_) + .assertResults("SELECT sum(d) AS s FROM tmp"); +} + +TEST_F(CudfDecimalTest, decimalSumGlobalSingle) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"d"}, + {makeFlatVector( + {12345, -2500, 10000, 200, -300}, DECIMAL(12, 2))}); + + std::vector vectors = {input}; + createDuckDbTable(vectors); + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .singleAggregation({}, {"sum(d) AS s"}) + .planNode(); + + facebook::velox::exec::test::AssertQueryBuilder(plan, duckDbQueryRunner_) + .assertResults("SELECT sum(d) AS s FROM tmp"); +} + +TEST_F(CudfDecimalTest, decimalSumGroupbySingleDecimal64Overflow) { + // One group of 12 values of 9e17 (DECIMAL(18,0)) sums to 1.08e19, past 2^63. + // sum(decimal(18,0)) -> decimal(38,0), computed in 128 bits, no wrap. + constexpr int64_t kBig = 900'000'000'000'000'000; + constexpr int kNumRows = 12; + std::vector keys(kNumRows, 1); + std::vector values(kNumRows, kBig); + + auto input = makeRowVector( + {"k", "d"}, + { + makeFlatVector(keys), + makeFlatVector(values, DECIMAL(18, 0)), + }); + + std::vector vectors = {input}; + + const int128_t expectedSum = static_cast(kBig) * kNumRows; + auto plan = exec::test::PlanBuilder() + .values(vectors) + .singleAggregation({"k"}, {"sum(d) AS s"}) + .planNode(); + + auto expected = makeRowVector( + {"k", "s"}, + { + makeFlatVector({1}), + makeFlatVector({expectedSum}, DECIMAL(38, 0)), + }); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalSumGlobalPartialFinalDecimal64Overflow) { + // Global SUM whose total overflows DECIMAL64; exercises the partial raw sum + // (serialized to VARBINARY) and the final merge, both in 128 bits. + constexpr int64_t kBig = 900'000'000'000'000'000; + constexpr int kNumRows = 12; + std::vector values(kNumRows, kBig); + + auto input = + makeRowVector({"d"}, {makeFlatVector(values, DECIMAL(18, 0))}); + + std::vector vectors = {input}; + + const int128_t expectedSum = static_cast(kBig) * kNumRows; + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"sum(d) AS s"}) + .finalAggregation() + .planNode(); + + auto expected = makeRowVector( + {"s"}, {makeFlatVector({expectedSum}, DECIMAL(38, 0))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalSumPartialFinalVarbinaryNullGroup) { + auto rowType = ROW({ + {"k", INTEGER()}, + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 1, 2, 2, 3, 3}), + makeNullableFlatVector( + {100, 200, std::nullopt, std::nullopt, 400, std::nullopt}, + DECIMAL(12, 2)), + }); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({"k"}, {"sum(d) AS s"}) + .finalAggregation() + .orderBy({"k"}, false) + .planNode(); + + auto expected = makeRowVector( + {"k", "s"}, + { + makeFlatVector({1, 2, 3}), + makeNullableFlatVector( + {static_cast(300), + std::nullopt, + static_cast(400)}, + DECIMAL(38, 2)), + }); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalSumIntermediateVarbinaryNullGroup) { + auto rowType = ROW({ + {"k", INTEGER()}, + {"d", DECIMAL(12, 2)}, + }); + + auto input1 = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 2, 3}), + makeNullableFlatVector( + {100, std::nullopt, 400}, DECIMAL(12, 2)), + }); + auto input2 = makeRowVector( + {"k", "d"}, + { + makeFlatVector({1, 2, 3}), + makeNullableFlatVector( + {200, std::nullopt, std::nullopt}, DECIMAL(12, 2)), + }); + + std::vector vectors = {input1, input2}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({"k"}, {"sum(d) AS s"}) + .intermediateAggregation() + .finalAggregation() + .orderBy({"k"}, false) + .planNode(); + + auto expected = makeRowVector( + {"k", "s"}, + { + makeFlatVector({1, 2, 3}), + makeNullableFlatVector( + {static_cast(300), + std::nullopt, + static_cast(400)}, + DECIMAL(38, 2)), + }); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalSumGlobalPartialFinalVarbinaryAllNulls) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"d"}, + {makeNullableFlatVector( + {std::nullopt, std::nullopt, std::nullopt}, DECIMAL(12, 2))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"sum(d) AS s"}) + .finalAggregation() + .planNode(); + + auto expected = makeRowVector( + {"s"}, + {makeNullableFlatVector({std::nullopt}, DECIMAL(38, 2))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalSumGlobalIntermediateVarbinaryAllNulls) { + auto rowType = ROW({ + {"d", DECIMAL(12, 2)}, + }); + + auto input = makeRowVector( + {"d"}, + {makeNullableFlatVector( + {std::nullopt, std::nullopt, std::nullopt}, DECIMAL(12, 2))}); + + std::vector vectors = {input}; + + auto plan = exec::test::PlanBuilder() + .values(vectors) + .partialAggregation({}, {"sum(d) AS s"}) + .intermediateAggregation() + .finalAggregation() + .planNode(); + + auto expected = makeRowVector( + {"s"}, + {makeNullableFlatVector({std::nullopt}, DECIMAL(38, 2))}); + + auto result = + facebook::velox::exec::test::AssertQueryBuilder(plan).copyResults(pool()); + facebook::velox::test::assertEqualVectors(expected, result); +} + +TEST_F(CudfDecimalTest, decimalDeserializeSumStateDecimal64) { + 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 = makeDecimalColumn(sums, 2, &sumValid, stream); + auto countCol = makeInt64Column(counts, &countValid, stream); + auto stateCol = + serializeDecimalSumState(sumCol->view(), countCol->view(), 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, decimalDeserializeSumStateDecimal128) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector<__int128_t> sums = { + static_cast<__int128_t>(123450), + static_cast<__int128_t>(-25000), + static_cast<__int128_t>(100000), + }; + std::vector counts = {2, 1, 0}; + std::vector sumValid = {true, true, true}; + std::vector countValid = {true, false, true}; + + auto sumCol = makeDecimalColumn<__int128_t>(sums, 3, &sumValid, stream); + auto countCol = makeInt64Column(counts, &countValid, stream); + auto stateCol = + serializeDecimalSumState(sumCol->view(), countCol->view(), stream, mr); + auto sumAndCount = deserializeDecimalSumState(stateCol->view(), 3, 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], sums[i]); + } + } +} + +TEST_F(CudfDecimalTest, decimalDeserializeSumStateAllNull) { + auto stream = cudf::get_default_stream(); + constexpr cudf::size_type numRows = 4; + + auto offsetsCol = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::INT32}, + numRows + 1, + cudf::mask_state::UNALLOCATED, + stream); + auto* offsetsPtr = offsetsCol->mutable_view().data(); + auto status = cudaMemsetAsync( + offsetsPtr, + 0, + static_cast(numRows + 1) * sizeof(int32_t), + stream.value()); + VELOX_CHECK_EQ(0, static_cast(status)); + stream.synchronize(); + + std::vector valid(numRows, false); + auto [nullMask, nullCount] = makeNullMask(valid, stream); + rmm::device_buffer charsBuf(0, stream); + auto stateCol = cudf::make_strings_column( + numRows, + std::move(offsetsCol), + std::move(charsBuf), + nullCount, + std::move(nullMask)); + + auto sumAndCount = deserializeDecimalSumState(stateCol->view(), 2, stream); + auto outSumView = sumAndCount.sum->view(); + auto outCountView = sumAndCount.count->view(); + + EXPECT_EQ(outSumView.size(), numRows); + EXPECT_EQ(outCountView.size(), numRows); + EXPECT_EQ(outSumView.null_count(), numRows); + EXPECT_EQ(outCountView.null_count(), numRows); + + auto outSumMask = copyNullMask(outSumView, stream); + auto outCountMask = copyNullMask(outCountView, stream); + for (size_t i = 0; i < static_cast(numRows); ++i) { + EXPECT_FALSE(isValidAt(outSumMask, i)); + EXPECT_FALSE(isValidAt(outCountMask, i)); + } +} + +TEST_F(CudfDecimalTest, decimalSerializeSumStateUsesInt64OffsetsWhenEnabled) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + ScopedEnvVar enableLargeStrings("LIBCUDF_LARGE_STRINGS_ENABLED", "1"); + ScopedEnvVar threshold("LIBCUDF_LARGE_STRINGS_THRESHOLD", "1"); + + std::vector sums = {100, -200}; + std::vector counts = {1, 1}; + + auto sumCol = makeDecimalColumn(sums, 2, nullptr, stream); + auto countCol = makeInt64Column(counts, nullptr, stream); + auto stateCol = + serializeDecimalSumState(sumCol->view(), countCol->view(), stream, mr); + + cudf::strings_column_view strings(stateCol->view()); + EXPECT_EQ(strings.offsets().type().id(), cudf::type_id::INT64); +} + +TEST_F(CudfDecimalTest, decimalSumStateRoundTripUsesInt64Offsets) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + ScopedEnvVar enableLargeStrings("LIBCUDF_LARGE_STRINGS_ENABLED", "1"); + ScopedEnvVar threshold("LIBCUDF_LARGE_STRINGS_THRESHOLD", "1"); + + 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); + + cudf::strings_column_view strings(stateCol->view()); + EXPECT_EQ(strings.offsets().type().id(), cudf::type_id::INT64); + + auto sumAndCount = deserializeDecimalSumState(stateCol->view(), 2, stream); + auto outSumView = sumAndCount.sum->view(); + auto outCountView = sumAndCount.count->view(); + auto outSum = copyColumnData<__int128_t>(outSumView, stream); + auto outCount = copyColumnData(outCountView, stream); + auto outSumMask = copyNullMask(outSumView, stream); + auto outCountMask = copyNullMask(outCountView, stream); + + EXPECT_EQ(outSumView.size(), sums.size()); + EXPECT_EQ(outCountView.size(), counts.size()); + EXPECT_EQ(outSumMask, outCountMask); + + for (size_t i = 0; i < sums.size(); ++i) { + bool expectedValid = sumValid[i] && countValid[i] && counts[i] != 0; + EXPECT_EQ(isValidAt(outSumMask, i), expectedValid); + EXPECT_EQ(isValidAt(outCountMask, i), expectedValid); + if (expectedValid) { + EXPECT_EQ(outSum[i], static_cast<__int128_t>(sums[i])); + EXPECT_EQ(outCount[i], counts[i]); + } + } +} + +TEST_F(CudfDecimalTest, decimalComputeAverageDecimal64) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector sums = {100, 105, 250, -125}; + std::vector counts = {4, 2, 0, 2}; + std::vector sumValid = {true, true, true, true}; + std::vector countValid = {true, false, true, true}; + + auto sumCol = makeDecimalColumn(sums, 2, &sumValid, stream); + auto countCol = makeInt64Column(counts, &countValid, stream); + auto avgCol = + computeDecimalAverage(sumCol->view(), countCol->view(), stream, mr); + + auto avgMask = copyNullMask(avgCol->view(), stream); + auto outAvg = copyColumnData(avgCol->view(), stream); + + auto avgUnscaled = [](int128_t sum, int64_t count) { + __int128_t out = 0; + facebook::velox::DecimalUtil:: + divideWithRoundUp<__int128_t, __int128_t, int64_t>( + out, sum, count, false, 0, 0); + return static_cast(out); + }; + + for (size_t i = 0; i < sums.size(); ++i) { + bool expectedValid = sumValid[i] && countValid[i] && counts[i] != 0; + EXPECT_EQ(isValidAt(avgMask, i), expectedValid); + if (expectedValid) { + EXPECT_EQ(outAvg[i], avgUnscaled(sums[i], counts[i])); + } + } +} + +TEST_F(CudfDecimalTest, decimalComputeAverageDecimal128) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector<__int128_t> sums = { + static_cast<__int128_t>(123450), + static_cast<__int128_t>(-25000), + static_cast<__int128_t>(100000), + }; + std::vector counts = {3, 2, 0}; + std::vector sumValid = {true, true, true}; + std::vector countValid = {true, true, true}; + + auto sumCol = makeDecimalColumn<__int128_t>(sums, 3, &sumValid, stream); + auto countCol = makeInt64Column(counts, &countValid, stream); + auto avgCol = + computeDecimalAverage(sumCol->view(), countCol->view(), stream, mr); + + auto avgMask = copyNullMask(avgCol->view(), stream); + auto outAvg = copyColumnData<__int128_t>(avgCol->view(), stream); + + auto avgUnscaled = [](int128_t sum, int64_t count) { + __int128_t out = 0; + facebook::velox::DecimalUtil:: + divideWithRoundUp<__int128_t, __int128_t, int64_t>( + out, sum, count, false, 0, 0); + return out; + }; + + for (size_t i = 0; i < sums.size(); ++i) { + bool expectedValid = sumValid[i] && countValid[i] && counts[i] != 0; + EXPECT_EQ(isValidAt(avgMask, i), expectedValid); + if (expectedValid) { + EXPECT_EQ(outAvg[i], avgUnscaled(sums[i], counts[i])); + } + } +} + +TEST_F(CudfDecimalTest, decimalComputeAverageDecimal64MostNegativeSum) { + // Negating INT64_MIN in the signed type is overflow UB; the magnitude and + // sign must be handled in the unsigned domain. avg of one INT64_MIN is + // itself. + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + constexpr int64_t kMin = std::numeric_limits::min(); + std::vector sums = {kMin}; + std::vector counts = {1}; + std::vector valid = {true}; + + auto sumCol = makeDecimalColumn(sums, 0, &valid, stream); + auto countCol = makeInt64Column(counts, &valid, stream); + auto avgCol = + computeDecimalAverage(sumCol->view(), countCol->view(), stream, mr); + + auto outAvg = copyColumnData(avgCol->view(), stream); + EXPECT_EQ(outAvg[0], kMin); +} + +TEST_F(CudfDecimalTest, decimalComputeAverageDecimal128MostNegativeSum) { + // Same regression at the __int128 boundary. avg of one -2^127 is itself. + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + const __int128_t kMin = + static_cast<__int128_t>(static_cast(1) << 127); + std::vector<__int128_t> sums = {kMin}; + std::vector counts = {1}; + std::vector valid = {true}; + + auto sumCol = makeDecimalColumn<__int128_t>(sums, 0, &valid, stream); + auto countCol = makeInt64Column(counts, &valid, stream); + auto avgCol = + computeDecimalAverage(sumCol->view(), countCol->view(), stream, mr); + + auto outAvg = copyColumnData<__int128_t>(avgCol->view(), stream); + EXPECT_EQ(outAvg[0], kMin); +} + +TEST_F(CudfDecimalTest, decimalComputeAverageDecimal64AllValid) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector sums = {100, 200, -150}; + std::vector counts = {4, 5, 3}; + std::vector sumValid = {true, true, true}; + std::vector countValid = {true, true, true}; + + auto sumCol = makeDecimalColumn(sums, 2, &sumValid, stream); + auto countCol = makeInt64Column(counts, &countValid, stream); + auto avgCol = + computeDecimalAverage(sumCol->view(), countCol->view(), stream, mr); + + EXPECT_EQ(avgCol->view().null_count(), 0); + auto avgMask = copyNullMask(avgCol->view(), stream); + auto outAvg = copyColumnData(avgCol->view(), stream); + + auto avgUnscaled = [](int128_t sum, int64_t count) { + __int128_t out = 0; + facebook::velox::DecimalUtil:: + divideWithRoundUp<__int128_t, __int128_t, int64_t>( + out, sum, count, false, 0, 0); + return static_cast(out); + }; + + for (size_t i = 0; i < sums.size(); ++i) { + if (avgCol->view().nullable()) { + EXPECT_TRUE(isValidAt(avgMask, i)) << "row " << i; + } + EXPECT_EQ(outAvg[i], avgUnscaled(sums[i], counts[i])) << "row " << i; + } +} + +TEST_F(CudfDecimalTest, decimalComputeAverageDecimal128AllValid) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector<__int128_t> sums = { + static_cast<__int128_t>(90000), + static_cast<__int128_t>(-5000), + static_cast<__int128_t>(1), + }; + std::vector counts = {3, 5, 1}; + std::vector sumValid = {true, true, true}; + std::vector countValid = {true, true, true}; + + auto sumCol = makeDecimalColumn<__int128_t>(sums, 3, &sumValid, stream); + auto countCol = makeInt64Column(counts, &countValid, stream); + auto avgCol = + computeDecimalAverage(sumCol->view(), countCol->view(), stream, mr); + + EXPECT_EQ(avgCol->view().null_count(), 0); + auto avgMask = copyNullMask(avgCol->view(), stream); + auto outAvg = copyColumnData<__int128_t>(avgCol->view(), stream); + + auto avgUnscaled = [](int128_t sum, int64_t count) { + __int128_t out = 0; + facebook::velox::DecimalUtil:: + divideWithRoundUp<__int128_t, __int128_t, int64_t>( + out, sum, count, false, 0, 0); + return out; + }; + + for (size_t i = 0; i < sums.size(); ++i) { + if (avgCol->view().nullable()) { + EXPECT_TRUE(isValidAt(avgMask, i)) << "row " << i; + } + EXPECT_EQ(outAvg[i], avgUnscaled(sums[i], counts[i])) << "row " << i; + } +} + +TEST_F(CudfDecimalTest, decimalComputeAverageDecimal64NonNullableInputs) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector sums = {80, -40, 1000}; + std::vector counts = {2, 4, 10}; + + auto sumCol = makeDecimalColumn(sums, 2, nullptr, stream); + auto countCol = makeInt64Column(counts, nullptr, stream); + ASSERT_FALSE(sumCol->nullable()); + ASSERT_FALSE(countCol->nullable()); + + auto avgCol = + computeDecimalAverage(sumCol->view(), countCol->view(), stream, mr); + + EXPECT_EQ(avgCol->view().null_count(), 0); + auto avgMask = copyNullMask(avgCol->view(), stream); + auto outAvg = copyColumnData(avgCol->view(), stream); + + auto avgUnscaled = [](int128_t sum, int64_t count) { + __int128_t out = 0; + facebook::velox::DecimalUtil:: + divideWithRoundUp<__int128_t, __int128_t, int64_t>( + out, sum, count, false, 0, 0); + return static_cast(out); + }; + + for (size_t i = 0; i < sums.size(); ++i) { + if (avgCol->view().nullable()) { + EXPECT_TRUE(isValidAt(avgMask, i)) << "row " << i; + } + EXPECT_EQ(outAvg[i], avgUnscaled(sums[i], counts[i])) << "row " << i; + } +} + +TEST_F(CudfDecimalTest, decimalComputeAverageDecimal128NonNullableInputs) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector<__int128_t> sums = { + static_cast<__int128_t>(600), + static_cast<__int128_t>(-99), + }; + std::vector counts = {3, 9}; + + auto sumCol = makeDecimalColumn<__int128_t>(sums, 4, nullptr, stream); + auto countCol = makeInt64Column(counts, nullptr, stream); + ASSERT_FALSE(sumCol->nullable()); + ASSERT_FALSE(countCol->nullable()); + + auto avgCol = + computeDecimalAverage(sumCol->view(), countCol->view(), stream, mr); + + EXPECT_EQ(avgCol->view().null_count(), 0); + auto avgMask = copyNullMask(avgCol->view(), stream); + auto outAvg = copyColumnData<__int128_t>(avgCol->view(), stream); + + auto avgUnscaled = [](int128_t sum, int64_t count) { + __int128_t out = 0; + facebook::velox::DecimalUtil:: + divideWithRoundUp<__int128_t, __int128_t, int64_t>( + out, sum, count, false, 0, 0); + return out; + }; + + for (size_t i = 0; i < sums.size(); ++i) { + if (avgCol->view().nullable()) { + EXPECT_TRUE(isValidAt(avgMask, i)) << "row " << i; + } + EXPECT_EQ(outAvg[i], avgUnscaled(sums[i], counts[i])) << "row " << i; + } +} + +TEST_F(CudfDecimalTest, decimalSumStateRoundTripDecimal64) { + 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 stateMask = copyNullMask(stateCol->view(), stream); + + auto sumAndCount = deserializeDecimalSumState(stateCol->view(), 2, stream); + auto outSumView = sumAndCount.sum->view(); + auto outCountView = sumAndCount.count->view(); + auto outSum = copyColumnData<__int128_t>(outSumView, stream); + auto outCount = copyColumnData(outCountView, stream); + auto outSumMask = copyNullMask(outSumView, stream); + auto outCountMask = copyNullMask(outCountView, stream); + + EXPECT_EQ(stateMask, outSumMask); + EXPECT_EQ(stateMask, outCountMask); + + for (size_t i = 0; i < sums.size(); ++i) { + bool expectedValid = sumValid[i] && countValid[i] && counts[i] != 0; + EXPECT_EQ(isValidAt(stateMask, i), expectedValid); + if (expectedValid) { + EXPECT_EQ(outSum[i], static_cast<__int128_t>(sums[i])); + EXPECT_EQ(outCount[i], counts[i]); + } + } +} + +TEST_F(CudfDecimalTest, decimalSumStateRoundTripDecimal128) { + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + std::vector<__int128_t> sums = { + static_cast<__int128_t>(123450), + static_cast<__int128_t>(-25000), + static_cast<__int128_t>(100000), + }; + std::vector counts = {2, 1, 0}; + std::vector sumValid = {true, false, true}; + std::vector countValid = {true, true, true}; + + auto sumCol = makeDecimalColumn<__int128_t>(sums, 3, &sumValid, stream); + auto countCol = makeInt64Column(counts, &countValid, stream); + auto stateCol = + serializeDecimalSumState(sumCol->view(), countCol->view(), stream, mr); + auto stateMask = copyNullMask(stateCol->view(), stream); + + auto sumAndCount = deserializeDecimalSumState(stateCol->view(), 3, stream); + auto outSumView = sumAndCount.sum->view(); + auto outCountView = sumAndCount.count->view(); + auto outSum = copyColumnData<__int128_t>(outSumView, stream); + auto outCount = copyColumnData(outCountView, stream); + auto outSumMask = copyNullMask(outSumView, stream); + auto outCountMask = copyNullMask(outCountView, stream); + + EXPECT_EQ(stateMask, outSumMask); + EXPECT_EQ(stateMask, outCountMask); + + for (size_t i = 0; i < sums.size(); ++i) { + bool expectedValid = sumValid[i] && countValid[i] && counts[i] != 0; + EXPECT_EQ(isValidAt(stateMask, i), expectedValid); + if (expectedValid) { + EXPECT_EQ(outSum[i], sums[i]); + EXPECT_EQ(outCount[i], counts[i]); + } + } +} + +TEST_F(CudfDecimalTest, cudfVarbinaryArrowRoundTrip) { + auto input = makeRowVector( + {"bin"}, + {makeNullableFlatVector( + {std::string("abc"), std::nullopt, std::string("xyz")}, + VARBINARY())}); + + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + auto cudfTable = with_arrow::toCudfTable(input, pool(), stream, mr); + auto roundTrip = + with_arrow::toVeloxColumn(cudfTable->view(), pool(), "rt_", stream, mr); + + ASSERT_EQ(roundTrip->childAt(0)->type()->kind(), TypeKind::VARCHAR); + VELOX_ASSERT_THROW( + roundTrip->setType(ROW({{"rt_0", VARBINARY()}})), + "Cannot change vector type"); + + auto expected = makeRowVector( + {"rt_0"}, + {makeNullableFlatVector( + {std::string("abc"), std::nullopt, std::string("xyz")}, VARCHAR())}); + + facebook::velox::test::assertEqualVectors(expected, roundTrip); +} + +TEST_F(CudfDecimalTest, cudfVarbinaryArrowRoundTripWithExpectedType) { + auto input = makeRowVector( + {"bin"}, + {makeNullableFlatVector( + {std::string("abc"), std::nullopt, std::string("xyz")}, + VARBINARY())}); + + auto expectedType = ROW({{"bin", VARBINARY()}}); + + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + auto cudfTable = with_arrow::toCudfTable(input, pool(), stream, mr); + auto roundTrip = with_arrow::toVeloxColumn( + cudfTable->view(), pool(), expectedType, "rt_", stream, mr); + + ASSERT_EQ(roundTrip->childAt(0)->type()->kind(), TypeKind::VARBINARY); + + auto expected = makeRowVector( + {"rt_0"}, + {makeNullableFlatVector( + {std::string("abc"), std::nullopt, std::string("xyz")}, + VARBINARY())}); + + facebook::velox::test::assertEqualVectors(expected, roundTrip); +} + +TEST_F(CudfDecimalTest, cudfVarbinaryRowTypeMismatch) { + auto input = makeRowVector( + {"l_returnflag", + "l_linestatus", + "avg_51", + "avg_52", + "avg_53", + "count_54", + "sum_47", + "sum_48", + "sum_49", + "sum_50"}, + {makeFlatVector({"A", "B"}, VARCHAR()), + makeFlatVector({"F", "O"}, VARCHAR()), + makeFlatVector({"x", "y"}, VARBINARY()), + makeFlatVector({"p", "q"}, VARBINARY()), + makeFlatVector({"m", "n"}, VARBINARY()), + makeFlatVector({10, 20}, BIGINT()), + makeFlatVector({"u", "v"}, VARBINARY()), + makeFlatVector({"r", "s"}, VARBINARY()), + makeFlatVector({"t", "w"}, VARBINARY()), + makeFlatVector({"c", "d"}, VARBINARY())}); + + auto expectedType = ROW({ + {"l_returnflag", VARCHAR()}, + {"l_linestatus", VARCHAR()}, + {"avg_51", VARBINARY()}, + {"avg_52", VARBINARY()}, + {"avg_53", VARBINARY()}, + {"count_54", BIGINT()}, + {"sum_47", VARBINARY()}, + {"sum_48", VARBINARY()}, + {"sum_49", VARBINARY()}, + {"sum_50", VARBINARY()}, + }); + + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + auto cudfTable = with_arrow::toCudfTable(input, pool(), stream, mr); + auto roundTrip = + with_arrow::toVeloxColumn(cudfTable->view(), pool(), "rt_", stream, mr); + + ASSERT_EQ(roundTrip->childAt(2)->type()->kind(), TypeKind::VARCHAR); + ASSERT_EQ(roundTrip->childAt(3)->type()->kind(), TypeKind::VARCHAR); + ASSERT_EQ(roundTrip->childAt(4)->type()->kind(), TypeKind::VARCHAR); + ASSERT_EQ(roundTrip->childAt(6)->type()->kind(), TypeKind::VARCHAR); + ASSERT_EQ(roundTrip->childAt(7)->type()->kind(), TypeKind::VARCHAR); + ASSERT_EQ(roundTrip->childAt(8)->type()->kind(), TypeKind::VARCHAR); + ASSERT_EQ(roundTrip->childAt(9)->type()->kind(), TypeKind::VARCHAR); + + VELOX_ASSERT_THROW( + roundTrip->setType(expectedType), "Cannot change vector type"); +} + +} // namespace +} // namespace facebook::velox::cudf_velox