Merge pull request #189 from yandex/has-column-in-table

METR-23248 Has column in table

dbms/include/DB/Functions/FunctionsMiscellaneous.h

@@ -1396,4 +1396,32 @@ public:
 	}
 };
 
+
+/** Usage:
+ *  hasColumnInTable('database', 'table', 'column')
+ */
+class FunctionHasColumnInTable : public IFunction
+{
+public:
+	static constexpr auto name = "hasColumnInTable";
+	static constexpr size_t number_of_arguments = 3;
+
+	static FunctionPtr create(const Context & context) { return std::make_shared<FunctionHasColumnInTable>(context.getGlobalContext()); }
+
+	FunctionHasColumnInTable(const Context & global_context_)
+	: global_context(global_context_)
+	{
+	}
+
+	String getName() const override { return name; }
+
+	void getReturnTypeAndPrerequisites(const ColumnsWithTypeAndName & arguments,
+										DataTypePtr & out_return_type,
+										ExpressionActions::Actions & out_prerequisites) override;
+	void execute(Block & block, const ColumnNumbers & arguments, size_t result) override;
+
+private:
+	const Context & global_context;
+};
+
 }

dbms/include/DB/Interpreters/Context.h

@@ -204,7 +204,10 @@ public:
 	/// Для методов ниже может быть необходимо захватывать блокировку самостоятельно.
 	std::unique_lock<Poco::Mutex> getLock() const;
 
+	const Context & getSessionContext() const;
 	Context & getSessionContext();
+
+	const Context & getGlobalContext() const;
 	Context & getGlobalContext();
 
 	void setSessionContext(Context & context_)								{ session_context = &context_; }

dbms/src/Functions/FunctionsMiscellaneous.cpp

@@ -380,6 +380,54 @@ void FunctionVisibleWidth::execute(Block & block, const ColumnNumbers & argument
 }
 
 
+void FunctionHasColumnInTable::getReturnTypeAndPrerequisites(
+	const ColumnsWithTypeAndName & arguments,
+	DataTypePtr & out_return_type,
+	ExpressionActions::Actions & out_prerequisites)
+{
+	if (arguments.size() != number_of_arguments)
+		throw Exception("Function " + getName() + " requires exactly three arguments.", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+
+
+	static const std::string arg_pos_description[] = {"First", "Second", "Third"};
+	for (size_t i = 0; i < number_of_arguments; ++i)
+	{
+		const ColumnWithTypeAndName & argument = arguments[i];
+
+		const ColumnConstString * column = typeid_cast<const ColumnConstString *>(argument.column.get());
+		if (!column)
+		{
+			throw Exception(
+				arg_pos_description[i] + " argument for function " + getName() + " must be const String.",
+				ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+		}
+	}
+
+	out_return_type = std::make_shared<DataTypeUInt8>();
+}
+
+
+void FunctionHasColumnInTable::execute(Block & block, const ColumnNumbers & arguments, size_t result)
+{
+	auto get_string_from_block =
+		[&](size_t column_pos) -> const String &
+		{
+			DB::ColumnPtr column = block.getByPosition(column_pos).column;
+			const ColumnConstString * const_column = typeid_cast<const ColumnConstString *>(column.get());
+			return const_column->getData();
+		};
+
+	const DB::String & database_name = get_string_from_block(arguments[0]);
+	const DB::String & table_name = get_string_from_block(arguments[1]);
+	const DB::String & column_name = get_string_from_block(arguments[2]);
+
+	const DB::StoragePtr & table = global_context.getTable(database_name, table_name);
+	const bool has_column = table->hasColumn(column_name);
+
+	block.getByPosition(result).column = std::make_shared<ColumnConstUInt8>(
+		block.rowsInFirstColumn(), has_column);
+}
+
 }
 
 
@@ -405,6 +453,7 @@ void registerFunctionsMiscellaneous(FunctionFactory & factory)
 	factory.registerFunction<FunctionArrayJoin>();
 	factory.registerFunction<FunctionReplicate>();
 	factory.registerFunction<FunctionBar>();
+	factory.registerFunction<FunctionHasColumnInTable>();
 
 	factory.registerFunction<FunctionTuple>();
 	factory.registerFunction<FunctionTupleElement>();

dbms/src/Interpreters/Context.cpp

@@ -725,6 +725,12 @@ void Context::setMacros(Macros && macros)
 	shared->macros = macros;
 }
 
+const Context & Context::getSessionContext() const
+{
+	if (!session_context)
+		throw Exception("There is no session", ErrorCodes::THERE_IS_NO_SESSION);
+	return *session_context;
+}
 
 Context & Context::getSessionContext()
 {
@@ -733,6 +739,12 @@ Context & Context::getSessionContext()
 	return *session_context;
 }
 
