Read Iceberg catalog once during cluster requests.

src/Storages/IStorageCluster.cpp

@@ -3,6 +3,7 @@
 #include <pcg_random.hpp>
 #include <Common/randomSeed.h>
 
+#include <Analyzer/createUniqueAliasesIfNecessary.h>
 #include <Functions/FunctionsLogical.h>
 #include <Functions/IFunctionAdaptors.h>
 
@@ -441,7 +442,8 @@ or (object_storage_cluster_join_mode='global')
 void IStorageCluster::updateQueryWithJoinToSendIfNeeded(
     ASTPtr & query_to_send,
     SelectQueryInfo query_info,
-    const ContextPtr & context)
+    const ContextPtr & context,
+    QueryTreeNodePtr * resolved_query_tree)
 {
     auto object_storage_cluster_join_mode = context->getSettingsRef()[Setting::object_storage_cluster_join_mode];
     switch (object_storage_cluster_join_mode)
@@ -453,7 +455,7 @@ void IStorageCluster::updateQueryWithJoinToSendIfNeeded(
         if (!needsInitiatorLocalJoin(info))
             return;
 
-        rewriteQueryForInitiatorLocalJoin(query_to_send, query_info, info, context);
+        rewriteQueryForInitiatorLocalJoin(query_to_send, query_info, info, context, resolved_query_tree);
         return;
     }
     case ObjectStorageClusterJoinMode::GLOBAL:
@@ -464,7 +466,7 @@ void IStorageCluster::updateQueryWithJoinToSendIfNeeded(
             if (queryHasRightOrFullJoin(query_to_send))
             {
                 auto info = getQueryJoinInfoFromAST(query_to_send);
-                rewriteQueryForInitiatorLocalJoin(query_to_send, query_info, info, context);
+                rewriteQueryForInitiatorLocalJoin(query_to_send, query_info, info, context, resolved_query_tree);
             }
             return;
         }
@@ -726,11 +728,26 @@ void IStorageCluster::read(
 
     SharedHeader sample_block;
 
-    updateQueryWithJoinToSendIfNeeded(query_to_send, query_info, context);
+    /// A join kept on the initiator replaces `query_to_send` but already has a resolved left-scan tree.
+    QueryTreeNodePtr resolved_query_tree;
+    updateQueryWithJoinToSendIfNeeded(query_to_send, query_info, context, &resolved_query_tree);
 
     if (settings[Setting::allow_experimental_analyzer])
     {
-        sample_block = InterpreterSelectQueryAnalyzer::getSampleBlock(query_to_send, context, SelectQueryOptions(processed_stage));
+        /// Re-analyzing the shard AST resolves a `DataLakeCatalog` table again and repeats a remote catalog lookup.
+        if (canReuseResolvedQueryTreeForHeader() && resolved_query_tree)
+            sample_block = InterpreterSelectQueryAnalyzer::getSampleBlock(resolved_query_tree, context, SelectQueryOptions(processed_stage));
+        else if (canReuseResolvedQueryTreeForHeader() && query_to_send == query_info.query && query_info.query_tree)
+        {
+            /// A shard re-analyzes this query and `createUniqueAliasesIfNecessary` restarts `__tableN` at 1.
+            /// The resolved tree still has the outer alias (`__table2` when this table is not leftmost).
+            auto query_info_for_header = query_info;
+            query_info_for_header.query_tree = query_info.query_tree->clone();
+            createUniqueAliasesIfNecessary(query_info_for_header.query_tree, context);
+            sample_block = getHeaderForProcessingStage(column_names, storage_snapshot, query_info_for_header, context, processed_stage);
+        }
+        else
+            sample_block = InterpreterSelectQueryAnalyzer::getSampleBlock(query_to_send, context, SelectQueryOptions(processed_stage));
     }
     else
     {
@@ -1060,7 +1077,8 @@ void IStorageCluster::rewriteQueryForInitiatorLocalJoin(
     ASTPtr & query_to_send,
     const SelectQueryInfo & query_info,
     const QueryTreeInfo & info,
-    const ContextPtr & context)
+    const ContextPtr & context,
+    QueryTreeNodePtr * resolved_query_tree)
 {
     /// Analyzer path: reuse `extractLeftTableExpression` and `buildQueryToReadColumnsFromTableExpression`
     /// (same helpers the planner uses when wrapping `IStorageCluster` in a subquery).
@@ -1102,6 +1120,9 @@ void IStorageCluster::rewriteQueryForInitiatorLocalJoin(
         rewritten_query_node.getPrewhere() = query_node.getPrewhere();
         rewritten_query_node.getWhere() = query_node.getWhere();
 
+        if (resolved_query_tree)
+            *resolved_query_tree = rewritten_query_tree;
+
         query_to_send = queryNodeToDistributedSelectQuery(rewritten_query_tree);
         return;
     }

src/Storages/IStorageCluster.h

@@ -80,7 +80,13 @@ class IStorageCluster : public IStorage
 
 protected:
     virtual void updateBeforeRead(const ContextPtr &) {}
-    void updateQueryWithJoinToSendIfNeeded(ASTPtr & query_to_send, SelectQueryInfo query_info, const ContextPtr & context);
+    void updateQueryWithJoinToSendIfNeeded(
+        ASTPtr & query_to_send,
+        SelectQueryInfo query_info,
+        const ContextPtr & context,
+        QueryTreeNodePtr * resolved_query_tree = nullptr);
+    /// `DataLakeCatalog` tables need this to avoid resolving the catalog table a second time for the header.
+    virtual bool canReuseResolvedQueryTreeForHeader() const { return false; }
 
     virtual void updateConfigurationIfNeeded(ContextPtr /* context */) {}
 
@@ -184,7 +190,8 @@ class IStorageCluster : public IStorage
         ASTPtr & query_to_send,
         const SelectQueryInfo & query_info,
         const QueryTreeInfo & info,
-        const ContextPtr & context);
+        const ContextPtr & context,
+        QueryTreeNodePtr * resolved_query_tree);
 };
 
 

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -126,6 +126,7 @@ StorageObjectStorageCluster::StorageObjectStorageCluster(
         cluster_name_, table_id_, getLogger(fmt::format("{}({})", configuration_->getEngineName(), table_id_.table_name)))
     , configuration{configuration_}
     , object_storage(object_storage_)
