Fix decimal agg signature on partial companion function (oap-project#465

)
GlutenPerfBot · Dec 22, 2023 · c3ba99c · c3ba99c
1 parent d48e9a2
commit c3ba99c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/velox/functions/sparksql/aggregates/AverageAggregate.cpp b/velox/functions/sparksql/aggregates/AverageAggregate.cpp
@@ -429,7 +429,9 @@ exec::AggregateRegistrationResult registerAverage(
                 auto inputScale = inputType->asShortDecimal().scale();
                 auto sumType =
                     DECIMAL(std::min(38, inputPrecision + 10), inputScale);
-                if (exec::isPartialOutput(step)) {
+                if (exec::isPartialOutput(step) ||
+                    (step == core::AggregationNode::Step::kSingle &&
+                     resultType->isRow())) {
                   return std::make_unique<
                       DecimalAverageAggregate<int64_t, int64_t>>(
                       resultType, sumType);

diff --git a/velox/functions/sparksql/aggregates/SumAggregate.cpp b/velox/functions/sparksql/aggregates/SumAggregate.cpp
@@ -29,7 +29,13 @@ using SumAggregate = SumAggregateBase<TInput, TAccumulator, ResultType, true>;
 TypePtr getDecimalSumType(
     const TypePtr& resultType,
     core::AggregationNode::Step step) {
-  return exec::isPartialOutput(step) ? resultType->childAt(0) : resultType;
+  if (exec::isPartialOutput(step)) {
+    return resultType->childAt(0);
+  }
+  if (step == core::AggregationNode::Step::kSingle && resultType->isRow()) {
+    return resultType->childAt(0);
+  }
+  return resultType;
 }
 } // namespace