Rewrite CROSS/COMMA to INNER JOIN using table's columns knowledge

dbms/src/Interpreters/CrossToInnerJoinVisitor.cpp

@@ -4,8 +4,10 @@
 #include <Interpreters/CrossToInnerJoinVisitor.h>
 #include <Interpreters/DatabaseAndTableWithAlias.h>
 #include <Interpreters/IdentifierSemantic.h>
+#include <Interpreters/QueryAliasesVisitor.h>
 #include <Interpreters/misc.h>
 #include <Parsers/ASTSelectQuery.h>
+#include <Parsers/ASTSubquery.h>
 #include <Parsers/ASTTablesInSelectQuery.h>
 #include <Parsers/ASTIdentifier.h>
 #include <Parsers/ASTFunction.h>
@@ -27,41 +29,26 @@ namespace ErrorCodes
 namespace
 {
 
-struct JoinedTable
+struct JoinedElement
 {
-    DatabaseAndTableWithAlias table;
-    ASTTablesInSelectQueryElement * element = nullptr;
-    ASTTableJoin * join = nullptr;
-    ASTPtr array_join = nullptr;
-    bool has_using = false;
-
-    JoinedTable(ASTPtr table_element)
+    JoinedElement(const ASTTablesInSelectQueryElement & table_element)
+        : element(table_element)
     {
-        element = table_element->as<ASTTablesInSelectQueryElement>();
-        if (!element)
-            throw Exception("Logical error: TablesInSelectQueryElement expected", ErrorCodes::LOGICAL_ERROR);
+        if (element.table_join)
+            join = element.table_join->as<ASTTableJoin>();
+    }
 
-        if (element->table_join)
-        {
-            join = element->table_join->as<ASTTableJoin>();
-            if (join->kind == ASTTableJoin::Kind::Cross ||
-                join->kind == ASTTableJoin::Kind::Comma)
-            {
-                if (!join->children.empty())
-                    throw Exception("Logical error: CROSS JOIN has expressions", ErrorCodes::LOGICAL_ERROR);
-            }
+    void checkTableName(const DatabaseAndTableWithAlias & table, const String & current_database) const
+    {
+        if (!element.table_expression)
+            throw Exception("Not a table expression in JOIN (ARRAY JOIN?)", ErrorCodes::LOGICAL_ERROR);
 
-            if (join->using_expression_list)
-                has_using = true;
-        }
+        ASTTableExpression * table_expression = element.table_expression->as<ASTTableExpression>();
+        if (!table_expression)
+            throw Exception("Wrong table expression in JOIN", ErrorCodes::LOGICAL_ERROR);
 
-        if (element->table_expression)
-        {
-            const auto & expr = element->table_expression->as<ASTTableExpression &>();
-            table = DatabaseAndTableWithAlias(expr);
-        }
-
-        array_join = element->array_join;
+        if (!table.same(DatabaseAndTableWithAlias(*table_expression, current_database)))
+            throw Exception("Inconsistent table names", ErrorCodes::LOGICAL_ERROR);
     }
 
     void rewriteCommaToCross()
@@ -70,7 +57,24 @@ struct JoinedTable
             join->kind = ASTTableJoin::Kind::Cross;
     }
 
+    void rewriteCrossToInner(ASTPtr on_expression)
+    {
+        join->kind = ASTTableJoin::Kind::Inner;
+        join->strictness = ASTTableJoin::Strictness::All;
+
+        join->on_expression = on_expression;
+        join->children.push_back(join->on_expression);
+    }
+
+    ASTPtr arrayJoin() const { return element.array_join; }
+    const ASTTableJoin * tableJoin() const { return join; }
+
     bool canAttachOnExpression() const { return join && !join->on_expression; }
+    bool hasUsing() const { return join && join->using_expression_list; }
+
+private:
+    const ASTTablesInSelectQueryElement & element;
+    ASTTableJoin * join = nullptr;
 };
 
 bool isComparison(const String & name)
@@ -89,13 +93,14 @@ class CheckExpressionVisitorData
 public:
     using TypeToVisit = const ASTFunction;
 
-    CheckExpressionVisitorData(const std::vector<JoinedTable> & tables_)
+    CheckExpressionVisitorData(const std::vector<JoinedElement> & tables_,
+                               const std::vector<TableWithColumnNamesAndTypes> & tables_with_columns,
+                               Aliases && aliases_)
         : joined_tables(tables_)
+        , tables(tables_with_columns)
+        , aliases(aliases_)
         , ands_only(true)
-    {
-        for (auto & joined : joined_tables)
-            tables.push_back(joined.table);
-    }
+    {}
 
     void visit(const ASTFunction & node, const ASTPtr & ast)
     {
@@ -160,9 +165,10 @@ class CheckExpressionVisitorData
     }
 
 private:
-    const std::vector<JoinedTable> & joined_tables;
-    std::vector<DatabaseAndTableWithAlias> tables;
+    const std::vector<JoinedElement> & joined_tables;
+    const std::vector<TableWithColumnNamesAndTypes> & tables;
     std::map<size_t, std::vector<ASTPtr>> asts_to_join_on;
