Correct number of columns is returned from an element of UNION ALL inside a...

Correct number of columns is returned from an element of UNION ALL inside a subquery if some column was queried more than once #2786

dbms/src/Interpreters/ExpressionAnalyzer.cpp

@@ -203,7 +203,7 @@ ExpressionAnalyzer::ExpressionAnalyzer(
     const SubqueriesForSets & subqueries_for_set_)
     : ast(ast_), context(context_), settings(context.getSettings()),
     subquery_depth(subquery_depth_),
-    source_columns(source_columns_), required_result_columns(required_result_columns_.begin(), required_result_columns_.end()),
+    source_columns(source_columns_), required_result_columns(required_result_columns_),
     storage(storage_),
     do_global(do_global_), subqueries_for_sets(subqueries_for_set_)
 {
@@ -2847,7 +2847,8 @@ void ExpressionAnalyzer::appendProjectResult(ExpressionActionsChain & chain) con
     for (size_t i = 0; i < asts.size(); ++i)
     {
         String result_name = asts[i]->getAliasOrColumnName();
-        if (required_result_columns.empty() || required_result_columns.count(result_name))
+        if (required_result_columns.empty()
+            || std::find(required_result_columns.begin(), required_result_columns.end(), result_name) !=  required_result_columns.end())
         {
             result_columns.emplace_back(asts[i]->getColumnName(), result_name);
             step.required_output.push_back(result_columns.back().second);
@@ -3393,15 +3394,37 @@ void ExpressionAnalyzer::removeUnneededColumnsFromSelectClause()
     if (!select_query)
         return;
 
-    if (required_result_columns.empty() || select_query->distinct)
+    if (required_result_columns.empty())
         return;
 
     ASTs & elements = select_query->select_expression_list->children;
 
-    elements.erase(std::remove_if(elements.begin(), elements.end(), [this](const auto & node)
+    ASTs new_elements;
+    new_elements.reserve(elements.size());
+
+    /// Some columns may be queried multiple times, like SELECT x, y, y FROM table.
+    /// In that case we keep them exactly same number of times.
+    std::map<String, size_t> required_columns_with_duplicate_count;
+    for (const auto & name : required_result_columns)
+        ++required_columns_with_duplicate_count[name];
+
+    for (const auto & elem : elements)
     {
-        return !required_result_columns.count(node->getAliasOrColumnName()) && !hasArrayJoin(node);
-    }), elements.end());
+        String name = elem->getAliasOrColumnName();
+
+        auto it = required_columns_with_duplicate_count.find(name);
+        if (required_columns_with_duplicate_count.end() != it && it->second)
+        {
+            new_elements.push_back(elem);
+            --it->second;
+        }
+        else if (select_query->distinct || hasArrayJoin(elem))
+        {
+            new_elements.push_back(elem);
+        }
+    }
+
+    elements = std::move(new_elements);
 }
 
 }

dbms/src/Interpreters/ExpressionAnalyzer.h

@@ -204,7 +204,7 @@ private:
 
     /** If non-empty, ignore all expressions in  not from this list.
       */
-    NameSet required_result_columns;
+    Names required_result_columns;
 
     /// Columns after ARRAY JOIN, JOIN, and/or aggregation.
     NamesAndTypesList aggregated_columns;

dbms/src/Interpreters/InterpreterSelectWithUnionQuery.cpp

@@ -92,7 +92,7 @@ InterpreterSelectWithUnionQuery::InterpreterSelectWithUnionQuery(
             ast.list_of_selects->children.at(query_num), context, current_required_result_column_names, to_stage, subquery_depth, only_analyze));
     }
 
-    /// Determine structure of result.
+    /// Determine structure of the result.
 
     if (num_selects == 1)
     {

dbms/tests/queries/0_stateless/00680_duplicate_columns_inside_union_all.reference

+1	2
+3	3
+1	2
+4	4
+1	2
+3	4
+1	2
+3	3

dbms/tests/queries/0_stateless/00680_duplicate_columns_inside_union_all.sql

+SELECT x, y FROM (SELECT x, y FROM (SELECT 1 AS x, 2 AS y) UNION ALL SELECT x, x FROM (SELECT 3 AS x, 4 AS y)) ORDER BY x, y;
+SELECT x, y FROM (SELECT x, y FROM (SELECT 1 AS x, 2 AS y) UNION ALL SELECT y, y FROM (SELECT 3 AS x, 4 AS y)) ORDER BY x, y;
+SELECT x, y FROM (SELECT x, x, y FROM (SELECT 1 AS x, 2 AS y) UNION ALL SELECT x, y, y FROM (SELECT 3 AS x, 4 AS y)) ORDER BY x, y;
+SELECT x, y FROM (SELECT x, y, y FROM (SELECT 1 AS x, 2 AS y) UNION ALL SELECT x, x, y FROM (SELECT 3 AS x, 4 AS y)) ORDER BY x, y;