feat(cudf): GPU Decimal (Part 3)

velox/experimental/cudf/exec/AggregationRegistry.cpp

@@ -82,6 +82,37 @@ void registerCommonAggregationFunctions(
           .argumentType("double")
           .build()};
 
+  auto decimalSumSingle = std::vector<exec::FunctionSignaturePtr>{
+      FunctionSignatureBuilder()
+          .integerVariable("a_precision")
+          .integerVariable("a_scale")
+          .returnType("decimal(38, a_scale)")
+          .argumentType("decimal(a_precision, a_scale)")
+          .build()};
+  auto decimalSumPartial = std::vector<exec::FunctionSignaturePtr>{
+      FunctionSignatureBuilder()
+          .integerVariable("a_precision")
+          .integerVariable("a_scale")
+          .returnType("varbinary")
+          .argumentType("decimal(a_precision, a_scale)")
+          .build()};
+  auto decimalSumFinal = std::vector<exec::FunctionSignaturePtr>{
+      FunctionSignatureBuilder()
+          .integerVariable("a_scale")
+          .returnType("decimal(38, a_scale)")
+          .argumentType("varbinary")
+          .build()};
+  auto decimalSumIntermediate =
+      std::vector<exec::FunctionSignaturePtr>{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<exec::FunctionSignaturePtr>{
       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<exec::FunctionSignaturePtr>{
+      FunctionSignatureBuilder()
+          .integerVariable("a_precision")
+          .integerVariable("a_scale")
+          .returnType("decimal(a_precision, a_scale)")
+          .argumentType("decimal(a_precision, a_scale)")
+          .build()};
+  auto decimalAvgPartial = std::vector<exec::FunctionSignaturePtr>{
+      FunctionSignatureBuilder()
+          .integerVariable("a_precision")
+          .integerVariable("a_scale")
+          .returnType("varbinary")
+          .argumentType("decimal(a_precision, a_scale)")
+          .build()};
+  auto decimalAvgFinal = std::vector<exec::FunctionSignaturePtr>{
+      FunctionSignatureBuilder()
+          .integerVariable("a_precision")
+          .integerVariable("a_scale")
+          .returnType("decimal(a_precision, a_scale)")
+          .argumentType("varbinary")
+          .build()};
+  auto decimalAvgIntermediate =
+      std::vector<exec::FunctionSignaturePtr>{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<exec::FunctionSignaturePtr>{
+  auto avgFinalIntermediateSignatures = std::vector<exec::FunctionSignaturePtr>{
       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",

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

velox/experimental/cudf/exec/CudfAggregation.cpp

@@ -136,6 +136,8 @@ std::vector<ResolvedAggregateInfo> 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<ResolvedAggregateInfo> resolveAggregateInfos(
         resultType,
         isCountFunctionName(aggregate.call->name())
             ? std::make_optional(getCountInputKind(aggregate, constants[i]))
-            : std::nullopt);
+            : std::nullopt,
+        isDecimalAggregate);
   }
   return params;
 }

velox/experimental/cudf/exec/CudfAggregation.h

@@ -62,6 +62,10 @@ struct ResolvedAggregateInfo {
   VectorPtr constant;
   TypePtr resultType;
   std::optional<CountInputKind> 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,