From 3145ca999db42f6014d979c49f959193a94dc38b Mon Sep 17 00:00:00 2001 From: Amos Bird Date: Mon, 4 Jan 2021 00:45:27 +0800 Subject: [PATCH] column excepts with regex --- src/Parsers/ASTColumnsTransformers.cpp | 62 +++++++++++++++---- src/Parsers/ASTColumnsTransformers.h | 9 +++ src/Parsers/ExpressionElementParsers.cpp | 11 +++- .../01470_columns_transformers2.reference | 1 + .../01470_columns_transformers2.sql | 7 +++ 5 files changed, 74 insertions(+), 16 deletions(-) create mode 100644 tests/queries/0_stateless/01470_columns_transformers2.reference create mode 100644 tests/queries/0_stateless/01470_columns_transformers2.sql diff --git a/src/Parsers/ASTColumnsTransformers.cpp b/src/Parsers/ASTColumnsTransformers.cpp index 9ec50a09cf1..024c9c3b767 100644 --- a/src/Parsers/ASTColumnsTransformers.cpp +++ b/src/Parsers/ASTColumnsTransformers.cpp @@ -6,6 +6,7 @@ #include #include #include +#include namespace DB @@ -14,6 +15,7 @@ namespace ErrorCodes { extern const int ILLEGAL_TYPE_OF_ARGUMENT; extern const int NO_SUCH_COLUMN_IN_TABLE; + extern const int CANNOT_COMPILE_REGEXP; } void IASTColumnsTransformer::transform(const ASTPtr & transformer, ASTs & nodes) @@ -86,6 +88,9 @@ void ASTColumnsExceptTransformer::formatImpl(const FormatSettings & settings, Fo (*it)->formatImpl(settings, state, frame); } + if (!original_pattern.empty()) + settings.ostr << quoteString(original_pattern); + if (children.size() > 1) settings.ostr << ")"; } @@ -93,24 +98,40 @@ void ASTColumnsExceptTransformer::formatImpl(const FormatSettings & settings, Fo void ASTColumnsExceptTransformer::transform(ASTs & nodes) const { std::set expected_columns; - for (const auto & child : children) - expected_columns.insert(child->as().name()); - - for (auto it = nodes.begin(); it != nodes.end();) + if (original_pattern.empty()) { - if (const auto * id = it->get()->as()) + for (const auto & child : children) + expected_columns.insert(child->as().name()); + + for (auto it = nodes.begin(); it != nodes.end();) { - auto expected_column = expected_columns.find(id->shortName()); - if (expected_column != expected_columns.end()) + if (const auto * id = it->get()->as()) { - expected_columns.erase(expected_column); - it = nodes.erase(it); + auto expected_column = expected_columns.find(id->shortName()); + if (expected_column != expected_columns.end()) + { + expected_columns.erase(expected_column); + it = nodes.erase(it); + continue; + } } - else - ++it; - } - else ++it; + } + } + else + { + for (auto it = nodes.begin(); it != nodes.end();) + { + if (const auto * id = it->get()->as()) + { + if (isColumnMatching(id->shortName())) + { + it = nodes.erase(it); + continue; + } + } + ++it; + } } if (is_strict && !expected_columns.empty()) @@ -125,6 +146,21 @@ void ASTColumnsExceptTransformer::transform(ASTs & nodes) const } } +void ASTColumnsExceptTransformer::setPattern(String pattern) +{ + original_pattern = std::move(pattern); + column_matcher = std::make_shared(original_pattern, RE2::Quiet); + if (!column_matcher->ok()) + throw DB::Exception( + "COLUMNS pattern " + original_pattern + " cannot be compiled: " + column_matcher->error(), + DB::ErrorCodes::CANNOT_COMPILE_REGEXP); +} + +bool ASTColumnsExceptTransformer::isColumnMatching(const String & column_name) const +{ + return RE2::PartialMatch(column_name, *column_matcher); +} + void ASTColumnsReplaceTransformer::Replacement::formatImpl( const FormatSettings & settings, FormatState & state, FormatStateStacked frame) const { diff --git a/src/Parsers/ASTColumnsTransformers.h b/src/Parsers/ASTColumnsTransformers.h index ec1fd3eee0a..49d29222f02 100644 --- a/src/Parsers/ASTColumnsTransformers.h +++ b/src/Parsers/ASTColumnsTransformers.h @@ -2,6 +2,11 @@ #include +namespace re2 +{ + class RE2; +} + namespace DB { class IASTColumnsTransformer : public IAST @@ -43,9 +48,13 @@ public: return clone; } void transform(ASTs & nodes) const override; + void setPattern(String pattern); + bool isColumnMatching(const String & column_name) const; protected: void formatImpl(const FormatSettings & settings, FormatState &, FormatStateStacked) const override; + std::shared_ptr column_matcher; + String original_pattern; }; class ASTColumnsReplaceTransformer : public IASTColumnsTransformer diff --git a/src/Parsers/ExpressionElementParsers.cpp b/src/Parsers/ExpressionElementParsers.cpp index 7c82c4aca1e..b421a7230db 100644 --- a/src/Parsers/ExpressionElementParsers.cpp +++ b/src/Parsers/ExpressionElementParsers.cpp @@ -1427,6 +1427,8 @@ bool ParserColumnsTransformers::parseImpl(Pos & pos, ASTPtr & node, Expected & e is_strict = true; ASTs identifiers; + ASTPtr regex_node; + ParserStringLiteral regex; auto parse_id = [&identifiers, &pos, &expected] { ASTPtr identifier; @@ -1441,7 +1443,7 @@ bool ParserColumnsTransformers::parseImpl(Pos & pos, ASTPtr & node, Expected & e { // support one or more parameter ++pos; - if (!ParserList::parseUtil(pos, expected, parse_id, false)) + if (!ParserList::parseUtil(pos, expected, parse_id, false) && !regex.parse(pos, regex_node, expected)) return false; if (pos->type != TokenType::ClosingRoundBracket) @@ -1451,12 +1453,15 @@ bool ParserColumnsTransformers::parseImpl(Pos & pos, ASTPtr & node, Expected & e else { // only one parameter - if (!parse_id()) + if (!parse_id() && !regex.parse(pos, regex_node, expected)) return false; } auto res = std::make_shared(); - res->children = std::move(identifiers); + if (regex_node) + res->setPattern(regex_node->as().value.get()); + else + res->children = std::move(identifiers); res->is_strict = is_strict; node = std::move(res); return true; diff --git a/tests/queries/0_stateless/01470_columns_transformers2.reference b/tests/queries/0_stateless/01470_columns_transformers2.reference new file mode 100644 index 00000000000..18c0f5c7e89 --- /dev/null +++ b/tests/queries/0_stateless/01470_columns_transformers2.reference @@ -0,0 +1 @@ +100 10 324 120.00 B 8.00 B 23.00 B diff --git a/tests/queries/0_stateless/01470_columns_transformers2.sql b/tests/queries/0_stateless/01470_columns_transformers2.sql new file mode 100644 index 00000000000..3691ef1e65d --- /dev/null +++ b/tests/queries/0_stateless/01470_columns_transformers2.sql @@ -0,0 +1,7 @@ +DROP TABLE IF EXISTS columns_transformers; + +CREATE TABLE columns_transformers (i int, j int, k int, a_bytes int, b_bytes int, c_bytes int) Engine=TinyLog; +INSERT INTO columns_transformers VALUES (100, 10, 324, 120, 8, 23); +SELECT * EXCEPT 'bytes', COLUMNS('bytes') APPLY formatReadableSize FROM columns_transformers; + +DROP TABLE IF EXISTS columns_transformers;