[WIP] read identify columns from iceberg metada

src/Storages/ObjectStorage/DataLakes/DataLakeConfiguration.h

@@ -323,6 +323,15 @@ class DataLakeConfiguration : public BaseStorageConfiguration, public std::enabl
         return current_metadata->getColumnMapperForCurrentSchema(storage_metadata_snapshot, context);
     }
 
+    Names getIdentityPartitionColumnNames(ContextPtr context) const override
+    {
+        if (!current_metadata)
+        {
+            return {};
+        }
+        return current_metadata->getIdentityPartitionColumnNames(context);
+    }
+
     void drop(ContextPtr local_context) override
     {
         if (current_metadata)
@@ -806,6 +815,8 @@ class StorageIcebergConfiguration : public StorageObjectStorageConfiguration, pu
     ColumnMapperPtr getColumnMapperForCurrentSchema(StorageMetadataPtr storage_metadata_snapshot, ContextPtr context) const override
         { return getImpl().getColumnMapperForCurrentSchema(storage_metadata_snapshot, context); }
 
+    Names getIdentityPartitionColumnNames(ContextPtr context) const override { return getImpl().getIdentityPartitionColumnNames(context); }
+
     std::shared_ptr<DataLake::ICatalog> getCatalog(ContextPtr context, const StorageID & table_id) const override
         { return getImpl().getCatalog(context, table_id); }
 

src/Storages/ObjectStorage/DataLakes/IDataLakeMetadata.h

@@ -192,6 +192,8 @@ class IDataLakeMetadata : boost::noncopyable
     virtual ColumnMapperPtr getColumnMapperForObject(ObjectInfoPtr /**/) const { return nullptr; }
     virtual ColumnMapperPtr getColumnMapperForCurrentSchema(StorageMetadataPtr, ContextPtr) const { return nullptr; }
 
+    virtual Names getIdentityPartitionColumnNames(ContextPtr) const { return {}; }
+
     virtual SinkToStoragePtr write(
         SharedHeader /*sample_block*/,
         const StorageID & /*table_id*/,

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

@@ -1436,6 +1436,21 @@ ColumnMapperPtr IcebergMetadata::getColumnMapperForCurrentSchema(StorageMetadata
     return persistent_components.schema_processor->getColumnMapperById(iceberg_table_state->schema_id);
 }
 
+Names IcebergMetadata::getIdentityPartitionColumnNames(ContextPtr local_context) const
+{
+    // check getPartitionKey
+    auto [data_snapshot, table_state_snapshot] = getRelevantState(local_context);
+    auto metadata_object = getMetadataJSONObject(
+        table_state_snapshot.metadata_file_path,
+        object_storage,
+        persistent_components.metadata_cache,
+        local_context,
+        log,
+        persistent_components.metadata_compression_method,
+        persistent_components.table_uuid);
+    return getIdentityPartitionColumnsFromMetadata(metadata_object);
+}
+
 std::optional<String> IcebergMetadata::getPartitionKey(ContextPtr local_context, TableStateSnapshot actual_table_state_snapshot) const
 {
     auto metadata_object = getMetadataJSONObject(

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

@@ -106,6 +106,8 @@ class IcebergMetadata : public IDataLakeMetadata
 
     ColumnMapperPtr getColumnMapperForCurrentSchema(StorageMetadataPtr storage_metadata_snapshot, ContextPtr context) const override;
 
+    Names getIdentityPartitionColumnNames(ContextPtr context) const override;
+
     SinkToStoragePtr write(
         SharedHeader sample_block,
         const StorageID & table_id,

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

@@ -1471,6 +1471,52 @@ std::optional<String> getPartitionKeyStringFromMetadata(Poco::JSON::Object::Ptr
     return result;
 }
 
+Names getIdentityPartitionColumnsFromMetadata(Poco::JSON::Object::Ptr metadata_object)
+{
+    // @todo refactor dupli
+    if (!metadata_object->has(f_partition_specs) || !metadata_object->has(f_default_spec_id))
+        return {};
+
+    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 {};
+
+    auto fields = partition_spec->getArray(f_fields);
+    if (fields->size() == 0)
+        return {};
+
+    Names result;
+    std::vector<String> part_exprs;
+    for (UInt32 i = 0; i < fields->size(); ++i)
+    {
+        auto field = fields->getObject(i);
+        if (field->getValue<String>(f_transform) != "identity")
+            continue;
+
+        if (auto it = source_id_to_column_name.find(field->getValue<Int64>(f_source_id)); it != 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

@@ -3,6 +3,7 @@
 
 #include <Common/Exception.h>
 #include <Common/Logger.h>
+#include <Common/UniqueLock.h>
 #include <Common/logger_useful.h>
 #include <Core/Settings.h>
 #include <Formats/FormatFactory.h>
@@ -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);
+    auto exclude = hive_partition_columns_to_read_from_file_path;
+
+    const auto & cols = getInMemoryMetadataPtr()->getColumns();
+    {
+        SharedLockGuard lock(mutex_file_constant_columns);
+        for (const auto & col : file_constant_columns)
+        {
+            if (getInMemoryMetadataPtr()->getColumns().has(col))
+                // tryGetColumn
+                exclude.emplace_back(col, cols.get(col).type);
+        }
+    }
+
+    LOG_DEBUG(log, "Prewhere exclude list: [{}]", exclude.toString());
+    return getInMemoryMetadataPtr()->getColumnsWithoutDefaultExpressions(/*exclude=*/exclude);
 }
 
 IStorage::ColumnSizeByName StorageObjectStorage::getColumnSizes() const
@@ -354,6 +369,12 @@ configuration->update(object_storage, query_context);
     return configuration->getExternalMetadata();
 }
 
+void StorageObjectStorage::updateFileConstantColumns(ContextPtr query_context)
+{
+    UniqueLock lock(mutex_file_constant_columns);
+    file_constant_columns = configuration->getIdentityPartitionColumnNames(query_context);
+}
+
 void StorageObjectStorage::updateExternalDynamicMetadataIfExists(ContextPtr query_context)
 {
     if (!configuration->isDataLakeConfiguration())
@@ -381,6 +402,8 @@ void StorageObjectStorage::updateExternalDynamicMetadataIfExists(ContextPtr quer
             new_metadata = *metadata_snapshot;
     }
 
+    updateFileConstantColumns(query_context);
+
     setInMemoryMetadata(new_metadata);
 }
 

src/Storages/ObjectStorage/StorageObjectStorage.h

@@ -155,6 +155,7 @@ class StorageObjectStorage : public IStorage
 
     void addInferredEngineArgsToCreateQuery(ASTs & args, const ContextPtr & context) const override;
 
+    void updateFileConstantColumns(ContextPtr query_context);
     void updateExternalDynamicMetadataIfExists(ContextPtr query_context) override;
 
     IDataLakeMetadata * getExternalMetadata(ContextPtr query_context);
@@ -214,6 +215,9 @@ class StorageObjectStorage : public IStorage
     NamesAndTypesList hive_partition_columns_to_read_from_file_path;
     NamesAndTypesList file_columns;
 
+    mutable SharedMutex mutex_file_constant_columns;
+    Names file_constant_columns TSA_GUARDED_BY(mutex_file_constant_columns);
+
     LoggerPtr log;
 
     std::shared_ptr<DataLake::ICatalog> catalog;

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -569,7 +569,10 @@ void StorageObjectStorageCluster::updateExternalDynamicMetadataIfExists(ContextP
     setInMemoryMetadata(new_metadata);
 
     if (pure_storage)
+    {
         pure_storage->setInMemoryMetadata(IStorageCluster::getInMemoryMetadata());
+        pure_storage->updateFileConstantColumns(query_context);
+    }
 }
 
 class TaskDistributor : public TaskIterator

src/Storages/ObjectStorage/StorageObjectStorageConfiguration.h

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