dbms: development [#METR-10498]

Signed-off-by: Pavel Kartavyy <kartavyy@yandex-team.ru>

dbms/include/DB/Interpreters/ExpressionAnalyzer.h

@@ -101,6 +101,10 @@ public:
 
 	/// Все новые временные таблицы, полученные при выполнении подзапросов GLOBAL IN.
 	Tables external_tables;
+
+	/// создает Set, заданные явно
+	void makeExplicitSets();
+
 private:
 	typedef std::set<String> NamesSet;
 	
@@ -313,6 +317,9 @@ private:
 	void assertSelect();
 	void assertAggregation();
 	void assertArrayJoin();
+
+	void makeExplicitSet(ASTFunction * node, const Block & sample_block);
+	void makeExplicitSetsRecursively(ASTPtr & node, const Block & sample_block);
 };
 
 }

dbms/include/DB/Interpreters/Set.h

@@ -21,6 +21,7 @@
 #include <DB/Columns/ColumnArray.h>
 
 #include <DB/Storages/MergeTree/BoolMask.h>
+#include <DB/Storages/MergeTree/PKCondition.h>
 
 namespace DB
 {

dbms/include/DB/Storages/MergeTree/PKCondition.h

@@ -8,7 +8,6 @@
 #include <DB/Parsers/ASTSelectQuery.h>
 #include <DB/Parsers/ASTFunction.h>
 #include <DB/Parsers/ASTLiteral.h>
-#include <DB/Columns/ColumnSet.h>
 #include <DB/Storages/MergeTree/BoolMask.h>
 
 
@@ -347,7 +346,7 @@ private:
 		Range range;
 		size_t key_column;
 		/// Для FUNCTION_IN_SET
-		SetPtr set;
+		ASTPtr in_function;
 	};
 	
 	typedef std::vector<RPNElement> RPN;

dbms/src/Interpreters/ExpressionAnalyzer.cpp

@@ -498,6 +498,26 @@ void ExpressionAnalyzer::normalizeTreeImpl(ASTPtr & ast, MapOfASTs & finished_as
 	finished_asts[initial_ast] = ast;
 }
 
+void ExpressionAnalyzer::makeExplicitSets()
+{
+	makeExplicitSetsRecursively(ast, storage->getSampleBlock());
+}
+
+void ExpressionAnalyzer::makeExplicitSetsRecursively(ASTPtr & node, const Block & sample_block)
+{
+	for (auto & child : node->children)
+		makeExplicitSetsRecursively(child, sample_block);
+
+	ASTFunction * func = dynamic_cast<ASTFunction *>(node.get());
+	if (func && func->kind == ASTFunction::FUNCTION && (func->name == "in" || func->name == "notIn"))
+	{
+		IAST & args = *func->arguments;
+		ASTPtr & arg = args.children[1];
+
+		if (!dynamic_cast<ASTSet *>(&*arg) && !dynamic_cast<ASTSubquery *>(&*arg))
+			makeExplicitSet(func, sample_block);
+	}
+}
 
 void ExpressionAnalyzer::findGlobalFunctions(ASTPtr & ast, std::vector<ASTPtr> & global_nodes)
 {
@@ -679,7 +699,15 @@ void ExpressionAnalyzer::makeSet(ASTFunction * node, const Block & sample_block)
 	}
 	else
 	{
-		/// Случай явного перечисления значений.
+		makeExplicitSet(node, sample_block);
+	}
+}
+
+/// Случай явного перечисления значений.
+void ExpressionAnalyzer::makeExplicitSet(ASTFunction * node, const Block & sample_block)
+{
+		IAST & args = *node->arguments;
+		ASTPtr & arg = args.children[1];
 
 		DataTypes set_element_types;
 		ASTPtr & left_arg = args.children[0];
@@ -741,7 +769,6 @@ void ExpressionAnalyzer::makeSet(ASTFunction * node, const Block & sample_block)
 		ast_set->set = new Set(settings.limits);
 		ast_set->set->createFromAST(set_element_types, elements_ast);
 		arg = ast_set_ptr;
-	}
 }
 
 

