Merge pull request #16935 from ucasFL/agg-func-setting-null-for-empty

`aggregate_functions_null_for_empty` move OrNull ahead

src/Interpreters/TreeRewriter.cpp

@@ -125,12 +125,60 @@ struct CustomizeAggregateFunctionsSuffixData
         {
             auto properties = instance.tryGetProperties(func.name);
             if (properties && !properties->returns_default_when_only_null)
-                func.name = func.name + customized_func_suffix;
+            {
+                func.name += customized_func_suffix;
+            }
+        }
+    }
+};
+
+// Used to rewrite aggregate functions with -OrNull suffix in some cases, such as sumIfOrNull, we shoule rewrite to sumOrNullIf
+struct CustomizeAggregateFunctionsMoveSuffixData
+{
+    using TypeToVisit = ASTFunction;
+
+    const String & customized_func_suffix;
+
+    String moveSuffixAhead(const String & name) const
+    {
+        auto prefix = name.substr(0, name.size() - customized_func_suffix.size());
+
+        auto prefix_size = prefix.size();
+
+        if (endsWith(prefix, "MergeState"))
+            return prefix.substr(0, prefix_size - 10) + customized_func_suffix + "MergeState";
+
+        if (endsWith(prefix, "Merge"))
+            return prefix.substr(0, prefix_size - 5) + customized_func_suffix + "Merge";
+
+        if (endsWith(prefix, "State"))
+            return prefix.substr(0, prefix_size - 5) + customized_func_suffix + "State";
+
+        if (endsWith(prefix, "If"))
+            return prefix.substr(0, prefix_size - 2) + customized_func_suffix + "If";
+
+        return name;
+    }
+
+    void visit(ASTFunction & func, ASTPtr &) const
+    {
+        const auto & instance = AggregateFunctionFactory::instance();
+        if (instance.isAggregateFunctionName(func.name))
+        {
+            if (endsWith(func.name, customized_func_suffix))
+            {
+                auto properties = instance.tryGetProperties(func.name);
+                if (properties && !properties->returns_default_when_only_null)
+                {
+                    func.name = moveSuffixAhead(func.name);
+                }
+            }
         }
     }
 };
 
 using CustomizeAggregateFunctionsOrNullVisitor = InDepthNodeVisitor<OneTypeMatcher<CustomizeAggregateFunctionsSuffixData>, true>;
+using CustomizeAggregateFunctionsMoveOrNullVisitor = InDepthNodeVisitor<OneTypeMatcher<CustomizeAggregateFunctionsMoveSuffixData>, true>;
 
 /// Translate qualified names such as db.table.column, table.column, table_alias.column to names' normal form.
 /// Expand asterisks and qualified asterisks with column names.
@@ -753,6 +801,10 @@ void TreeRewriter::normalize(ASTPtr & query, Aliases & aliases, const Settings &
         CustomizeAggregateFunctionsOrNullVisitor(data_or_null).visit(query);
     }
 
+    /// Move -OrNull suffix ahead, this should execute after add -OrNull suffix
+    CustomizeAggregateFunctionsMoveOrNullVisitor::Data data_or_null{"OrNull"};
+    CustomizeAggregateFunctionsMoveOrNullVisitor(data_or_null).visit(query);
+
     /// Creates a dictionary `aliases`: alias -> ASTPtr
     QueryAliasesVisitor(aliases).visit(query);
 

tests/queries/0_stateless/01562_agg_null_for_empty_ahead.reference

@@ -0,0 +1,20 @@
+0
+0
+0
+0
+0
+1
+\N
+\N
+1
+\N
+\N
+0
+\N
+0
+\N
+1
+\N
+\N
+1
+\N

tests/queries/0_stateless/01562_agg_null_for_empty_ahead.sql

@@ -0,0 +1,36 @@
+SELECT sumMerge(s) FROM (SELECT sumState(number) s FROM numbers(0));
+SELECT sumMerge(s) FROM (SELECT sumState(number) s FROM numbers(1));
+
+SELECT sumMerge(s) FROM (SELECT sumMergeState(n) s FROM (SELECT sumState(number) n FROM numbers(0)));
+SELECT sumMerge(s) FROM (SELECT sumMergeState(n) s FROM (SELECT sumState(number) n FROM numbers(1)));
+
+SELECT sumIf(1, 0);
+
+SELECT sumIf(1, 1);
+
+-- should return Null even if we donn't set aggregate_functions_null_for_empty
+SELECT sumIfOrNull(1, 0);
+SELECT sumOrNullIf(1, 0);
+
+SELECT nullIf(1, 0);
+
+SELECT nullIf(1, 1);
+
+SET aggregate_functions_null_for_empty=1;
+
+SELECT sumMerge(s) FROM (SELECT sumState(number) s FROM numbers(0));
+SELECT sumMerge(s) FROM (SELECT sumState(number) s FROM numbers(1));
+
+SELECT sumMerge(s) FROM (SELECT sumMergeState(n) s FROM (SELECT sumState(number) n FROM numbers(0)));
+SELECT sumMerge(s) FROM (SELECT sumMergeState(n) s FROM (SELECT sumState(number) n FROM numbers(1)));
+
+SELECT sumIf(1, 0);
+
+SELECT sumIf(1, 1);
+
+SELECT sumIfOrNull(1, 0);
+SELECT sumOrNullIf(1, 0);
+
+SELECT nullIf(1, 0);
+
+SELECT nullIf(1, 1);