Merge pull request ClickHouse#76078 from ClickHouse/analyzer-scalar-correlated-subqueries

Support correlated subqueries in WHERE clause

Author: KochetovNicolai

src/Analyzer/HashUtils.h

@@ -54,6 +54,11 @@ using QueryTreeNodeConstRawPtrWithHashSet = std::unordered_set<QueryTreeNodeCons
 template <typename Value>
 using QueryTreeNodePtrWithHashMap = std::unordered_map<QueryTreeNodePtrWithHash, Value>;
 
+class ColumnNode;
+using ColumnNodePtr = std::shared_ptr<ColumnNode>;
+using ColumnNodePtrWithHash = QueryTreeNodeWithHash<ColumnNodePtr>;
+using ColumnNodePtrWithHashSet = std::unordered_set<ColumnNodePtrWithHash>;
+
 }
 
 template <typename T, bool compare_aliases, bool compare_types>

src/Analyzer/IQueryTreeNode.h

@@ -47,7 +47,7 @@ enum class QueryTreeNodeType : uint8_t
     ARRAY_JOIN,
     CROSS_JOIN,
     JOIN,
-    UNION
+    UNION,
 };
 
 /// Convert query tree node type to string
@@ -91,12 +91,12 @@ class IQueryTreeNode : public TypePromotion<IQueryTreeNode>
       */
     virtual DataTypePtr getResultType() const
     {
-        throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "Method getResultType is not supported for {} query node", getNodeTypeName());
+        throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "Method getResultType is not supported for {} query tree node", getNodeTypeName());
     }
 
     virtual void convertToNullable()
     {
-        throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "Method convertToNullable is not supported for {} query node", getNodeTypeName());
+        throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "Method convertToNullable is not supported for {} query tree node", getNodeTypeName());
     }
 
     struct CompareOptions

src/Analyzer/Passes/RemoveUnusedProjectionColumnsPass.cpp

@@ -2,12 +2,13 @@
 
 #include <Functions/FunctionFactory.h>
 
-#include <Analyzer/InDepthQueryTreeVisitor.h>
+#include <Analyzer/AggregationUtils.h>
+#include <Analyzer/ColumnNode.h>
 #include <Analyzer/FunctionNode.h>
+#include <Analyzer/InDepthQueryTreeVisitor.h>
 #include <Analyzer/QueryNode.h>
-#include <Analyzer/ColumnNode.h>
 #include <Analyzer/SortNode.h>
-#include <Analyzer/AggregationUtils.h>
+#include <Analyzer/UnionNode.h>
 #include <Analyzer/Utils.h>
 
 namespace DB
@@ -48,6 +49,29 @@ class CollectUsedColumnsVisitor : public InDepthQueryTreeVisitorWithContext<Coll
             return;
         }
 
+        if (node_type == QueryTreeNodeType::FUNCTION)
+        {
+            auto & function_node = node->as<FunctionNode &>();
+
+            if (function_node.getFunctionName() != "exists")
+                return;
+
+            const auto & subquery_argument = function_node.getArguments().getNodes().front();
+            auto * query_node = subquery_argument->as<QueryNode>();
+            auto * union_node = subquery_argument->as<UnionNode>();
+
+            const auto & correlated_columns = query_node != nullptr ? query_node->getCorrelatedColumns() : union_node->getCorrelatedColumns();
+            for (const auto & correlated_column : correlated_columns)
+            {
+                auto * column_node = correlated_column->as<ColumnNode>();
+                auto column_source_node = column_node->getColumnSource();
+                auto column_source_node_type = column_source_node->getNodeType();
+                if (column_source_node_type == QueryTreeNodeType::QUERY || column_source_node_type == QueryTreeNodeType::UNION)
+                    query_or_union_node_to_used_columns[column_source_node].insert(column_node->getColumnName());
+            }
+            return;
+        }
+
         if (node_type != QueryTreeNodeType::COLUMN)
             return;
 
