Merge pull request #52645 from JackyWoo/optimize_uniq_to_count

Optimize uniq to count

src/Analyzer/Passes/UniqToCountPass.cpp

@@ -0,0 +1,195 @@
+#include "UniqToCountPass.h"
+
+#include <AggregateFunctions/AggregateFunctionFactory.h>
+#include <AggregateFunctions/IAggregateFunction.h>
+
+#include <Analyzer/ColumnNode.h>
+#include <Analyzer/FunctionNode.h>
+#include <Analyzer/InDepthQueryTreeVisitor.h>
+#include <Analyzer/QueryNode.h>
+
+namespace DB
+{
+
+namespace
+{
+
+bool matchFnUniq(String func_name)
+{
+    auto name = Poco::toLower(func_name);
+    return name == "uniq" || name == "uniqHLL12" || name == "uniqExact" || name == "uniqTheta" || name == "uniqCombined"
+        || name == "uniqCombined64";
+}
+
+/// Extract the corresponding projection columns for group by node list.
+/// For example:
+///     SELECT a as aa, any(b) FROM table group by a;  ->  aa(ColumnNode)
+NamesAndTypes extractProjectionColumnsForGroupBy(const QueryNode * query_node)
+{
+    if (!query_node->hasGroupBy())
+        return {};
+
+    NamesAndTypes result;
+    for (const auto & group_by_ele : query_node->getGroupByNode()->getChildren())
+    {
+        const auto & projection_columns = query_node->getProjectionColumns();
+        const auto & projection_nodes = query_node->getProjection().getNodes();
+
+        assert(projection_columns.size() == projection_nodes.size());
+
+        for (size_t i = 0; i < projection_columns.size(); i++)
+        {
+            if (projection_nodes[i]->isEqual(*group_by_ele))
+                result.push_back(projection_columns[i]);
+        }
+    }
+    return result;
+}
+
+/// Whether query_columns equals subquery_columns.
+///     query_columns: query columns from query
+///     subquery_columns: projection columns from subquery
+bool nodeListEquals(const QueryTreeNodes & query_columns, const NamesAndTypes & subquery_columns)
+{
+    if (query_columns.size() != subquery_columns.size())
+        return false;
+
+    for (const auto & query_column : query_columns)
+    {
+        auto find = std::find_if(
+            subquery_columns.begin(),
+            subquery_columns.end(),
+            [&](const auto & subquery_column) -> bool
+            {
+                if (auto * column_node = query_column->as<ColumnNode>())
+                {
+                    return subquery_column == column_node->getColumn();
+                }
+                return false;
+            });
+
+        if (find == subquery_columns.end())
+            return false;
+    }
+    return true;
+}
+
+/// Whether subquery_columns contains all columns in subquery_columns.
+///     query_columns: query columns from query
+///     subquery_columns: projection columns from subquery
+bool nodeListContainsAll(const QueryTreeNodes & query_columns, const NamesAndTypes & subquery_columns)
+{
+    if (query_columns.size() > subquery_columns.size())
+        return false;
+
+    for (const auto & query_column : query_columns)
+    {
+        auto find = std::find_if(
+            subquery_columns.begin(),
+            subquery_columns.end(),
+            [&](const auto & subquery_column) -> bool
+            {
+                if (auto * column_node = query_column->as<ColumnNode>())
+                {
+                    return subquery_column == column_node->getColumn();
+                }
+                return false;
+            });
+
+        if (find == subquery_columns.end())
+            return false;
+    }
+    return true;
+}
+
+}
+
+class UniqToCountVisitor : public InDepthQueryTreeVisitorWithContext<UniqToCountVisitor>
+{
+public:
+    using Base = InDepthQueryTreeVisitorWithContext<UniqToCountVisitor>;
+    using Base::Base;
+
+    void enterImpl(QueryTreeNodePtr & node)
+    {
+        if (!getSettings().optimize_uniq_to_count)
+            return;
+
+        auto * query_node = node->as<QueryNode>();
+        if (!query_node)
+            return;
+
+        /// Check that query has only single table expression which is subquery
+        auto * subquery_node = query_node->getJoinTree()->as<QueryNode>();
+        if (!subquery_node)
+            return;
+
+        /// Check that query has only single node in projection
+        auto & projection_nodes = query_node->getProjection().getNodes();
+        if (projection_nodes.size() != 1)
+            return;
+
+        /// Check that projection_node is a function
+        auto & projection_node = projection_nodes[0];
+        auto * function_node = projection_node->as<FunctionNode>();
+        if (!function_node)
+            return;
+
+        /// Check that query single projection node is `uniq` or its variants
+        if (!matchFnUniq(function_node->getFunctionName()))
+            return;
+
+        auto & uniq_arguments_nodes = function_node->getArguments().getNodes();
+
+        /// Whether query matches 'SELECT uniq(x ...) FROM (SELECT DISTINCT x ...)'
+        auto match_subquery_with_distinct = [&]() -> bool
+        {
+            if (!subquery_node->isDistinct())
+                return false;
+
+            /// uniq expression list == subquery projection columns
+            if (!nodeListEquals(uniq_arguments_nodes, subquery_node->getProjectionColumns()))
+                return false;
+
+            return true;
+        };
+
+        /// Whether query matches 'SELECT uniq(x ...) FROM (SELECT x ... GROUP BY x ...)'
+        auto match_subquery_with_group_by = [&]() -> bool
+        {
+            if (!subquery_node->hasGroupBy())
+                return false;
+
+            /// uniq argument node list == subquery group by node list
+            auto group_by_columns = extractProjectionColumnsForGroupBy(subquery_node);
+
+            if (!nodeListEquals(uniq_arguments_nodes, group_by_columns))
+                return false;
+
+            /// subquery projection columns must contain all columns in uniq argument node list
+            if (!nodeListContainsAll(uniq_arguments_nodes, subquery_node->getProjectionColumns()))
+                return false;
+
+            return true;
+        };
+
+        /// Replace uniq of initial query to count
+        if (match_subquery_with_distinct() || match_subquery_with_group_by())
+        {
+            AggregateFunctionProperties properties;
+            auto aggregate_function = AggregateFunctionFactory::instance().get("count", {}, {}, properties);
+
+            function_node->getArguments().getNodes().clear();
+            function_node->resolveAsAggregateFunction(std::move(aggregate_function));
+        }
+    }
+};
+
+
+void UniqToCountPass::run(QueryTreeNodePtr query_tree_node, ContextPtr context)
+{
+    UniqToCountVisitor visitor(context);
+    visitor.visit(query_tree_node);
+}
+
+}

