Add config for registration (7110)

oap-project · Dec 29, 2023 · 3d82641 · 3d82641
1 parent b7baa7c
commit 3d82641
Show file tree

Hide file tree

Showing 18 changed files with 182 additions and 71 deletions.
diff --git a/velox/functions/lib/aggregates/BitwiseAggregateBase.h b/velox/functions/lib/aggregates/BitwiseAggregateBase.h
@@ -70,7 +70,10 @@ class BitwiseAggregateBase : public SimpleNumericAggregate<T, T, T> {
 };
 
 template <template <typename U> class T>
-exec::AggregateRegistrationResult registerBitwise(const std::string& name) {
+exec::AggregateRegistrationResult registerBitwise(
+    const std::string& name,
+    bool withCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures;
   for (const auto& inputType : {"tinyint", "smallint", "integer", "bigint"}) {
     signatures.push_back(exec::AggregateFunctionSignatureBuilder()
@@ -106,7 +109,9 @@ exec::AggregateRegistrationResult registerBitwise(const std::string& name) {
                 name,
                 inputType->kindName());
         }
-      });
+      },
+      withCompanionFunctions,
+      overwrite);
 }
 
 } // namespace facebook::velox::functions::aggregate
diff --git a/velox/functions/prestosql/aggregates/BitwiseAggregates.cpp b/velox/functions/prestosql/aggregates/BitwiseAggregates.cpp
@@ -101,9 +101,14 @@ class BitwiseAndAggregate : public BitwiseAggregateBase<T> {
 
 } // namespace
 
-void registerBitwiseAggregates(const std::string& prefix) {
-  registerBitwise<BitwiseOrAggregate>(prefix + kBitwiseOr);
-  registerBitwise<BitwiseAndAggregate>(prefix + kBitwiseAnd);
+void registerBitwiseAggregates(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite) {
+  registerBitwise<BitwiseOrAggregate>(
+      prefix + kBitwiseOr, withCompanionFunctions, overwrite);
+  registerBitwise<BitwiseAndAggregate>(
+      prefix + kBitwiseAnd, withCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/CountAggregate.cpp b/velox/functions/prestosql/aggregates/CountAggregate.cpp
@@ -150,7 +150,10 @@ class CountAggregate : public SimpleNumericAggregate<bool, int64_t, int64_t> {
 
 } // namespace
 
-void registerCountAggregate(const std::string& prefix) {
+void registerCountAggregate(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures{
       exec::AggregateFunctionSignatureBuilder()
           .returnType("bigint")
@@ -177,7 +180,9 @@ void registerCountAggregate(const std::string& prefix) {
         VELOX_CHECK_LE(
             argTypes.size(), 1, "{} takes at most one argument", name);
         return std::make_unique<CountAggregate>();
-      });
+      },
+      withCompanionFunctions,
+      overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/CovarianceAggregates.cpp b/velox/functions/prestosql/aggregates/CovarianceAggregates.cpp
@@ -571,7 +571,10 @@ template <
     typename TIntermediateInput,
     typename TIntermediateResult,
     typename TResultAccessor>
-exec::AggregateRegistrationResult registerCovariance(const std::string& name) {
+exec::AggregateRegistrationResult registerCovariance(
+    const std::string& name,
+    bool withCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures = {
       // (double, double) -> double
       exec::AggregateFunctionSignatureBuilder()
@@ -620,37 +623,46 @@ exec::AggregateRegistrationResult registerCovariance(const std::string& name) {
                 "Unsupported raw input type: {}. Expected DOUBLE or REAL.",
                 rawInputType->toString())
         }
-      });
+      },
+      withCompanionFunctions,
+      overwrite);
 }
 
 } // namespace
 