@@ -56,10 +80,15 @@ class CollectUsedColumnsVisitor : public InDepthQueryTreeVisitorWithContext<Coll
             return;
 
         auto column_source_node = column_node.getColumnSource();
-        auto column_source_node_type = column_source_node->getNodeType();
 
-        if (column_source_node_type == QueryTreeNodeType::QUERY || column_source_node_type == QueryTreeNodeType::UNION)
-            query_or_union_node_to_used_columns[column_source_node].insert(column_node.getColumnName());
+        auto it = query_or_union_node_to_used_columns.find(column_source_node);
+        /// If the source node is not found in the map then:
+        /// 1. Tt's either not a Query or Union node.
+        /// 2. It's a correlated column and it comes from the outer scope.
+        if (it != query_or_union_node_to_used_columns.end())
+        {
+            it->second.insert(column_node.getColumnName());
+        }
     }
 
     void reset()

src/Analyzer/QueryNode.cpp

@@ -22,6 +22,7 @@
 #include <Parsers/ASTSelectWithUnionQuery.h>
 #include <Parsers/ASTSetQuery.h>
 
+#include <Analyzer/ColumnNode.h>
 #include <Analyzer/InterpolateNode.h>
 #include <Analyzer/UnionNode.h>
 #include <Analyzer/Utils.h>
@@ -46,6 +47,7 @@ QueryNode::QueryNode(ContextMutablePtr context_, SettingsChanges settings_change
     children[window_child_index] = std::make_shared<ListNode>();
     children[order_by_child_index] = std::make_shared<ListNode>();
     children[limit_by_child_index] = std::make_shared<ListNode>();
+    children[correlated_columns_list_index] = std::make_shared<ListNode>();
 }
 
 QueryNode::QueryNode(ContextMutablePtr context_)
@@ -108,6 +110,31 @@ void QueryNode::removeUnusedProjectionColumns(const std::unordered_set<size_t> &
     }
 }
 
