Some fixes

ianton-ru · ianton-ru · commit 43825ba92394 · 2025-10-15T18:18:07.000+02:00
diff --git a/src/Storages/IStorageCluster.cpp b/src/Storages/IStorageCluster.cpp
@@ -33,6 +33,7 @@
 #include <Analyzer/ColumnNode.h>
 #include <Analyzer/JoinNode.h>
 #include <Analyzer/InDepthQueryTreeVisitor.h>
+#include <Analyzer/Utils.h>
 #include <Storages/StorageDistributed.h>
 #include <TableFunctions/TableFunctionFactory.h>
 
@@ -222,37 +223,40 @@ void IStorageCluster::updateQueryWithJoinToSendIfNeeded(
     {
     case ObjectStorageClusterJoinMode::LOCAL:
     {
-        auto modified_query_tree = query_tree->clone();
-        bool need_modify = false;
+        if (has_join || has_local_columns_in_where)
+        {
+            auto modified_query_tree = query_tree->clone();
 
-        SearcherVisitor table_function_searcher({QueryTreeNodeType::TABLE, QueryTreeNodeType::TABLE_FUNCTION}, context);
-        table_function_searcher.visit(query_tree);
-        auto table_function_node = table_function_searcher.getNode();
-        if (!table_function_node)
-            throw Exception(ErrorCodes::LOGICAL_ERROR, "Can't find table function node");
+            SearcherVisitor table_function_searcher({QueryTreeNodeType::TABLE, QueryTreeNodeType::TABLE_FUNCTION}, context);
+            table_function_searcher.visit(modified_query_tree);
+            auto table_function_node = table_function_searcher.getNode();
+            if (!table_function_node)
+                throw Exception(ErrorCodes::LOGICAL_ERROR, "Can't find table function node");
 
-        if (has_join)
-        {
             QueryTreeNodePtr query_tree_distributed;
 
             auto & query_node = modified_query_tree->as<QueryNode &>();
 
-            if (table_function_searcher.getType().value() == QueryTreeNodeType::TABLE_FUNCTION)
-            {
-                auto table_function = extractTableFunctionASTPtrFromSelectQuery(query_to_send);
-                query_tree_distributed = buildTableFunctionQueryTree(table_function, context);
-                auto & table_function_ast = table_function->as<ASTFunction &>();
-                query_tree_distributed->setAlias(table_function_ast.alias);
-            }
-            else
+            if (has_join)
             {
-                auto join_node = query_node.getJoinTree();
-                query_tree_distributed = join_node->as<JoinNode>()->getLeftTableExpression()->clone();
+                if (table_function_searcher.getType().value() == QueryTreeNodeType::TABLE_FUNCTION)
+                {
+                    auto table_function = extractTableFunctionASTPtrFromSelectQuery(query_to_send);
+                    query_tree_distributed = buildTableFunctionQueryTree(table_function, context);
+                    auto & table_function_ast = table_function->as<ASTFunction &>();
+                    query_tree_distributed->setAlias(table_function_ast.alias);
+                }
+                else
+                {
+                    auto join_node = query_node.getJoinTree();
+                    query_tree_distributed = join_node->as<JoinNode>()->getLeftTableExpression()->clone();
+                }
             }
 
             // Find add used columns from table function to make proper projection list
+            // Need to do before changing WHERE condition
             CollectUsedColumnsForSourceVisitor collector(table_function_node, context);
-            collector.visit(query_tree);
+            collector.visit(modified_query_tree);
             const auto & columns = collector.getColumns();
 
             query_node.resolveProjectionColumns(columns);
@@ -262,20 +266,26 @@ void IStorageCluster::updateQueryWithJoinToSendIfNeeded(
                 column_nodes_to_select->getNodes().emplace_back(std::make_shared<ColumnNode>(column, table_function_node));
             query_node.getProjectionNode() = column_nodes_to_select;
 
-            // Left only table function to send on cluster nodes
-            modified_query_tree = modified_query_tree->cloneAndReplace(query_node.getJoinTree(), query_tree_distributed);
+            if (has_local_columns_in_where)
+            {
+                if (query_node.getPrewhere())
+                    removeExpressionsThatDoNotDependOnTableIdentifiers(query_node.getPrewhere(), table_function_node, context);
+                if (query_node.getWhere())
+                    removeExpressionsThatDoNotDependOnTableIdentifiers(query_node.getWhere(), table_function_node, context);
+            }
+
+            query_node.getOrderByNode() = std::make_shared<ListNode>();
+            query_node.getGroupByNode() = std::make_shared<ListNode>();
 
-            need_modify = true;
-        }
+            if (query_tree_distributed)
+            {
+                // Left only table function to send on cluster nodes
+                modified_query_tree = modified_query_tree->cloneAndReplace(query_node.getJoinTree(), query_tree_distributed);
+            }
 
-        if (has_local_columns_in_where)
-        {
-            auto & query_node = modified_query_tree->as<QueryNode &>();
-            query_node.getWhere() = {};
+            query_to_send = queryNodeToDistributedSelectQuery(modified_query_tree);
         }
 
-        if (need_modify)
-            query_to_send = queryNodeToDistributedSelectQuery(modified_query_tree);
         return;
     }
     case ObjectStorageClusterJoinMode::GLOBAL:
@@ -530,12 +540,16 @@ QueryProcessingStage::Enum IStorageCluster::getQueryProcessingStage(
             throw Exception(ErrorCodes::LOGICAL_ERROR, "Can't find table or table function node");
 
         auto & query_node = query_info.query_tree->as<QueryNode &>();
-        if (query_node.hasWhere())
+        if (query_node.hasWhere() || query_node.hasPrewhere())
         {
             CollectUsedColumnsForSourceVisitor collector_where(table_function_node, context, true);
-            collector_where.visit(query_node.getWhere());
+            if (query_node.hasPrewhere())
+                collector_where.visit(query_node.getPrewhere());
+            if (query_node.hasWhere())
+                collector_where.visit(query_node.getWhere());
 
-            // Can't use 'WHERE' on remote node if it contains columns from other sources
+            // SELECT x FROM datalake.table WHERE x IN local.table
+            // Need to modify 'WHERE' on remote node if it contains columns from other sources
             if (!collector_where.getColumns().empty())
                 has_local_columns_in_where = true;
         }
diff --git a/tests/integration/test_s3_cluster/test.py b/tests/integration/test_s3_cluster/test.py
@@ -1163,6 +1163,20 @@ def test_joins(started_cluster):
     res = list(map(str.split, result5.splitlines()))
     assert len(res) == 6
 
+    result6 = node.query(
+        f"""
+        SELECT name FROM
+            s3Cluster('cluster_simple',
+                'http://minio1:9001/root/data/{{clickhouse,database}}/*', 'minio', '{minio_secret_key}', 'CSV',
+                'name String, value UInt32, polygon Array(Array(Tuple(Float64, Float64)))')
+        WHERE value IN (SELECT id FROM join_table)
+        ORDER BY name
+        SETTINGS object_storage_cluster_join_mode='local';
+        """
+    )
+    res = list(map(str.split, result6.splitlines()))
+    assert len(res) == 25
+
 
 def test_graceful_shutdown(started_cluster):
     node = started_cluster.instances["s0_0_0"]
diff --git a/tests/integration/test_storage_iceberg/test.py b/tests/integration/test_storage_iceberg/test.py
@@ -3492,3 +3492,94 @@ def compare_selects(query):
     compare_selects(f"SELECT _path,* FROM {creation_expression} ORDER BY ALL")
     compare_selects(f"SELECT _path,* FROM {creation_expression} WHERE name_old='vasily' ORDER BY ALL")
     compare_selects(f"SELECT _path,* FROM {creation_expression} WHERE ((tag + length(name_old)) % 2 = 1) ORDER BY ALL")
+
+
+@pytest.mark.parametrize("storage_type", ["s3", "azure"])
+def test_cluster_joins(started_cluster, storage_type):
+    instance = started_cluster.instances["node1"]
+    spark = started_cluster.spark_session
+    TABLE_NAME = "test_cluster_joins_" + storage_type + "_" + get_uuid_str()
+    TABLE_NAME_2 = "test_cluster_joins_2_" + storage_type + "_" + get_uuid_str()
+
+    def execute_spark_query(query: str, table_name):
+        return execute_spark_query_general(
+            spark,
+            started_cluster,
+            storage_type,
+            table_name,
+            query,
+        )
+
+    execute_spark_query(
+        f"""
+            CREATE TABLE {TABLE_NAME} (
+                tag INT,
+                name VARCHAR(50)
+            )
+            USING iceberg
+            OPTIONS('format-version'='2')
+        """, TABLE_NAME
+    )
+
+    execute_spark_query(
+        f"""
+        INSERT INTO {TABLE_NAME} VALUES
+        (1, 'john'),
+        (2, 'jack')
+    """, TABLE_NAME
+    )
+
+    execute_spark_query(
+        f"""
+            CREATE TABLE {TABLE_NAME_2} (
+                id INT,
+                second_name VARCHAR(50)
+            )
+            USING iceberg
+            OPTIONS('format-version'='2')
+        """, TABLE_NAME_2
+    )
+
+    execute_spark_query(
+        f"""
+        INSERT INTO {TABLE_NAME_2} VALUES
+        (1, 'dow'),
+        (2, 'sparrow')
+    """, TABLE_NAME_2
+    )
+
+    creation_expression = get_creation_expression(
+        storage_type, TABLE_NAME, started_cluster, table_function=True, run_on_cluster=True
+    )
+
+    creation_expression_2 = get_creation_expression(
+        storage_type, TABLE_NAME_2, started_cluster, table_function=True, run_on_cluster=True
+    )
+
+    res = instance.query(
+        f"""
+            SELECT t1.name,t2.second_name
+            FROM {creation_expression} AS t1
+                JOIN {creation_expression_2} AS t2
+                ON t1.tag=t2.id
+            ORDER BY ALL
+            SETTINGS object_storage_cluster_join_mode='local'
+        """
+    )
+
+    assert res == "jack\tsparrow\njohn\tdow\n"
+
+    res = instance.query(
+        f"""
+            SELECT name
+            FROM {creation_expression}
+            WHERE tag in (
+                SELECT id
+                FROM {creation_expression_2}
+            )
+            ORDER BY ALL
+            SETTINGS object_storage_cluster_join_mode='local'
+        """
+    )
+
+    assert res == "jack\njohn\n"