src/Analyzer/Passes/UniqToCountPass.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <Analyzer/IQueryTreePass.h>
+
+namespace DB
+{
+
+/** Optimize `uniq` and its variants(except uniqUpTo) into `count` over subquery.
+ *     Example: 'SELECT uniq(x ...) FROM (SELECT DISTINCT x ...)' to
+ *     Result: 'SELECT count() FROM (SELECT DISTINCT x ...)'
+ *
+ *     Example: 'SELECT uniq(x ...) FROM (SELECT x ... GROUP BY x ...)' to
+ *     Result: 'SELECT count() FROM (SELECT x ... GROUP BY x ...)'
+ *
+ *     Note that we can rewrite all uniq variants except uniqUpTo.
+ */
+class UniqToCountPass final : public IQueryTreePass
+{
+public:
+    String getName() override { return "UniqToCount"; }
+
+    String getDescription() override
+    {
+        return "Rewrite uniq and its variants(except uniqUpTo) to count if subquery has distinct or group by clause.";
+    }
+
+    void run(QueryTreeNodePtr query_tree_node, ContextPtr context) override;
+};
+
+}

src/Analyzer/QueryTreePassManager.cpp

@@ -18,6 +18,7 @@
 #include <Analyzer/Utils.h>
 #include <Analyzer/Passes/QueryAnalysisPass.h>
 #include <Analyzer/Passes/CountDistinctPass.h>
+#include <Analyzer/Passes/UniqToCountPass.h>
 #include <Analyzer/Passes/FunctionToSubcolumnsPass.h>
 #include <Analyzer/Passes/RewriteAggregateFunctionWithIfPass.h>
 #include <Analyzer/Passes/SumIfToCountIfPass.h>
@@ -247,6 +248,7 @@ void addQueryTreePasses(QueryTreePassManager & manager)
     manager.addPass(std::make_unique<ConvertLogicalExpressionToCNFPass>());
 
     manager.addPass(std::make_unique<CountDistinctPass>());