-void registerCovarianceAggregates(const std::string& prefix) {
+void registerCovarianceAggregates(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite) {
   registerCovariance<
       CovarAccumulator,
       CovarIntermediateInput,
       CovarIntermediateResult,
-      CovarPopResultAccessor>(prefix + kCovarPop);
+      CovarPopResultAccessor>(
+      prefix + kCovarPop, withCompanionFunctions, overwrite);
   registerCovariance<
       CovarAccumulator,
       CovarIntermediateInput,
       CovarIntermediateResult,
-      CovarSampResultAccessor>(prefix + kCovarSamp);
+      CovarSampResultAccessor>(
+      prefix + kCovarSamp, withCompanionFunctions, overwrite);
   registerCovariance<
       CorrAccumulator,
       CorrIntermediateInput,
       CorrIntermediateResult,
-      CorrResultAccessor>(prefix + kCorr);
+      CorrResultAccessor>(prefix + kCorr, withCompanionFunctions, overwrite);
   registerCovariance<
       RegrAccumulator,
       RegrIntermediateInput,
       RegrIntermediateResult,
-      RegrInterceptResultAccessor>(prefix + kRegrIntercept);
+      RegrInterceptResultAccessor>(
+      prefix + kRegrIntercept, withCompanionFunctions, overwrite);
   registerCovariance<
       RegrAccumulator,
       RegrIntermediateInput,
       RegrIntermediateResult,
-      RegrSlopeResultAccessor>(prefix + kRegrSlop);
+      RegrSlopeResultAccessor>(
+      prefix + kRegrSlop, withCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/MinMaxAggregates.cpp b/velox/functions/prestosql/aggregates/MinMaxAggregates.cpp
@@ -904,7 +904,10 @@ template <
     typename TNonNumeric,
     template <typename T>
     class TNumericN>
-exec::AggregateRegistrationResult registerMinMax(const std::string& name) {
+exec::AggregateRegistrationResult registerMinMax(
+    const std::string& name,
+    bool withCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures;
   signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                            .orderableTypeVariable("T")
@@ -1008,16 +1011,21 @@ exec::AggregateRegistrationResult registerMinMax(const std::string& name) {
                   inputType->kindName());
           }
         }
-      });
+      },
+      withCompanionFunctions,
+      overwrite);
 }
 
 } // namespace
 
-void registerMinMaxAggregates(const std::string& prefix) {
+void registerMinMaxAggregates(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite) {
   registerMinMax<MinAggregate, NonNumericMinAggregate, MinNAggregate>(
-      prefix + kMin);
+      prefix + kMin, withCompanionFunctions, overwrite);
   registerMinMax<MaxAggregate, NonNumericMaxAggregate, MaxNAggregate>(
-      prefix + kMax);
+      prefix + kMax, withCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp
@@ -31,7 +31,10 @@ extern void registerAverageAggregate(
     bool withCompanionFunctions);
 extern void registerBitwiseXorAggregate(const std::string& prefix);
 extern void registerChecksumAggregate(const std::string& prefix);
-extern void registerCountAggregate(const std::string& prefix);
+extern void registerCountAggregate(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite);
 extern void registerCountIfAggregate(const std::string& prefix);
 extern void registerEntropyAggregate(const std::string& prefix);
 extern void registerGeometricMeanAggregate(const std::string& prefix);
@@ -49,32 +52,45 @@ extern void registerSetUnionAggregate(const std::string& prefix);
 extern void registerApproxDistinctAggregates(
     const std::string& prefix,
     bool withCompanionFunctions);
-extern void registerBitwiseAggregates(const std::string& prefix);
+extern void registerBitwiseAggregates(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite);
 extern void registerBoolAggregates(const std::string& prefix);
 extern void registerCentralMomentsAggregates(const std::string& prefix);
-extern void registerCovarianceAggregates(const std::string& prefix);
-extern void registerMinMaxAggregates(const std::string& prefix);
+extern void registerCovarianceAggregates(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite);
+extern void registerMinMaxAggregates(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite);
 extern void registerMinMaxByAggregates(const std::string& prefix);
 extern void registerSumAggregate(const std::string& prefix);
-extern void registerVarianceAggregates(const std::string& prefix);
+extern void registerVarianceAggregates(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite);
 
 void registerAllAggregateFunctions(
     const std::string& prefix,
-    bool withCompanionFunctions) {
+    bool withCompanionFunctions,
+    bool overwrite) {
   registerApproxDistinctAggregates(prefix, withCompanionFunctions);
   registerApproxMostFrequentAggregate(prefix);
   registerApproxPercentileAggregate(prefix, withCompanionFunctions);
   registerArbitraryAggregate(prefix);
   registerArrayAggAggregate(prefix, withCompanionFunctions);
   registerAverageAggregate(prefix, withCompanionFunctions);
-  registerBitwiseAggregates(prefix);
+  registerBitwiseAggregates(prefix, withCompanionFunctions, overwrite);
   registerBitwiseXorAggregate(prefix);
   registerBoolAggregates(prefix);
   registerCentralMomentsAggregates(prefix);
   registerChecksumAggregate(prefix);
-  registerCountAggregate(prefix);
+  registerCountAggregate(prefix, withCompanionFunctions, overwrite);
   registerCountIfAggregate(prefix);
-  registerCovarianceAggregates(prefix);
+  registerCovarianceAggregates(prefix, withCompanionFunctions, overwrite);
   registerEntropyAggregate(prefix);
   registerGeometricMeanAggregate(prefix);
   registerHistogramAggregate(prefix);
@@ -84,13 +100,13 @@ void registerAllAggregateFunctions(
   registerMaxDataSizeForStatsAggregate(prefix);
   registerMultiMapAggAggregate(prefix);
   registerSumDataSizeForStatsAggregate(prefix);
-  registerMinMaxAggregates(prefix);
+  registerMinMaxAggregates(prefix, withCompanionFunctions, overwrite);
   registerMinMaxByAggregates(prefix);
   registerReduceAgg(prefix);
   registerSetAggAggregate(prefix);
   registerSetUnionAggregate(prefix);
   registerSumAggregate(prefix);
-  registerVarianceAggregates(prefix);
+  registerVarianceAggregates(prefix, withCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h
@@ -21,6 +21,7 @@ namespace facebook::velox::aggregate::prestosql {
 
 void registerAllAggregateFunctions(
     const std::string& prefix = "",
-    bool withCompanionFunctions = true);
+    bool withCompanionFunctions = true,
+    bool overwrite = false);
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/VarianceAggregates.cpp b/velox/functions/prestosql/aggregates/VarianceAggregates.cpp
@@ -465,7 +465,10 @@ void checkSumCountRowType(
 }
 
 template <template <typename TInput> class TClass>
-exec::AggregateRegistrationResult registerVariance(const std::string& name) {
+exec::AggregateRegistrationResult registerVariance(
+    const std::string& name,
+    bool withCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures;
   std::vector<std::string> inputTypes = {
       "smallint", "integer", "bigint", "real", "double"};
@@ -514,18 +517,29 @@ exec::AggregateRegistrationResult registerVariance(const std::string& name) {
               "(count:bigint, mean:double, m2:double) struct");
           return std::make_unique<TClass<int64_t>>(resultType);
         }
-      });
+      },
+      withCompanionFunctions,
+      overwrite);
 }
 
 } // namespace
 
