Merge pull request ClickHouse#79600 from ClickHouse/analyzer-correlated-scalar

Support scalar correlated subqueries in WHERE clause

Author: KochetovNicolai

src/Analyzer/QueryNode.cpp

@@ -1,3 +1,4 @@
+#include <memory>
 #include <Analyzer/QueryNode.h>
 
 #include <fmt/core.h>
@@ -34,6 +35,7 @@ namespace ErrorCodes
 {
     extern const int LOGICAL_ERROR;
     extern const int BAD_ARGUMENTS;
+    extern const int UNSUPPORTED_METHOD;
 }
 
 QueryNode::QueryNode(ContextMutablePtr context_, SettingsChanges settings_changes_)
@@ -135,6 +137,22 @@ void QueryNode::addCorrelatedColumn(const QueryTreeNodePtr & correlated_column)
     correlated_columns.push_back(correlated_column);
 }
 
+DataTypePtr QueryNode::getResultType() const
+{
+    if (isCorrelated())
+    {
+        if (projection_columns.size() == 1)
+        {
+            return projection_columns[0].type;
+        }
+        else
+            throw Exception(ErrorCodes::UNSUPPORTED_METHOD,
+                "Method getResultType is supported only for correlated query node with 1 column, but got {}",
+                projection_columns.size());
+    }
+    throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "Method getResultType is supported only for correlated query node");
+}
+
 void QueryNode::dumpTreeImpl(WriteBuffer & buffer, FormatState & format_state, size_t indent) const
 {
     buffer << std::string(indent, ' ') << "QUERY id: " << format_state.getNodeId(this);

src/Analyzer/QueryNode.h

@@ -647,6 +647,10 @@ class QueryNode final : public IQueryTreeNode
 
     void addCorrelatedColumn(const QueryTreeNodePtr & correlated_column);
 
+    /// Returns result type of projection expression if query is correlated
+    /// or throws an exception otherwise.
+    DataTypePtr getResultType() const override;
+
     QueryTreeNodeType getNodeType() const override
     {
         return QueryTreeNodeType::QUERY;

src/Analyzer/Resolve/QueryAnalyzer.cpp

@@ -547,6 +547,11 @@ void QueryAnalyzer::evaluateScalarSubqueryIfNeeded(QueryTreeNodePtr & node, Iden
             node->getNodeTypeName(),
             node->formatASTForErrorMessage());
 
+    bool is_correlated_subquery = (query_node != nullptr && query_node->isCorrelated())
+                                || (union_node != nullptr && union_node->isCorrelated());
+    if (is_correlated_subquery)
+        throw Exception(ErrorCodes::BAD_ARGUMENTS, "Cannot evaluate correlated scalar subquery");
+
     auto & context = scope.context;
 
     Block scalar_block;
@@ -2533,7 +2538,8 @@ ProjectionName QueryAnalyzer::resolveWindow(QueryTreeNodePtr & node, IdentifierR
     auto & window_node = node->as<WindowNode &>();
     window_node.setParentWindowName({});
 
-    ProjectionNames partition_by_projection_names = resolveExpressionNodeList(window_node.getPartitionByNode(),
+    ProjectionNames partition_by_projection_names = resolveExpressionNodeList(
+        window_node.getPartitionByNode(),
         scope,
         false /*allow_lambda_expression*/,
         false /*allow_table_expression*/);
@@ -2743,7 +2749,8 @@ ProjectionNames QueryAnalyzer::resolveFunction(QueryTreeNodePtr & node, Identifi
 
     /// Resolve function parameters
 
-    auto parameters_projection_names = resolveExpressionNodeList(function_node_ptr->getParametersNode(),
+    auto parameters_projection_names = resolveExpressionNodeList(
+        function_node_ptr->getParametersNode(),
         scope,
         false /*allow_lambda_expression*/,
         false /*allow_table_expression*/);
@@ -2924,7 +2931,8 @@ ProjectionNames QueryAnalyzer::resolveFunction(QueryTreeNodePtr & node, Identifi
 
     /// Resolve function arguments
     bool allow_table_expressions = is_special_function_in || is_special_function_exists;
-    auto arguments_projection_names = resolveExpressionNodeList(function_node_ptr->getArgumentsNode(),
+    auto arguments_projection_names = resolveExpressionNodeList(
+        function_node_ptr->getArgumentsNode(),
         scope,
         true /*allow_lambda_expression*/,
         allow_table_expressions /*allow_table_expression*/);
@@ -3929,7 +3937,11 @@ ProjectionNames QueryAnalyzer::resolveExpressionNode(
             else
                 resolveUnion(node, subquery_scope);
 
-            if (!allow_table_expression)
+            bool is_correlated_subquery = node_type == QueryTreeNodeType::QUERY
+                ? node->as<QueryNode>()->isCorrelated()
+                : node->as<UnionNode>()->isCorrelated();
+
+            if (!allow_table_expression && !is_correlated_subquery)
                 evaluateScalarSubqueryIfNeeded(node, subquery_scope);
 
             if (result_projection_names.empty())
@@ -4022,7 +4034,12 @@ ProjectionNames QueryAnalyzer::resolveExpressionNode(
   * Example: CREATE TABLE test_table (id UInt64, value UInt64) ENGINE=TinyLog; SELECT plus(*) FROM test_table;
   * Example: SELECT *** FROM system.one;
   */
-ProjectionNames QueryAnalyzer::resolveExpressionNodeList(QueryTreeNodePtr & node_list, IdentifierResolveScope & scope, bool allow_lambda_expression, bool allow_table_expression)
+ProjectionNames QueryAnalyzer::resolveExpressionNodeList(
+    QueryTreeNodePtr & node_list,
+    IdentifierResolveScope & scope,
+    bool allow_lambda_expression,
+    bool allow_table_expression
+)
 {
     auto & node_list_typed = node_list->as<ListNode &>();
     size_t node_list_size = node_list_typed.getNodes().size();

src/Interpreters/Context.cpp

@@ -5529,10 +5529,10 @@ void Context::setGoogleProtosPath(const String & path)
     shared->google_protos_path = path;
 }
 
-Context::SampleBlockCache & Context::getSampleBlockCache() const
+std::pair<Context::SampleBlockCache *, std::unique_lock<std::mutex>> Context::getSampleBlockCache() const
 {
     assert(hasQueryContext());
-    return getQueryContext()->sample_block_cache;
+    return std::make_pair(&getQueryContext()->sample_block_cache, std::unique_lock(getQueryContext()->sample_block_cache_mutex));
 }
 
 

src/Interpreters/Context.h

@@ -534,6 +534,7 @@ class ContextData
 protected:
     using SampleBlockCache = std::unordered_map<std::string, Block>;
     mutable SampleBlockCache sample_block_cache;
+    mutable std::mutex sample_block_cache_mutex;
 
     PartUUIDsPtr part_uuids; /// set of parts' uuids, is used for query parts deduplication
     PartUUIDsPtr ignored_part_uuids; /// set of parts' uuids are meant to be excluded from query processing
@@ -1405,7 +1406,7 @@ class Context: public ContextData, public std::enable_shared_from_this<Context>
     String getGoogleProtosPath() const;
     void setGoogleProtosPath(const String & path);
 
-    SampleBlockCache & getSampleBlockCache() const;
+    std::pair<Context::SampleBlockCache *, std::unique_lock<std::mutex>> getSampleBlockCache() const;
 
     /// Query parameters for prepared statements.
     bool hasQueryParameters() const;

src/Interpreters/InterpreterSelectWithUnionQuery.cpp

@@ -283,20 +283,23 @@ Block InterpreterSelectWithUnionQuery::getSampleBlock(const ASTPtr & query_ptr_,
         return InterpreterSelectWithUnionQuery(query_ptr_, context_, std::move(options.analyze())).getSampleBlock();
     }
 
-    auto & cache = context_->getSampleBlockCache();
     /// Using query string because query_ptr changes for every internal SELECT
     auto key = query_ptr_->formatWithSecretsOneLine();
-    if (cache.find(key) != cache.end())
     {
-        return cache[key];
+        auto [cache, lock] = context_->getSampleBlockCache();
+        if (cache->find(key) != cache->end())
+            return cache->at(key);
     }
 
     SelectQueryOptions options;
     if (is_subquery)
         options = options.subquery();
     if (is_create_parameterized_view)
         options = options.createParameterizedView();
-    return cache[key] = InterpreterSelectWithUnionQuery(query_ptr_, context_, std::move(options.analyze())).getSampleBlock();
+
+    auto sample_block = InterpreterSelectWithUnionQuery(query_ptr_, context_, std::move(options.analyze())).getSampleBlock();
+    auto [cache, lock] = context_->getSampleBlockCache();
+    return (*cache)[key] = sample_block;
 }
 
 

src/Interpreters/JoinInfo.cpp

@@ -356,6 +356,16 @@ JoinActionRef JoinActionRef::deserialize(ReadBuffer & in, const ActionsDAGRawPtr
     return res;
 }
 
+JoinActionRef JoinActionRef::clone(const ActionsDAG * actions_dag_) const
+{
+    return JoinActionRef{actions_dag_, column_name};
+}
+
+JoinActionRef::JoinActionRef(const ActionsDAG * actions_dag_, const String & column_name_)
+    : actions_dag(actions_dag_)
+    , column_name(column_name_)
+{}
+
 void JoinPredicate::serialize(WriteBuffer & out, const JoinActionRef::ActionsDAGRawPtrs & dags) const
 {
     serializePredicateOperator(op, out);
@@ -435,6 +445,40 @@ JoinCondition JoinCondition::deserialize(ReadBuffer & in, const JoinActionRef::A
     };
 }
 
+JoinCondition JoinCondition::clone(const JoinExpressionActions & expression_actions) const
+{
+    JoinCondition copy;
+
+    copy.predicates.reserve(predicates.size());
+    for (const auto & predicate : predicates)
+    {
+        copy.predicates.emplace_back(
+            predicate.left_node.clone(expression_actions.left_pre_join_actions.get()),
+            predicate.right_node.clone(expression_actions.right_pre_join_actions.get()),
+            predicate.op);
+    }
+
+    copy.left_filter_conditions.reserve(left_filter_conditions.size());
+    for (const auto & condition: left_filter_conditions)
+    {
+        copy.left_filter_conditions.emplace_back(condition.clone(expression_actions.left_pre_join_actions.get()));
+    }
+
+    copy.right_filter_conditions.reserve(right_filter_conditions.size());
+    for (const auto & condition: right_filter_conditions)
+    {
+        copy.right_filter_conditions.emplace_back(condition.clone(expression_actions.right_pre_join_actions.get()));
+    }
+
+    copy.residual_conditions.reserve(residual_conditions.size());
+    for (const auto & condition: residual_conditions)
+    {
+        copy.residual_conditions.emplace_back(condition.clone(expression_actions.post_join_actions.get()));
+    }
+
+    return copy;
+}
+
 void JoinExpression::serialize(WriteBuffer & out, const JoinActionRef::ActionsDAGRawPtrs & dags) const
 {
     UInt8 is_using_flag = is_using ? 1 : 0;
@@ -473,6 +517,20 @@ JoinExpression JoinExpression::deserialize(ReadBuffer & in, const JoinActionRef:
     return {std::move(condition), std::move(disjunctive_conditions), bool(is_using_flag)};
 }
 
+JoinExpression JoinExpression::clone(const JoinExpressionActions & expression_copy) const
+{
+    JoinExpression copy;
+    copy.condition = condition.clone(expression_copy);
+
+    copy.disjunctive_conditions.reserve(disjunctive_conditions.size());
+    for (const auto & disjunctive_condition : disjunctive_conditions)
+        copy.disjunctive_conditions.emplace_back(disjunctive_condition.clone(expression_copy));
+
+    copy.is_using = is_using;
+
+    return copy;
+}
+
 void JoinInfo::serialize(WriteBuffer & out, const JoinActionRef::ActionsDAGRawPtrs & dags) const
 {
     expression.serialize(out, dags);

src/Interpreters/JoinInfo.h

@@ -75,6 +75,20 @@ struct JoinExpressionActions
     {
     }
 
+    JoinExpressionActions clone() const
+    {
+        return  JoinExpressionActions(
+            std::make_unique<ActionsDAG>(left_pre_join_actions->clone()),
+            std::make_unique<ActionsDAG>(right_pre_join_actions->clone()),
+            std::make_unique<ActionsDAG>(post_join_actions->clone())
+        );
+    }
+
+    bool hasCorrelatedExpressions() const noexcept
+    {
+        return left_pre_join_actions->hasCorrelatedColumns() || right_pre_join_actions->hasCorrelatedColumns() || post_join_actions->hasCorrelatedColumns();
+    }
+
     ActionsDAGPtr left_pre_join_actions;
     ActionsDAGPtr right_pre_join_actions;
     ActionsDAGPtr post_join_actions;
@@ -103,7 +117,11 @@ class JoinActionRef
     void serialize(WriteBuffer & out, const ActionsDAGRawPtrs & dags) const;
     static JoinActionRef deserialize(ReadBuffer & in, const ActionsDAGRawPtrs & dags);
 
+    JoinActionRef clone(const ActionsDAG * actions_dag_) const;
+
 private:
+    JoinActionRef(const ActionsDAG * actions_dag_, const String & column_name_);
+
     const ActionsDAG * actions_dag = nullptr;
     String column_name;
 };
@@ -136,6 +154,8 @@ struct JoinCondition
 
     void serialize(WriteBuffer & out, const JoinActionRef::ActionsDAGRawPtrs & dags) const;
     static JoinCondition deserialize(ReadBuffer & in, const JoinActionRef::ActionsDAGRawPtrs & dags);
+
+    JoinCondition clone(const JoinExpressionActions & expression_actions) const;
 };
 
 struct JoinExpression
@@ -154,6 +174,8 @@ struct JoinExpression
 
     void serialize(WriteBuffer & out, const JoinActionRef::ActionsDAGRawPtrs & dags) const;
     static JoinExpression deserialize(ReadBuffer & in, const JoinActionRef::ActionsDAGRawPtrs & dags);
+
+    JoinExpression clone(const JoinExpressionActions & expression_copy) const;
 };
 
 struct JoinInfo
@@ -170,6 +192,11 @@ struct JoinInfo
     /// The locality of the join (e.g., LOCAL, GLOBAL)
     JoinLocality locality;
 
+    JoinInfo clone(const JoinExpressionActions & expression_actions) const
+    {
+        return JoinInfo{ expression.clone(expression_actions), kind, strictness, locality};
+    }
+
     void serialize(WriteBuffer & out, const JoinActionRef::ActionsDAGRawPtrs & dags) const;
     static JoinInfo deserialize(ReadBuffer & in, const JoinActionRef::ActionsDAGRawPtrs & dags);
 };

src/Planner/CollectColumnIdentifiers.cpp

@@ -1,7 +1,11 @@
 #include <Planner/CollectColumnIdentifiers.h>
 
-#include <Analyzer/InDepthQueryTreeVisitor.h>
 #include <Analyzer/ColumnNode.h>
+#include <Analyzer/FunctionNode.h>
+#include <Analyzer/InDepthQueryTreeVisitor.h>
+#include <Analyzer/QueryNode.h>
+#include <Analyzer/UnionNode.h>
+#include <Analyzer/Utils.h>
 
 #include <Planner/PlannerContext.h>
 
@@ -35,6 +39,30 @@ class CollectTopLevelColumnIdentifiersVisitor : public ConstInDepthQueryTreeVisi
 
     void visitImpl(const QueryTreeNodePtr & node)
     {
+        if (node->getNodeType() == QueryTreeNodeType::FUNCTION)
+        {
+            auto * function_node = node->as<FunctionNode>();
+            for (const auto & argument : function_node->getArguments().getNodes())
+            {
+                if (!isCorrelatedQueryOrUnionNode(argument))
+                    continue;
+
+                auto * query_node = argument->as<QueryNode>();
+                auto * union_node = argument->as<UnionNode>();
+
+                const auto & correlated_columns = query_node != nullptr ? query_node->getCorrelatedColumns() : union_node->getCorrelatedColumns();
+                for (const auto & column : correlated_columns)
+                {
+                    const auto * column_identifier = planner_context->getColumnNodeIdentifierOrNull(column);
+                    if (!column_identifier)
+                        return;
+
+                    used_identifiers.insert(*column_identifier);
+                }
+            }
+            return;
+        }
+
         if (node->getNodeType() != QueryTreeNodeType::COLUMN)
             return;