[WIP] read identify columns from iceberg metada

src/Interpreters/InterpreterSelectQuery.cpp

@@ -831,7 +831,7 @@ InterpreterSelectQuery::InterpreterSelectQuery(
                 current_info.syntax_analyzer_result = syntax_analyzer_result;
 
                 Names queried_columns = syntax_analyzer_result->requiredSourceColumns();
-                const auto & supported_prewhere_columns = storage->supportedPrewhereColumns();
+                const auto & supported_prewhere_columns = storage->supportedAutomaticPrewhereColumns(storage_snapshot->metadata);
 
                 RangesInDataParts parts_for_estimator;
                 if (storage_snapshot->data)

src/Processors/QueryPlan/Optimizations/optimizePrewhere.cpp

@@ -186,7 +186,7 @@ void optimizePrewhere(QueryPlan::Node & parent_node, const bool remove_unused_co
         storage_snapshot,
         read_from_merge_tree_step ? read_from_merge_tree_step->getConditionSelectivityEstimator(queried_columns) : nullptr,
         queried_columns,
-        storage.supportedPrewhereColumns(),
+        storage.supportedAutomaticPrewhereColumns(storage_snapshot->metadata),
         getLogger("QueryPlanOptimizePrewhere")};
 
     auto optimize_result = where_optimizer.optimize(filter_step->getExpression(),

src/Storages/IStorage.h

@@ -146,6 +146,8 @@ class IStorage : public std::enable_shared_from_this<IStorage>, public TypePromo
     /// This is needed for engines whose aggregates data from multiple tables, like Merge.
     virtual std::optional<NameSet> supportedPrewhereColumns() const { return std::nullopt; }
 
+    virtual std::optional<NameSet> supportedAutomaticPrewhereColumns(const StorageMetadataPtr & /* metadata */) const { return supportedPrewhereColumns(); }
+
     /// Returns true if the storage supports optimization of moving conditions to PREWHERE section.
     virtual bool canMoveConditionsToPrewhere() const { return supportsPrewhere(); }
 

src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergMetadata.cpp

@@ -134,6 +134,7 @@ extern const SettingsBool allow_insert_into_iceberg;
 extern const SettingsBool allow_experimental_iceberg_compaction;
 extern const SettingsBool allow_experimental_expire_snapshots;
 extern const SettingsBool iceberg_delete_data_on_drop;
+extern const SettingsBool allow_experimental_iceberg_read_optimization;
 }
 
 static constexpr size_t MAX_TRANSACTION_RETRIES = 100;
@@ -1236,7 +1237,21 @@ std::unique_ptr<StorageInMemoryMetadata> IcebergMetadata::buildStorageMetadataFr
     result->setColumns(
         ColumnsDescription{*persistent_components.schema_processor->getClickhouseTableSchemaById(iceberg_state.schema_id)});
     result->setDataLakeTableState(state);
-    result->sorting_key = getSortingKey(local_context, iceberg_state);
+
+    auto metadata_object = getMetadataJSONObject(
+        iceberg_state.metadata_file_path,
+        object_storage,
+        persistent_components.metadata_cache,
+        local_context,
+        log,
+        persistent_components.metadata_compression_method,
+        persistent_components.table_uuid);
+
+    result->sorting_key = getSortingKeyFromMetadata(metadata_object, local_context);
+
+    if (local_context->getSettingsRef()[Setting::allow_experimental_iceberg_read_optimization])
+        result->setIdentityPartitionColumns(getIdentityPartitionColumnsFromMetadata(metadata_object));
+
     return result;
 }
 
@@ -1464,10 +1479,14 @@ KeyDescription IcebergMetadata::getSortingKey(ContextPtr local_context, TableSta
         persistent_components.metadata_compression_method,
         persistent_components.table_uuid);
 
+    return getSortingKeyFromMetadata(metadata_object, local_context);
+}
+
+KeyDescription IcebergMetadata::getSortingKeyFromMetadata(const Poco::JSON::Object::Ptr & metadata_object, ContextPtr local_context) const
+{
     auto [schema, current_schema_id] = parseTableSchemaV2Method(metadata_object);
     auto result = getSortingKeyDescriptionFromMetadata(metadata_object, *persistent_components.schema_processor->getClickhouseTableSchemaById(current_schema_id), local_context);
-    auto sort_order_id = metadata_object->getValue<Int64>(f_default_sort_order_id);
-    result.sort_order_id = sort_order_id;
+    result.sort_order_id = metadata_object->getValue<Int64>(f_default_sort_order_id);
     return result;
 }
 

src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergMetadata.h

@@ -209,6 +209,7 @@ class IcebergMetadata : public IDataLakeMetadata
 
     std::optional<String> getPartitionKey(ContextPtr local_context, Iceberg::TableStateSnapshot actual_table_state_snapshot) const;
     KeyDescription getSortingKey(ContextPtr local_context, Iceberg::TableStateSnapshot actual_table_state_snapshot) const;
