From 47e19a4cd96fa50cd69848a398717b8b469e22e3 Mon Sep 17 00:00:00 2001 From: Anton Popov Date: Wed, 10 Nov 2021 22:19:18 +0300 Subject: [PATCH] optimize tupleElement to reading of subcolumn --- src/DataTypes/DataTypeTuple.cpp | 9 ++++ src/DataTypes/DataTypeTuple.h | 1 + .../RewriteFunctionToSubcolumnVisitor.cpp | 38 ++++++++++++++--- .../0_stateless/02116_tuple_element.reference | 25 +++++++++++ .../0_stateless/02116_tuple_element.sql | 42 +++++++++++++++++++ 5 files changed, 109 insertions(+), 6 deletions(-) create mode 100644 tests/queries/0_stateless/02116_tuple_element.reference create mode 100644 tests/queries/0_stateless/02116_tuple_element.sql diff --git a/src/DataTypes/DataTypeTuple.cpp b/src/DataTypes/DataTypeTuple.cpp index 4e1a5a05d45..0660f371258 100644 --- a/src/DataTypes/DataTypeTuple.cpp +++ b/src/DataTypes/DataTypeTuple.cpp @@ -29,6 +29,7 @@ namespace ErrorCodes extern const int NOT_FOUND_COLUMN_IN_BLOCK; extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; extern const int SIZES_OF_COLUMNS_IN_TUPLE_DOESNT_MATCH; + extern const int ILLEGAL_INDEX; } @@ -193,6 +194,14 @@ size_t DataTypeTuple::getPositionByName(const String & name) const throw Exception("Tuple doesn't have element with name '" + name + "'", ErrorCodes::NOT_FOUND_COLUMN_IN_BLOCK); } +String DataTypeTuple::getNameByPosition(size_t i) const +{ + if (i == 0 || i > names.size()) + throw Exception(ErrorCodes::ILLEGAL_INDEX, "Index of tuple element ({}) if out range ([1, {}])", i, names.size()); + + return names[i - 1]; +} + bool DataTypeTuple::textCanContainOnlyValidUTF8() const { diff --git a/src/DataTypes/DataTypeTuple.h b/src/DataTypes/DataTypeTuple.h index 8dae8b7765b..d168d73efbf 100644 --- a/src/DataTypes/DataTypeTuple.h +++ b/src/DataTypes/DataTypeTuple.h @@ -61,6 +61,7 @@ public: const Strings & getElementNames() const { return names; } size_t getPositionByName(const String & name) const; + String getNameByPosition(size_t i) const; bool haveExplicitNames() const { return have_explicit_names; } bool serializeNames() const { return serialize_names; } diff --git a/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp b/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp index debbd9ae8fe..5128bb136c3 100644 --- a/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp +++ b/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include @@ -79,7 +80,8 @@ void RewriteFunctionToSubcolumnData::visit(ASTFunction & function, ASTPtr & ast) if (!columns.has(name_in_storage)) return; - TypeIndex column_type_id = columns.get(name_in_storage).type->getTypeId(); + const auto & column_type = columns.get(name_in_storage).type; + TypeIndex column_type_id = column_type->getTypeId(); if (arguments.size() == 1) { @@ -93,12 +95,36 @@ void RewriteFunctionToSubcolumnData::visit(ASTFunction & function, ASTPtr & ast) } else { - auto it = binary_function_to_subcolumn.find(function.name); - if (it != binary_function_to_subcolumn.end()) + if (function.name == "tupleElement" && column_type_id == TypeIndex::Tuple) { - const auto & [type_id, subcolumn_name, transformer] = it->second; - if (column_type_id == type_id) - ast = transformer(name_in_storage, subcolumn_name, arguments[1]); + const auto * literal = arguments[1]->as(); + if (!literal) + return; + + String subcolumn_name; + auto value_type = literal->value.getType(); + if (value_type == Field::Types::UInt64) + { + const auto & type_tuple = assert_cast(*column_type); + auto index = get(literal->value); + subcolumn_name = type_tuple.getNameByPosition(index); + } + else if (value_type == Field::Types::String) + subcolumn_name = get(literal->value); + else + return; + + ast = transformToSubcolumn(name_in_storage, subcolumn_name); + } + else + { + auto it = binary_function_to_subcolumn.find(function.name); + if (it != binary_function_to_subcolumn.end()) + { + const auto & [type_id, subcolumn_name, transformer] = it->second; + if (column_type_id == type_id) + ast = transformer(name_in_storage, subcolumn_name, arguments[1]); + } } } } diff --git a/tests/queries/0_stateless/02116_tuple_element.reference b/tests/queries/0_stateless/02116_tuple_element.reference new file mode 100644 index 00000000000..121b08d02f1 --- /dev/null +++ b/tests/queries/0_stateless/02116_tuple_element.reference @@ -0,0 +1,25 @@ +1 +SELECT `t1.a` +FROM t_tuple_element +a +SELECT `t1.s` +FROM t_tuple_element +1 +SELECT `t1.a` +FROM t_tuple_element +2 +SELECT `t2.1` +FROM t_tuple_element +2 +SELECT `t2.1` +FROM t_tuple_element +1 2 +WITH (1, 2) AS t +SELECT + t.1, + t.2 +1 2 +WITH CAST(\'(1, 2)\', \'Tuple(a UInt32, b UInt32)\') AS t +SELECT + t.1, + tupleElement(t, \'b\') diff --git a/tests/queries/0_stateless/02116_tuple_element.sql b/tests/queries/0_stateless/02116_tuple_element.sql new file mode 100644 index 00000000000..4ce6e5cf136 --- /dev/null +++ b/tests/queries/0_stateless/02116_tuple_element.sql @@ -0,0 +1,42 @@ +DROP TABLE IF EXISTS t_tuple_element; + +CREATE TABLE t_tuple_element(t1 Tuple(a UInt32, s String), t2 Tuple(UInt32, String)) ENGINE = Memory; +INSERT INTO t_tuple_element VALUES ((1, 'a'), (2, 'b')); + +SET optimize_functions_to_subcolumns = 1; + +SELECT t1.1 FROM t_tuple_element; +EXPLAIN SYNTAX SELECT t1.1 FROM t_tuple_element; + +SELECT tupleElement(t1, 2) FROM t_tuple_element; +EXPLAIN SYNTAX SELECT tupleElement(t1, 2) FROM t_tuple_element; + +SELECT tupleElement(t1, 'a') FROM t_tuple_element; +EXPLAIN SYNTAX SELECT tupleElement(t1, 'a') FROM t_tuple_element; + +SELECT tupleElement(number, 1) FROM numbers(1); -- { serverError 43 } +SELECT tupleElement(t1) FROM t_tuple_element; -- { serverError 42 } +SELECT tupleElement(t1, 'b') FROM t_tuple_element; -- { serverError 47 } +SELECT tupleElement(t1, 0) FROM t_tuple_element; -- { serverError 127 } +SELECT tupleElement(t1, 3) FROM t_tuple_element; -- { serverError 127 } +SELECT tupleElement(t1, materialize('a')) FROM t_tuple_element; -- { serverError 43 } + +SELECT t2.1 FROM t_tuple_element; +EXPLAIN SYNTAX SELECT t2.1 FROM t_tuple_element; + +SELECT tupleElement(t2, 1) FROM t_tuple_element; +EXPLAIN SYNTAX SELECT tupleElement(t2, 1) FROM t_tuple_element; + +SELECT tupleElement(t2) FROM t_tuple_element; -- { serverError 42 } +SELECT tupleElement(t2, 'a') FROM t_tuple_element; -- { serverError 47 } +SELECT tupleElement(t2, 0) FROM t_tuple_element; -- { serverError 127 } +SELECT tupleElement(t2, 3) FROM t_tuple_element; -- { serverError 127 } +SELECT tupleElement(t2, materialize(1)) FROM t_tuple_element; -- { serverError 43 } + +DROP TABLE t_tuple_element; + +WITH (1, 2) AS t SELECT t.1, t.2; +EXPLAIN SYNTAX WITH (1, 2) AS t SELECT t.1, t.2; + +WITH (1, 2)::Tuple(a UInt32, b UInt32) AS t SELECT t.1, tupleElement(t, 'b'); +EXPLAIN SYNTAX WITH (1, 2)::Tuple(a UInt32, b UInt32) AS t SELECT t.1, tupleElement(t, 'b');