dbms: allowed constant expressions in IN [#METR-2944].

dbms/include/DB/Interpreters/Set.h

@@ -283,7 +283,7 @@ public:
 	  * node - это список значений: 1, 2, 3 или список tuple-ов: (1, 2), (3, 4), (5, 6).
 	  * create_ordered_set - создавать ли вектор упорядоченных элементов. Нужен для работы индекса
 	  */
-	void createFromAST(DataTypes & types, ASTPtr node, bool create_ordered_set);
+	void createFromAST(DataTypes & types, ASTPtr node, const Context & context, bool create_ordered_set);
 
 	// Возвращает false, если превышено какое-нибудь ограничение, и больше не нужно вставлять.
 	bool insertFromBlock(const Block & block, bool create_ordered_set = false);

dbms/src/Interpreters/ExpressionAnalyzer.cpp

@@ -954,7 +954,7 @@ void ExpressionAnalyzer::makeExplicitSet(ASTFunction * node, const Block & sampl
 	ASTPtr ast_set_ptr = ast_set;
 	ast_set->set = new Set(settings.limits);
 	ast_set->is_explicit = true;
-	ast_set->set->createFromAST(set_element_types, elements_ast, create_ordered_set);
+	ast_set->set->createFromAST(set_element_types, elements_ast, context, create_ordered_set);
 	arg = ast_set_ptr;
 }
 

dbms/src/Interpreters/Set.cpp

@@ -10,8 +10,12 @@
 #include <DB/Parsers/ASTExpressionList.h>
 #include <DB/Parsers/ASTFunction.h>
 #include <DB/Parsers/ASTLiteral.h>
+#include <DB/Parsers/formatAST.h>
 
 #include <DB/Interpreters/Set.h>
+#include <DB/Interpreters/ExpressionAnalyzer.h>
+#include <DB/Interpreters/ExpressionActions.h>
+
 #include <DB/DataTypes/DataTypeArray.h>
 #include <DB/DataTypes/DataTypesNumberFixed.h>
 #include <DB/DataTypes/DataTypeString.h>
@@ -259,10 +263,10 @@ static Field convertToType(const Field & src, const IDataType & type)
 		if (is_uint8 || is_uint16 || is_uint32 || is_uint64)
 		{
 			if (src.getType() == Field::Types::Int64)
-				throw Exception("Type mismatch in IN section: " + type.getName() + " at left, signed literal at right");
+				throw Exception("Type mismatch in IN section: " + type.getName() + " at left, signed at right");
 
 			if (src.getType() == Field::Types::Float64)
-				throw Exception("Type mismatch in IN section: " + type.getName() + " at left, floating point literal at right");
+				throw Exception("Type mismatch in IN section: " + type.getName() + " at left, floating point at right");
 
 			if (src.getType() == Field::Types::UInt64)
 			{
@@ -276,12 +280,12 @@ static Field convertToType(const Field & src, const IDataType & type)
 			}
 
 			throw Exception("Type mismatch in IN section: " + type.getName() + " at left, "
-				+ Field::Types::toString(src.getType()) + " literal at right");
+				+ Field::Types::toString(src.getType()) + " at right");
 		}
 		else if (is_int8 || is_int16 || is_int32 || is_int64)
 		{
 			if (src.getType() == Field::Types::Float64)
-				throw Exception("Type mismatch in IN section: " + type.getName() + " at left, floating point literal at right");
+				throw Exception("Type mismatch in IN section: " + type.getName() + " at left, floating point at right");
 
 			if (src.getType() == Field::Types::UInt64)
 			{
@@ -308,7 +312,7 @@ static Field convertToType(const Field & src, const IDataType & type)
 			}
 
 			throw Exception("Type mismatch in IN section: " + type.getName() + " at left, "
-				+ Field::Types::toString(src.getType()) + " literal at right");
+				+ Field::Types::toString(src.getType()) + " at right");
 		}
 		else if (is_float32 || is_float64)
 		{
@@ -322,7 +326,7 @@ static Field convertToType(const Field & src, const IDataType & type)
 				return src;
 
 			throw Exception("Type mismatch in IN section: " + type.getName() + " at left, "
-				+ Field::Types::toString(src.getType()) + " literal at right");
+				+ Field::Types::toString(src.getType()) + " at right");
 		}
 	}
 	else
@@ -337,22 +341,54 @@ static Field convertToType(const Field & src, const IDataType & type)
 			|| (src.getType() == Field::Types::Array
 				&& !typeid_cast<const DataTypeArray *>(&type)))
 			throw Exception("Type mismatch in IN section: " + type.getName() + " at left, "
-				+ Field::Types::toString(src.getType()) + " literal at right");
+				+ Field::Types::toString(src.getType()) + " at right");
 	}
 
 	return src;
 }
 
 