+    KeyDescription getSortingKeyFromMetadata(const Poco::JSON::Object::Ptr & metadata_object, ContextPtr local_context) const;
 
     bool commitImportPartitionTransactionImpl(
         FileNamesGenerator & filename_generator,

src/Storages/ObjectStorage/DataLakes/Iceberg/Utils.cpp

@@ -108,6 +108,53 @@ static constexpr auto MAX_TRANSACTION_RETRIES = 100;
 namespace DB::Iceberg
 {
 using namespace DB;
+
+namespace
+{
+struct ResolvedPartitionSpec
+{
+    Poco::JSON::Array::Ptr fields;
+    std::unordered_map<Int64, String> source_id_to_column_name;
+};
+
+std::optional<ResolvedPartitionSpec> resolveDefaultPartitionSpec(const Poco::JSON::Object::Ptr & metadata_object)
+{
+    if (!metadata_object->has(f_partition_specs) || !metadata_object->has(f_default_spec_id))
+        return std::nullopt;
+
+    auto partition_spec_id = metadata_object->getValue<Int64>(f_default_spec_id);
+    Poco::JSON::Array::Ptr partition_specs = metadata_object->getArray(f_partition_specs);
+    if (!partition_specs)
+        return std::nullopt;
+
+    std::unordered_map<Int64, String> source_id_to_column_name;
+    auto [schema, current_schema_id] = parseTableSchemaV2Method(metadata_object);
+    auto mapper = createColumnMapper(schema)->getStorageColumnEncoding();
+    for (const auto & [col_name, source_id] : mapper)
+        source_id_to_column_name[source_id] = col_name;
+
+    Poco::JSON::Object::Ptr partition_spec;
+    for (size_t i = 0; i < partition_specs->size(); ++i)
+    {
+        auto spec = partition_specs->getObject(static_cast<UInt32>(i));
+        if (spec && spec->getValue<Int64>(f_spec_id) == partition_spec_id)
+        {
+            partition_spec = spec;
+            break;
+        }
+    }
+
+    if (!partition_spec || !partition_spec->has(f_fields))
+        return std::nullopt;
+
+    auto fields = partition_spec->getArray(f_fields);
+    if (!fields || fields->size() == 0)
+        return std::nullopt;
+
+    return ResolvedPartitionSpec{fields, std::move(source_id_to_column_name)};
+}
+}
+
 static CompressionMethod getCompressionMethodFromMetadataFile(const String & path)
 {
     constexpr std::string_view metadata_suffix = ".metadata.json";
@@ -1423,44 +1470,20 @@ static String formatPartitionFieldDisplay(const String & iceberg_transform_name,
 
 std::optional<String> getPartitionKeyStringFromMetadata(Poco::JSON::Object::Ptr metadata_object, const NamesAndTypesList & /* ch_schema */, ContextPtr /* local_context */)
 {
-    if (!metadata_object->has(f_partition_specs) || !metadata_object->has(f_default_spec_id))
-        return std::nullopt;
-    auto partition_spec_id = metadata_object->getValue<Int64>(f_default_spec_id);
-    Poco::JSON::Array::Ptr partition_specs = metadata_object->getArray(f_partition_specs);
-    std::unordered_map<Int64, String> source_id_to_column_name;
-    auto [schema, current_schema_id] = parseTableSchemaV2Method(metadata_object);
-    auto mapper = createColumnMapper(schema)->getStorageColumnEncoding();
-    for (const auto & [col_name, source_id] : mapper)
-        source_id_to_column_name[source_id] = col_name;
-
-    Poco::JSON::Object::Ptr partition_spec;
-    for (size_t i = 0; i < partition_specs->size(); ++i)
-    {
-        auto spec = partition_specs->getObject(static_cast<UInt32>(i));
-        if (spec->getValue<Int64>(f_spec_id) == partition_spec_id)
-        {
-            partition_spec = spec;
-            break;
-        }
-    }
-    if (!partition_spec || !partition_spec->has(f_fields))
-        return std::nullopt;
-    auto fields = partition_spec->getArray(f_fields);
-    if (fields->size() == 0)
+    auto resolved = resolveDefaultPartitionSpec(metadata_object);
+    if (!resolved)
         return std::nullopt;
 
     std::vector<String> part_exprs;
-    for (UInt32 i = 0; i < fields->size(); ++i)
+    for (UInt32 i = 0; i < resolved->fields->size(); ++i)
     {
-        auto field = fields->getObject(i);
-        auto source_id = field->getValue<Int64>(f_source_id);
-        auto it = source_id_to_column_name.find(source_id);
-        if (it == source_id_to_column_name.end())
+        auto field = resolved->fields->getObject(i);
+        auto it = resolved->source_id_to_column_name.find(field->getValue<Int64>(f_source_id));
+        if (it == resolved->source_id_to_column_name.end())
             return std::nullopt;
-        String column_name = it->second;
-        auto iceberg_transform_name = field->getValue<String>(f_transform);
-        part_exprs.push_back(formatPartitionFieldDisplay(iceberg_transform_name, column_name));
+        part_exprs.push_back(formatPartitionFieldDisplay(field->getValue<String>(f_transform), it->second));
     }
+
     String result;
     for (size_t i = 0; i < part_exprs.size(); ++i)
     {
@@ -1471,6 +1494,27 @@ std::optional<String> getPartitionKeyStringFromMetadata(Poco::JSON::Object::Ptr
     return result;
 }
 
+Names getIdentityPartitionColumnsFromMetadata(Poco::JSON::Object::Ptr metadata_object)
+{
+    auto resolved = resolveDefaultPartitionSpec(metadata_object);
+    if (!resolved)
+        return {};
+
+    Names result;
+    for (UInt32 i = 0; i < resolved->fields->size(); ++i)
+    {
+        auto field = resolved->fields->getObject(i);
+
+        if (Poco::toLower(field->getValue<String>(f_transform)) != "identity")
+            continue;
+
+        if (auto it = resolved->source_id_to_column_name.find(field->getValue<Int64>(f_source_id));
+            it != resolved->source_id_to_column_name.end())
+            result.push_back(it->second);
+    }
+    return result;
+}
+
 std::optional<String> getSortingKeyDisplayStringFromMetadata(Poco::JSON::Object::Ptr metadata_object, const NamesAndTypesList & /* ch_schema */)
 {
     if (!metadata_object->has(f_sort_orders) || !metadata_object->has(f_default_sort_order_id))

src/Storages/ObjectStorage/DataLakes/Iceberg/Utils.h

@@ -133,6 +133,7 @@ std::optional<String> getSortingKeyDisplayStringFromMetadata(
     Poco::JSON::Object::Ptr metadata_object, const NamesAndTypesList & ch_schema);
 std::optional<String> getPartitionKeyStringFromMetadata(
     Poco::JSON::Object::Ptr metadata_object, const NamesAndTypesList & ch_schema, ContextPtr local_context);
+Names getIdentityPartitionColumnsFromMetadata(Poco::JSON::Object::Ptr metadata_object);
 void sortBlockByKeyDescription(Block & block, const KeyDescription & sort_description, ContextPtr context);
 }
 

src/Storages/ObjectStorage/StorageObjectStorage.cpp

@@ -51,6 +51,7 @@ namespace Setting
     extern const SettingsInt64 delta_lake_snapshot_start_version;
     extern const SettingsInt64 delta_lake_snapshot_end_version;
     extern const SettingsUInt64 max_streams_for_files_processing_in_cluster_functions;
+    extern const SettingsBool allow_experimental_iceberg_read_optimization;
 }
 
 namespace ErrorCodes
@@ -336,7 +337,21 @@ bool StorageObjectStorage::canMoveConditionsToPrewhere() const
 
 std::optional<NameSet> StorageObjectStorage::supportedPrewhereColumns() const
 {
-    return getInMemoryMetadataPtr()->getColumnsWithoutDefaultExpressions(/*exclude=*/ hive_partition_columns_to_read_from_file_path);
+    return getInMemoryMetadataPtr()->getColumnsWithoutDefaultExpressions(/*exclude=*/hive_partition_columns_to_read_from_file_path);
+}
+
+std::optional<NameSet> StorageObjectStorage::supportedAutomaticPrewhereColumns(const StorageMetadataPtr & metadata) const
+{
+    auto exclude = hive_partition_columns_to_read_from_file_path;
+
+    for (const auto & identity_partition_column : metadata->identity_partition_columns)
+    {
+        if (metadata->getColumns().has(identity_partition_column))
+            exclude.emplace_back(identity_partition_column, metadata->getColumns().get(identity_partition_column).type);
+    }
+
+    LOG_DEBUG(log, "Automatic-PREWHERE exclude list: [{}]", exclude.toString());
+    return metadata->getColumnsWithoutDefaultExpressions(/*exclude=*/exclude);
 }
 
 IStorage::ColumnSizeByName StorageObjectStorage::getColumnSizes() const

src/Storages/ObjectStorage/StorageObjectStorage.h

@@ -124,6 +124,7 @@ class StorageObjectStorage : public IStorage
     bool supportsPrewhere() const override;
     bool canMoveConditionsToPrewhere() const override;
     std::optional<NameSet> supportedPrewhereColumns() const override;
+    std::optional<NameSet> supportedAutomaticPrewhereColumns(const StorageMetadataPtr & metadata) const override;
     ColumnSizeByName getColumnSizes() const override;
 
     bool prefersLargeBlocks() const override;

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -1062,6 +1062,13 @@ std::optional<NameSet> StorageObjectStorageCluster::supportedPrewhereColumns() c
     return IStorageCluster::supportedPrewhereColumns();
 }
 
+std::optional<NameSet> StorageObjectStorageCluster::supportedAutomaticPrewhereColumns(const StorageMetadataPtr & metadata) const
+{
+    if (pure_storage)
+        return pure_storage->supportedAutomaticPrewhereColumns(metadata);
+    return IStorageCluster::supportedAutomaticPrewhereColumns(metadata);
+}
+
 IStorageCluster::ColumnSizeByName StorageObjectStorageCluster::getColumnSizes() const
 {
     if (pure_storage)

src/Storages/ObjectStorage/StorageObjectStorageCluster.h

@@ -163,6 +163,7 @@ class StorageObjectStorageCluster : public IStorageCluster
     bool supportsPrewhere() const override;
     bool canMoveConditionsToPrewhere() const override;
     std::optional<NameSet> supportedPrewhereColumns() const override;
+    std::optional<NameSet> supportedAutomaticPrewhereColumns(const StorageMetadataPtr & metadata) const override;
     ColumnSizeByName getColumnSizes() const override;
 
     bool parallelizeOutputAfterReading(ContextPtr context) const override;

src/Storages/ObjectStorage/StorageObjectStorageConfiguration.h

@@ -312,7 +312,6 @@ class StorageObjectStorageConfiguration
 
     virtual ColumnMapperPtr getColumnMapperForCurrentSchema(StorageMetadataPtr /**/, ContextPtr /**/) const { return nullptr; }
 
-
     virtual std::shared_ptr<DataLake::ICatalog> getCatalog(ContextPtr /*context*/, const StorageID & /*table_id*/) const
     {
         return nullptr;

src/Storages/ObjectStorage/StorageObjectStorageSource.cpp

@@ -955,8 +955,8 @@ StorageObjectStorageSource::ReaderHolder StorageObjectStorageSource::createReade
             if (requested_columns_copy.empty()
                 && (!format_filter_info || (!format_filter_info->row_level_filter && !format_filter_info->prewhere_info)))
                 need_only_count = true;
+            }
         }
-    }
 
     std::optional<size_t> num_rows_from_cache
         = need_only_count && context_->getSettingsRef()[Setting::use_cache_for_count_from_files] ? try_get_num_rows_from_cache() : std::nullopt;

src/Storages/StorageInMemoryMetadata.cpp

@@ -62,6 +62,7 @@ StorageInMemoryMetadata::StorageInMemoryMetadata(const StorageInMemoryMetadata &
     , comment(other.comment)
     , metadata_version(other.metadata_version)
     , datalake_table_state(other.datalake_table_state)
+    , identity_partition_columns(other.identity_partition_columns)
 {
 }
 
@@ -99,6 +100,7 @@ StorageInMemoryMetadata & StorageInMemoryMetadata::operator=(const StorageInMemo
     comment = other.comment;
     metadata_version = other.metadata_version;
     datalake_table_state = other.datalake_table_state;
+    identity_partition_columns = other.identity_partition_columns;
 
     return *this;
 }
@@ -245,6 +247,11 @@ void StorageInMemoryMetadata::setDataLakeTableState(const DataLakeTableStateSnap
     datalake_table_state = datalake_table_state_;
 }
 
+void StorageInMemoryMetadata::setIdentityPartitionColumns(const Names & identity_partition_columns_)
+{
+    identity_partition_columns = identity_partition_columns_;
+}
+
 StorageInMemoryMetadata StorageInMemoryMetadata::withMetadataVersion(int32_t metadata_version_) const
 {
     StorageInMemoryMetadata copy(*this);

src/Storages/StorageInMemoryMetadata.h

@@ -78,6 +78,10 @@ struct StorageInMemoryMetadata
     ///  Current state of a datalake table.
     std::optional<DataLakeTableStateSnapshot> datalake_table_state;
 
+    /// Names of identity-partition columns (constant within every data file),
+    /// resolved for the same pinned `datalake_table_state` snapshot above.
+    Names identity_partition_columns;
+
     StorageInMemoryMetadata() = default;
 
     StorageInMemoryMetadata(const StorageInMemoryMetadata & other);
@@ -130,6 +134,7 @@ struct StorageInMemoryMetadata
     void setSQLSecurity(const ASTSQLSecurity & sql_security);
 
     void setDataLakeTableState(const DataLakeTableStateSnapshot & datalake_table_state_);
+    void setIdentityPartitionColumns(const Names & identity_partition_columns_);
     UUID getDefinerID(ContextPtr context) const;
 
     /// Returns a copy of the context with the correct user from SQL security options.