diff --git a/src/Storages/IStorageCluster.cpp b/src/Storages/IStorageCluster.cpp index 7e972f2802038..c8e1633adc790 100644 --- a/src/Storages/IStorageCluster.cpp +++ b/src/Storages/IStorageCluster.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -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; } diff --git a/src/Storages/IStorageCluster.h b/src/Storages/IStorageCluster.h index c592a259e12b0..5bbca00713c77 100644 --- a/src/Storages/IStorageCluster.h +++ b/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); }; diff --git a/src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp b/src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp index 545122e68f0f4..f686b67d37725 100644 --- a/src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp +++ b/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_); diff --git a/src/Storages/ObjectStorage/StorageObjectStorageCluster.h b/src/Storages/ObjectStorage/StorageObjectStorageCluster.h index 9a4a4fc53122e..cac884a204ce6 100644 --- a/src/Storages/ObjectStorage/StorageObjectStorageCluster.h +++ b/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 pure_storage; diff --git a/tests/integration/test_database_iceberg/test.py b/tests/integration/test_database_iceberg/test.py index 430f3d593f656..6f4834209be8f 100644 --- a/tests/integration/test_database_iceberg/test.py +++ b/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"]