Merge pull request ClickHouse#79399 from ClickHouse/improve-correlated-subqueries

Follow up to ClickHouse#76078

Author: novikd

src/Analyzer/QueryNode.cpp

@@ -124,7 +124,7 @@ ColumnNodePtrWithHashSet QueryNode::getCorrelatedColumnsSet() const
     return result;
 }
 
-void QueryNode::addCorrelatedColumn(ColumnNodePtr correlated_column)
+void QueryNode::addCorrelatedColumn(const QueryTreeNodePtr & correlated_column)
 {
     auto & correlated_columns = getCorrelatedColumns().getNodes();
     for (const auto & column : correlated_columns)

src/Analyzer/QueryNode.h

@@ -645,7 +645,7 @@ class QueryNode final : public IQueryTreeNode
 
     ColumnNodePtrWithHashSet getCorrelatedColumnsSet() const;
 
-    void addCorrelatedColumn(ColumnNodePtr correlated_column);
+    void addCorrelatedColumn(const QueryTreeNodePtr & correlated_column);
 
     QueryTreeNodeType getNodeType() const override
     {

src/Analyzer/Resolve/QueryAnalyzer.cpp

@@ -1382,17 +1382,17 @@ IdentifierResolveResult QueryAnalyzer::tryResolveIdentifierInParentScopes(const
     {
         auto current = nodes_to_process.back();
         nodes_to_process.pop_back();
-        if (ColumnNodePtr current_column = std::dynamic_pointer_cast<ColumnNode>(current))
+        if (current->getNodeType() == QueryTreeNodeType::COLUMN)
         {
-            auto is_correlated_column = checkCorrelatedColumn(&scope, current_column);
+            auto is_correlated_column = checkCorrelatedColumn(&scope, current);
             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 '{}'"
                     " (Enable 'allow_experimental_correlated_subqueries' setting to allow correlated subqueries execution). In scope {}",
                     identifier_lookup.identifier.getFullName(),
                     resolved_identifier->formatASTForErrorMessage(),
-                    current_column->getColumnName(),
+                    current->as<ColumnNode>()->getColumnName(),
                     scope.scope_node->formatASTForErrorMessage());
             }
         }
@@ -2932,13 +2932,13 @@ 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);
         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)
         {
+            /// Rewrite EXISTS (subquery) into 1 IN (SELECT 1 FROM (subquery) LIMIT 1).
             auto constant_data_type = std::make_shared<DataTypeUInt64>();
 
             auto in_subquery = std::make_shared<QueryNode>(Context::createCopy(scope.context));

src/Analyzer/UnionNode.cpp

@@ -132,7 +132,7 @@ void UnionNode::removeUnusedProjectionColumns(const std::unordered_set<size_t> &
     }
 }
 
-void UnionNode::addCorrelatedColumn(ColumnNodePtr correlated_column)
+void UnionNode::addCorrelatedColumn(const QueryTreeNodePtr & correlated_column)
 {
     auto & correlated_columns = getCorrelatedColumns().getNodes();
     for (const auto & column : correlated_columns)

src/Analyzer/UnionNode.h

@@ -193,7 +193,7 @@ class UnionNode final : public IQueryTreeNode
         return children[correlated_columns_list_index]->as<ListNode &>();
     }
 
-    void addCorrelatedColumn(ColumnNodePtr correlated_column);
+    void addCorrelatedColumn(const QueryTreeNodePtr & correlated_column);
 
     QueryTreeNodeType getNodeType() const override
     {

src/Analyzer/Utils.cpp

@@ -252,10 +252,13 @@ bool isCorrelatedQueryOrUnionNode(const QueryTreeNodePtr & node)
 
 bool checkCorrelatedColumn(
     IdentifierResolveScope * scope_to_check,
-    const ColumnNodePtr & column
+    const QueryTreeNodePtr & column
 )
 {
-    auto column_source = column->getColumnSource();
+    auto * current_scope = scope_to_check;
+    chassert(column->getNodeType() == QueryTreeNodeType::COLUMN);
+    auto * column_node = column->as<ColumnNode>();
+    auto column_source = column_node->getColumnSource();
 
     /// The case of lambda argument. Example:
     /// arrayMap(X -> X + Y, [0])
@@ -283,7 +286,11 @@ bool checkCorrelatedColumn(
         scope_to_check = scope_to_check->parent_scope;
     }
     if (!scope_to_check)
-        throw Exception(ErrorCodes::LOGICAL_ERROR, "Cannot find the original scope of the column");
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Cannot find the original scope of the column '{}'. Current scope: {}",
+            column_node->getColumnName(),
+            current_scope->scope_node->formatASTForErrorMessage());
 
     return is_correlated;
 }

src/Analyzer/Utils.h

@@ -72,7 +72,7 @@ bool isCorrelatedQueryOrUnionNode(const QueryTreeNodePtr & node);
  */
 bool checkCorrelatedColumn(
     IdentifierResolveScope * scope_to_check,
-    const ColumnNodePtr & column
+    const QueryTreeNodePtr & column
 );
 
 DataTypePtr getExpressionNodeResultTypeOrNull(const QueryTreeNodePtr & query_tree_node);

src/Planner/CollectTableExpressionData.cpp

@@ -49,6 +49,11 @@ class CollectSourceColumnsVisitor : public InDepthQueryTreeVisitorWithContext<Co
 
         if (isExistsFunction(node))
         {
+            /// Add used in correlated subquery columns to the table expression data.
+            /// These columns can be used only by correlated subquery, but still they
+            /// must be read by query plan for current query.
+            ///
+            /// Example: SELECT 1 FROM table as t WHERE EXISTS (SELECT * FROM numbers(10) WHERE t.id = number);
             auto * function_node = node->as<FunctionNode>();
             const auto & subquery_argument = function_node->getArguments().getNodes().front();
             auto * query_node = subquery_argument->as<QueryNode>();