-void registerVarianceAggregates(const std::string& prefix) {
-  registerVariance<StdDevSampAggregate>(prefix + kStdDev);
-  registerVariance<StdDevPopAggregate>(prefix + kStdDevPop);
-  registerVariance<StdDevSampAggregate>(prefix + kStdDevSamp);
-  registerVariance<VarSampAggregate>(prefix + kVariance);
-  registerVariance<VarPopAggregate>(prefix + kVarPop);
-  registerVariance<VarSampAggregate>(prefix + kVarSamp);
+void registerVarianceAggregates(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite) {
+  registerVariance<StdDevSampAggregate>(
+      prefix + kStdDev, withCompanionFunctions, overwrite);
+  registerVariance<StdDevPopAggregate>(
+      prefix + kStdDevPop, withCompanionFunctions, overwrite);
+  registerVariance<StdDevSampAggregate>(
+      prefix + kStdDevSamp, withCompanionFunctions, overwrite);
+  registerVariance<VarSampAggregate>(
+      prefix + kVariance, withCompanionFunctions, overwrite);
+  registerVariance<VarPopAggregate>(
+      prefix + kVarPop, withCompanionFunctions, overwrite);
+  registerVariance<VarSampAggregate>(
+      prefix + kVarSamp, withCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/sparksql/aggregates/AverageAggregate.cpp b/velox/functions/sparksql/aggregates/AverageAggregate.cpp
@@ -370,7 +370,8 @@ class DecimalAverageAggregate : public DecimalAggregate<TInputType> {
 ///     DECIMAL         |     DECIMAL         |    DECIMAL
 exec::AggregateRegistrationResult registerAverage(
     const std::string& name,
-    bool withCompanionFunctions) {
+    bool withCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures;
 
   for (const auto& inputType :
@@ -494,7 +495,8 @@ exec::AggregateRegistrationResult registerAverage(
           }
         }
       },
-      withCompanionFunctions);
+      withCompanionFunctions,
+      overwrite);
 }
 
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/AverageAggregate.h b/velox/functions/sparksql/aggregates/AverageAggregate.h
@@ -24,6 +24,7 @@ namespace facebook::velox::functions::aggregate::sparksql {
 
 exec::AggregateRegistrationResult registerAverage(
     const std::string& name,
-    bool withCompanionFunctions);
+    bool withCompanionFunctions,
+    bool overwrite);
 
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/BitwiseXorAggregate.cpp b/velox/functions/sparksql/aggregates/BitwiseXorAggregate.cpp
@@ -66,9 +66,11 @@ class BitwiseXorAggregate : public BitwiseAggregateBase<T> {
 } // namespace
 
 exec::AggregateRegistrationResult registerBitwiseXorAggregate(
-    const std::string& prefix) {
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite) {
   return functions::aggregate::registerBitwise<BitwiseXorAggregate>(
-      prefix + "bit_xor");
+      prefix + "bit_xor", withCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/BitwiseXorAggregate.h b/velox/functions/sparksql/aggregates/BitwiseXorAggregate.h
@@ -23,6 +23,8 @@
 namespace facebook::velox::functions::aggregate::sparksql {
 
 exec::AggregateRegistrationResult registerBitwiseXorAggregate(
-    const std::string& name);
+    const std::string& name,
+    bool registerCompanionFunctions,
+    bool overwrite);
 
 } // namespace facebook::velox::functions::aggregate::sparksql