visit.cpp

struct CustomizeASTSelectWithUnionQueryNormalize
{
    using TypeToVisit = ASTSelectWithUnionQuery;

    const UnionMode & union_default_mode;

    static void getSelectsFromUnionListNode(ASTPtr & ast_select, ASTs & selects)
    {
        if (auto * inner_union = ast_select->as<ASTSelectWithUnionQuery>())
        {
            /// We need flatten from last to first
            for (auto child = inner_union->list_of_selects->children.rbegin(); child != inner_union->list_of_selects->children.rend();
                 ++child)
                getSelectsFromUnionListNode(*child, selects);

            return;
        }

        selects.push_back(std::move(ast_select));
    }
  
void visit(ASTSelectWithUnionQuery & ast, ASTPtr & ast_ptr)
    {
        auto & union_modes = ast.list_of_modes;
        ASTs selects;
        auto & select_list = ast.list_of_selects->children;

        int i;
        for (i = union_modes.size() - 1; i >= 0; --i)
        {
            /// Rewrite UNION Mode
            if (union_modes[i] == ASTSelectWithUnionQuery::Mode::Unspecified)
            {
                if (union_default_mode == UnionMode::ALL)
                    union_modes[i] = ASTSelectWithUnionQuery::Mode::ALL;
                else if (union_default_mode == UnionMode::DISTINCT)
                    union_modes[i] = ASTSelectWithUnionQuery::Mode::DISTINCT;
                else
                    throw Exception(
                        "Expected ALL or DISTINCT in SelectWithUnion query, because setting (union_default_mode) is empty",
                        DB::ErrorCodes::EXPECTED_ALL_OR_DISTINCT);
            }

            if (union_modes[i] == ASTSelectWithUnionQuery::Mode::ALL)
            {
                if (auto * inner_union = select_list[i + 1]->as<ASTSelectWithUnionQuery>())
                {
                    /// If inner_union is an UNION ALL list, just lift up
                    if (inner_union->union_mode == ASTSelectWithUnionQuery::Mode::ALL)
                    {
                        for (auto child = inner_union->list_of_selects->children.rbegin();
                             child != inner_union->list_of_selects->children.rend();
                             ++child)
                            selects.push_back(std::move(*child));
                    }
                    /// inner_union is an UNION DISTINCT list,
                    // we cann't lift up
                    else
                        selects.push_back(std::move(select_list[i + 1]));
                }
                else
                  selects.push_back(std::move(select_list[i + 1]));
            }
            /// flatten all left nodes and current node to a UNION DISTINCT list
            else if (union_modes[i] == ASTSelectWithUnionQuery::Mode::DISTINCT)
            {
                auto distinct_list = std::make_shared<ASTSelectWithUnionQuery>();
                distinct_list->list_of_selects = std::make_shared<ASTExpressionList>();
                distinct_list->children.push_back(distinct_list->list_of_selects);
                for (int j = i + 1; j >= 0; j--)
                {
                    getSelectsFromUnionListNode(select_list[j], distinct_list->list_of_selects->children);
                }
                distinct_list->union_mode = ASTSelectWithUnionQuery::Mode::DISTINCT;
                // Reverse children list
                std::reverse(distinct_list->list_of_selects->children.begin(), distinct_list->list_of_selects->children.end());
                distinct_list->is_normalized = true;
                selects.push_back(std::move(distinct_list));
                break;
            }
        }

        /// No UNION DISTINCT or only one SELECT in select_list
        if (i == -1)
        {
            if (auto * inner_union = select_list[0]->as<ASTSelectWithUnionQuery>())
            {
                /// If inner_union is an UNION ALL list, just lift it up
                if (inner_union->union_mode == ASTSelectWithUnionQuery::Mode::ALL)
                {
                    for (auto child = inner_union->list_of_selects->children.rbegin();
                         child != inner_union->list_of_selects->children.rend();
                         ++child)
                        selects.push_back(std::move(*child));
            }
            /// inner_union is an UNION DISTINCT list,
            // we cann't lift it up
            else
                selects.push_back(std::move(select_list[i + 1]));
            }
        else
            selects.push_back(std::move(select_list[0]));
        }
        // reverse children list
        std::reverse(selects.begin(), selects.end());

        ast.is_normalized = true;
        ast.union_mode = ASTSelectWithUnionQuery::Mode::ALL;

        /// After normalize, if we only have one ASTSelectWithUnionQuery child, lift if up
        if (selects.size() == 1)
        {
            if (selects.at(0)->as<ASTSelectWithUnionQuery>())
            {
                ast_ptr = std::move(selects.at(0));
                return;
            }
        }
        ast.list_of_selects->children = std::move(selects);
    }
};

using CustomizeASTSelectWithUnionQueryNormalizeVisitor
    = InDepthNodeVisitor<OneTypeMatcher<CustomizeASTSelectWithUnionQueryNormalize>, false>;