dbms/src/Interpreters/InterpreterSelectQuery.cpp

@@ -501,6 +501,7 @@ QueryProcessingStage::Enum InterpreterSelectQuery::executeFetchColumns(BlockInpu
 
 	QueryProcessingStage::Enum from_stage = QueryProcessingStage::FetchColumns;
 	
+	query_analyzer->makeExplicitSets();
 	/// Инициализируем изначальные потоки данных, на которые накладываются преобразования запроса. Таблица или подзапрос?
 	if (!interpreter_subquery)
 	{

dbms/src/Storages/MergeTree/PKCondition.cpp

@@ -3,6 +3,7 @@
 #include <DB/Interpreters/ExpressionAnalyzer.h>
 #include <DB/Columns/ColumnSet.h>
 #include <DB/Columns/ColumnTuple.h>
+#include <DB/Parsers/ASTSet.h>
 
 namespace DB
 {
@@ -139,15 +140,11 @@ bool PKCondition::atomFromAST(ASTPtr & node, Block & block_with_constants, RPNEl
 			inverted = true;
 			column = pk_columns[args[1]->getColumnName()];
 		}
-		/// для In, notIn
+		else if (pk_columns.count(args[0]->getColumnName()) && dynamic_cast<ASTSet *>(args[1].get()))
-		else if (pk_columns.count(args[0]->getColumnName()) && dynamic_cast<DB::ColumnSet *>(args[1].get()))
 		{
-			/// для in не бывает inverted
 			inverted = false;
-			/// не поддерживаем Primary Key, если аргумент функции in tuple
-			if (dynamic_cast<DB::ColumnTuple*>(args[0].get()))
-				return false;
 			column = pk_columns[args[0]->getColumnName()];
+			dynamic_cast<ASTSet *>(args[1].get())->set->createOrderedSet();
 		}
 		else
 			return false;
@@ -188,8 +185,7 @@ bool PKCondition::atomFromAST(ASTPtr & node, Block & block_with_constants, RPNEl
 		else if (func->name == "in" || func->name == "notIn")
 		{
 			out.function = func->name == "in" ? RPNElement::FUNCTION_IN_SET : RPNElement::FUNCTION_NOT_IN_SET;
-			out.set = (dynamic_cast<DB::ColumnSet *>(args[1].get())->getData());
+			out.in_function = node;
-			out.set->createOrderedSet();
 		}
 		else
 			return false;
@@ -282,11 +278,21 @@ bool PKCondition::mayBeTrueInRange(const Field * left_pk, const Field * right_pk
 		}
 		else if (element.function == RPNElement::FUNCTION_IN_SET || element.function == RPNElement::FUNCTION_NOT_IN_SET)
 		{
-			const Range & key_range = key_ranges[element.key_column];
+			ASTFunction * in_func = dynamic_cast<ASTFunction *>(element.in_function.get());
+			ASTs & args = dynamic_cast<ASTExpressionList &>(*in_func->arguments).children;
+			ASTSet * ast_set = dynamic_cast<ASTSet *>(args[1].get());
+			if (in_func && ast_set)
+			{
+				const Range & key_range = key_ranges[element.key_column];
 
-			rpn_stack.push_back(element.set->mayBeTrueInRange(key_range.left, key_range.right));
+				rpn_stack.push_back(ast_set->set->mayBeTrueInRange(key_range.left, key_range.right));
-			if (element.function == RPNElement::FUNCTION_NOT_IN_SET)
+				if (element.function == RPNElement::FUNCTION_NOT_IN_SET)
-				rpn_stack.back() = !rpn_stack.back();
+					rpn_stack.back() = !rpn_stack.back();
+			}
+			else
+			{
+				throw DB::Exception("Set for IN is not created yet!");
+			}
 		}
 		else if (element.function == RPNElement::FUNCTION_NOT)
 		{