Merge pull request #17620 from ClickHouse/backport/20.12/16183

Backport #16183 to 20.12: Fix: throw error when column transformer use non-exist column

src/Parsers/ASTColumnsTransformers.cpp

@@ -13,6 +13,7 @@ namespace DB
 namespace ErrorCodes
 {
     extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+    extern const int NO_SUCH_COLUMN_IN_TABLE;
 }
 
 void IASTColumnsTransformer::transform(const ASTPtr & transformer, ASTs & nodes)
@@ -33,25 +34,48 @@ void IASTColumnsTransformer::transform(const ASTPtr & transformer, ASTs & nodes)
 
 void ASTColumnsApplyTransformer::formatImpl(const FormatSettings & settings, FormatState & state, FormatStateStacked frame) const
 {
-    settings.ostr << (settings.hilite ? hilite_keyword : "") << "APPLY" << (settings.hilite ? hilite_none : "") << "(" << func_name;
+    settings.ostr << (settings.hilite ? hilite_keyword : "") << "APPLY" << (settings.hilite ? hilite_none : "") << " ";
+
+    if (!column_name_prefix.empty())
+        settings.ostr << "(";
+    settings.ostr << func_name;
+
     if (parameters)
         parameters->formatImpl(settings, state, frame);
-    settings.ostr << ")";
+
+    if (!column_name_prefix.empty())
+        settings.ostr << ", '" << column_name_prefix << "')";
 }
 
 void ASTColumnsApplyTransformer::transform(ASTs & nodes) const
 {
     for (auto & column : nodes)
     {
+        String name;
+        auto alias = column->tryGetAlias();
+        if (!alias.empty())
+            name = alias;
+        else
+        {
+            if (const auto * id = column->as<ASTIdentifier>())
+                name = id->shortName();
+            else
+                name = column->getColumnName();
+        }
         auto function = makeASTFunction(func_name, column);
         function->parameters = parameters;
         column = function;
+        if (!column_name_prefix.empty())
+            column->setAlias(column_name_prefix + name);
     }
 }
 
 void ASTColumnsExceptTransformer::formatImpl(const FormatSettings & settings, FormatState & state, FormatStateStacked frame) const
 {
-    settings.ostr << (settings.hilite ? hilite_keyword : "") << "EXCEPT" << (settings.hilite ? hilite_none : "") << "(";
+    settings.ostr << (settings.hilite ? hilite_keyword : "") << "EXCEPT" << (is_strict ? " STRICT " : " ") << (settings.hilite ? hilite_none : "");
+
+    if (children.size() > 1)
+        settings.ostr << "(";
 
     for (ASTs::const_iterator it = children.begin(); it != children.end(); ++it)
     {
@@ -62,28 +86,43 @@ void ASTColumnsExceptTransformer::formatImpl(const FormatSettings & settings, Fo
         (*it)->formatImpl(settings, state, frame);
     }
 
-    settings.ostr << ")";
+    if (children.size() > 1)
+        settings.ostr << ")";
 }
 
 void ASTColumnsExceptTransformer::transform(ASTs & nodes) const
 {
-    nodes.erase(
-        std::remove_if(
-            nodes.begin(),
-            nodes.end(),
-            [this](const ASTPtr & node_child)
+    std::set<String> expected_columns;
+    for (const auto & child : children)
+        expected_columns.insert(child->as<const ASTIdentifier &>().name());
+
+    for (auto it = nodes.begin(); it != nodes.end();)
+    {
+        if (const auto * id = it->get()->as<ASTIdentifier>())
+        {
+            auto expected_column = expected_columns.find(id->shortName());
+            if (expected_column != expected_columns.end())
             {
-                if (const auto * id = node_child->as<ASTIdentifier>())
-                {
-                    for (const auto & except_child : children)
-                    {
-                        if (except_child->as<const ASTIdentifier &>().name() == id->shortName())
-                            return true;
-                    }
-                }
-                return false;
-            }),
-        nodes.end());
+                expected_columns.erase(expected_column);
+                it = nodes.erase(it);
+            }
+            else
+                ++it;
+        }
+        else
+            ++it;
+    }
+
+    if (is_strict && !expected_columns.empty())
+    {
+        String expected_columns_str;
+        std::for_each(expected_columns.begin(), expected_columns.end(),
+            [&](String x) { expected_columns_str += (" " + x) ; });
+
+        throw Exception(
+            "Columns transformer EXCEPT expects following column(s) :" + expected_columns_str,
+            ErrorCodes::NO_SUCH_COLUMN_IN_TABLE);
+    }
 }
 
 void ASTColumnsReplaceTransformer::Replacement::formatImpl(
@@ -95,7 +134,10 @@ void ASTColumnsReplaceTransformer::Replacement::formatImpl(
 
 void ASTColumnsReplaceTransformer::formatImpl(const FormatSettings & settings, FormatState & state, FormatStateStacked frame) const
 {
-    settings.ostr << (settings.hilite ? hilite_keyword : "") << "REPLACE" << (settings.hilite ? hilite_none : "") << "(";
+    settings.ostr << (settings.hilite ? hilite_keyword : "") << "REPLACE" << (is_strict ? " STRICT " : " ") << (settings.hilite ? hilite_none : "");
+
+    if (children.size() > 1)
+        settings.ostr << "(";
 
     for (ASTs::const_iterator it = children.begin(); it != children.end(); ++it)
     {
@@ -106,7 +148,8 @@ void ASTColumnsReplaceTransformer::formatImpl(const FormatSettings & settings, F
         (*it)->formatImpl(settings, state, frame);
     }
 
-    settings.ostr << ")";
+    if (children.size() > 1)
+        settings.ostr << ")";
 }
 
 void ASTColumnsReplaceTransformer::replaceChildren(ASTPtr & node, const ASTPtr & replacement, const String & name)
@@ -145,6 +188,7 @@ void ASTColumnsReplaceTransformer::transform(ASTs & nodes) const
             {
                 column = replace_it->second;
                 column->setAlias(replace_it->first);
+                replace_map.erase(replace_it);
             }
         }
         else if (auto * ast_with_alias = dynamic_cast<ASTWithAlias *>(column.get()))
@@ -157,9 +201,25 @@ void ASTColumnsReplaceTransformer::transform(ASTs & nodes) const
                 replaceChildren(new_ast, column, replace_it->first);
                 column = new_ast;
                 column->setAlias(replace_it->first);
+                replace_map.erase(replace_it);
             }
         }
     }