+    , is_datalake_catalog_table(is_datalake_query)
 {
     configuration->initPartitionStrategy(partition_by, columns_in_table_or_function_definition, context_);
     configuration->check(context_);

src/Storages/ObjectStorage/StorageObjectStorageCluster.h

@@ -203,6 +203,7 @@ class StorageObjectStorageCluster : public IStorageCluster
 
     /// Whether this query is served by `pure_storage` instead of being distributed over a cluster.
     bool readsFromPureStorage(ContextPtr context) const;
+    bool canReuseResolvedQueryTreeForHeader() const override { return is_datalake_catalog_table; }
 
     void readFallBackToPure(
         QueryPlan & query_plan,
@@ -243,6 +244,8 @@ class StorageObjectStorageCluster : public IStorageCluster
     const ObjectStoragePtr object_storage;
     /// True for explicit *Cluster(...); false for SETTINGS / ENGINE / DataLake paths.
     bool cluster_name_from_function_argument = false;
+    /// True only for tables resolved through a `DataLakeCatalog` database.
+    bool is_datalake_catalog_table = false;
 
     /// non-clustered storage to fall back on pure realisation if needed
     std::shared_ptr<StorageObjectStorage> pure_storage;

tests/integration/test_database_iceberg/test.py

@@ -1441,6 +1441,119 @@ def test_cluster_select(started_cluster):
     assert node2.query(f"SELECT * FROM {CATALOG_NAME}.`{root_namespace}.{table_name}`", settings={"parallel_replicas_for_cluster_engines": 1, "enable_parallel_replicas": 2, "cluster_for_parallel_replicas": "cluster_simple"}) == 'pablo\n'
 
 
+def test_object_storage_cluster_loads_rest_catalog_table_once(started_cluster):
+    node = started_cluster.instances["node1"]
+
+    test_ref = f"test_object_storage_cluster_{uuid.uuid4().hex}"
+    table_name = f"{test_ref}_table"
+    namespace = f"{test_ref}_namespace"
+
+    catalog = load_catalog_impl(started_cluster)
+    catalog.create_namespace(namespace)
+    table = create_table(catalog, namespace, table_name)
+    table.append(pa.Table.from_pylist([generate_record()]))
+
+    create_clickhouse_iceberg_database(started_cluster, node, CATALOG_NAME)
+
+    query_id = f"object_storage_cluster_catalog_{uuid.uuid4().hex}"
+    assert (
+        node.query(
+            f"SELECT * FROM {CATALOG_NAME}.`{namespace}.{table_name}` "
+            "SETTINGS object_storage_cluster = 'cluster_simple' FORMAT Null",
+            query_id=query_id,
+        )
+        == ""
+    )
+
+    node.query("SYSTEM FLUSH LOGS system.query_log")
+    assert (
+        node.query(
+            "SELECT ProfileEvents['DataLakeRestCatalogGetTableMetadata'] "
+            "FROM system.query_log "
+            f"WHERE query_id = '{query_id}' AND type = 'QueryFinish' AND is_initial_query"
+        )
+        == "1\n"
+    )
+
+
+def test_object_storage_cluster_join_loads_rest_catalog_table_once(started_cluster):
+    node = started_cluster.instances["node1"]
+
+    test_ref = f"test_object_storage_cluster_join_{uuid.uuid4().hex}"
+    table_name = f"{test_ref}_table"
+    namespace = f"{test_ref}_namespace"
+    memory_table = f"{test_ref}_memory"
+
+    catalog = load_catalog_impl(started_cluster)
+    catalog.create_namespace(namespace)
+    table = create_table(catalog, namespace, table_name)
+    table.append(pa.Table.from_pylist([generate_record()]))
+
+    create_clickhouse_iceberg_database(started_cluster, node, CATALOG_NAME)
+    node.query(f"CREATE TABLE {memory_table} (symbol String) ENGINE = Memory")
+    node.query(f"INSERT INTO {memory_table} VALUES ('kek')")
+
+    query_id = f"object_storage_cluster_join_catalog_{uuid.uuid4().hex}"
+    assert (
+        node.query(
+            f"SELECT t.symbol FROM {CATALOG_NAME}.`{namespace}.{table_name}` AS t "
+            f"INNER JOIN {memory_table} AS m ON t.symbol = m.symbol "
+            "SETTINGS object_storage_cluster = 'cluster_simple'",
+            query_id=query_id,
+        )
+        == "kek\n"
+    )
+
+    node.query("SYSTEM FLUSH LOGS system.query_log")
+    assert (
+        node.query(
+            "SELECT ProfileEvents['DataLakeRestCatalogGetTableMetadata'] "
+            "FROM system.query_log "
+            f"WHERE query_id = '{query_id}' AND type = 'QueryFinish' AND is_initial_query"
+        )
+        == "1\n"
+    )
+
+
+def test_object_storage_cluster_right_join_loads_rest_catalog_table_once(started_cluster):
+    node = started_cluster.instances["node1"]
+
+    test_ref = f"test_object_storage_cluster_right_join_{uuid.uuid4().hex}"
+    table_name = f"{test_ref}_table"
+    namespace = f"{test_ref}_namespace"
+    memory_table = f"{test_ref}_memory"
+
+    catalog = load_catalog_impl(started_cluster)
+    catalog.create_namespace(namespace)
+    table = create_table(catalog, namespace, table_name)
+    table.append(pa.Table.from_pylist([generate_record()]))
+
+    create_clickhouse_iceberg_database(started_cluster, node, CATALOG_NAME)
+    node.query(f"CREATE TABLE {memory_table} (symbol String) ENGINE = Memory")
+    node.query(f"INSERT INTO {memory_table} VALUES ('kek')")
+
+    query_id = f"object_storage_cluster_right_join_catalog_{uuid.uuid4().hex}"
+    assert (
+        node.query(
+            f"SELECT t.symbol FROM {memory_table} AS m "
+            f"INNER JOIN {CATALOG_NAME}.`{namespace}.{table_name}` AS t ON m.symbol = t.symbol "
+            "SETTINGS object_storage_cluster = 'cluster_simple'",
+            query_id=query_id,
+        )
+        == "kek\n"
+    )
+
+    node.query("SYSTEM FLUSH LOGS system.query_log")
+    assert (
+        node.query(
+            "SELECT ProfileEvents['DataLakeRestCatalogGetTableMetadata'] "
+            "FROM system.query_log "
+            f"WHERE query_id = '{query_id}' AND type = 'QueryFinish' AND is_initial_query"
+        )
+        == "1\n"
+    )
+
+
 def test_on_cluster_ddl_rejected_for_datalake_catalog(started_cluster):
     node1 = started_cluster.instances["node1"]
     node2 = started_cluster.instances["node2"]