+const Context & Context::getGlobalContext() const
+{
+	if (!global_context)
+		throw Exception("Logical error: there is no global context", ErrorCodes::LOGICAL_ERROR);
+	return *global_context;
+}
 
 Context & Context::getGlobalContext()
 {

dbms/tests/clickhouse-test

@@ -15,28 +15,14 @@ from subprocess import CalledProcessError
 from datetime import datetime
 from time import sleep
 from errno import ESRCH
+from termcolor import colored
 
-if sys.stdout.isatty():
-	COLORS = {
-		"COLOR_RESET": "\033[0m",
-		"COLOR_WHITE": "\033[1;37m",
-		"COLOR_FAIL": "\033[1;31m",
-		"COLOR_UNKNOWN": "\033[1;30m",
-		"COLOR_OK": "\033[1;32m",
-		"COLOR_SKIPPED": "\033[1;34m" }
-else:
-	COLORS = {
-		"COLOR_RESET": "",
-		"COLOR_WHITE": "",
-		"COLOR_FAIL": "",
-		"COLOR_UNKNOWN": "",
-		"COLOR_OK": "",
-		"COLOR_SKIPPED": "" }
 
-MSG_FAIL = "{COLOR_WHITE}[ {COLOR_FAIL}FAIL{COLOR_WHITE} ]{COLOR_RESET}".format(**COLORS)
-MSG_UNKNOWN = "{COLOR_WHITE}[ {COLOR_UNKNOWN}UNKNOWN{COLOR_WHITE} ]{COLOR_RESET}".format(**COLORS)
-MSG_OK = "{COLOR_WHITE}[ {COLOR_OK}OK{COLOR_WHITE} ]{COLOR_RESET}".format(**COLORS)
-MSG_SKIPPED = "{COLOR_WHITE}[ {COLOR_SKIPPED}SKIPPED{COLOR_WHITE} ]{COLOR_RESET}".format(**COLORS)
+MSG_FAIL = "[ " + colored("FAIL", "red") + " ]"
+MSG_UNKNOWN = "[ " + colored("UNKNOWN", "yellow") + " ]"
+MSG_OK = "[ " + colored("OK", "green") + " ]"
+MSG_SKIPPED = "[ " + colored("SKIPPED", "cyan") + " ]"
+
 
 def main(args):
 
@@ -107,7 +93,7 @@ def main(args):
 					
 					if not args.zookeeper and 'zookeeper' in name:
 						report_testcase.append(et.Element("skipped", attrib = {"message": "no zookeeper"}))
-						print(MSG_SKIPPED, " - no zookeeper")
+						print(MSG_SKIPPED + " - no zookeeper")
 					else:
 						reference_file = os.path.join(suite_dir, name) + '.reference'
 						stdout_file = os.path.join(suite_dir, name) + '.stdout'
@@ -137,7 +123,9 @@ def main(args):
 							print("{0} - Timeout!".format(MSG_FAIL))
 						else:
 							stdout = open(stdout_file, 'r').read() if os.path.exists(stdout_file) else ''
+							stdout = unicode(stdout, errors='replace', encoding='utf-8')
 							stderr = open(stderr_file, 'r').read() if os.path.exists(stderr_file) else ''
+							stderr = unicode(stderr, errors='replace', encoding='utf-8')
 							
 							if proc.returncode != 0:
 								failure = et.Element("failure", attrib = {"message": "return code {}".format(proc.returncode)})
@@ -208,10 +196,10 @@ def main(args):
 			failures_total = failures_total + failures
 	
 	if failures_total > 0:
-		print("\n{COLOR_FAIL}Having {0} errors!{COLOR_RESET}".format(failures_total, **COLORS))
+		print(colored("\nHaving {0} errors!".format(failures_total), "red"))
 		sys.exit(1)
 	else:
-		print("\n{COLOR_OK}All tests passed.{COLOR_RESET}".format(**COLORS))
+		print(colored("\nAll tests passed.", "green"))
 		sys.exit(0)
 			
 

dbms/tests/queries/0_stateless/00386_has_column_in_table.reference

@@ -0,0 +1,7 @@
+1
+1
+1
+1
+0
+0
+0

dbms/tests/queries/0_stateless/00386_has_column_in_table.sql

@@ -0,0 +1,16 @@
+CREATE DATABASE IF NOT EXISTS test;
+DROP TABLE IF EXISTS test.has_column_in_table;
+CREATE TABLE test.has_column_in_table (i Int64, s String, nest Nested(x UInt8, y UInt32)) ENGINE = Memory;
+
+/* existing column */
+SELECT hasColumnInTable('test', 'has_column_in_table', 'i');
+SELECT hasColumnInTable('test', 'has_column_in_table', 's');
+SELECT hasColumnInTable('test', 'has_column_in_table', 'nest.x');
+SELECT hasColumnInTable('test', 'has_column_in_table', 'nest.y');
+
+/* not existing column */
+SELECT hasColumnInTable('test', 'has_column_in_table', 'nest');
+SELECT hasColumnInTable('test', 'has_column_in_table', 'nest.not_existing');
+SELECT hasColumnInTable('test', 'has_column_in_table', 'not_existing');
+
+DROP TABLE test.has_column_in_table;

doc/reference_en.html

@@ -5727,6 +5727,12 @@ Note that 0 is returned for a NaN.
 ===isNaN(x)===
 Accepts Float32 and Float64 and returns UInt8 equal to 1 if the argument is a NaN, otherwise 0.
 
+===hasColumnInTable('database', 'table', 'column')===
+Accepts constant String columns - database name, table name and column name. Returns constant UInt8 value, equal to 1 if column exists,
+otherwise 0.
+If table doesn't exist than exception is thrown.
+For elements of nested data structure function checks existence of column. For nested data structure 0 is returned.
+
 ===bar===
 Allows building a unicode-art diagram.
 

doc/reference_ru.html

@@ -5831,6 +5831,12 @@ N - индекс столбца начиная с 1. N должно быть к
 ===isNaN(x)===
 Принимает Float32 или Float64 и возвращает UInt8, равный 1, если аргумент является NaN, иначе 0.
 
+===hasColumnInTable('database', 'table', 'column')===
+Принимает константные строки - имя базы данных, имя таблицы и название столбца. Возвращает константное выражение типа UInt8, равное 1,
+если есть столбец, иначе 0.
+Функция кидает исключение, если таблица не существует.
+Для элементов вложенной структуры данных функция проверяет сушествование столбца. Для самой же вложенной структуры данных функция возвращает 0.
+
 ===bar===
 Позволяет построить unicode-art диаграмму.