+    manager.addPass(std::make_unique<UniqToCountPass>());
     manager.addPass(std::make_unique<RewriteAggregateFunctionWithIfPass>());
     manager.addPass(std::make_unique<SumIfToCountIfPass>());
     manager.addPass(std::make_unique<RewriteArrayExistsToHasPass>());

src/Core/Settings.h

@@ -775,6 +775,7 @@ class IColumn;
     M(Bool, function_json_value_return_type_allow_nullable, false, "Allow function JSON_VALUE to return nullable type.", 0) \
     M(Bool, function_json_value_return_type_allow_complex, false, "Allow function JSON_VALUE to return complex type, such as: struct, array, map.", 0) \
     M(Bool, use_with_fill_by_sorting_prefix, true, "Columns preceding WITH FILL columns in ORDER BY clause form sorting prefix. Rows with different values in sorting prefix are filled independently", 0) \
+    M(Bool, optimize_uniq_to_count, true, "Rewrite uniq and its variants(except uniqUpTo) to count if subquery has distinct or group by clause.", 0) \
     \
     /** Experimental functions */ \
     M(Bool, allow_experimental_funnel_functions, false, "Enable experimental functions for funnel analysis.", 0) \

src/Interpreters/InterpreterSelectQuery.cpp

@@ -39,6 +39,7 @@
 #include <Interpreters/QueryLog.h>
 #include <Interpreters/replaceAliasColumnsInQuery.h>
 #include <Interpreters/RewriteCountDistinctVisitor.h>
+#include <Interpreters/RewriteUniqToCountVisitor.h>
 #include <Interpreters/getCustomKeyFilterForParallelReplicas.h>
 
 #include <QueryPipeline/Pipe.h>
@@ -421,6 +422,12 @@ InterpreterSelectQuery::InterpreterSelectQuery(
         RewriteCountDistinctFunctionVisitor(data_rewrite_countdistinct).visit(query_ptr);
     }
 
+    if (settings.optimize_uniq_to_count)
+    {
+        RewriteUniqToCountMatcher::Data data_rewrite_uniq_count;
+        RewriteUniqToCountVisitor(data_rewrite_uniq_count).visit(query_ptr);
+    }
+
     JoinedTables joined_tables(getSubqueryContext(context), getSelectQuery(), options.with_all_cols, options_.is_create_parameterized_view);
 
     bool got_storage_from_query = false;

src/Interpreters/RewriteUniqToCountVisitor.cpp

