Iceberg partition aggregate metadata

src/Storages/IStorage.h

@@ -25,6 +25,8 @@
 #include <expected>
 #include <optional>
 #include <list>
+#include <utility>
+#include <vector>
 
 
 namespace DB
@@ -81,6 +83,14 @@ using ObjectStoragePtr = std::shared_ptr<IObjectStorage>;
 
 class ActionsDAG;
 
+/// One live data file whose identity-partition values and `record_count` are known exactly.
+/// A null Field is SQL NULL. A column that is absent is not an identity partition of this file.
+struct IdentityPartitionFile
+{
+    std::vector<std::pair<String, Field>> identity_values;
+    UInt64 record_count = 0;
+};
+
 /** Storage. Describes the table. Responsible for
   * - storage of the table data;
   * - the definition in which files (or not in files) the data is stored;
@@ -870,6 +880,11 @@ It is currently only implemented in StorageObjectStorage.
     /// Same as above but also take partition predicate into account.
     virtual std::optional<UInt64> totalRowsByPartitionPredicate(const ActionsDAG &, ContextPtr) const { return {}; }
 
+    /// Live data files with identity partition values and `record_count`.
+    /// `storage_metadata` is the snapshot pinned at analysis. The files must come from that
+    /// state, not from a newer commit. nullopt means the storage cannot answer exactly, so the caller must scan.
+    virtual std::optional<std::vector<IdentityPartitionFile>> getIdentityPartitionFiles(const StorageMetadataPtr & /*storage_metadata*/, ContextPtr) const { return {}; }
+
     /// Aggregated `(num_rows, num_defaults)` for `column_name` across all visible parts,
     /// taken from per-part `SerializationInfo`. Returns nullopt when the storage cannot
     /// supply an exact count -- see `Storages/MergeTree/SparsityFilter.h` for the precise

src/Storages/ObjectStorage/DataLakes/DataLakeConfiguration.h

@@ -653,6 +653,9 @@ class StorageIcebergConfiguration : public StorageObjectStorageConfiguration, pu
     bool isStaticConfiguration() const override { return getImpl().isStaticConfiguration(); }
 
     bool isDataLakeConfiguration() const override { return getImpl().isDataLakeConfiguration(); }
+    /// `StorageIcebergConfiguration` selects S3, Azure, HDFS, or local at runtime.
+    /// `system.iceberg_files` and partition-column aggregates recognize Iceberg through this flag.
+    bool isIcebergConfiguration() const override { return getImpl().isIcebergConfiguration(); }
 
     bool supportsTotalRows(ContextPtr context, ObjectStorageType storage_type) const override { return getImpl().supportsTotalRows(context, storage_type); }
     std::optional<size_t> totalRows(ContextPtr context) override { return getImpl().totalRows(context); }

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

@@ -1630,6 +1630,84 @@ std::optional<size_t> IcebergMetadata::totalRows(ContextPtr local_context) const
     return static_cast<size_t>(result);
 }
 
