dbms: Server: Fixed the bug in a much nicer way: InterpreterSelectQuery does not split anymore its AST. [#METR-16778]

dbms/include/DB/Parsers/ASTSelectQuery.h

@@ -32,6 +32,8 @@ public:
 	/// Переписывает select_expression_list, чтобы вернуть только необходимые столбцы в правильном порядке.
 	void rewriteSelectExpressionList(const Names & column_names);
 
+	bool isUnionAllHead() const { return prev_union_all.isNull() && !next_union_all.isNull(); }
+
 	ASTPtr clone() const override;
 
 	/// Возвращает указатель на формат из последнего SELECT'а цепочки UNION ALL.
@@ -55,6 +57,7 @@ public:
 	ASTPtr limit_offset;
 	ASTPtr limit_length;
 	ASTPtr settings;
+	ASTPtr prev_union_all;
 	ASTPtr next_union_all; /// Следующий запрос SELECT в цепочке UNION ALL, если такой есть
 };
 

dbms/src/Interpreters/InterpreterSelectQuery.cpp

@@ -55,7 +55,6 @@ void InterpreterSelectQuery::init(BlockInputStreamPtr input, const Names & requi
 		/// Создать цепочку запросов SELECT.
 		InterpreterSelectQuery * interpreter = this;
 		ASTPtr tail = query.next_union_all;
-		query.next_union_all = nullptr;
 
 		while (!tail.isNull())
 		{
@@ -63,7 +62,6 @@ void InterpreterSelectQuery::init(BlockInputStreamPtr input, const Names & requi
 
 			ASTSelectQuery & head_query = static_cast<ASTSelectQuery &>(*head);
 			tail = head_query.next_union_all;
-			head_query.next_union_all = nullptr;
 
 			interpreter->next_select_in_union_all.reset(new InterpreterSelectQuery(head, context, to_stage, subquery_depth, nullptr, false));
 			interpreter = interpreter->next_select_in_union_all.get();
@@ -178,7 +176,7 @@ InterpreterSelectQuery::InterpreterSelectQuery(ASTPtr query_ptr_, const Context
 	size_t subquery_depth_, BlockInputStreamPtr input_, bool is_union_all_head_)
 	: query_ptr(query_ptr_), query(typeid_cast<ASTSelectQuery &>(*query_ptr)),
 	context(context_), to_stage(to_stage_), subquery_depth(subquery_depth_),
-	is_first_select_inside_union_all(is_union_all_head_ && !query.next_union_all.isNull()),
+	is_first_select_inside_union_all(is_union_all_head_ && query.isUnionAllHead()),
 	log(&Logger::get("InterpreterSelectQuery"))
 {
 	init(input_);
@@ -189,7 +187,7 @@ InterpreterSelectQuery::InterpreterSelectQuery(ASTPtr query_ptr_, const Context
 	QueryProcessingStage::Enum to_stage_, size_t subquery_depth_, BlockInputStreamPtr input_)
 	: query_ptr(query_ptr_), query(typeid_cast<ASTSelectQuery &>(*query_ptr)),
 	context(context_), to_stage(to_stage_), subquery_depth(subquery_depth_),
-	is_first_select_inside_union_all(!query.next_union_all.isNull()),
+	is_first_select_inside_union_all(query.isUnionAllHead()),
 	log(&Logger::get("InterpreterSelectQuery"))
 {
 	init(input_, required_column_names_);
@@ -200,7 +198,7 @@ InterpreterSelectQuery::InterpreterSelectQuery(ASTPtr query_ptr_, const Context
 	const NamesAndTypesList & table_column_names, QueryProcessingStage::Enum to_stage_, size_t subquery_depth_, BlockInputStreamPtr input_)
 	: query_ptr(query_ptr_), query(typeid_cast<ASTSelectQuery &>(*query_ptr)),
 	context(context_), to_stage(to_stage_), subquery_depth(subquery_depth_),
-	is_first_select_inside_union_all(!query.next_union_all.isNull()),
+	is_first_select_inside_union_all(query.isUnionAllHead()),
 	log(&Logger::get("InterpreterSelectQuery"))
 {
 	init(input_, required_column_names_, table_column_names);
@@ -212,11 +210,13 @@ bool InterpreterSelectQuery::hasAsterisk() const
 		return true;
 
 	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())
 		{
 			if (p->query.hasAsterisk())
 				return true;
 		}
+	}
 
 	return false;
 }
@@ -224,8 +224,10 @@ bool InterpreterSelectQuery::hasAsterisk() const
 void InterpreterSelectQuery::renameColumns()
 {
 	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.renameColumns(query);
+	}
 }
 
 void InterpreterSelectQuery::rewriteExpressionList(const Names & required_column_names)
@@ -234,17 +236,21 @@ void InterpreterSelectQuery::rewriteExpressionList(const Names & required_column
 		return;
 
 	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())
 		{
 			if (p->query.distinct)
 				return;
 		}
+	}
 
 	query.rewriteSelectExpressionList(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.rewriteSelectExpressionList(required_column_names);
+	}
 }
 
 void InterpreterSelectQuery::getDatabaseAndTableNames(String & database_name, String & table_name)
