From d31d92162168b78665752d0f16b42783901be84c Mon Sep 17 00:00:00 2001
From: PHILO-HE <feilong.he@intel.com>
Date: Fri, 3 Nov 2023 10:11:35 +0800
Subject: [PATCH] Add registerCompanionFunctions and overwrite as parameters in
 agg registration (7110)

---
 .../lib/aggregates/BitwiseAggregateBase.h     |  9 ++++-
 .../aggregates/BitwiseAggregates.cpp          | 11 ++++--
 .../prestosql/aggregates/CountAggregate.cpp   |  8 +++-
 .../aggregates/CovarianceAggregates.cpp       | 29 ++++++++++----
 .../prestosql/aggregates/MinMaxAggregates.cpp | 18 ++++++---
 .../aggregates/RegisterAggregateFunctions.cpp | 39 +++++++++++++------
 .../aggregates/RegisterAggregateFunctions.h   |  5 ++-
 .../aggregates/VarianceAggregates.cpp         | 32 ++++++++++-----
 .../sparksql/aggregates/AverageAggregate.cpp  |  8 +++-
 .../sparksql/aggregates/AverageAggregate.h    |  5 ++-
 .../aggregates/BitwiseXorAggregate.cpp        |  6 ++-
 .../sparksql/aggregates/BitwiseXorAggregate.h |  4 +-
 .../aggregates/FirstLastAggregate.cpp         | 26 +++++++++----
 .../sparksql/aggregates/MinMaxByAggregate.cpp | 20 +++++++---
 .../sparksql/aggregates/Register.cpp          | 26 ++++++++-----
 .../functions/sparksql/aggregates/Register.h  |  5 ++-
 .../sparksql/aggregates/SumAggregate.cpp      |  9 ++++-
 .../sparksql/aggregates/SumAggregate.h        |  5 ++-
 18 files changed, 193 insertions(+), 72 deletions(-)

diff --git a/velox/functions/lib/aggregates/BitwiseAggregateBase.h b/velox/functions/lib/aggregates/BitwiseAggregateBase.h
index 428b905ea83c..d0b0efe298bd 100644
--- 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 registerCompanionFunctions,
+    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());
         }
-      });
+      },
+      registerCompanionFunctions,
+      overwrite);
 }
 
 } // namespace facebook::velox::functions::aggregate
diff --git a/velox/functions/prestosql/aggregates/BitwiseAggregates.cpp b/velox/functions/prestosql/aggregates/BitwiseAggregates.cpp
index a5127e37bed2..943060e47524 100644
--- 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 registerCompanionFunctions,
+    bool overwrite) {
+  registerBitwise<BitwiseOrAggregate>(
+      prefix + kBitwiseOr, registerCompanionFunctions, overwrite);
+  registerBitwise<BitwiseAndAggregate>(
+      prefix + kBitwiseAnd, registerCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/CountAggregate.cpp b/velox/functions/prestosql/aggregates/CountAggregate.cpp
index edf0622e69ac..1fb19c414405 100644
--- a/velox/functions/prestosql/aggregates/CountAggregate.cpp
+++ b/velox/functions/prestosql/aggregates/CountAggregate.cpp
@@ -151,7 +151,9 @@ class CountAggregate : public SimpleNumericAggregate<bool, int64_t, int64_t> {
 } // namespace
 
 exec::AggregateRegistrationResult registerCountAggregate(
-    const std::string& prefix) {
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures{
       exec::AggregateFunctionSignatureBuilder()
           .returnType("bigint")
@@ -178,7 +180,9 @@ exec::AggregateRegistrationResult registerCountAggregate(
         VELOX_CHECK_LE(
             argTypes.size(), 1, "{} takes at most one argument", name);
         return std::make_unique<CountAggregate>();
-      });
+      },
+      registerCompanionFunctions,
+      overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/CovarianceAggregates.cpp b/velox/functions/prestosql/aggregates/CovarianceAggregates.cpp
index f318c4263495..3cdbdbfae5a6 100644
--- a/velox/functions/prestosql/aggregates/CovarianceAggregates.cpp
+++ b/velox/functions/prestosql/aggregates/CovarianceAggregates.cpp
@@ -558,7 +558,10 @@ template <
     typename TIntermediateInput,
     typename TIntermediateResult,
     typename TResultAccessor>
-exec::AggregateRegistrationResult registerCovariance(const std::string& name) {
+exec::AggregateRegistrationResult registerCovariance(
+    const std::string& name,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures = {
       // (double, double) -> double
       exec::AggregateFunctionSignatureBuilder()
@@ -607,37 +610,47 @@ exec::AggregateRegistrationResult registerCovariance(const std::string& name) {
                 "Unsupported raw input type: {}. Expected DOUBLE or REAL.",
                 rawInputType->toString())
         }
-      });
+      },
+      registerCompanionFunctions,
+      overwrite);
 }
 
 } // namespace
 
