Combinator -MergeState now returns AggregateFunction(nested_func). [#CLICKHOUSE-2891]

dbms/src/AggregateFunctions/AggregateFunctionMerge.h

@@ -34,6 +34,11 @@ public:
         return nested_func->getReturnType();
     }
 
+    AggregateFunctionPtr getNestedFunction() const
+    {
+        return nested_func_owner;
+    }
+
     void setArguments(const DataTypes & arguments) override
     {
         if (arguments.size() != 1)

dbms/src/AggregateFunctions/AggregateFunctionState.cpp

 #include <AggregateFunctions/AggregateFunctionState.h>
+#include <AggregateFunctions/AggregateFunctionMerge.h>
 
 namespace DB
 {
 
+namespace ErrorCodes
+{
+    extern const int BAD_ARGUMENTS;
+}
+
+DataTypePtr AggregateFunctionState::getReturnType() const
+{
+    auto ptr = std::make_shared<DataTypeAggregateFunction>(nested_func_owner, arguments, params);
+
+    /// Special case: it is -MergeState combinator
+    if (typeid_cast<const AggregateFunctionMerge *>(ptr->getFunction().get()))
+    {
+        if (arguments.size() != 1)
+            throw Exception("Combinator -MergeState expects only one argument", ErrorCodes::BAD_ARGUMENTS);
+
+        if (!typeid_cast<const DataTypeAggregateFunction *>(arguments[0].get()))
+            throw Exception("Combinator -MergeState expects argument with AggregateFunction type", ErrorCodes::BAD_ARGUMENTS);
+
+        return arguments[0];
+    }
+
+    return ptr;
+}
+
+
 AggregateFunctionPtr createAggregateFunctionState(AggregateFunctionPtr & nested)
 {
     return std::make_shared<AggregateFunctionState>(nested);

dbms/src/AggregateFunctions/AggregateFunctionState.h

+
 #pragma once
 
 #include <DataTypes/DataTypeAggregateFunction.h>
@@ -30,10 +31,7 @@ public:
         return nested_func->getName() + "State";
     }
 
-    DataTypePtr getReturnType() const override
-    {
-        return std::make_shared<DataTypeAggregateFunction>(nested_func_owner, arguments, params);
-    }
+    DataTypePtr getReturnType() const override;
 
     void setArguments(const DataTypes & arguments_) override
     {

dbms/tests/queries/0_stateless/00432_aggregate_function_scalars_and_constants.reference

@@ -27,3 +27,6 @@
 
 1
 0
+
+1
+2

dbms/tests/queries/0_stateless/00432_aggregate_function_scalars_and_constants.sql

@@ -38,3 +38,20 @@ SELECT arrayReduce('groupUniqArrayMergeIf',
 SELECT '';
 SELECT arrayReduce('avgState', [0]) IN (arrayReduce('avgState', [0, 1]), arrayReduce('avgState', [0]));
 SELECT arrayReduce('avgState', [0]) IN (arrayReduce('avgState', [0, 1]), arrayReduce('avgState', [1]));
+
+SELECT '';
+SELECT arrayReduce('uniqExactMerge',
+    [arrayReduce('uniqExactMergeState',
+        [
+            arrayReduce('uniqExactState', [12345678901]),
+            arrayReduce('uniqExactState', [12345678901])
+        ])
+    ]);
+
+SELECT arrayReduce('uniqExactMerge',
+    [arrayReduce('uniqExactMergeState',
+        [
+            arrayReduce('uniqExactState', [12345678901]),
+            arrayReduce('uniqExactState', [12345678902])
+        ])
+    ]);