+
+    if (is_strict && !replace_map.empty())
+    {
+        String expected_columns;
+        for (auto & elem: replace_map)
+        {
+            if (!expected_columns.empty())
+                expected_columns += ", ";
+            expected_columns += elem.first;
+        }
+        throw Exception(
+            "Columns transformer REPLACE expects following column(s) : " + expected_columns,
+            ErrorCodes::NO_SUCH_COLUMN_IN_TABLE);
+    }
+
 }
 
 }

src/Parsers/ASTColumnsTransformers.h

@@ -24,6 +24,7 @@ public:
     }
     void transform(ASTs & nodes) const override;
     String func_name;
+    String column_name_prefix;
     ASTPtr parameters;
 
 protected:
@@ -33,6 +34,7 @@ protected:
 class ASTColumnsExceptTransformer : public IASTColumnsTransformer
 {
 public:
+    bool is_strict = false;
     String getID(char) const override { return "ColumnsExceptTransformer"; }
     ASTPtr clone() const override
     {
@@ -69,6 +71,7 @@ public:
         void formatImpl(const FormatSettings & settings, FormatState &, FormatStateStacked) const override;
     };
 
+    bool is_strict = false;
     String getID(char) const override { return "ColumnsReplaceTransformer"; }
     ASTPtr clone() const override
     {

src/Parsers/ExpressionElementParsers.cpp

@@ -1271,12 +1271,17 @@ bool ParserColumnsTransformers::parseImpl(Pos & pos, ASTPtr & node, Expected & e
     ParserKeyword except("EXCEPT");
     ParserKeyword replace("REPLACE");
     ParserKeyword as("AS");
+    ParserKeyword strict("STRICT");
 
     if (apply.ignore(pos, expected))
     {
-        if (pos->type != TokenType::OpeningRoundBracket)
-            return false;
-        ++pos;
+        bool with_open_round_bracket = false;
+
+        if (pos->type == TokenType::OpeningRoundBracket)
+        {
+            ++pos;
+            with_open_round_bracket = true;
+        }
 
         ASTPtr func_name;
         if (!ParserIdentifier().parse(pos, func_name, expected))
@@ -1294,21 +1299,37 @@ bool ParserColumnsTransformers::parseImpl(Pos & pos, ASTPtr & node, Expected & e
             ++pos;
         }
 
-        if (pos->type != TokenType::ClosingRoundBracket)
-            return false;
-        ++pos;
+        String column_name_prefix;
+        if (with_open_round_bracket && pos->type == TokenType::Comma)
+        {
+            ++pos;
+
+            ParserStringLiteral parser_string_literal;
+            ASTPtr ast_prefix_name;
+            if (!parser_string_literal.parse(pos, ast_prefix_name, expected))
+                return false;
+
+            column_name_prefix = ast_prefix_name->as<ASTLiteral &>().value.get<const String &>();
+        }
+
+        if (with_open_round_bracket)
+        {
+            if (pos->type != TokenType::ClosingRoundBracket)
+                return false;
+            ++pos;
+        }
 
         auto res = std::make_shared<ASTColumnsApplyTransformer>();
         res->func_name = getIdentifierName(func_name);
         res->parameters = expr_list_args;
+        res->column_name_prefix = column_name_prefix;
         node = std::move(res);
         return true;
     }
     else if (except.ignore(pos, expected))
     {
-        if (pos->type != TokenType::OpeningRoundBracket)
-            return false;
-        ++pos;
+        if (strict.ignore(pos, expected))
+            is_strict = true;
 
         ASTs identifiers;
         auto parse_id = [&identifiers, &pos, &expected]
@@ -1321,23 +1342,34 @@ bool ParserColumnsTransformers::parseImpl(Pos & pos, ASTPtr & node, Expected & e
             return true;
         };
 
-        if (!ParserList::parseUtil(pos, expected, parse_id, false))
-            return false;
+        if (pos->type == TokenType::OpeningRoundBracket)
+        {
+            // support one or more parameter
+            ++pos;
+            if (!ParserList::parseUtil(pos, expected, parse_id, false))
+                return false;
 
-        if (pos->type != TokenType::ClosingRoundBracket)
-            return false;
-        ++pos;
+            if (pos->type != TokenType::ClosingRoundBracket)
+                return false;
+            ++pos;
+        }
+        else
+        {
+            // only one parameter
+            if (!parse_id())
+                return false;
+        }
 
         auto res = std::make_shared<ASTColumnsExceptTransformer>();
         res->children = std::move(identifiers);
+        res->is_strict = is_strict;
         node = std::move(res);
         return true;
     }
     else if (replace.ignore(pos, expected))
     {
-        if (pos->type != TokenType::OpeningRoundBracket)
-            return false;
-        ++pos;
+        if (strict.ignore(pos, expected))
+            is_strict = true;
 
         ASTs replacements;
         ParserExpression element_p;
@@ -1362,15 +1394,27 @@ bool ParserColumnsTransformers::parseImpl(Pos & pos, ASTPtr & node, Expected & e
             return true;
         };
 
-        if (!ParserList::parseUtil(pos, expected, parse_id, false))
-            return false;
+        if (pos->type == TokenType::OpeningRoundBracket)
+        {
+            ++pos;
 
-        if (pos->type != TokenType::ClosingRoundBracket)
-            return false;
-        ++pos;
+            if (!ParserList::parseUtil(pos, expected, parse_id, false))
+                return false;
+
+            if (pos->type != TokenType::ClosingRoundBracket)
+                return false;
+            ++pos;
+        }
+        else
+        {
+            // only one parameter
+            if (!parse_id())
+                return false;
+        }
 
         auto res = std::make_shared<ASTColumnsReplaceTransformer>();
         res->children = std::move(replacements);
+        res->is_strict = is_strict;
         node = std::move(res);
         return true;
     }

src/Parsers/ExpressionElementParsers.h

@@ -100,9 +100,12 @@ protected:
   */
 class ParserColumnsTransformers : public IParserBase
 {
+public:
+    ParserColumnsTransformers(bool is_strict_ = false): is_strict(is_strict_) {}
 protected:
     const char * getName() const override { return "COLUMNS transformers"; }
     bool parseImpl(Pos & pos, ASTPtr & node, Expected & expected) override;
+    bool is_strict;
 };
 
 /** A function, for example, f(x, y + 1, g(z)).

tests/queries/0_stateless/01470_columns_transformers.reference

 220	18	347
+220	18	347
 110	9	173.5
 1970-04-11	1970-01-11	1970-11-21
 2	3
@@ -6,6 +7,12 @@
 18	347
 110	173.5
 1970-04-11	1970-01-11	1970-11-21
+10	324
+8	23
+324
+23
+101	10	324
+121	8	23
 222	18	347
 111	11	173.5
 1970-04-11	1970-01-11	1970-11-21

tests/queries/0_stateless/01470_columns_transformers.sql

@@ -4,6 +4,7 @@ CREATE TABLE columns_transformers (i Int64, j Int16, k Int64) Engine=TinyLog;
 INSERT INTO columns_transformers VALUES (100, 10, 324), (120, 8, 23);
 
 SELECT * APPLY(sum) from columns_transformers;
+SELECT * APPLY sum from columns_transformers;
 SELECT columns_transformers.* APPLY(avg) from columns_transformers;
 SELECT a.* APPLY(toDate) APPLY(any) from columns_transformers a;
 SELECT COLUMNS('[jk]') APPLY(toString) APPLY(length) from columns_transformers;
@@ -13,11 +14,18 @@ SELECT columns_transformers.* EXCEPT(j) APPLY(avg) from columns_transformers;
 -- EXCEPT after APPLY will not match anything
 SELECT a.* APPLY(toDate) EXCEPT(i, j) APPLY(any) from columns_transformers a;
 
+SELECT * EXCEPT STRICT i from columns_transformers;
+SELECT * EXCEPT STRICT (i, j) from columns_transformers;
+SELECT * EXCEPT STRICT i, j1 from columns_transformers; -- { serverError 47 }
+SELECT * EXCEPT STRICT(i, j1) from columns_transformers; -- { serverError 16 }
+SELECT * REPLACE STRICT i + 1 AS i from columns_transformers;
+SELECT * REPLACE STRICT(i + 1 AS col) from columns_transformers; -- { serverError 16 }
 SELECT * REPLACE(i + 1 AS i) APPLY(sum) from columns_transformers;
 SELECT columns_transformers.* REPLACE(j + 2 AS j, i + 1 AS i) APPLY(avg) from columns_transformers;
 SELECT columns_transformers.* REPLACE(j + 1 AS j, j + 2 AS j) APPLY(avg) from columns_transformers; -- { serverError 43 }
 -- REPLACE after APPLY will not match anything
 SELECT a.* APPLY(toDate) REPLACE(i + 1 AS i) APPLY(any) from columns_transformers a;
+SELECT a.* APPLY(toDate) REPLACE STRICT(i + 1 AS i) APPLY(any) from columns_transformers a; -- { serverError 16 }
 
 EXPLAIN SYNTAX SELECT * APPLY(sum) from columns_transformers;
 EXPLAIN SYNTAX SELECT columns_transformers.* APPLY(avg) from columns_transformers;