+    Aliases aliases;
     bool ands_only;
 
     size_t canMoveEqualsToJoinOn(const ASTFunction & node)
@@ -177,6 +183,12 @@ class CheckExpressionVisitorData
         if (!left || !right)
             return false;
 
+        /// Moving expressions that use column aliases is not supported.
+        if (left->isShort() && aliases.count(left->shortName()))
+            return false;
+        if (right->isShort() && aliases.count(right->shortName()))
+            return false;
+
         return checkIdentifiers(*left, *right);
     }
 
@@ -185,15 +197,17 @@ class CheckExpressionVisitorData
     /// @return table position to attach expression to or 0.
     size_t checkIdentifiers(const ASTIdentifier & left, const ASTIdentifier & right)
     {
-        size_t left_table_pos = 0;
-        bool left_match = IdentifierSemantic::chooseTable(left, tables, left_table_pos);
+        std::optional<size_t> left_table_pos = IdentifierSemantic::getMembership(left);
+        if (!left_table_pos)
+            left_table_pos = IdentifierSemantic::chooseTable(left, tables);
 
-        size_t right_table_pos = 0;
-        bool right_match = IdentifierSemantic::chooseTable(right, tables, right_table_pos);
+        std::optional<size_t> right_table_pos = IdentifierSemantic::getMembership(right);
+        if (!right_table_pos)
+            right_table_pos = IdentifierSemantic::chooseTable(right, tables);
 
-        if (left_match && right_match && (left_table_pos != right_table_pos))
+        if (left_table_pos && right_table_pos && (*left_table_pos != *right_table_pos))
         {
-            size_t table_pos = std::max(left_table_pos, right_table_pos);
+            size_t table_pos = std::max(*left_table_pos, *right_table_pos);
             if (joined_tables[table_pos].canAttachOnExpression())
                 return table_pos;
         }
@@ -205,7 +219,7 @@ using CheckExpressionMatcher = ConstOneTypeMatcher<CheckExpressionVisitorData, f
 using CheckExpressionVisitor = ConstInDepthNodeVisitor<CheckExpressionMatcher, true>;
 
 
-bool getTables(ASTSelectQuery & select, std::vector<JoinedTable> & joined_tables, size_t & num_comma)
+bool getTables(ASTSelectQuery & select, std::vector<JoinedElement> & joined_tables, size_t & num_comma)
 {
     if (!select.tables())
         return false;
@@ -224,23 +238,37 @@ bool getTables(ASTSelectQuery & select, std::vector<JoinedTable> & joined_tables
 
     for (auto & child : tables->children)
     {
-        joined_tables.emplace_back(JoinedTable(child));
-        JoinedTable & t = joined_tables.back();
-        if (t.array_join)
+        auto table_element = child->as<ASTTablesInSelectQueryElement>();
+        if (!table_element)
+            throw Exception("Logical error: TablesInSelectQueryElement expected", ErrorCodes::LOGICAL_ERROR);
+
+        joined_tables.emplace_back(JoinedElement(*table_element));
+        JoinedElement & t = joined_tables.back();
+
+        if (t.arrayJoin())
         {
             ++num_array_join;
             continue;
         }
 
-        if (t.has_using)
+        if (t.hasUsing())
         {
             ++num_using;
             continue;
         }
 
-        if (auto * join = t.join)
+        if (auto * join = t.tableJoin())
+        {
+            if (join->kind == ASTTableJoin::Kind::Cross ||
+                join->kind == ASTTableJoin::Kind::Comma)
+            {
+                if (!join->children.empty())
+                    throw Exception("Logical error: CROSS JOIN has expressions", ErrorCodes::LOGICAL_ERROR);
+            }
+
             if (join->kind == ASTTableJoin::Kind::Comma)
                 ++num_comma;
+        }
     }
 
     if (num_using && (num_tables - num_array_join) > 2)
@@ -251,12 +279,20 @@ bool getTables(ASTSelectQuery & select, std::vector<JoinedTable> & joined_tables
 
     if (num_array_join || num_using)
         return false;
+
     return true;
 }
 
 }
 
 
+bool CrossToInnerJoinMatcher::needChildVisit(ASTPtr & node, const ASTPtr &)
+{
+    if (node->as<ASTSubquery>())
+        return false;
+    return true;
+}
+
 void CrossToInnerJoinMatcher::visit(ASTPtr & ast, Data & data)
 {
     if (auto * t = ast->as<ASTSelectQuery>())
@@ -266,10 +302,19 @@ void CrossToInnerJoinMatcher::visit(ASTPtr & ast, Data & data)
 void CrossToInnerJoinMatcher::visit(ASTSelectQuery & select, ASTPtr &, Data & data)
 {
     size_t num_comma = 0;
-    std::vector<JoinedTable> joined_tables;
+    std::vector<JoinedElement> joined_tables;
     if (!getTables(select, joined_tables, num_comma))
         return;
 
+    /// Check if joined_tables are consistent with known tables_with_columns
+    {
+        if (joined_tables.size() != data.tables_with_columns.size())
+            throw Exception("Logical error: inconsistent number of tables", ErrorCodes::LOGICAL_ERROR);
+
+        for (size_t i = 0; i < joined_tables.size(); ++i)
+            joined_tables[i].checkTableName(data.tables_with_columns[i].table, data.current_database);
+    }
+
     /// COMMA to CROSS
 
     if (num_comma)
@@ -283,7 +328,13 @@ void CrossToInnerJoinMatcher::visit(ASTSelectQuery & select, ASTPtr &, Data & da
     if (!select.where())
         return;
 
-    CheckExpressionVisitor::Data visitor_data{joined_tables};
+    Aliases aliases;
+    QueryAliasesVisitor::Data query_aliases_data{aliases};
+    if (ASTPtr with = select.with())
+        QueryAliasesVisitor(query_aliases_data).visit(with);
+    QueryAliasesVisitor(query_aliases_data).visit(select.select());
+
+    CheckExpressionVisitor::Data visitor_data{joined_tables, data.tables_with_columns, std::move(aliases)};
     CheckExpressionVisitor(visitor_data).visit(select.where());
 
     if (visitor_data.complex())
@@ -293,12 +344,7 @@ void CrossToInnerJoinMatcher::visit(ASTSelectQuery & select, ASTPtr &, Data & da
     {
         if (visitor_data.matchAny(i))
         {
-            ASTTableJoin & join = *joined_tables[i].join;
-            join.kind = ASTTableJoin::Kind::Inner;
-            join.strictness = ASTTableJoin::Strictness::All;
-
-            join.on_expression = visitor_data.makeOnExpression(i);
-            join.children.push_back(join.on_expression);
+            joined_tables[i].rewriteCrossToInner(visitor_data.makeOnExpression(i));
             data.done = true;
         }
     }

dbms/src/Interpreters/CrossToInnerJoinVisitor.h

@@ -6,17 +6,20 @@ namespace DB
 {
 
 class ASTSelectQuery;
+struct TableWithColumnNamesAndTypes;
 
 /// AST transformer. It replaces cross joins with equivalented inner join if possible.
 class CrossToInnerJoinMatcher
 {
 public:
     struct Data
     {
+        const std::vector<TableWithColumnNamesAndTypes> & tables_with_columns;
+        const String current_database;
         bool done = false;
     };
 
-    static bool needChildVisit(ASTPtr &, const ASTPtr &) { return true; }
+    static bool needChildVisit(ASTPtr &, const ASTPtr &);
     static void visit(ASTPtr & ast, Data & data);
 
 private:

dbms/src/Interpreters/DatabaseAndTableWithAlias.h

@@ -35,6 +35,12 @@ struct DatabaseAndTableWithAlias
 
     /// Check if it satisfies another db_table name. @note opterion is not symmetric.
     bool satisfies(const DatabaseAndTableWithAlias & table, bool table_may_be_an_alias);
+
+    /// Exactly the same table name
+    bool same(const DatabaseAndTableWithAlias & db_table) const
+    {
+        return database == db_table.database && table == db_table.table && alias == db_table.alias;
+    }
 };
 
 struct TableWithColumnNames
@@ -80,6 +86,19 @@ struct TableWithColumnNamesAndTypes
         , columns(columns_)
     {}
 
+    bool hasColumn(const String & name) const
+    {
+        if (names.empty())
+        {
+            for (auto & col : columns)
+                names.insert(col.name);
+            for (auto & col : hidden_columns)
+                names.insert(col.name);
+        }
+
+        return names.count(name);
+    }
+
     void addHiddenColumns(const NamesAndTypesList & addition)
     {
         hidden_columns.insert(hidden_columns.end(), addition.begin(), addition.end());
@@ -99,6 +118,9 @@ struct TableWithColumnNamesAndTypes
 
         return TableWithColumnNames(table, std::move(out_columns), std::move(out_hidden_columns));
     }
+
+private:
+    mutable NameSet names;
 };
 
 std::vector<DatabaseAndTableWithAlias> getDatabaseAndTables(const ASTSelectQuery & select_query, const String & current_database);

dbms/src/Interpreters/ExtractExpressionInfoVisitor.cpp

@@ -50,9 +50,8 @@ void ExpressionInfoMatcher::visit(const ASTIdentifier & identifier, const ASTPtr
     }
     else
     {
-        size_t best_table_pos = 0;
-        if (IdentifierSemantic::chooseTable(identifier, data.tables, best_table_pos))
-            data.unique_reference_tables_pos.emplace(best_table_pos);
+        if (auto best_table_pos = IdentifierSemantic::chooseTable(identifier, data.tables))
+            data.unique_reference_tables_pos.emplace(*best_table_pos);
     }
 }