[fix](decimal256) support decimal256 for many functions (#42136)

## Proposed changes Issue Number: close #xxx Support decimal256 for the following functions: ``` multi_distinct_sum multi_distinct_count array_sum array_avg array_product array_cum_sum ```
apache · Oct 23, 2024 · bf1b47d · bf1b47d
1 parent efc0621
commit bf1b47d
Show file tree

Hide file tree

Showing 67 changed files with 956 additions and 217 deletions.
diff --git a/be/src/vec/aggregate_functions/aggregate_function.h b/be/src/vec/aggregate_functions/aggregate_function.h
@@ -38,6 +38,10 @@ class Arena;
 class IColumn;
 class IDataType;
 
+struct AggregateFunctionAttr {
+    bool enable_decimal256 {false};
+};
+
 template <bool nullable, typename ColVecType>
 class AggregateFunctionBitmapCount;
 template <typename Op>

diff --git a/be/src/vec/aggregate_functions/aggregate_function_approx_count_distinct.cpp b/be/src/vec/aggregate_functions/aggregate_function_approx_count_distinct.cpp
@@ -31,7 +31,8 @@
 namespace doris::vectorized {
 
 AggregateFunctionPtr create_aggregate_function_approx_count_distinct(
-        const std::string& name, const DataTypes& argument_types, const bool result_is_nullable) {
+        const std::string& name, const DataTypes& argument_types, const bool result_is_nullable,
+        const AggregateFunctionAttr& attr) {
     WhichDataType which(remove_nullable(argument_types[0]));
 
 #define DISPATCH(TYPE, COLUMN_TYPE)                                                             \

diff --git a/be/src/vec/aggregate_functions/aggregate_function_avg.cpp b/be/src/vec/aggregate_functions/aggregate_function_avg.cpp
@@ -45,8 +45,17 @@ template <typename T>
 using AggregateFuncAvgDecimal256 = typename AvgDecimal256<T>::Function;
 
 void register_aggregate_function_avg(AggregateFunctionSimpleFactory& factory) {
-    factory.register_function_both("avg", creator_with_type::creator<AggregateFuncAvg>);
-    factory.register_function_both("avg_decimal256",
-                                   creator_with_type::creator<AggregateFuncAvgDecimal256>);
+    AggregateFunctionCreator creator = [&](const std::string& name, const DataTypes& types,
+                                           const bool result_is_nullable,
+                                           const AggregateFunctionAttr& attr) {
+        if (attr.enable_decimal256) {
+            return creator_with_type::creator<AggregateFuncAvgDecimal256>(name, types,
+                                                                          result_is_nullable, attr);
+        } else {
+            return creator_with_type::creator<AggregateFuncAvg>(name, types, result_is_nullable,
+                                                                attr);
+        }
+    };
+    factory.register_function_both("avg", creator);
 }
 } // namespace doris::vectorized
diff --git a/be/src/vec/aggregate_functions/aggregate_function_bitmap.cpp b/be/src/vec/aggregate_functions/aggregate_function_bitmap.cpp
@@ -40,9 +40,9 @@ AggregateFunctionPtr create_with_int_data_type(const DataTypes& argument_type) {
     return nullptr;
 }
 
-AggregateFunctionPtr create_aggregate_function_bitmap_union_count(const std::string& name,
-                                                                  const DataTypes& argument_types,
-                                                                  const bool result_is_nullable) {
+AggregateFunctionPtr create_aggregate_function_bitmap_union_count(
+        const std::string& name, const DataTypes& argument_types, const bool result_is_nullable,
+        const AggregateFunctionAttr& attr) {
     const bool arg_is_nullable = argument_types[0]->is_nullable();
     if (arg_is_nullable) {
         return std::make_shared<AggregateFunctionBitmapCount<true, ColumnBitmap>>(argument_types);
@@ -53,7 +53,8 @@ AggregateFunctionPtr create_aggregate_function_bitmap_union_count(const std::str
 
 AggregateFunctionPtr create_aggregate_function_bitmap_union_int(const std::string& name,
                                                                 const DataTypes& argument_types,
-                                                                const bool result_is_nullable) {
+                                                                const bool result_is_nullable,
+                                                                const AggregateFunctionAttr& attr) {
     const bool arg_is_nullable = argument_types[0]->is_nullable();
     if (arg_is_nullable) {
         return AggregateFunctionPtr(

diff --git a/be/src/vec/aggregate_functions/aggregate_function_bitmap_agg.cpp b/be/src/vec/aggregate_functions/aggregate_function_bitmap_agg.cpp
@@ -41,7 +41,8 @@ AggregateFunctionPtr create_with_int_data_type(const DataTypes& argument_types)
 
 AggregateFunctionPtr create_aggregate_function_bitmap_agg(const std::string& name,
                                                           const DataTypes& argument_types,
-                                                          const bool result_is_nullable) {
+                                                          const bool result_is_nullable,
+                                                          const AggregateFunctionAttr& attr) {
     const bool arg_is_nullable = argument_types[0]->is_nullable();
     if (arg_is_nullable) {
         return AggregateFunctionPtr(create_with_int_data_type<true>(argument_types));

diff --git a/be/src/vec/aggregate_functions/aggregate_function_collect.cpp b/be/src/vec/aggregate_functions/aggregate_function_collect.cpp
@@ -96,7 +96,8 @@ AggregateFunctionPtr create_aggregate_function_collect_impl(const std::string& n
 
 AggregateFunctionPtr create_aggregate_function_collect(const std::string& name,
                                                        const DataTypes& argument_types,
-                                                       const bool result_is_nullable) {
+                                                       const bool result_is_nullable,
+                                                       const AggregateFunctionAttr& attr) {
     if (argument_types.size() == 1) {
         if (name == "array_agg") {
             return create_aggregate_function_collect_impl<std::false_type, std::true_type>(

diff --git a/be/src/vec/aggregate_functions/aggregate_function_corr.cpp b/be/src/vec/aggregate_functions/aggregate_function_corr.cpp
@@ -89,7 +89,8 @@ struct CorrMoment {
 
 AggregateFunctionPtr create_aggregate_corr_function(const std::string& name,
                                                     const DataTypes& argument_types,
-                                                    const bool result_is_nullable) {
+                                                    const bool result_is_nullable,
+                                                    const AggregateFunctionAttr& attr) {
     assert_binary(name, argument_types);
     return create_with_two_basic_numeric_types<CorrMoment>(argument_types[0], argument_types[1],
                                                            argument_types, result_is_nullable);

diff --git a/be/src/vec/aggregate_functions/aggregate_function_count.cpp b/be/src/vec/aggregate_functions/aggregate_function_count.cpp
@@ -29,15 +29,16 @@ namespace doris::vectorized {
 
 AggregateFunctionPtr create_aggregate_function_count(const std::string& name,
                                                      const DataTypes& argument_types,
-                                                     const bool result_is_nullable) {
+                                                     const bool result_is_nullable,
+                                                     const AggregateFunctionAttr& attr) {
     assert_arity_at_most<1>(name, argument_types);
 
     return std::make_shared<AggregateFunctionCount>(argument_types);
 }
 
-AggregateFunctionPtr create_aggregate_function_count_not_null_unary(const std::string& name,
-                                                                    const DataTypes& argument_types,
-                                                                    const bool result_is_nullable) {
+AggregateFunctionPtr create_aggregate_function_count_not_null_unary(
+        const std::string& name, const DataTypes& argument_types, const bool result_is_nullable,
+        const AggregateFunctionAttr& attr) {
     assert_arity_at_most<1>(name, argument_types);
 
     return std::make_shared<AggregateFunctionCountNotNullUnary>(argument_types);

diff --git a/be/src/vec/aggregate_functions/aggregate_function_count_by_enum.cpp b/be/src/vec/aggregate_functions/aggregate_function_count_by_enum.cpp
@@ -29,7 +29,8 @@ namespace doris::vectorized {
 
 AggregateFunctionPtr create_aggregate_function_count_by_enum(const std::string& name,
                                                              const DataTypes& argument_types,
-                                                             const bool result_is_nullable) {
+                                                             const bool result_is_nullable,
+                                                             const AggregateFunctionAttr& attr) {
     if (argument_types.size() < 1) {
         LOG(WARNING) << fmt::format("Illegal number {} of argument for aggregate function {}",
                                     argument_types.size(), name);

diff --git a/be/src/vec/aggregate_functions/aggregate_function_covar.cpp b/be/src/vec/aggregate_functions/aggregate_function_covar.cpp
@@ -53,14 +53,16 @@ AggregateFunctionPtr create_function_single_value(const String& name,
 
 AggregateFunctionPtr create_aggregate_function_covariance_samp(const std::string& name,
                                                                const DataTypes& argument_types,
-                                                               const bool result_is_nullable) {
+                                                               const bool result_is_nullable,
+                                                               const AggregateFunctionAttr& attr) {
     return create_function_single_value<AggregateFunctionSamp, CovarSampName, SampData>(
             name, argument_types, result_is_nullable, NOTNULLABLE);
 }
 
 AggregateFunctionPtr create_aggregate_function_covariance_pop(const std::string& name,
                                                               const DataTypes& argument_types,
-                                                              const bool result_is_nullable) {
+                                                              const bool result_is_nullable,
+                                                              const AggregateFunctionAttr& attr) {
     return create_function_single_value<AggregateFunctionPop, CovarName, PopData>(
             name, argument_types, result_is_nullable, NOTNULLABLE);
 }

diff --git a/be/src/vec/aggregate_functions/aggregate_function_distinct.cpp b/be/src/vec/aggregate_functions/aggregate_function_distinct.cpp
@@ -83,7 +83,8 @@ const std::string DISTINCT_FUNCTION_PREFIX = "multi_distinct_";
 
 void register_aggregate_function_combinator_distinct(AggregateFunctionSimpleFactory& factory) {
     AggregateFunctionCreator creator = [&](const std::string& name, const DataTypes& types,
-                                           const bool result_is_nullable) {
+                                           const bool result_is_nullable,
+                                           const AggregateFunctionAttr& attr) {
         // 1. we should get not nullable types;
         DataTypes nested_types(types.size());
         std::transform(types.begin(), types.end(), nested_types.begin(),
@@ -92,7 +93,7 @@ void register_aggregate_function_combinator_distinct(AggregateFunctionSimpleFact
         auto transform_arguments = function_combinator->transform_arguments(nested_types);
         auto nested_function_name = name.substr(DISTINCT_FUNCTION_PREFIX.size());
         auto nested_function = factory.get(nested_function_name, transform_arguments, false,
-                                           BeExecVersionManager::get_newest_version());
+                                           BeExecVersionManager::get_newest_version(), attr);
         return function_combinator->transform_aggregate_function(nested_function, types,
                                                                  result_is_nullable);
     };

diff --git a/be/src/vec/aggregate_functions/aggregate_function_foreach.cpp b/be/src/vec/aggregate_functions/aggregate_function_foreach.cpp
@@ -34,8 +34,9 @@
 namespace doris::vectorized {
 
 void register_aggregate_function_combinator_foreach(AggregateFunctionSimpleFactory& factory) {
-    AggregateFunctionCreator creator = [&](const std::string& name, const DataTypes& types,
-                                           const bool result_is_nullable) -> AggregateFunctionPtr {
+    AggregateFunctionCreator creator =
+            [&](const std::string& name, const DataTypes& types, const bool result_is_nullable,
+                const AggregateFunctionAttr& attr) -> AggregateFunctionPtr {
         const std::string& suffix = AggregateFunctionForEach::AGG_FOREACH_SUFFIX;
         DataTypes transform_arguments;
         for (const auto& t : types) {
@@ -46,7 +47,7 @@ void register_aggregate_function_combinator_foreach(AggregateFunctionSimpleFacto
         auto nested_function_name = name.substr(0, name.size() - suffix.size());
         auto nested_function =
                 factory.get(nested_function_name, transform_arguments, result_is_nullable,
-                            BeExecVersionManager::get_newest_version(), false);
+                            BeExecVersionManager::get_newest_version(), attr);
         if (!nested_function) {
             throw Exception(
                     ErrorCode::INTERNAL_ERROR,

diff --git a/be/src/vec/aggregate_functions/aggregate_function_group_array_intersect.cpp b/be/src/vec/aggregate_functions/aggregate_function_group_array_intersect.cpp
@@ -70,7 +70,8 @@ inline AggregateFunctionPtr create_aggregate_function_group_array_intersect_impl
 }
 
 AggregateFunctionPtr create_aggregate_function_group_array_intersect(
-        const std::string& name, const DataTypes& argument_types, const bool result_is_nullable) {
+        const std::string& name, const DataTypes& argument_types, const bool result_is_nullable,
+        const AggregateFunctionAttr& attr) {
     assert_unary(name, argument_types);
     const DataTypePtr& argument_type = remove_nullable(argument_types[0]);
 

diff --git a/be/src/vec/aggregate_functions/aggregate_function_group_concat.cpp b/be/src/vec/aggregate_functions/aggregate_function_group_concat.cpp
@@ -28,7 +28,8 @@ const std::string AggregateFunctionGroupConcatImplStr::separator = ",";
 
 AggregateFunctionPtr create_aggregate_function_group_concat(const std::string& name,
                                                             const DataTypes& argument_types,
-                                                            const bool result_is_nullable) {
+                                                            const bool result_is_nullable,
+                                                            const AggregateFunctionAttr& attr) {
     if (argument_types.size() == 1) {
         return creator_without_type::create<
                 AggregateFunctionGroupConcat<AggregateFunctionGroupConcatImplStr>>(

diff --git a/be/src/vec/aggregate_functions/aggregate_function_histogram.cpp b/be/src/vec/aggregate_functions/aggregate_function_histogram.cpp
@@ -47,7 +47,8 @@ AggregateFunctionPtr create_agg_function_histogram(const DataTypes& argument_typ
 
 AggregateFunctionPtr create_aggregate_function_histogram(const std::string& name,
                                                          const DataTypes& argument_types,
-                                                         const bool result_is_nullable) {
+                                                         const bool result_is_nullable,
+                                                         const AggregateFunctionAttr& attr) {
     WhichDataType type(remove_nullable(argument_types[0]));
 
 #define DISPATCH(TYPE)               \

diff --git a/be/src/vec/aggregate_functions/aggregate_function_kurtosis.cpp b/be/src/vec/aggregate_functions/aggregate_function_kurtosis.cpp
@@ -45,7 +45,8 @@ AggregateFunctionPtr type_dispatch_for_aggregate_function_kurt(const DataTypes&
 
 AggregateFunctionPtr create_aggregate_function_kurt(const std::string& name,
                                                     const DataTypes& argument_types,
-                                                    const bool result_is_nullable) {
+                                                    const bool result_is_nullable,
+                                                    const AggregateFunctionAttr& attr) {
     if (argument_types.size() != 1) {
         LOG(WARNING) << "aggregate function " << name << " requires exactly 1 argument";
         return nullptr;

diff --git a/be/src/vec/aggregate_functions/aggregate_function_linear_histogram.cpp b/be/src/vec/aggregate_functions/aggregate_function_linear_histogram.cpp
@@ -41,7 +41,8 @@ AggregateFunctionPtr create_agg_function_linear_histogram(const DataTypes& argum
 
 AggregateFunctionPtr create_aggregate_function_linear_histogram(const std::string& name,
                                                                 const DataTypes& argument_types,
-                                                                const bool result_is_nullable) {
+                                                                const bool result_is_nullable,
+                                                                const AggregateFunctionAttr& attr) {
     WhichDataType type(remove_nullable(argument_types[0]));
 
 #define DISPATCH(TYPE)               \

diff --git a/be/src/vec/aggregate_functions/aggregate_function_map.cpp b/be/src/vec/aggregate_functions/aggregate_function_map.cpp
@@ -32,7 +32,8 @@ AggregateFunctionPtr create_agg_function_map_agg(const DataTypes& argument_types
 
 AggregateFunctionPtr create_aggregate_function_map_agg(const std::string& name,
                                                        const DataTypes& argument_types,
-                                                       const bool result_is_nullable) {
+                                                       const bool result_is_nullable,
+                                                       const AggregateFunctionAttr& attr) {
     WhichDataType type(remove_nullable(argument_types[0]));
 
 #define DISPATCH(TYPE)               \

diff --git a/be/src/vec/aggregate_functions/aggregate_function_min_max.cpp b/be/src/vec/aggregate_functions/aggregate_function_min_max.cpp
@@ -30,7 +30,8 @@ namespace doris::vectorized {
 template <template <typename> class Data>
 AggregateFunctionPtr create_aggregate_function_single_value(const String& name,
                                                             const DataTypes& argument_types,
-                                                            const bool result_is_nullable) {
+                                                            const bool result_is_nullable,
+                                                            const AggregateFunctionAttr& attr) {
     assert_unary(name, argument_types);
 
     AggregateFunctionPtr res(creator_with_numeric_type::create<AggregateFunctionsSingleValue, Data,

diff --git a/be/src/vec/aggregate_functions/aggregate_function_min_max.h b/be/src/vec/aggregate_functions/aggregate_function_min_max.h
@@ -714,5 +714,6 @@ class AggregateFunctionsSingleValue final
 template <template <typename> class Data>
 AggregateFunctionPtr create_aggregate_function_single_value(const String& name,
                                                             const DataTypes& argument_types,
-                                                            const bool result_is_nullable);
+                                                            const bool result_is_nullable,
+                                                            const AggregateFunctionAttr& attr = {});
 } // namespace doris::vectorized
diff --git a/be/src/vec/aggregate_functions/aggregate_function_min_max_by.h b/be/src/vec/aggregate_functions/aggregate_function_min_max_by.h
@@ -243,7 +243,8 @@ template <template <typename> class AggregateFunctionTemplate,
           template <typename, typename> class Data>
 AggregateFunctionPtr create_aggregate_function_min_max_by(const String& name,
                                                           const DataTypes& argument_types,
-                                                          const bool result_is_nullable) {
+                                                          const bool result_is_nullable,
+                                                          const AggregateFunctionAttr& attr) {
     if (argument_types.size() != 2) {
         return nullptr;
     }

diff --git a/be/src/vec/aggregate_functions/aggregate_function_orthogonal_bitmap.cpp b/be/src/vec/aggregate_functions/aggregate_function_orthogonal_bitmap.cpp
@@ -35,7 +35,8 @@ template <template <typename> class Impl>
 AggregateFunctionPtr create_aggregate_function_orthogonal(const std::string& name,
                                                           const DataTypes& argument_types,
 
-                                                          const bool result_is_nullable) {
+                                                          const bool result_is_nullable,
+                                                          const AggregateFunctionAttr& attr) {
     if (argument_types.empty()) {
         LOG(WARNING) << "Incorrect number of arguments for aggregate function " << name;
         return nullptr;

diff --git a/be/src/vec/aggregate_functions/aggregate_function_percentile.cpp b/be/src/vec/aggregate_functions/aggregate_function_percentile.cpp
@@ -23,9 +23,9 @@
 
 namespace doris::vectorized {
 
-AggregateFunctionPtr create_aggregate_function_percentile_approx(const std::string& name,
-                                                                 const DataTypes& argument_types,
-                                                                 const bool result_is_nullable) {
+AggregateFunctionPtr create_aggregate_function_percentile_approx(
+        const std::string& name, const DataTypes& argument_types, const bool result_is_nullable,
+        const AggregateFunctionAttr& attr) {
     const DataTypePtr& argument_type = remove_nullable(argument_types[0]);
     WhichDataType which(argument_type);
     if (which.idx != TypeIndex::Float64) {
@@ -43,7 +43,8 @@ AggregateFunctionPtr create_aggregate_function_percentile_approx(const std::stri
 }
 
 AggregateFunctionPtr create_aggregate_function_percentile_approx_weighted(
-        const std::string& name, const DataTypes& argument_types, const bool result_is_nullable) {
+        const std::string& name, const DataTypes& argument_types, const bool result_is_nullable,
+        const AggregateFunctionAttr& attr) {
     const DataTypePtr& argument_type = remove_nullable(argument_types[0]);
     WhichDataType which(argument_type);
     if (which.idx != TypeIndex::Float64) {