Merge pull request #1763 from zhang2014/fix/ISSUES-320

ISSUES-320 Fix rewrite expression

dbms/src/Interpreters/ExpressionAnalyzer.cpp

@@ -158,11 +158,11 @@ ExpressionAnalyzer::ExpressionAnalyzer(
     const StoragePtr & storage_,
     const NamesAndTypesList & columns_,
     size_t subquery_depth_,
-    bool do_global_)
+    bool do_global_, SubqueriesForSets subquery_for_set_)
     : ast(ast_), context(context_), settings(context.getSettings()),
     subquery_depth(subquery_depth_), columns(columns_),
     storage(storage_ ? storage_ : getTable()),
-    do_global(do_global_)
+    do_global(do_global_), subqueries_for_sets(subquery_for_set_)
 {
     init();
 }

dbms/src/Interpreters/ExpressionAnalyzer.h

@@ -72,7 +72,7 @@ public:
         const StoragePtr & storage_,
         const NamesAndTypesList & columns_,
         size_t subquery_depth_ = 0,
-        bool do_global_ = false);
+        bool do_global_ = false, SubqueriesForSets subquery_for_set_ = {});
 
     /// Does the expression have aggregate functions or a GROUP BY or HAVING section.
     bool hasAggregation() const { return has_aggregation; }
@@ -166,6 +166,9 @@ private:
     NamesAndTypesList aggregation_keys;
     AggregateDescriptions aggregate_descriptions;
 
+    /// Do I need to prepare for execution global subqueries when analyzing the query.
+    bool do_global;
+
     SubqueriesForSets subqueries_for_sets;
 
     PreparedSets prepared_sets;
@@ -201,14 +204,11 @@ private:
     /// The backward mapping for array_join_alias_to_name.
     NameToNameMap array_join_name_to_alias;
 
-    /// Do I need to prepare for execution global subqueries when analyzing the query.
-    bool do_global;
 
     /// All new temporary tables obtained by performing the GLOBAL IN/JOIN subqueries.
     Tables external_tables;
     size_t external_table_id = 1;
 
-
     void init();
 
     static NamesAndTypesList::iterator findColumn(const String & name, NamesAndTypesList & cols);

dbms/src/Interpreters/InterpreterSelectQuery.cpp

@@ -95,7 +95,8 @@ void InterpreterSelectQuery::init(const BlockInputStreamPtr & input, const Names
             ASTSelectQuery & head_query = static_cast<ASTSelectQuery &>(*head);
             tail = head_query.next_union_all;
 
-            interpreter->next_select_in_union_all = std::make_unique<InterpreterSelectQuery>(head, context, to_stage, subquery_depth);
+            interpreter->next_select_in_union_all =
+                std::make_unique<InterpreterSelectQuery>(head, context, to_stage, subquery_depth);
             interpreter = interpreter->next_select_in_union_all.get();
         }
     }
@@ -119,9 +120,17 @@ void InterpreterSelectQuery::init(const BlockInputStreamPtr & input, const Names
     else
     {
         renameColumns();
-        if (!required_column_names.empty())
+        if (!required_column_names.empty()) {
             rewriteExpressionList(required_column_names);
 
+            if (is_first_select_inside_union_all) {
+                for (auto p = next_select_in_union_all.get(); p != nullptr; p = p->next_select_in_union_all.get())
+                    p->query_analyzer.reset(new ExpressionAnalyzer(
+                        p->query_ptr, p->context, p->storage, p->table_column_names, p->subquery_depth,
+                        false, p->query_analyzer->getSubqueriesForSets()));
+            }
+        }
+
         basicInit(input);
     }
 }

dbms/tests/queries/0_stateless/00562_rewrite_select_expression_with_union.reference

@@ -0,0 +1,4 @@
+test_string	2
+test_string	2
+test_string
+test_string

dbms/tests/queries/0_stateless/00562_rewrite_select_expression_with_union.sql

@@ -0,0 +1,10 @@
+DROP TABLE IF EXISTS test.test;
+
+CREATE TABLE test.test ( s String,  i Int64) ENGINE = Memory;
+
+INSERT INTO test.test VALUES('test_string', 1);
+
+SELECT s, SUM(i*2) AS i FROM test.test GROUP BY s  UNION ALL  SELECT s, SUM(i*2) AS i FROM test.test GROUP BY s;
+SELECT s FROM (SELECT s, SUM(i*2) AS i FROM test.test GROUP BY s  UNION ALL  SELECT s, SUM(i*2) AS i FROM test.test GROUP BY s);
+
+DROP TABLE IF EXISTS test.test;