@@ -0,0 +1,163 @@
+#include <Interpreters/RewriteUniqToCountVisitor.h>
+
+#include <AggregateFunctions/AggregateFunctionFactory.h>
+#include <Parsers/ASTExpressionList.h>
+#include <Parsers/ASTFunction.h>
+#include <Parsers/ASTIdentifier.h>
+#include <Parsers/ASTSelectQuery.h>
+#include <Parsers/ASTSelectWithUnionQuery.h>
+#include <Parsers/ASTSubquery.h>
+#include <Parsers/ASTTablesInSelectQuery.h>
+#include <Parsers/parseQuery.h>
+
+
+namespace DB
+{
+
+using Aliases = std::unordered_map<String, ASTPtr>;
+
+namespace
+{
+
+bool matchFnUniq(String func_name)
+{
+    auto name = Poco::toLower(func_name);
+    return name == "uniq" || name == "uniqHLL12" || name == "uniqExact" || name == "uniqTheta" || name == "uniqCombined"
+        || name == "uniqCombined64";
+}
+
+bool expressionEquals(const ASTPtr & lhs, const ASTPtr & rhs, const Aliases & alias)
+{
+    if (lhs->getTreeHash() == rhs->getTreeHash())
+    {
+        return true;
+    }
+    else
+    {
+        auto * lhs_idf = lhs->as<ASTIdentifier>();
+        auto * rhs_idf = rhs->as<ASTIdentifier>();
+        if (lhs_idf && rhs_idf)
+        {
+            /// compound identifiers, such as: <t.name, name>
+            if (lhs_idf->shortName() == rhs_idf->shortName())
+                return true;
+
+            /// translate alias
+            if (alias.find(lhs_idf->shortName()) != alias.end())
+                lhs_idf = alias.find(lhs_idf->shortName())->second->as<ASTIdentifier>();
+
+            if (alias.find(rhs_idf->shortName()) != alias.end())
+                rhs_idf = alias.find(rhs_idf->shortName())->second->as<ASTIdentifier>();
+
+            if (lhs_idf->shortName() == rhs_idf->shortName())
+                return true;
+        }
+    }
+    return false;
+}
+
+bool expressionListEquals(ASTExpressionList * lhs, ASTExpressionList * rhs, const Aliases & alias)
+{
+    if (!lhs || !rhs)
+        return false;
+    if (lhs->children.size() != rhs->children.size())
+        return false;
+    for (size_t i = 0; i < lhs->children.size(); i++)
+    {
+        if (!expressionEquals(lhs->children[i], rhs->children[i], alias))
+            return false;
+    }
+    return true;
+}
+
+/// Test whether lhs contains all expressions in rhs.
+bool expressionListContainsAll(ASTExpressionList * lhs, ASTExpressionList * rhs, const Aliases & alias)
+{
+    if (!lhs || !rhs)
+        return false;
+    if (lhs->children.size() < rhs->children.size())
+        return false;
+    for (const auto & re : rhs->children)
+    {
+        auto predicate = [&re, &alias](ASTPtr & le) { return expressionEquals(le, re, alias); };
+        if (std::find_if(lhs->children.begin(), lhs->children.end(), predicate) == lhs->children.end())
+            return false;
+    }
+    return true;
+}
+
+}
+
+void RewriteUniqToCountMatcher::visit(ASTPtr & ast, Data & /*data*/)
+{
+    auto * selectq = ast->as<ASTSelectQuery>();
+    if (!selectq || !selectq->tables() || selectq->tables()->children.size() != 1)
+        return;
+    auto expr_list = selectq->select();
+    if (!expr_list || expr_list->children.size() != 1)
+        return;
+    auto * func = expr_list->children[0]->as<ASTFunction>();
+    if (!func || !matchFnUniq(func->name))
+        return;
+    if (selectq->tables()->as<ASTTablesInSelectQuery>()->children[0]->as<ASTTablesInSelectQueryElement>()->children.size() != 1)
+        return;
+    auto * table_expr = selectq->tables()
+                            ->as<ASTTablesInSelectQuery>()
+                            ->children[0]
+                            ->as<ASTTablesInSelectQueryElement>()
+                            ->children[0]
+                            ->as<ASTTableExpression>();
+    if (!table_expr || table_expr->children.size() != 1 || !table_expr->subquery)
+        return;
+    auto * subquery = table_expr->subquery->as<ASTSubquery>();
+    if (!subquery)
+        return;
+    auto * sub_selectq = subquery->children[0]
+                             ->as<ASTSelectWithUnionQuery>()->children[0]
+                             ->as<ASTExpressionList>()->children[0]
+                             ->as<ASTSelectQuery>();
+    if (!sub_selectq)
+        return;
+    auto sub_expr_list = sub_selectq->select();
+    if (!sub_expr_list)
+        return;
+
+    /// collect subquery select expressions alias
+    Aliases alias;
+    for (const auto & expr : sub_expr_list->children)
+    {
+        if (!expr->tryGetAlias().empty())
+            alias.insert({expr->tryGetAlias(), expr});
+    }
+
+    /// Whether query matches 'SELECT uniq(x ...) FROM (SELECT DISTINCT x ...)'
+    auto match_subquery_with_distinct = [&]() -> bool
+    {
+        if (!sub_selectq->distinct)
+            return false;
+        /// uniq expression list == subquery group by expression list
+        if (!expressionListEquals(func->children[0]->as<ASTExpressionList>(), sub_expr_list->as<ASTExpressionList>(), alias))
+            return false;
+        return true;
+    };
+
+    /// Whether query matches 'SELECT uniq(x ...) FROM (SELECT x ... GROUP BY x ...)'
+    auto match_subquery_with_group_by = [&]() -> bool
+    {
+        auto group_by = sub_selectq->groupBy();
+        if (!group_by)
+            return false;
+        /// uniq expression list == subquery group by expression list
+        if (!expressionListEquals(func->children[0]->as<ASTExpressionList>(), group_by->as<ASTExpressionList>(), alias))
+            return false;
+        /// subquery select expression list must contain all columns in uniq expression list
+        if (!expressionListContainsAll(sub_expr_list->as<ASTExpressionList>(), func->children[0]->as<ASTExpressionList>(), alias))
+            return false;
+        return true;
+    };
+
+    if (match_subquery_with_distinct() || match_subquery_with_group_by())
+        expr_list->children[0] = makeASTFunction("count");
+}
+
+}

