Add withCompanionFunctions for SetAgg

oap-project · Feb 29, 2024 · 39b0f11 · 39b0f11
1 parent 96448fe
commit 39b0f11
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp
@@ -48,7 +48,9 @@ extern void registerMaxDataSizeForStatsAggregate(const std::string& prefix);
 extern void registerMultiMapAggAggregate(const std::string& prefix);
 extern void registerSumDataSizeForStatsAggregate(const std::string& prefix);
 extern void registerReduceAgg(const std::string& prefix);
-extern void registerSetAggAggregate(const std::string& prefix);
+extern void registerSetAggAggregate(
+    const std::string& prefix,
+    bool withCompanionFunctions);
 extern void registerSetUnionAggregate(const std::string& prefix);
 
 extern void registerApproxDistinctAggregates(
@@ -108,7 +110,7 @@ void registerAllAggregateFunctions(
   registerMinMaxAggregates(prefix, withCompanionFunctions, overwrite);
   registerMinMaxByAggregates(prefix);
   registerReduceAgg(prefix);
-  registerSetAggAggregate(prefix);
+  registerSetAggAggregate(prefix, withCompanionFunctions);
   registerSetUnionAggregate(prefix);
   registerSumAggregate(prefix);
   registerVarianceAggregates(prefix, withCompanionFunctions, overwrite);

diff --git a/velox/functions/prestosql/aggregates/SetAggregates.cpp b/velox/functions/prestosql/aggregates/SetAggregates.cpp
@@ -427,7 +427,9 @@ std::unique_ptr<exec::Aggregate> create(
 
 } // namespace
 
-void registerSetAggAggregate(const std::string& prefix) {
+void registerSetAggAggregate(
+    const std::string& prefix,
+    bool withCompanionFunctions) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures = {
       exec::AggregateFunctionSignatureBuilder()
           .typeVariable("T")
@@ -491,7 +493,8 @@ void registerSetAggAggregate(const std::string& prefix) {
             VELOX_UNREACHABLE(
                 "Unexpected type {}", mapTypeKindToName(typeKind));
         }
-      });
+      },
+      withCompanionFunctions);
 }
 
 void registerSetUnionAggregate(const std::string& prefix) {