+std::optional<std::vector<IdentityPartitionFile>> IcebergMetadata::getDataFileIdentityPartitions(
+    StorageMetadataPtr storage_metadata, ContextPtr local_context) const
+{
+    /// Same snapshot the scan uses. `getRelevantState` would read the latest metadata file instead.
+    auto table_state = extractIcebergSnapshotIdFromMetadataObject(storage_metadata);
+    if (!table_state)
+        return std::nullopt;
+
+    auto actual_data_snapshot = getRelevantDataSnapshotFromTableStateSnapshot(*table_state, local_context);
+
+    if (!actual_data_snapshot)
+        return std::vector<IdentityPartitionFile>{};
+
+    /// Same fail-closed conditions as `totalRows`: deletes make `record_count` not equal to
+    /// the number of surviving rows, and an unreadable path must not be answered from metadata.
+    if (actual_data_snapshot->total_equality_delete_rows.has_value()
+        && *actual_data_snapshot->total_equality_delete_rows > 0)
+        return std::nullopt;
+
+    if (!persistent_components.schema_processor)
+        return std::nullopt;
+
+    std::vector<IdentityPartitionFile> result;
+    for (const auto & manifest_list_entry : actual_data_snapshot->manifest_list_entries)
+    {
+        auto manifest_file_ptr = getManifestFileEntriesHandle(
+            object_storage, persistent_components, local_context, log, manifest_list_entry, table_state->schema_id, *external_storages);
+
+        for (const auto & processed : manifest_file_ptr.getFilesWithoutDeleted(FileContentType::DATA))
+            if (!isPathReadable(
+                    persistent_components.table_location,
+                    processed->parsed_entry->file_path_key.serialize(),
+                    object_storage,
+                    local_context))
+                return std::nullopt;
+
+        if (!manifest_file_ptr.getFilesWithoutDeleted(FileContentType::EQUALITY_DELETE).empty()
+            || !manifest_file_ptr.getFilesWithoutDeleted(FileContentType::POSITION_DELETE).empty())
+            return std::nullopt;
+
+        for (const auto & processed : manifest_file_ptr.getFilesWithoutDeleted(FileContentType::DATA))
+        {
+            const auto & entry = *processed->parsed_entry;
+            /// A deletion vector on the data file itself, or a negative count, is not an exact row count.
+            if (entry.isDeletionVector() || entry.record_count < 0)
+                return std::nullopt;
+            /// An empty file contributes no rows and no group.
+            if (entry.record_count == 0)
+                continue;
+            if (!processed->common_partition_specification)
+                return std::nullopt;
+
+            auto partition_values = Iceberg::getIdentityPartitionColumnValues(
+                *processed, *persistent_components.schema_processor, /* skip_nulls */ false);
+
+            IdentityPartitionFile file;
+            file.record_count = static_cast<UInt64>(entry.record_count);
+            for (auto & [column_name, column_value] : partition_values)
+            {
+                auto it = std::find_if(file.identity_values.begin(), file.identity_values.end(), [&](const auto & existing)
+                {
+                    return existing.first == column_name;
+                });
+                if (it != file.identity_values.end())
+                {
+                    if (!(it->second == column_value))
+                        return std::nullopt;
+                    continue;
+                }
+                file.identity_values.emplace_back(std::move(column_name), std::move(column_value));
+            }
+            result.push_back(std::move(file));
+        }
+    }
+
+    return result;
+}
+
 std::optional<size_t> IcebergMetadata::totalBytes(ContextPtr local_context) const
 {
     auto [actual_data_snapshot, actual_table_state_snapshot] = getRelevantState(local_context);

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

@@ -129,6 +129,12 @@ class IcebergMetadata : public IDataLakeMetadata
 
     static bool supportsTotalRows(ContextPtr, ObjectStorageType) { return true; }
     std::optional<size_t> totalRows(ContextPtr Local_context) const override;
+
+    /// Live DATA files with identity partition values (including SQL NULL) and `record_count`.
+    /// Resolved from `datalake_table_state` on `storage_metadata`, the snapshot pinned at analysis.
+    /// nullopt when that state is missing or an exact answer is impossible: equality or position deletes,
+    /// deletion vectors, an unreadable path, or a negative `record_count`. An empty snapshot is an empty vector.
+    std::optional<std::vector<IdentityPartitionFile>> getDataFileIdentityPartitions(StorageMetadataPtr storage_metadata, ContextPtr local_context) const;
     static bool supportsTotalBytes(ContextPtr, ObjectStorageType) { return true; }
     std::optional<size_t> totalBytes(ContextPtr Local_context) const override;
 

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

@@ -2187,7 +2187,7 @@ void forEachAvroEntry(
 }
 
 PartitionColumnValues getIdentityPartitionColumnValues(
-    const ProcessedManifestFileEntry & manifest_file_entry, const IcebergSchemaProcessor & schema_processor)
+    const ProcessedManifestFileEntry & manifest_file_entry, const IcebergSchemaProcessor & schema_processor, bool skip_nulls)
 {
     const auto & partition_key_value = manifest_file_entry.parsed_entry->partition_key_value;
     if (partition_key_value.empty())
@@ -2228,7 +2228,9 @@ PartitionColumnValues getIdentityPartitionColumnValues(
         {
             value = convertFieldToTypeOrThrow(raw_value, *name_and_type->type);
         }
-        if (value.isNull())
+        /// A null identity value is a real SQL NULL group. Callers that only need the
+        /// non-null values (pruning, compaction) skip them; aggregation must keep them.
+        if (value.isNull() && skip_nulls)
             continue;
 
         result.emplace_back(name_and_type->name, std::move(value));

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

@@ -189,7 +189,9 @@ void forEachAvroEntry(
 using PartitionColumnValues = std::vector<std::pair<String, DB::Field>>;
 
 PartitionColumnValues getIdentityPartitionColumnValues(
-    const ProcessedManifestFileEntry & manifest_file_entry, const IcebergSchemaProcessor & schema_processor);
+    const ProcessedManifestFileEntry & manifest_file_entry,
+    const IcebergSchemaProcessor & schema_processor,
+    bool skip_nulls = true);
 }
 
 #endif

src/Storages/ObjectStorage/StorageObjectStorage.cpp

@@ -1,4 +1,8 @@
 #include <Storages/ObjectStorage/StorageObjectStorage.h>
+#include "config.h"
+#if USE_AVRO
+#include <Storages/ObjectStorage/DataLakes/Iceberg/IcebergMetadata.h>
+#endif
 #include <Storages/MergeTree/MergeTreePartInfo.h>
 
 #include <Common/CurrentThread.h>
@@ -584,6 +588,33 @@ std::optional<UInt64> StorageObjectStorage::totalBytes(ContextPtr query_context)
     return configuration->totalBytes(query_context);
 }
 