-void registerCovarianceAggregates(const std::string& prefix) {
+void registerCovarianceAggregates(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   registerCovariance<
       CovarAccumulator,
       CovarIntermediateInput,
       CovarIntermediateResult,
-      CovarPopResultAccessor>(prefix + kCovarPop);
+      CovarPopResultAccessor>(
+      prefix + kCovarPop, registerCompanionFunctions, overwrite);
   registerCovariance<
       CovarAccumulator,
       CovarIntermediateInput,
       CovarIntermediateResult,
-      CovarSampResultAccessor>(prefix + kCovarSamp);
+      CovarSampResultAccessor>(
+      prefix + kCovarSamp, registerCompanionFunctions, overwrite);
   registerCovariance<
       CorrAccumulator,
       CorrIntermediateInput,
       CorrIntermediateResult,
-      CorrResultAccessor>(prefix + kCorr);
+      CorrResultAccessor>(
+      prefix + kCorr, registerCompanionFunctions, overwrite);
   registerCovariance<
       RegrAccumulator,
       RegrIntermediateInput,
       RegrIntermediateResult,
-      RegrInterceptResultAccessor>(prefix + kRegrIntercept);
+      RegrInterceptResultAccessor>(
+      prefix + kRegrIntercept, registerCompanionFunctions, overwrite);
   registerCovariance<
       RegrAccumulator,
       RegrIntermediateInput,
       RegrIntermediateResult,
-      RegrSlopeResultAccessor>(prefix + kRegrSlop);
+      RegrSlopeResultAccessor>(
+      prefix + kRegrSlop, registerCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/MinMaxAggregates.cpp b/velox/functions/prestosql/aggregates/MinMaxAggregates.cpp
index 2aa40e497a61..d9897d13f7ff 100644
--- 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 registerCompanionFunctions,
+    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());
           }
         }
-      });
+      },
+      registerCompanionFunctions,
+      overwrite);
 }
 
 } // namespace
 