src/Interpreters/RewriteUniqToCountVisitor.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <Parsers/IAST.h>
+#include <Interpreters/InDepthNodeVisitor.h>
+#include "Interpreters/TreeRewriter.h"
+
+namespace DB
+{
+
+class ASTFunction;
+
+/** Optimize `uniq` into `count` over subquery.
+ *     Example: 'SELECT uniq(x ...) FROM (SELECT DISTINCT x ...)' to
+ *     Result: 'SELECT count() FROM (SELECT DISTINCT x ...)'
+ *
+ *     Example: 'SELECT uniq(x ...) FROM (SELECT x ... GROUP BY x ...)' to
+ *     Result: 'SELECT count() FROM (SELECT x ... GROUP BY x ...)'
+ *
+ *     Note that we can rewrite all uniq variants except uniqUpTo.
+ */
+class RewriteUniqToCountMatcher
+{
+public:
+    struct Data {};
+    static void visit(ASTPtr & ast, Data &);
+    static bool needChildVisit(const ASTPtr &, const ASTPtr &) { return true; }
+};
+
+using RewriteUniqToCountVisitor = InDepthNodeVisitor<RewriteUniqToCountMatcher, true>;
+}

tests/performance/uniq_to_count.xml

@@ -0,0 +1,8 @@
+<test>
+    <query>select uniq(number) from (select DISTINCT number from numbers(1000000))</query>
+    <query>select uniq(number) from (select number from numbers(1000000) group by number)</query>
+
+    <!--For new analyzer-->
+    <query>select uniq(number) from (select DISTINCT number from numbers(1000000)) SETTINGS allow_experimental_analyzer=1</query>
+    <query>select uniq(number) from (select number from numbers(1000000) group by number) SETTINGS allow_experimental_analyzer=1</query>
+</test>

tests/queries/0_stateless/02868_distinct_to_count_optimization.reference