+std::optional<std::vector<IdentityPartitionFile>> StorageObjectStorage::getIdentityPartitionFiles(
+    const StorageMetadataPtr & storage_metadata, ContextPtr query_context) const
+{
+#if USE_AVRO
+    /// A replica with `distributed_processing` sees only its own tasks. Answering from the whole
+    /// snapshot here would count every shard's rows on every shard.
+    if (distributed_processing || !configuration->isIcebergConfiguration())
+        return std::nullopt;
+
+    /// Analysis pinned `datalake_table_state`. `configuration->update` would load a newer commit
+    /// and disagree with the column types of this query. A missing pin scans, and the read path pins one snapshot.
+    if (!storage_metadata || !storage_metadata->datalake_table_state.has_value())
+        return std::nullopt;
+
+    configuration->lazyInitializeIfNeeded(object_storage, query_context);
+
+    auto iceberg = std::dynamic_pointer_cast<IcebergMetadata>(configuration->getExternalMetadata());
+    if (!iceberg)
+        return std::nullopt;
+    return iceberg->getDataFileIdentityPartitions(storage_metadata, query_context);
+#else
+    (void)storage_metadata;
+    (void)query_context;
+    return std::nullopt;
+#endif
+}
+
 void StorageObjectStorage::read(
     QueryPlan & query_plan,
     const Names & column_names,

src/Storages/ObjectStorage/StorageObjectStorage.h

@@ -212,6 +212,7 @@ class StorageObjectStorage : public IStorage, public IBackgroundOperation
 
     std::optional<UInt64> totalRows(ContextPtr query_context) const override;
     std::optional<UInt64> totalBytes(ContextPtr query_context) const override;
+    std::optional<std::vector<IdentityPartitionFile>> getIdentityPartitionFiles(const StorageMetadataPtr & storage_metadata, ContextPtr query_context) const override;
 
     bool optimize(
         const ASTPtr & /*query*/,

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -331,6 +331,17 @@ std::optional<UInt64> StorageObjectStorageCluster::totalBytes(ContextPtr query_c
     return configuration->totalBytes(query_context);
 }
 
+std::optional<std::vector<IdentityPartitionFile>> StorageObjectStorageCluster::getIdentityPartitionFiles(
+    const StorageMetadataPtr & storage_metadata, ContextPtr query_context) const
+{
+    /// The initiator's `pure_storage` reads the snapshot once. Shards are not asked.
+    /// `pure_storage` itself returns nullopt when `distributed_processing` is set.
+    /// `storage_metadata` is the cluster snapshot pinned at analysis. `pure_storage` shares this configuration.
+    if (pure_storage)
+        return pure_storage->getIdentityPartitionFiles(storage_metadata, query_context);
+    return std::nullopt;
+}
+
 bool StorageObjectStorageCluster::updateQueryForDistributedEngineIfNeeded(ASTPtr & query, ContextPtr context, bool make_cluster_function)
 {
     // Change table engine on table function for distributed request

src/Storages/ObjectStorage/StorageObjectStorageCluster.h

@@ -65,6 +65,7 @@ class StorageObjectStorageCluster : public IStorageCluster
 
     std::optional<UInt64> totalRows(ContextPtr query_context) const override;
     std::optional<UInt64> totalBytes(ContextPtr query_context) const override;
+    std::optional<std::vector<IdentityPartitionFile>> getIdentityPartitionFiles(const StorageMetadataPtr & storage_metadata, ContextPtr query_context) const override;
 
     void setClusterNameFromFunctionArgument(bool cluster_name_from_function_argument_)
     {