dbms: development [#METR-10498]

dbms/include/DB/Core/Field.h

@@ -352,7 +352,36 @@ private:
 	}
 };
 
-
+static std::ostream & operator<< (std::ostream & os, const Field & x)
+{
+	bool first = false;
+	switch (x.getType())
+	{
+		case Field::Types::Null: 			os << "Null";						break;
+		case Field::Types::UInt64: 			os << x.get<UInt64>();				break;
+		case Field::Types::Int64: 			os << x.get<Int64>();				break;
+		case Field::Types::Float64: 		os << x.get<Float64>();				break;
+		case Field::Types::String: 			os << x.get<String>();				break;
+		case Field::Types::Array: 				
+			for(const Field & f : x.get<Array>())
+			{
+				os << "[";
+				if (first)
+					os << f;
+				else
+					os << ", " << f;
+				os << "]";
+			}
+			break;
+		default: 
+		{
+			std::stringstream ss;
+			ss << "Unsupported type " << x.getTypeName();
+			throw DB::Exception(ss.str());
+		}
+	}
+	return os;
+}
 template <> struct Field::TypeToEnum<Null> 								{ static const Types::Which value = Types::Null; };
 template <> struct Field::TypeToEnum<UInt64> 							{ static const Types::Which value = Types::UInt64; };
 template <> struct Field::TypeToEnum<Int64> 							{ static const Types::Which value = Types::Int64; };

dbms/include/DB/Interpreters/ExpressionAnalyzer.h

@@ -102,9 +102,8 @@ public:
 	/// Все новые временные таблицы, полученные при выполнении подзапросов GLOBAL IN.
 	Tables external_tables;
 
-	/// создает Set, заданные явно
-	void makeExplicitSets();
-
+	/// ordered_set нужен если в In используется Primary Key
+	void makeExplicitSets(bool create_ordered_set);
 private:
 	typedef std::set<String> NamesSet;
 	
@@ -318,8 +317,8 @@ private:
 	void assertAggregation();
 	void assertArrayJoin();
 
-	void makeExplicitSet(ASTFunction * node, const Block & sample_block);
-	void makeExplicitSetsRecursively(ASTPtr & node, const Block & sample_block);
+	void makeExplicitSet(ASTFunction * node, const Block & sample_block, bool create_ordered_set);
+	void makeExplicitSetsRecursively(ASTPtr & node, const Block & sample_block, bool create_ordered_set);
 };
 
 }

dbms/include/DB/Interpreters/Set.h

@@ -46,7 +46,7 @@ public:
 	  * types - типы того, что стоит слева от IN.
 	  * node - это список значений: 1, 2, 3 или список tuple-ов: (1, 2), (3, 4), (5, 6).
 	  */
-	void createFromAST(DataTypes & types, ASTPtr node);
+	void createFromAST(DataTypes & types, ASTPtr node, bool create_ordered_set);
 
 	/** Запомнить поток блоков (для подзапросов), чтобы потом его можно было прочитать и создать множество.
 	  */
@@ -55,7 +55,7 @@ public:
 	BlockInputStreamPtr getSource() { return source; }
 
 	// Возвращает false, если превышено какое-нибудь ограничение, и больше не нужно вставлять.
-	bool insertFromBlock(Block & block);
+	bool insertFromBlock(Block & block, bool create_ordered_set = false);
 
 	size_t size() const { return getTotalRowCount(); }
 