-void registerMinMaxAggregates(const std::string& prefix) {
+void registerMinMaxAggregates(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   registerMinMax<MinAggregate, NonNumericMinAggregate, MinNAggregate>(
-      prefix + kMin);
+      prefix + kMin, registerCompanionFunctions, overwrite);
   registerMinMax<MaxAggregate, NonNumericMaxAggregate, MaxNAggregate>(
-      prefix + kMax);
+      prefix + kMax, registerCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp
index d9dc50277ed5..88e7f8901a68 100644
--- a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp
+++ b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp
@@ -33,7 +33,9 @@ extern exec::AggregateRegistrationResult registerBitwiseXorAggregate(
 extern exec::AggregateRegistrationResult registerChecksumAggregate(
     const std::string& prefix);
 extern exec::AggregateRegistrationResult registerCountAggregate(
-    const std::string& prefix);
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite);
 extern exec::AggregateRegistrationResult registerCountIfAggregate(
     const std::string& prefix);
 extern exec::AggregateRegistrationResult registerEntropyAggregate(
@@ -62,30 +64,45 @@ extern exec::AggregateRegistrationResult registerSetUnionAggregate(
     const std::string& prefix);
 
 extern void registerApproxDistinctAggregates(const std::string& prefix);
-extern void registerBitwiseAggregates(const std::string& prefix);
+extern void registerBitwiseAggregates(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    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 registerCompanionFunctions,
+    bool overwrite);
+extern void registerMinMaxAggregates(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    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 registerCompanionFunctions,
+    bool overwrite);
 
-void registerAllAggregateFunctions(const std::string& prefix) {
+void registerAllAggregateFunctions(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   registerApproxDistinctAggregates(prefix);
   registerApproxMostFrequentAggregate(prefix);
   registerApproxPercentileAggregate(prefix);
   registerArbitraryAggregate(prefix);
   registerArrayAggAggregate(prefix);
   registerAverageAggregate(prefix);
-  registerBitwiseAggregates(prefix);
+  registerBitwiseAggregates(prefix, registerCompanionFunctions, overwrite);
   registerBitwiseXorAggregate(prefix);
   registerBoolAggregates(prefix);
   registerCentralMomentsAggregates(prefix);
   registerChecksumAggregate(prefix);
-  registerCountAggregate(prefix);
+  registerCountAggregate(prefix, registerCompanionFunctions, overwrite);
   registerCountIfAggregate(prefix);
-  registerCovarianceAggregates(prefix);
+  registerCovarianceAggregates(prefix, registerCompanionFunctions, overwrite);
   registerEntropyAggregate(prefix);
   registerGeometricMeanAggregate(prefix);
   registerHistogramAggregate(prefix);
@@ -95,13 +112,13 @@ void registerAllAggregateFunctions(const std::string& prefix) {
   registerMaxDataSizeForStatsAggregate(prefix);
   registerMultiMapAggAggregate(prefix);
   registerSumDataSizeForStatsAggregate(prefix);
-  registerMinMaxAggregates(prefix);
+  registerMinMaxAggregates(prefix, registerCompanionFunctions, overwrite);
   registerMinMaxByAggregates(prefix);
   registerReduceAgg(prefix);
   registerSetAggAggregate(prefix);
   registerSetUnionAggregate(prefix);
   registerSumAggregate(prefix);
-  registerVarianceAggregates(prefix);
+  registerVarianceAggregates(prefix, registerCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h
index e6a91e0676aa..5ae9542cdcdf 100644
--- a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h
+++ b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h
@@ -19,6 +19,9 @@
 
 namespace facebook::velox::aggregate::prestosql {
 
-void registerAllAggregateFunctions(const std::string& prefix = "");
+void registerAllAggregateFunctions(
+    const std::string& prefix = "",
+    bool registerCompanionFunctions = false,
+    bool overwrite = false);
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/prestosql/aggregates/VarianceAggregates.cpp b/velox/functions/prestosql/aggregates/VarianceAggregates.cpp
index 4e0eb2c6ae68..f7cb59d1b20c 100644
--- a/velox/functions/prestosql/aggregates/VarianceAggregates.cpp
+++ b/velox/functions/prestosql/aggregates/VarianceAggregates.cpp
@@ -459,7 +459,10 @@ void checkSumCountRowType(
 }
 
 template <template <typename TInput> class TClass>
-exec::AggregateRegistrationResult registerVariance(const std::string& name) {
+exec::AggregateRegistrationResult registerVariance(
+    const std::string& name,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures;
   std::vector<std::string> inputTypes = {
       "smallint", "integer", "bigint", "real", "double"};
@@ -508,18 +511,29 @@ exec::AggregateRegistrationResult registerVariance(const std::string& name) {
               "(count:bigint, mean:double, m2:double) struct");
           return std::make_unique<TClass<int64_t>>(resultType);
         }
-      });
+      },
+      registerCompanionFunctions,
+      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 registerCompanionFunctions,
+    bool overwrite) {
+  registerVariance<StdDevSampAggregate>(
+      prefix + kStdDev, registerCompanionFunctions, overwrite);
+  registerVariance<StdDevPopAggregate>(
+      prefix + kStdDevPop, registerCompanionFunctions, overwrite);
+  registerVariance<StdDevSampAggregate>(
+      prefix + kStdDevSamp, registerCompanionFunctions, overwrite);
+  registerVariance<VarSampAggregate>(
+      prefix + kVariance, registerCompanionFunctions, overwrite);
+  registerVariance<VarPopAggregate>(
+      prefix + kVarPop, registerCompanionFunctions, overwrite);
+  registerVariance<VarSampAggregate>(
+      prefix + kVarSamp, registerCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::aggregate::prestosql
diff --git a/velox/functions/sparksql/aggregates/AverageAggregate.cpp b/velox/functions/sparksql/aggregates/AverageAggregate.cpp
index 58952a4dd42c..045f88e1c90d 100644
--- a/velox/functions/sparksql/aggregates/AverageAggregate.cpp
+++ b/velox/functions/sparksql/aggregates/AverageAggregate.cpp
@@ -362,7 +362,10 @@ class DecimalAverageAggregate : public DecimalAggregate<TInputType> {
 ///     REAL            |     DOUBLE          |    DOUBLE
 ///     ALL INTs        |     DOUBLE          |    DOUBLE
 ///     DECIMAL         |     DECIMAL         |    DECIMAL
-exec::AggregateRegistrationResult registerAverage(const std::string& name) {
+exec::AggregateRegistrationResult registerAverage(
+    const std::string& name,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures;
 
   for (const auto& inputType :
@@ -495,7 +498,8 @@ exec::AggregateRegistrationResult registerAverage(const std::string& name) {
           }
         }
       },
-      /*registerCompanionFunctions*/ true);
+      registerCompanionFunctions,
+      overwrite);
 }
 
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/AverageAggregate.h b/velox/functions/sparksql/aggregates/AverageAggregate.h
index 33bda31409be..e06e5604d84e 100644
--- a/velox/functions/sparksql/aggregates/AverageAggregate.h
+++ b/velox/functions/sparksql/aggregates/AverageAggregate.h
@@ -22,6 +22,9 @@
 
 namespace facebook::velox::functions::aggregate::sparksql {
 
-exec::AggregateRegistrationResult registerAverage(const std::string& name);
+exec::AggregateRegistrationResult registerAverage(
+    const std::string& name,
+    bool registerCompanionFunctions,
+    bool overwrite);
 
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/BitwiseXorAggregate.cpp b/velox/functions/sparksql/aggregates/BitwiseXorAggregate.cpp
index 2029e364f33e..6ad5c6b68f8e 100644
--- 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 registerCompanionFunctions,
+    bool overwrite) {
   return functions::aggregate::registerBitwise<BitwiseXorAggregate>(
-      prefix + "bit_xor");
+      prefix + "bit_xor", registerCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/BitwiseXorAggregate.h b/velox/functions/sparksql/aggregates/BitwiseXorAggregate.h
index f141c5e0b806..d1275d25049f 100644
--- 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
diff --git a/velox/functions/sparksql/aggregates/FirstLastAggregate.cpp b/velox/functions/sparksql/aggregates/FirstLastAggregate.cpp
index 1d5a579aed85..0273c569d8f5 100644
--- a/velox/functions/sparksql/aggregates/FirstLastAggregate.cpp
+++ b/velox/functions/sparksql/aggregates/FirstLastAggregate.cpp
@@ -429,7 +429,10 @@ class LastAggregate : public FirstLastAggregateBase<numeric, TData> {
 } // namespace
 
 template <template <bool B1, typename T, bool B2> class TClass, bool ignoreNull>
-AggregateRegistrationResult registerFirstLast(const std::string& name) {
+AggregateRegistrationResult registerFirstLast(
+    const std::string& name,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<AggregateFunctionSignature>> signatures = {
       AggregateFunctionSignatureBuilder()
           .typeVariable("T")
@@ -499,14 +502,23 @@ AggregateRegistrationResult registerFirstLast(const std::string& name) {
                 name,
                 inputType->toString());
         }
-      });
+      },
+      registerCompanionFunctions,
+      overwrite);
 }
 
-void registerFirstLastAggregates(const std::string& prefix) {
-  registerFirstLast<FirstAggregate, false>(prefix + "first");
-  registerFirstLast<FirstAggregate, true>(prefix + "first_ignore_null");
-  registerFirstLast<LastAggregate, false>(prefix + "last");
-  registerFirstLast<LastAggregate, true>(prefix + "last_ignore_null");
+void registerFirstLastAggregates(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite) {
+  registerFirstLast<FirstAggregate, false>(
+      prefix + "first", registerCompanionFunctions, overwrite);
+  registerFirstLast<FirstAggregate, true>(
+      prefix + "first_ignore_null", registerCompanionFunctions, overwrite);
+  registerFirstLast<LastAggregate, false>(
+      prefix + "last", registerCompanionFunctions, overwrite);
+  registerFirstLast<LastAggregate, true>(
+      prefix + "last_ignore_null", registerCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/MinMaxByAggregate.cpp b/velox/functions/sparksql/aggregates/MinMaxByAggregate.cpp
index 015dd86b2347..57e7d77780c6 100644
--- a/velox/functions/sparksql/aggregates/MinMaxByAggregate.cpp
+++ b/velox/functions/sparksql/aggregates/MinMaxByAggregate.cpp
@@ -86,7 +86,10 @@ template <
         class C>
     class Aggregate,
     bool isMaxFunc>
-exec::AggregateRegistrationResult registerMinMaxBy(const std::string& name) {
+exec::AggregateRegistrationResult registerMinMaxBy(
+    const std::string& name,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures;
   // V, C -> row(V, C) -> V.
   signatures.push_back(exec::AggregateFunctionSignatureBuilder()
@@ -137,14 +140,21 @@ exec::AggregateRegistrationResult registerMinMaxBy(const std::string& name) {
           return create<Aggregate, SparkComparator, isMaxFunc>(
               resultType, valueType, compareType, errorMessage);
         }
-      });
+      },
+      registerCompanionFunctions,
+      overwrite);
 }
 
 } // namespace
 
-void registerMinMaxByAggregates(const std::string& prefix) {
-  registerMinMaxBy<MinMaxByAggregateBase, true>(prefix + "max_by");
-  registerMinMaxBy<MinMaxByAggregateBase, false>(prefix + "min_by");
+void registerMinMaxByAggregates(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite) {
+  registerMinMaxBy<MinMaxByAggregateBase, true>(
+      prefix + "max_by", registerCompanionFunctions, overwrite);
+  registerMinMaxBy<MinMaxByAggregateBase, false>(
+      prefix + "min_by", registerCompanionFunctions, overwrite);
 }
 
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/Register.cpp b/velox/functions/sparksql/aggregates/Register.cpp
index eaa7df9a9390..3a6b2cc0e888 100644
--- a/velox/functions/sparksql/aggregates/Register.cpp
+++ b/velox/functions/sparksql/aggregates/Register.cpp
@@ -24,16 +24,24 @@
 
 namespace facebook::velox::functions::aggregate::sparksql {
 
-extern void registerFirstLastAggregates(const std::string& prefix);
-extern void registerMinMaxByAggregates(const std::string& prefix);
+extern void registerFirstLastAggregates(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite);
+extern void registerMinMaxByAggregates(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite);
 
-void registerAggregateFunctions(const std::string& prefix) {
-  registerFirstLastAggregates(prefix);
-  registerMinMaxByAggregates(prefix);
-  registerBitwiseXorAggregate(prefix);
+void registerAggregateFunctions(
+    const std::string& prefix,
+    bool registerCompanionFunctions,
+    bool overwrite) {
+  registerFirstLastAggregates(prefix, registerCompanionFunctions, overwrite);
+  registerMinMaxByAggregates(prefix, registerCompanionFunctions, overwrite);
+  registerBitwiseXorAggregate(prefix, registerCompanionFunctions, overwrite);
   registerBloomFilterAggAggregate(prefix + "bloom_filter_agg");
-  registerAverage(prefix + "avg");
-  registerSum(prefix + "sum");
-  registerDecimalSumAggregate(prefix + "sum");
+  registerAverage(prefix + "avg", registerCompanionFunctions, overwrite);
+  registerSum(prefix + "sum", registerCompanionFunctions, overwrite);
 }
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/Register.h b/velox/functions/sparksql/aggregates/Register.h
index 5bfea69916b3..4fd96a839e7f 100644
--- a/velox/functions/sparksql/aggregates/Register.h
+++ b/velox/functions/sparksql/aggregates/Register.h
@@ -18,5 +18,8 @@
 #include <string>
 
 namespace facebook::velox::functions::aggregate::sparksql {
-void registerAggregateFunctions(const std::string& prefix);
+void registerAggregateFunctions(
+    const std::string& prefix,
+    bool registerCompanionFunctions = false,
+    bool overwrite = false);
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/SumAggregate.cpp b/velox/functions/sparksql/aggregates/SumAggregate.cpp
index bd9a2c5b7c87..a9ae735f73f9 100644
--- a/velox/functions/sparksql/aggregates/SumAggregate.cpp
+++ b/velox/functions/sparksql/aggregates/SumAggregate.cpp
@@ -27,7 +27,10 @@ template <typename TInput, typename TAccumulator, typename ResultType>
 using SumAggregate = SumAggregateBase<TInput, TAccumulator, ResultType, true>;
 }
 
-exec::AggregateRegistrationResult registerSum(const std::string& name) {
+exec::AggregateRegistrationResult registerSum(
+    const std::string& name,
+    bool registerCompanionFunctions,
+    bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures{
       exec::AggregateFunctionSignatureBuilder()
           .returnType("real")
@@ -143,7 +146,9 @@ exec::AggregateRegistrationResult registerSum(const std::string& name) {
                 name,
                 inputType->kindName());
         }
-      });
+      },
+      registerCompanionFunctions,
+      overwrite);
 }
 
 } // namespace facebook::velox::functions::aggregate::sparksql
diff --git a/velox/functions/sparksql/aggregates/SumAggregate.h b/velox/functions/sparksql/aggregates/SumAggregate.h
index fbeb4198a9a3..f73d0f4fddb9 100644
--- a/velox/functions/sparksql/aggregates/SumAggregate.h
+++ b/velox/functions/sparksql/aggregates/SumAggregate.h
@@ -22,6 +22,9 @@
 
 namespace facebook::velox::functions::aggregate::sparksql {
 
-exec::AggregateRegistrationResult registerSum(const std::string& name);
+exec::AggregateRegistrationResult registerSum(
+    const std::string& name,
+    bool registerCompanionFunctions,
+    bool overwrite);
 
 } // namespace facebook::velox::functions::aggregate::sparksql