Register merge extract companion agg functions without suffix

Author: zhztheplayer

velox/exec/AggregateCompanionAdapter.cpp

@@ -266,10 +266,13 @@ bool CompanionFunctionsRegistrar::registerPartialFunction(
                  const core::QueryConfig& config)
                  -> std::unique_ptr<Aggregate> {
                if (auto func = getAggregateFunctionEntry(name)) {
+                 core::AggregationNode::Step usedStep{
+                     core::AggregationNode::Step::kPartial};
                  if (!exec::isRawInput(step)) {
-                   step = core::AggregationNode::Step::kIntermediate;
+                   usedStep = core::AggregationNode::Step::kIntermediate;
                  }
-                 auto fn = func->factory(step, argTypes, resultType, config);
+                 auto fn =
+                     func->factory(usedStep, argTypes, resultType, config);
                  VELOX_CHECK_NOT_NULL(fn);
                  return std::make_unique<
                      AggregateCompanionAdapter::PartialFunction>(
@@ -387,56 +390,50 @@ bool CompanionFunctionsRegistrar::registerMergeExtractFunction(
     const std::string& name,
     const std::vector<AggregateFunctionSignaturePtr>& signatures,
     bool overwrite) {
+  bool registered = false;
   if (CompanionSignatures::hasSameIntermediateTypesAcrossSignatures(
           signatures)) {
-    return registerMergeExtractFunctionWithSuffix(name, signatures, overwrite);
+    registered |=
+        registerMergeExtractFunctionWithSuffix(name, signatures, overwrite);
   }
 
   auto mergeExtractSignatures =
       CompanionSignatures::mergeExtractFunctionSignatures(signatures);
   if (mergeExtractSignatures.empty()) {
-    return false;
+    return registered;
   }
 
   auto mergeExtractFunctionName =
       CompanionSignatures::mergeExtractFunctionName(name);
-  return exec::registerAggregateFunction(
-             mergeExtractFunctionName,
-             std::move(mergeExtractSignatures),
-             [name, mergeExtractFunctionName](
-                 core::AggregationNode::Step /*step*/,
-                 const std::vector<TypePtr>& argTypes,
-                 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 resolveable 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,
-                     config);
-                 VELOX_CHECK_NOT_NULL(fn);
-                 return std::make_unique<
-                     AggregateCompanionAdapter::MergeExtractFunction>(
-                     std::move(fn), resultType);
-               }
-               VELOX_FAIL(
-                   "Original aggregation function {} not found: {}",
-                   name,
-                   mergeExtractFunctionName);
-             },
-             /*registerCompanionFunctions*/ false,
-             overwrite)
-      .mainFunction;
+  registered |=
+      exec::registerAggregateFunction(
+          mergeExtractFunctionName,
+          std::move(mergeExtractSignatures),
+          [name, mergeExtractFunctionName](
+              core::AggregationNode::Step /*step*/,
+              const std::vector<TypePtr>& argTypes,
+              const TypePtr& resultType,
+              const core::QueryConfig& config) -> std::unique_ptr<Aggregate> {
+            if (auto func = getAggregateFunctionEntry(name)) {
+              auto fn = func->factory(
+                  core::AggregationNode::Step::kFinal,
+                  argTypes,
+                  resultType,
+                  config);
+              VELOX_CHECK_NOT_NULL(fn);
+              return std::make_unique<
+                  AggregateCompanionAdapter::MergeExtractFunction>(
+                  std::move(fn), resultType);
+            }
+            VELOX_FAIL(
+                "Original aggregation function {} not found: {}",
+                name,
+                mergeExtractFunctionName);
+          },
+          /*registerCompanionFunctions*/ false,
+          overwrite)
+          .mainFunction;
+  return registered;
 }
 
 bool CompanionFunctionsRegistrar::registerExtractFunctionWithSuffix(