@@ -0,0 +1,252 @@
+1. test simple distinct
+3
+SELECT count()
+FROM
+(
+    SELECT DISTINCT a
+    FROM test_rewrite_uniq_to_count
+)
+SETTINGS allow_experimental_analyzer = 0
+3
+QUERY id: 0
+  PROJECTION COLUMNS
+    uniq(a) UInt64
+  PROJECTION
+    LIST id: 1, nodes: 1
+      FUNCTION id: 2, function_name: count, function_type: aggregate, result_type: UInt64
+  JOIN TREE
+    QUERY id: 3, is_subquery: 1, is_distinct: 1
+      PROJECTION COLUMNS
+        a UInt8
+      PROJECTION
+        LIST id: 4, nodes: 1
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+      JOIN TREE
+        TABLE id: 6, table_name: default.test_rewrite_uniq_to_count
+  SETTINGS allow_experimental_analyzer=1
+2. test distinct with subquery alias
+3
+SELECT count()
+FROM
+(
+    SELECT DISTINCT a
+    FROM test_rewrite_uniq_to_count
+) AS t
+SETTINGS allow_experimental_analyzer = 0
+3
+QUERY id: 0
+  PROJECTION COLUMNS
+    uniq(a) UInt64
+  PROJECTION
+    LIST id: 1, nodes: 1
+      FUNCTION id: 2, function_name: count, function_type: aggregate, result_type: UInt64
+  JOIN TREE
+    QUERY id: 3, alias: t, is_subquery: 1, is_distinct: 1
+      PROJECTION COLUMNS
+        a UInt8
+      PROJECTION
+        LIST id: 4, nodes: 1
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+      JOIN TREE
+        TABLE id: 6, table_name: default.test_rewrite_uniq_to_count
+  SETTINGS allow_experimental_analyzer=1
+3. test distinct with compound column name
+3
+SELECT count()
+FROM
+(
+    SELECT DISTINCT a
+    FROM test_rewrite_uniq_to_count
+) AS t
+SETTINGS allow_experimental_analyzer = 0
+3
+QUERY id: 0
+  PROJECTION COLUMNS
+    uniq(a) UInt64
+  PROJECTION
+    LIST id: 1, nodes: 1
+      FUNCTION id: 2, function_name: count, function_type: aggregate, result_type: UInt64
+  JOIN TREE
+    QUERY id: 3, alias: t, is_subquery: 1, is_distinct: 1
+      PROJECTION COLUMNS
+        a UInt8
+      PROJECTION
+        LIST id: 4, nodes: 1
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+      JOIN TREE
+        TABLE id: 6, table_name: default.test_rewrite_uniq_to_count
+  SETTINGS allow_experimental_analyzer=1
+4. test distinct with select expression alias
+3
+SELECT count()
+FROM
+(
+    SELECT DISTINCT a AS alias_of_a
+    FROM test_rewrite_uniq_to_count
+) AS t
+SETTINGS allow_experimental_analyzer = 0
+3
+QUERY id: 0
+  PROJECTION COLUMNS
+    uniq(alias_of_a) UInt64
+  PROJECTION
+    LIST id: 1, nodes: 1
+      FUNCTION id: 2, function_name: count, function_type: aggregate, result_type: UInt64
+  JOIN TREE
+    QUERY id: 3, alias: t, is_subquery: 1, is_distinct: 1
+      PROJECTION COLUMNS
+        alias_of_a UInt8
+      PROJECTION
+        LIST id: 4, nodes: 1
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+      JOIN TREE
+        TABLE id: 6, table_name: default.test_rewrite_uniq_to_count
+  SETTINGS allow_experimental_analyzer=1
+5. test simple group by
+3
+SELECT count()
+FROM
+(
+    SELECT
+        a,
+        sum(b)
+    FROM test_rewrite_uniq_to_count
+    GROUP BY a
+)
+SETTINGS allow_experimental_analyzer = 0
+3
+QUERY id: 0
+  PROJECTION COLUMNS
+    uniq(a) UInt64
+  PROJECTION
+    LIST id: 1, nodes: 1
+      FUNCTION id: 2, function_name: count, function_type: aggregate, result_type: UInt64
+  JOIN TREE
+    QUERY id: 3, is_subquery: 1
+      PROJECTION COLUMNS
+        a UInt8
+        sum(b) UInt64
+      PROJECTION
+        LIST id: 4, nodes: 2
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+          FUNCTION id: 7, function_name: sum, function_type: aggregate, result_type: UInt64
+            ARGUMENTS
+              LIST id: 8, nodes: 1
+                COLUMN id: 9, column_name: b, result_type: UInt8, source_id: 6
+      JOIN TREE
+        TABLE id: 6, table_name: default.test_rewrite_uniq_to_count
+      GROUP BY
+        LIST id: 10, nodes: 1
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+  SETTINGS allow_experimental_analyzer=1
+6. test group by with subquery alias
+3
+SELECT count()
+FROM
+(
+    SELECT
+        a,
+        sum(b)
+    FROM test_rewrite_uniq_to_count
+    GROUP BY a
+) AS t
+SETTINGS allow_experimental_analyzer = 0
+3
+QUERY id: 0
+  PROJECTION COLUMNS
+    uniq(a) UInt64
+  PROJECTION
+    LIST id: 1, nodes: 1
+      FUNCTION id: 2, function_name: count, function_type: aggregate, result_type: UInt64
+  JOIN TREE
+    QUERY id: 3, alias: t, is_subquery: 1
+      PROJECTION COLUMNS
+        a UInt8
+        sum(b) UInt64
+      PROJECTION
+        LIST id: 4, nodes: 2
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+          FUNCTION id: 7, function_name: sum, function_type: aggregate, result_type: UInt64
+            ARGUMENTS
+              LIST id: 8, nodes: 1
+                COLUMN id: 9, column_name: b, result_type: UInt8, source_id: 6
+      JOIN TREE
+        TABLE id: 6, table_name: default.test_rewrite_uniq_to_count
+      GROUP BY
+        LIST id: 10, nodes: 1
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+  SETTINGS allow_experimental_analyzer=1
+7. test group by with compound column name
+3
+SELECT count()
+FROM
+(
+    SELECT
+        a AS alias_of_a,
+        sum(b)
+    FROM test_rewrite_uniq_to_count
+    GROUP BY a
+) AS t
+SETTINGS allow_experimental_analyzer = 0
+3
+QUERY id: 0
+  PROJECTION COLUMNS
+    uniq(alias_of_a) UInt64
+  PROJECTION
+    LIST id: 1, nodes: 1
+      FUNCTION id: 2, function_name: count, function_type: aggregate, result_type: UInt64
+  JOIN TREE
+    QUERY id: 3, alias: t, is_subquery: 1
+      PROJECTION COLUMNS
+        alias_of_a UInt8
+        sum(b) UInt64
+      PROJECTION
+        LIST id: 4, nodes: 2
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+          FUNCTION id: 7, function_name: sum, function_type: aggregate, result_type: UInt64
+            ARGUMENTS
+              LIST id: 8, nodes: 1
+                COLUMN id: 9, column_name: b, result_type: UInt8, source_id: 6
+      JOIN TREE
+        TABLE id: 6, table_name: default.test_rewrite_uniq_to_count
+      GROUP BY
+        LIST id: 10, nodes: 1
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+  SETTINGS allow_experimental_analyzer=1
+8. test group by with select expression alias
+3
+SELECT count()
+FROM
+(
+    SELECT
+        a AS alias_of_a,
+        sum(b)
+    FROM test_rewrite_uniq_to_count
+    GROUP BY alias_of_a
+) AS t
+SETTINGS allow_experimental_analyzer = 0
+3
+QUERY id: 0
+  PROJECTION COLUMNS
+    uniq(alias_of_a) UInt64
+  PROJECTION
+    LIST id: 1, nodes: 1
+      FUNCTION id: 2, function_name: count, function_type: aggregate, result_type: UInt64
+  JOIN TREE
+    QUERY id: 3, alias: t, is_subquery: 1
+      PROJECTION COLUMNS
+        alias_of_a UInt8
+        sum(b) UInt64
+      PROJECTION
+        LIST id: 4, nodes: 2
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+          FUNCTION id: 7, function_name: sum, function_type: aggregate, result_type: UInt64
+            ARGUMENTS
+              LIST id: 8, nodes: 1
+                COLUMN id: 9, column_name: b, result_type: UInt8, source_id: 6
+      JOIN TREE
+        TABLE id: 6, table_name: default.test_rewrite_uniq_to_count
+      GROUP BY
+        LIST id: 10, nodes: 1
+          COLUMN id: 5, column_name: a, result_type: UInt8, source_id: 6
+  SETTINGS allow_experimental_analyzer=1