@@ -66,41 +66,25 @@ public:
 
 	std::string descibe()
 	{
-		if (type == KEY_64)
-			return setToString(key64);
-		else if (type == KEY_STRING)
-			return setToString(key_string);
-		else if (type == HASHED)
-			return "{hashed values}";
-		else if (type == EMPTY)
+		if (!ordered_set)
 			return "{}";
-		else
-			throw DB::Exception("Unknown type");
-	}
-	
-	void createOrderedSet()
-	{
-		for (auto & key : key64)
-			ordered_key64.push_back(key);
-		std::sort(ordered_key64.begin(), ordered_key64.end());
-
-		for (auto & key : key_string)
-			ordered_string.push_back(key);
-		std::sort(ordered_string.begin(), ordered_string.end());
-	}
-
-	BoolMask mayBeTrueInRange(const Field & left, const Field & right)
-	{
-		if (type == KEY_64)
-			return mayBeTrueInRangeImpl(left, right, ordered_key64);
-		else if (type == KEY_STRING)
-			return mayBeTrueInRangeImpl(left, right, ordered_string);
-		else{
-			std::stringstream ss;
-			ss << "Unsupported set of type " << type;
-			throw DB::Exception(ss.str());
+		
+		bool first = true;
+		std::stringstream ss;
+		
+		for (const Field & f : *ordered_set)
+		{
+			if (!first)
+				ss << ", " << f;
+			else
+				ss << f;
+			first = false;
 		}
+		return ss.str();
 	}
+
+	BoolMask mayBeTrueInRange(const Range & range);
+	
 private:
 	/** Разные структуры данных, которые могут использоваться для проверки принадлежности
 	  *  одного или нескольких столбцов значений множеству.
@@ -166,91 +150,9 @@ private:
 	/// Считает суммарный размер в байтах буфферов всех Set'ов + размер string_pool'а
 	size_t getTotalByteCount() const;
 
-	template <typename T>
-	std::string setToString(const T & set)
-	{
-		std::stringstream ss;
-		bool first = false;
-
-		for (auto it = set.begin(); it != set.end(); ++it)
-		{
-			if (first)
-			{
-				ss << *it;
-				first = false;
-			}
-			else
-			{
-				ss << ", " << *it;
-			}
-		}
-
-		ss << "}";
-		return ss.str();
-	}
-
-	/// несколько столбцов пока не поддерживаем
-	std::vector<UInt64> ordered_key64;
-	std::vector<StringRef> ordered_string;
-
-	template <class T>
-	BoolMask mayBeTrueInRangeImpl(const Field & field_left, const Field & field_right, const std::vector<T> & v)
-	{
-		T left = field_left.get<T>();
-		T right = field_right.get<T>();
-
-		bool can_be_true;
-		bool can_be_false = true;
-
-		/// Если во всем диапазоне одинаковый ключ и он есть в Set, то выбираем блок для in и не выбираем для notIn
-		if (left == right)
-		{
-			if (std::find(v.begin(), v.end(), left) != v.end())
-			{
-				can_be_false = false;
-				can_be_true = true;
-			}
-			else
-			{
-				can_be_true = false;
-				can_be_false = true;
-			}
-		}
-		else
-		{
-			auto left_it = std::lower_bound(v.begin(), v.end(), left);
-			/// если весь диапазон, правее in
-			if (left_it == v.end())
-			{
-				can_be_true = false;
-			}
-			else
-			{
-				auto right_it = std::upper_bound(v.begin(), v.end(), right);
-				/// весь диапазон, левее in
-				if (right_it == v.begin())
-				{
-					can_be_true = false;
-				}
-				else
-				{
-					--right_it;
-					/// в диапазон не попадает ни одного ключа из in
-					if (*right_it < *left_it)
-					{
-						can_be_true = false;
-					}
-					else
-					{
-						can_be_true = true;
-					}
-				}
-			}
-		}
-
-		return BoolMask(can_be_true, can_be_false);
-	}
-
+	typedef std::vector<Field> OrderedSet;
+	typedef std::unique_ptr<OrderedSet> OrderedSetPtr;
+	OrderedSetPtr ordered_set;
 };
 
 typedef Poco::SharedPtr<Set> SetPtr;

dbms/src/Interpreters/ExpressionAnalyzer.cpp

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

dbms/src/Interpreters/InterpreterSelectQuery.cpp

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

dbms/src/Interpreters/Set.cpp

@@ -82,13 +82,12 @@ Set::Type Set::chooseMethod(const ConstColumnPlainPtrs & key_columns, bool & key
 		|| dynamic_cast<const ColumnConstString *>(key_columns[0])
 		|| (dynamic_cast<const ColumnFixedString *>(key_columns[0]) && !keys_fit_128_bits)))
 		return KEY_STRING;
-
 	/// Если много ключей - будем строить множество хэшей от них
 	return HASHED;
 }
 
 
-bool Set::insertFromBlock(Block & block)
+bool Set::insertFromBlock(Block & block, bool create_ordered_set)
 {
 	size_t keys_size = block.columns();
 	Row key(keys_size);
@@ -119,6 +118,9 @@ bool Set::insertFromBlock(Block & block)
 			/// Строим ключ
 			UInt64 key = get<UInt64>(column[i]);
 			res.insert(key);
+
+			if(create_ordered_set)
+				ordered_set->push_back(column[i]);
 		}
 	}
 	else if (type == KEY_STRING)
@@ -143,6 +145,9 @@ bool Set::insertFromBlock(Block & block)
 
 				if (inserted)
 					it->data = string_pool.insert(ref.data, ref.size);
+
+				if(create_ordered_set)
+					ordered_set->push_back(std::string(ref.data, ref.size));
 			}
 		}
 		else if (const ColumnFixedString * column_string = dynamic_cast<const ColumnFixedString *>(&column))
@@ -162,6 +167,9 @@ bool Set::insertFromBlock(Block & block)
 
 				if (inserted)
 					it->data = string_pool.insert(ref.data, ref.size);
+
+				if(create_ordered_set)
+					ordered_set->push_back(std::string(ref.data, ref.size));
 			}
 		}
 		else
@@ -198,7 +206,7 @@ bool Set::insertFromBlock(Block & block)
 }
 
 
-void Set::createFromAST(DataTypes & types, ASTPtr node)
+void Set::createFromAST(DataTypes & types, ASTPtr node, bool create_ordered_set)
 {
 	data_types = types;
 
@@ -247,7 +255,13 @@ void Set::createFromAST(DataTypes & types, ASTPtr node)
 		/// NOTE: Потом можно реализовать возможность задавать константные выражения в множествах.
 	}
 
-	insertFromBlock(block);
+	if (create_ordered_set)
+		ordered_set = OrderedSetPtr(new std::vector<Field>());
+
+	insertFromBlock(block, create_ordered_set);
+
+	if (create_ordered_set)
+		std::sort(ordered_set->begin(), ordered_set->end());
 }
 
 
@@ -528,4 +542,73 @@ void Set::executeConstArray(const ColumnConstArray * key_column, ColumnUInt8::Co
 		vec_res[i] = res;
 }
 
+BoolMask Set::mayBeTrueInRange(const Range & range)
+{
+	if (!ordered_set)
+		throw DB::Exception("Ordered set in not created.");
+	
+	if (ordered_set->empty())
+		return BoolMask(false, true);
+
+	const Field & left = range.left;
+	const Field & right = range.right;
+
+	bool can_be_true;
+	bool can_be_false = true;
+
+	/// Если во всем диапазоне одинаковый ключ и он есть в Set, то выбираем блок для in и не выбираем для notIn
+	if (range.left_bounded && range.right_bounded && range.right_included && range.left_included && left == right)
+	{
+		if (std::find(ordered_set->begin(), ordered_set->end(), left) != ordered_set->end())
+		{
+			can_be_false = false;
+			can_be_true = true;
+		}
+		else
+		{
+			can_be_true = false;
+			can_be_false = true;
+		}
+	}
+	else
+	{
+		auto left_it = range.left_bounded ? std::lower_bound(ordered_set->begin(), ordered_set->end(), left) : ordered_set->begin();
+		if (!range.left_included && left_it != ordered_set->end() && *left_it == left)
+			++left_it;
+
+		/// если весь диапазон, правее in
+		if (left_it == ordered_set->end())
+		{
+			can_be_true = false;
+		}
+		else
+		{
+			auto right_it = range.right_bounded ? std::upper_bound(ordered_set->begin(), ordered_set->end(), right) : ordered_set->end();
+			if (!range.right_included && right_it != ordered_set->begin() && *(right_it--) == right)
+				--right_it;
+
+			/// весь диапазон, левее in
+			if (right_it == ordered_set->begin())
+			{
+				can_be_true = false;
+			}
+			else
+			{
+				--right_it;
+				/// в диапазон не попадает ни одного ключа из in
+				if (*right_it < *left_it)
+				{
+					can_be_true = false;
+				}
+				else
+				{
+					can_be_true = true;
+				}
+			}
+		}
+	}
+
+	return BoolMask(can_be_true, can_be_false);
+}
+
 }

dbms/src/Storages/MergeTree/PKCondition.cpp

@@ -144,7 +144,6 @@ bool PKCondition::atomFromAST(ASTPtr & node, Block & block_with_constants, RPNEl
 		{
 			inverted = false;
 			column = pk_columns[args[0]->getColumnName()];
-			dynamic_cast<ASTSet *>(args[1].get())->set->createOrderedSet();
 		}
 		else
 			return false;
@@ -285,7 +284,7 @@ bool PKCondition::mayBeTrueInRange(const Field * left_pk, const Field * right_pk
 			{
 				const Range & key_range = key_ranges[element.key_column];
 
-				rpn_stack.push_back(ast_set->set->mayBeTrueInRange(key_range.left, key_range.right));
+				rpn_stack.push_back(ast_set->set->mayBeTrueInRange(key_range));
 				if (element.function == RPNElement::FUNCTION_NOT_IN_SET)
 					rpn_stack.back() = !rpn_stack.back();
 			}