Allow partial, merge and merge_extract companion functions

facebookincubator · Jan 8, 2025 · 9911e8f · 9911e8f
1 parent 3f86559
commit 9911e8f
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 30 deletions.
diff --git a/velox/exec/AggregateCompanionAdapter.cpp b/velox/exec/AggregateCompanionAdapter.cpp
@@ -340,20 +340,11 @@ bool registerAggregateFunction(
                  const TypePtr& resultType,
                  const core::QueryConfig& config)
                  -> std::unique_ptr<Aggregate> {
-               const auto& [originalResultType, _] =
-                   resolveAggregateFunction(mergeExtractFunctionName, argTypes);
-               if (!originalResultType) {
-                 // TODO: limitation -- result type must be resolvable given
-                 // intermediate type of the original UDAF.
-                 VELOX_UNREACHABLE(
-                     "Signatures whose result types are not resolvable given intermediate types should have been excluded.");
-               }
-
                if (auto func = getAggregateFunctionEntry(name)) {
                  auto fn = func->factory(
                      core::AggregationNode::Step::kFinal,
                      argTypes,
-                     originalResultType,
+                     resultType,
                      config);
                  VELOX_CHECK_NOT_NULL(fn);
                  return std::make_unique<

diff --git a/velox/exec/AggregateCompanionSignatures.cpp b/velox/exec/AggregateCompanionSignatures.cpp
@@ -100,9 +100,6 @@ CompanionSignatures::partialFunctionSignatures(
     const std::vector<AggregateFunctionSignaturePtr>& signatures) {
   std::vector<AggregateFunctionSignaturePtr> partialSignatures;
   for (const auto& signature : signatures) {
-    if (!isResultTypeResolvableGivenIntermediateType(signature)) {
-      continue;
-    }
     std::vector<TypeSignature> usedTypes = signature->argumentTypes();
     usedTypes.push_back(signature->intermediateType());
     auto variables = usedTypeVariables(usedTypes, signature->variables());
@@ -124,10 +121,6 @@ std::string CompanionSignatures::partialFunctionName(const std::string& name) {
 
 AggregateFunctionSignaturePtr CompanionSignatures::mergeFunctionSignature(
     const AggregateFunctionSignaturePtr& signature) {
-  if (!isResultTypeResolvableGivenIntermediateType(signature)) {
-    return nullptr;
-  }
-
   std::vector<TypeSignature> usedTypes = {signature->intermediateType()};
   auto variables = usedTypeVariables(usedTypes, signature->variables());
   return std::make_shared<AggregateFunctionSignature>(
@@ -170,10 +163,6 @@ bool CompanionSignatures::hasSameIntermediateTypesAcrossSignatures(
 AggregateFunctionSignaturePtr
 CompanionSignatures::mergeExtractFunctionSignature(
     const AggregateFunctionSignaturePtr& signature) {
-  if (!isResultTypeResolvableGivenIntermediateType(signature)) {
-    return nullptr;
-  }
-
   std::vector<TypeSignature> usedTypes = {
       signature->intermediateType(), signature->returnType()};
   auto variables = usedTypeVariables(usedTypes, signature->variables());

diff --git a/velox/exec/tests/AggregateCompanionAdapterTest.cpp b/velox/exec/tests/AggregateCompanionAdapterTest.cpp
@@ -414,22 +414,27 @@ TEST_F(
 TEST_F(
     AggregateCompanionRegistryTest,
     resultTypeNotResolvableFromIntermediateType) {
-  // We only register companion functions for original signatures whose result
-  // type can be resolved from its intermediate type.
+  // We only register partial, merge and merge_extract companion functions for
+  // original signatures whose result type cannot be resolved from its
+  // intermediate type.
   std::vector<std::shared_ptr<AggregateFunctionSignature>> signatures{
       AggregateFunctionSignatureBuilder()
-          .typeVariable("T")
-          .returnType("array(T)")
-          .intermediateType("varbinary")
-          .argumentType("T")
+          .integerVariable("a_precision")
+          .integerVariable("a_scale")
+          .integerVariable("i_precision", "min(38, a_precision + 10)")
+          .integerVariable("r_precision", "min(38, a_precision + 4)")
+          .integerVariable("r_scale", "min(38, a_scale + 4)")
+          .returnType("DECIMAL(r_precision, r_scale)")
+          .intermediateType("ROW(DECIMAL(i_precision, a_scale), bigint)")
+          .argumentType("DECIMAL(a_precision, a_scale)")
           .build()};
   registerDummyAggregateFunction("aggregateFunc6", signatures);
 
-  checkAggregateSignaturesCount("aggregateFunc6_partial", 0);
+  checkAggregateSignaturesCount("aggregateFunc6_partial", 1);
 
-  checkAggregateSignaturesCount("aggregateFunc6_merge", 0);
+  checkAggregateSignaturesCount("aggregateFunc6_merge", 1);
 
-  checkAggregateSignaturesCount("aggregateFunc6_merge_extract", 0);
+  checkAggregateSignaturesCount("aggregateFunc6_merge_extract", 1);
 
   checkScalarSignaturesCount("aggregateFunc6_extract", 0);
 }