tests/queries/0_stateless/02868_distinct_to_count_optimization.sql

@@ -0,0 +1,69 @@
+drop table if exists test_rewrite_uniq_to_count;
+
+CREATE TABLE test_rewrite_uniq_to_count
+(
+    `a` UInt8,
+    `b` UInt8,
+    `c` UInt8
+) ENGINE = MergeTree ORDER BY `a`;
+
+
+INSERT INTO test_rewrite_uniq_to_count values ('1', '1', '1'), ('1', '1', '1');
+INSERT INTO test_rewrite_uniq_to_count values ('2', '2', '2'), ('2', '2', '2');
+INSERT INTO test_rewrite_uniq_to_count values ('3', '3', '3'), ('3', '3', '3');
+
+set optimize_uniq_to_count=true;
+
+
+SELECT '1. test simple distinct';
+SELECT uniq(a) FROM (SELECT DISTINCT a FROM test_rewrite_uniq_to_count) settings allow_experimental_analyzer=0;
+EXPLAIN SYNTAX SELECT uniq(a) FROM (SELECT DISTINCT a FROM test_rewrite_uniq_to_count) settings allow_experimental_analyzer=0;
+SELECT uniq(a) FROM (SELECT DISTINCT a FROM test_rewrite_uniq_to_count) settings allow_experimental_analyzer=1;
+EXPLAIN QUERY TREE SELECT uniq(a) FROM (SELECT DISTINCT a FROM test_rewrite_uniq_to_count) settings allow_experimental_analyzer=1;
+
+
+SELECT '2. test distinct with subquery alias';
+SELECT uniq(t.a) FROM (SELECT DISTINCT a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=0;
+EXPLAIN SYNTAX SELECT uniq(a) FROM (SELECT DISTINCT a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=0;
+SELECT uniq(t.a) FROM (SELECT DISTINCT a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=1;
+EXPLAIN QUERY TREE SELECT uniq(t.a) FROM (SELECT DISTINCT a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=1;
+
+SELECT '3. test distinct with compound column name';
+SELECT uniq(a) FROM (SELECT DISTINCT test_rewrite_uniq_to_count.a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=0;
+EXPLAIN SYNTAX SELECT uniq(a) FROM (SELECT DISTINCT test_rewrite_uniq_to_count.a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=0;
+SELECT uniq(a) FROM (SELECT DISTINCT test_rewrite_uniq_to_count.a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=1;
+EXPLAIN QUERY TREE SELECT uniq(a) FROM (SELECT DISTINCT test_rewrite_uniq_to_count.a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=1;
+
+SELECT '4. test distinct with select expression alias';
+SELECT uniq(alias_of_a) FROM (SELECT DISTINCT a as alias_of_a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=0;
+EXPLAIN SYNTAX SELECT uniq(alias_of_a) FROM (SELECT DISTINCT a as alias_of_a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=0;
+SELECT uniq(alias_of_a) FROM (SELECT DISTINCT a as alias_of_a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=1;
+EXPLAIN QUERY TREE SELECT uniq(alias_of_a) FROM (SELECT DISTINCT a as alias_of_a FROM test_rewrite_uniq_to_count) t settings allow_experimental_analyzer=1;
+
+
+SELECT '5. test simple group by';
+SELECT uniq(a) FROM (SELECT a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) settings allow_experimental_analyzer=0;
+EXPLAIN SYNTAX SELECT uniq(a) FROM (SELECT a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) settings allow_experimental_analyzer=0;
+SELECT uniq(a) FROM (SELECT a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) settings allow_experimental_analyzer=1;
+EXPLAIN QUERY TREE SELECT uniq(a) FROM (SELECT a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) settings allow_experimental_analyzer=1;
+
+SELECT '6. test group by with subquery alias';
+SELECT uniq(t.a) FROM (SELECT a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) t settings allow_experimental_analyzer=0;
+EXPLAIN SYNTAX SELECT uniq(t.a) FROM (SELECT a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) t settings allow_experimental_analyzer=0;
+SELECT uniq(t.a) FROM (SELECT a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) t settings allow_experimental_analyzer=1;
+EXPLAIN QUERY TREE SELECT uniq(t.a) FROM (SELECT a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) t settings allow_experimental_analyzer=1;
+
+SELECT '7. test group by with compound column name';
+SELECT uniq(t.alias_of_a) FROM (SELECT a as alias_of_a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) t settings allow_experimental_analyzer=0;
+EXPLAIN SYNTAX SELECT uniq(t.alias_of_a) FROM (SELECT a as alias_of_a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) t settings allow_experimental_analyzer=0;
+SELECT uniq(t.alias_of_a) FROM (SELECT a as alias_of_a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) t settings allow_experimental_analyzer=1;
+EXPLAIN QUERY TREE SELECT uniq(t.alias_of_a) FROM (SELECT a as alias_of_a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY a) t settings allow_experimental_analyzer=1;
+
+SELECT '8. test group by with select expression alias';
+SELECT uniq(t.alias_of_a) FROM (SELECT a as alias_of_a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY alias_of_a) t settings allow_experimental_analyzer=0;
+EXPLAIN SYNTAX SELECT uniq(t.alias_of_a) FROM (SELECT a as alias_of_a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY alias_of_a) t  settings allow_experimental_analyzer=0;
+SELECT uniq(t.alias_of_a) FROM (SELECT a as alias_of_a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY alias_of_a) t  settings allow_experimental_analyzer=1;
+EXPLAIN QUERY TREE SELECT uniq(t.alias_of_a) FROM (SELECT a as alias_of_a, sum(b) FROM test_rewrite_uniq_to_count GROUP BY alias_of_a) t  settings allow_experimental_analyzer=1;
+
+drop table if exists test_rewrite_uniq_to_count;
+