+ColumnNodePtrWithHashSet QueryNode::getCorrelatedColumnsSet() const
+{
+    ColumnNodePtrWithHashSet result;
+
+    const auto & correlated_columns = getCorrelatedColumns().getNodes();
+    result.reserve(correlated_columns.size());
+
+    for (const auto & column : correlated_columns)
+    {
+        result.insert(std::static_pointer_cast<ColumnNode>(column));
+    }
+    return result;
+}
+
+void QueryNode::addCorrelatedColumn(ColumnNodePtr correlated_column)
+{
+    auto & correlated_columns = getCorrelatedColumns().getNodes();
+    for (const auto & column : correlated_columns)
+    {
+        if (column->isEqual(*correlated_column))
+            return;
+    }
+    correlated_columns.push_back(correlated_column);
+}
+
 void QueryNode::dumpTreeImpl(WriteBuffer & buffer, FormatState & format_state, size_t indent) const
 {
     buffer << std::string(indent, ' ') << "QUERY id: " << format_state.getNodeId(this);
@@ -153,6 +180,12 @@ void QueryNode::dumpTreeImpl(WriteBuffer & buffer, FormatState & format_state, s
     if (!cte_name.empty())
         buffer << ", cte_name: " << cte_name;
 
+    if (isCorrelated())
+    {
+        buffer << ", is_correlated: 1\n" << std::string(indent + 2, ' ') << "CORRELATED COLUMNS\n";
+        getCorrelatedColumns().dumpTreeImpl(buffer, format_state, indent + 4);
+    }
+
     if (hasWith())
     {
         buffer << '\n' << std::string(indent + 2, ' ') << "WITH\n";

src/Analyzer/QueryNode.h

@@ -5,6 +5,7 @@
 #include <Core/NamesAndTypes.h>
 #include <Core/Field.h>
 
+#include <Analyzer/HashUtils.h>
 #include <Analyzer/IQueryTreeNode.h>
 #include <Analyzer/ListNode.h>
 #include <Analyzer/TableExpressionModifiers.h>
@@ -59,6 +60,9 @@ namespace DB
 class QueryNode;
 using QueryNodePtr = std::shared_ptr<QueryNode>;
 
+class ColumnNode;
+using ColumnNodePtr = std::shared_ptr<ColumnNode>;
+
 class QueryNode final : public IQueryTreeNode
 {
 public:
@@ -619,6 +623,30 @@ class QueryNode final : public IQueryTreeNode
     /// Remove unused projection columns
     void removeUnusedProjectionColumns(const std::unordered_set<size_t> & used_projection_columns_indexes);
 
+    bool isCorrelated() const
+    {
+        return !children[correlated_columns_list_index]->as<ListNode>()->getNodes().empty();
+    }
+
+    QueryTreeNodePtr & getCorrelatedColumnsNode()
+    {
+        return children[correlated_columns_list_index];
+    }
+
+    ListNode & getCorrelatedColumns()
+    {
+        return children[correlated_columns_list_index]->as<ListNode &>();
+    }
+
+    const ListNode & getCorrelatedColumns() const
+    {
+        return children[correlated_columns_list_index]->as<ListNode &>();
+    }
+
+    ColumnNodePtrWithHashSet getCorrelatedColumnsSet() const;
+
+    void addCorrelatedColumn(ColumnNodePtr correlated_column);
+
     QueryTreeNodeType getNodeType() const override
     {
         return QueryTreeNodeType::QUERY;
@@ -675,7 +703,8 @@ class QueryNode final : public IQueryTreeNode
     static constexpr size_t limit_by_child_index = 13;
     static constexpr size_t limit_child_index = 14;
     static constexpr size_t offset_child_index = 15;
-    static constexpr size_t children_size = offset_child_index + 1;
+    static constexpr size_t correlated_columns_list_index = 16;
+    static constexpr size_t children_size = correlated_columns_list_index + 1;
 };
 
 }

src/Analyzer/QueryTreePassManager.cpp

@@ -107,6 +107,9 @@ class ValidationChecker : public InDepthQueryTreeVisitor<ValidationChecker>
         if (isNameOfInFunction(function->getFunctionName()))
             return;
 
+        if (function->getFunctionName() == "exists")
+            return;
+
         const auto & expected_argument_types = function->getArgumentTypes();
         size_t expected_argument_types_size = expected_argument_types.size();
         auto actual_argument_columns = function->getArgumentColumns();

src/Analyzer/Resolve/QueryAnalyzer.cpp

@@ -18,6 +18,7 @@
 #include <Functions/IFunctionAdaptors.h>
 #include <Functions/UserDefined/UserDefinedExecutableFunctionFactory.h>
 #include <Functions/UserDefined/UserDefinedSQLFunctionFactory.h>
+#include <Functions/exists.h>
 #include <Functions/grouping.h>
 
 #include <TableFunctions/TableFunctionFactory.h>
@@ -113,6 +114,7 @@ namespace Setting
     extern const SettingsBool allow_suspicious_types_in_order_by;
     extern const SettingsBool allow_not_comparable_types_in_order_by;
     extern const SettingsBool use_concurrency_control;
+    extern const SettingsBool allow_experimental_correlated_subqueries;
     extern const SettingsString implicit_table_at_top_level;
 }
 
@@ -1380,12 +1382,14 @@ IdentifierResolveResult QueryAnalyzer::tryResolveIdentifierInParentScopes(const
     {
         auto current = nodes_to_process.back();
         nodes_to_process.pop_back();
-        if (auto * current_column = current->as<ColumnNode>())
+        if (ColumnNodePtr current_column = std::dynamic_pointer_cast<ColumnNode>(current))
         {
-            if (isDependentColumn(&scope, current_column->getColumnSource()))
+            auto is_correlated_column = checkCorrelatedColumn(&scope, current_column);
+            if (is_correlated_column && !scope.context->getSettingsRef()[Setting::allow_experimental_correlated_subqueries])
             {
                 throw Exception(ErrorCodes::UNSUPPORTED_METHOD,
-                    "Resolved identifier '{}' in parent scope to expression '{}' with correlated column '{}'. In scope {}",
+                    "Resolved identifier '{}' in parent scope to expression '{}' with correlated column '{}'"
+                    " (Enable 'allow_experimental_correlated_subqueries' setting to allow correlated subqueries execution). In scope {}",
                     identifier_lookup.identifier.getFullName(),
                     resolved_identifier->formatASTForErrorMessage(),
                     current_column->getColumnName(),
@@ -2848,27 +2852,6 @@ ProjectionNames QueryAnalyzer::resolveFunction(QueryTreeNodePtr & node, Identifi
         }
     }
 
-    if (is_special_function_exists)
-    {
-        checkFunctionNodeHasEmptyNullsAction(*function_node_ptr);
-        /// Rewrite EXISTS (subquery) into 1 IN (SELECT 1 FROM (subquery) LIMIT 1).
-        auto & exists_subquery_argument = function_node_ptr->getArguments().getNodes().at(0);
-
-        auto constant_data_type = std::make_shared<DataTypeUInt64>();
-
-        auto in_subquery = std::make_shared<QueryNode>(Context::createCopy(scope.context));
-        in_subquery->setIsSubquery(true);
-        in_subquery->getProjection().getNodes().push_back(std::make_shared<ConstantNode>(1UL, constant_data_type));
-        in_subquery->getJoinTree() = exists_subquery_argument;
-        in_subquery->getLimit() = std::make_shared<ConstantNode>(1UL, constant_data_type);
-
-        function_node_ptr = std::make_shared<FunctionNode>("in");
-        function_node_ptr->getArguments().getNodes() = {std::make_shared<ConstantNode>(1UL, constant_data_type), in_subquery};
-        node = function_node_ptr;
-        function_name = "in";
-        is_special_function_in = true;
-    }
-
     if (is_special_function_if && !function_node_ptr->getArguments().getNodes().empty())
     {
         checkFunctionNodeHasEmptyNullsAction(*function_node_ptr);
@@ -2925,12 +2908,61 @@ ProjectionNames QueryAnalyzer::resolveFunction(QueryTreeNodePtr & node, Identifi
     }
 
     /// Resolve function arguments
-    bool allow_table_expressions = is_special_function_in;
+    bool allow_table_expressions = is_special_function_in || is_special_function_exists;
     auto arguments_projection_names = resolveExpressionNodeList(function_node_ptr->getArgumentsNode(),
         scope,
         true /*allow_lambda_expression*/,
         allow_table_expressions /*allow_table_expression*/);
 
+    if (is_special_function_exists)
+    {
+        checkFunctionNodeHasEmptyNullsAction(*function_node_ptr);
+        /// Rewrite EXISTS (subquery) into 1 IN (SELECT 1 FROM (subquery) LIMIT 1).
+        auto & exists_subquery_argument = function_node_ptr->getArguments().getNodes().at(0);
+        bool correlated_exists_subquery = exists_subquery_argument->getNodeType() == QueryTreeNodeType::QUERY
+            ? exists_subquery_argument->as<QueryNode>()->isCorrelated()
+            : exists_subquery_argument->as<UnionNode>()->isCorrelated();
+        if (!correlated_exists_subquery)
+        {
+            auto constant_data_type = std::make_shared<DataTypeUInt64>();
+
+            auto in_subquery = std::make_shared<QueryNode>(Context::createCopy(scope.context));
+            in_subquery->setIsSubquery(true);
+            in_subquery->getProjection().getNodes().push_back(std::make_shared<ConstantNode>(1UL, constant_data_type));
+            in_subquery->getJoinTree() = exists_subquery_argument;
+            in_subquery->getLimit() = std::make_shared<ConstantNode>(1UL, constant_data_type);
+
+            function_node_ptr = std::make_shared<FunctionNode>("in");
+            function_node_ptr->getArguments().getNodes() = {
+                std::make_shared<ConstantNode>(1UL, constant_data_type),
+                std::move(in_subquery)
+            };
+
+            /// Resolve modified arguments
+            arguments_projection_names = resolveExpressionNodeList(function_node_ptr->getArgumentsNode(),
+                scope,
+                true /*allow_lambda_expression*/,
+                true /*allow_table_expression*/);
+
+            node = function_node_ptr;
+            function_name = "in";
+            is_special_function_in = true;
+        }
+        else
+        {
+            /// Subquery is correlated and EXISTS can not be replaced by IN function.
+            /// EXISTS function will be replated by JOIN during query planning.
+            auto function_exists = std::make_shared<FunctionExists>();
+            function_node_ptr->resolveAsFunction(
+                std::make_shared<FunctionToFunctionBaseAdaptor>(
+                    function_exists, DataTypes{}, function_exists->getReturnTypeImpl({})
+                )
+            );
+
+            return { calculateFunctionProjectionName(node, parameters_projection_names, arguments_projection_names) };
+        }
+    }
+
     /// Mask arguments if needed
     if (!scope.context->getSettingsRef()[Setting::format_display_secrets_in_show_and_select])
     {
@@ -2981,6 +3013,10 @@ ProjectionNames QueryAnalyzer::resolveFunction(QueryTreeNodePtr & node, Identifi
             throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function '{}' expects 2 arguments", function_name);
 
         auto & in_second_argument = function_in_arguments_nodes[1];
+        if (isCorrelatedQueryOrUnionNode(function_in_arguments_nodes[0]) || isCorrelatedQueryOrUnionNode(function_in_arguments_nodes[1]))
+            throw Exception(ErrorCodes::NOT_IMPLEMENTED,
+                "Correlated subqueries are not supported as IN function arguments yet, but found in expression: {}",
+                node->formatASTForErrorMessage());
         auto * table_node = in_second_argument->as<TableNode>();
         auto * table_function_node = in_second_argument->as<TableFunctionNode>();
 
@@ -5143,6 +5179,16 @@ void QueryAnalyzer::resolveJoin(QueryTreeNodePtr & join_node, IdentifierResolveS
     resolveQueryJoinTreeNode(join_node_typed.getRightTableExpression(), scope, expressions_visitor);
     validateJoinTableExpressionWithoutAlias(join_node, join_node_typed.getRightTableExpression(), scope);
 
+    if (isCorrelatedQueryOrUnionNode(join_node_typed.getLeftTableExpression()))
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED,
+            "Correlated subqueries are not supported in JOINs yet, but found in expression: {}",
+            join_node_typed.getLeftTableExpression()->formatASTForErrorMessage());
+
+    if (isCorrelatedQueryOrUnionNode(join_node_typed.getRightTableExpression()))
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED,
+            "Correlated subqueries are not supported in JOINs yet, but found in expression: {}",
+            join_node_typed.getRightTableExpression()->formatASTForErrorMessage());
+
     if (!join_node_typed.getLeftTableExpression()->hasAlias() && !join_node_typed.getRightTableExpression()->hasAlias())
         checkDuplicateTableNamesOrAliasForPasteJoin(join_node_typed, scope);
 

src/Analyzer/TableNode.cpp

@@ -85,7 +85,10 @@ void TableNode::updateTreeHashImpl(HashState & state, CompareOptions) const
     }
     else
     {
-        auto full_name = storage_id.getFullNameNotQuoted();
+        // In case of cross-replication we don't know what database is used for the table.
+        // `storage_id.hasDatabase()` can return false only on the initiator node.
+        // Each shard will use the default database (in the case of cross-replication shards may have different defaults).
+        auto full_name = storage_id.hasDatabase() ? storage_id.getFullNameNotQuoted() : storage_id.getTableName();
         state.update(full_name.size());
         state.update(full_name);
     }