-void Set::createFromAST(DataTypes & types, ASTPtr node, bool create_ordered_set)
+/** Выполнить константное выражение (для элемента множества в IN). Весьма неоптимально. */
+static Field evaluateConstantExpression(ASTPtr & node, const Context & context)
+{
+	ExpressionActionsPtr expr_for_constant_folding = ExpressionAnalyzer(
+		node, context, NamesAndTypesList{{ "_dummy", new DataTypeUInt8 }}).getConstActions();
+
+	/// В блоке должен быть хотя бы один столбец, чтобы у него было известно число строк.
+	Block block_with_constants{{ new ColumnConstUInt8(1, 0), new DataTypeUInt8, "_dummy" }};
+
+	expr_for_constant_folding->execute(block_with_constants);
+
+	if (!block_with_constants || block_with_constants.rows() == 0)
+		throw Exception("Logical error: empty block after evaluation constant expression for IN", ErrorCodes::LOGICAL_ERROR);
+
+	String name = node->getColumnName();
+
+	if (!block_with_constants.has(name))
+		throw Exception("Element of set in IN is not a constant expression: " + name, ErrorCodes::BAD_ARGUMENTS);
+
+	const IColumn & result_column = *block_with_constants.getByName(name).column;
+
+	if (!result_column.isConst())
+		throw Exception("Element of set in IN is not a constant expression: " + name, ErrorCodes::BAD_ARGUMENTS);
+
+	return result_column[0];
+}
+
+
+static Field extractValueFromNode(ASTPtr & node, const IDataType & type, const Context & context)
 {
-	/** NOTE:
-	  * На данный момент в секции IN не поддерживаются выражения (вызовы функций), кроме кортежей.
-	  * То есть, например, не поддерживаются массивы. А по хорошему, хотелось бы поддерживать.
-	  * Для этого можно сделать constant folding с помощью ExpressionAnalyzer/ExpressionActions.
-	  * Но при этом, конечно же, не забыть про производительность работы с крупными множествами.
-	  */
+	if (ASTLiteral * lit = typeid_cast<ASTLiteral *>(node.get()))
+		return convertToType(lit->value, type);
+	else if (typeid_cast<ASTFunction *>(node.get()))
+		return convertToType(evaluateConstantExpression(node, context), type);
+	else
+		throw Exception("Incorrect element of set. Must be literal or constant expression.", ErrorCodes::INCORRECT_ELEMENT_OF_SET);
+}
 
+
+void Set::createFromAST(DataTypes & types, ASTPtr node, const Context & context, bool create_ordered_set)
+{
 	data_types = types;
 
 	/// Засунем множество в блок.
@@ -372,10 +408,7 @@ void Set::createFromAST(DataTypes & types, ASTPtr node, bool create_ordered_set)
 	{
 		if (data_types.size() == 1)
 		{
-			if (ASTLiteral * lit = typeid_cast<ASTLiteral *>(&**it))
-				block.getByPosition(0).column->insert(convertToType(lit->value, *data_types[0]));
-			else
-				throw Exception("Incorrect element of set. Must be literal.", ErrorCodes::INCORRECT_ELEMENT_OF_SET);
+			block.getByPosition(0).column->insert(extractValueFromNode(*it, *data_types[0], context));
 		}
 		else if (ASTFunction * func = typeid_cast<ASTFunction *>(&**it))
 		{
@@ -388,16 +421,11 @@ void Set::createFromAST(DataTypes & types, ASTPtr node, bool create_ordered_set)
 
 			for (size_t j = 0; j < tuple_size; ++j)
 			{
-				if (ASTLiteral * lit = typeid_cast<ASTLiteral *>(&*func->arguments->children[j]))
-					block.getByPosition(j).column->insert(convertToType(lit->value, *data_types[j]));
-				else
-					throw Exception("Incorrect element of tuple in set. Must be literal.", ErrorCodes::INCORRECT_ELEMENT_OF_SET);
+				block.getByPosition(j).column->insert(extractValueFromNode(func->arguments->children[j], *data_types[j], context));
 			}
 		}
 		else
 			throw Exception("Incorrect element of set", ErrorCodes::INCORRECT_ELEMENT_OF_SET);
-
-		/// NOTE: Потом можно реализовать возможность задавать константные выражения в множествах.
 	}
 
 	if (create_ordered_set)