@@ -710,7 +716,18 @@ QueryProcessingStage::Enum InterpreterSelectQuery::executeFetchColumns(BlockInpu
 		if (max_streams > 1 && !is_remote)
 			max_streams *= settings.max_streams_to_max_threads_ratio;
 
-		streams = storage->read(required_columns, query_ptr,
+		ASTPtr actual_query_ptr;
+		if (storage->isRemote())
+		{
+			/// В случае удаленного запроса отправляем только SELECT, который выполнится.
+			actual_query_ptr = query_ptr->clone();
+			auto actual_select_query = static_cast<ASTSelectQuery *>(&*actual_query_ptr);
+			actual_select_query->next_union_all = nullptr;
+		}
+		else
+			actual_query_ptr = query_ptr;
+
+		streams = storage->read(required_columns, actual_query_ptr,
 			context, settings_for_storage, from_stage,
 			settings.max_block_size, max_streams);
 

dbms/src/Parsers/ParserSelectQuery.cpp

@@ -329,6 +329,8 @@ bool ParserSelectQuery::parseImpl(Pos & pos, Pos end, ASTPtr & node, Pos & max_p
 			ParserSelectQuery select_p;
 			if (!select_p.parse(pos, end, select_query->next_union_all, max_parsed_pos, expected))
 				return false;
+			auto next_select_query = static_cast<ASTSelectQuery *>(&*select_query->next_union_all);
+			next_select_query->prev_union_all = node;
 		}
 		else
 			return false;
@@ -367,6 +369,8 @@ bool ParserSelectQuery::parseImpl(Pos & pos, Pos end, ASTPtr & node, Pos & max_p
 		select_query->children.push_back(select_query->settings);
 	if (select_query->format)
 		select_query->children.push_back(select_query->format);
+	if (select_query->prev_union_all)
+		select_query->children.push_back(select_query->prev_union_all);
 	if (select_query->next_union_all)
 		select_query->children.push_back(select_query->next_union_all);
 

dbms/tests/queries/0_stateless/00098_l_union_all.sql

@@ -1,3 +1,3 @@
-SELECT a,b,c,d FROM (SELECT 1 AS a,2 AS b, 3 AS c UNION ALL SELECT 2,3,4 ) ANY inner JOIN (SELECT 1 AS a,2 AS b,4 AS d UNION ALL SELECT 2,3,5) USING (a) ORDER BY a,b,c,d ASC;
+SELECT a,b,c,d FROM (SELECT 1 AS a,2 AS b, 3 AS c UNION ALL SELECT 2,3,4 ) ANY INNER JOIN (SELECT 1 AS a,2 AS b,4 AS d UNION ALL SELECT 2,3,5) USING (a) ORDER BY a,b,c,d ASC;
-SELECT a,b,c,d FROM (SELECT 1 AS a,2 AS b, 3 AS c UNION ALL SELECT 2,3,4 ) ALL left JOIN (SELECT 1 AS a,2 AS b,4 AS d UNION ALL SELECT 2,3,5) USING (a) ORDER BY a,b,c,d ASC;
+SELECT a,b,c,d FROM (SELECT 1 AS a,2 AS b, 3 AS c UNION ALL SELECT 2,3,4 ) ALL LEFT JOIN (SELECT 1 AS a,2 AS b,4 AS d UNION ALL SELECT 2,3,5) USING (a) ORDER BY a,b,c,d ASC;
-SELECT a,b,c,d FROM (SELECT 1 AS a,2 AS b, 3 AS c UNION ALL SELECT 2,3,4 ) ALL left JOIN (SELECT 1 AS a,2 AS b,4 AS d UNION ALL SELECT 2,3,5) USING a,b ORDER BY a,b,c,d ASC;
+SELECT a,b,c,d FROM (SELECT 1 AS a,2 AS b, 3 AS c UNION ALL SELECT 2,3,4 ) ALL LEFT JOIN (SELECT 1 AS a,2 AS b,4 AS d UNION ALL SELECT 2,3,5) USING a,b ORDER BY a,b,c,d ASC;