Merge pull request JanKaul#266 from JanKaul/populate-snapshot-summary

Populate snapshot summary

Author: JanKaul

iceberg-rust/src/table/transaction/operation.rs

@@ -10,9 +10,11 @@ use std::{collections::HashMap, sync::Arc};
 use bytes::Bytes;
 use futures::future;
 use iceberg_rust_spec::manifest_list::{
-    manifest_list_schema_v1, manifest_list_schema_v2, Content, ManifestListEntry,
+    manifest_list_schema_v1, manifest_list_schema_v2, Content as ManifestListContent,
+    ManifestListEntry,
 };
 use iceberg_rust_spec::snapshot::{Operation as SnapshotOperation, Snapshot};
+use iceberg_rust_spec::spec::manifest::Content;
 use iceberg_rust_spec::spec::table_metadata::TableMetadata;
 use iceberg_rust_spec::spec::{
     manifest::{partition_value_schema, DataFile, ManifestEntry, Status},
@@ -201,12 +203,12 @@ impl Operation {
                     // Split manifest file if limit is exceeded
                     for (content, files, n_files) in [
                         (
-                            Content::Data,
+                            ManifestListContent::Data,
                             Either::Left(new_datafile_iter),
                             n_data_files_in_group,
                         ),
                         (
-                            Content::Deletes,
+                            ManifestListContent::Deletes,
                             Either::Right(new_deletefile_iter),
                             n_delete_files_in_group,
                         ),
@@ -230,14 +232,28 @@ impl Operation {
                     (_, _) => Ok(SnapshotOperation::Overwrite),
                 }?;
 
+                // Compute updated snapshot summary
+                // Separate data files from delete files using iterators
+                let data_files_iter = all_files.iter().filter(|f| *f.content() == Content::Data);
+                let delete_files_iter = all_files.iter().filter(|f| *f.content() != Content::Data);
+
+                let mut summary_fields = update_snapshot_summary(
+                    old_snapshot.map(|s| s.summary()),
+                    data_files_iter,
+                    delete_files_iter,
+                );
+
+                // Merge with any additional summary fields
+                summary_fields.extend(additional_summary.unwrap_or_default());
+
                 let mut snapshot_builder = SnapshotBuilder::default();
                 snapshot_builder
                     .with_snapshot_id(snapshot_id)
                     .with_manifest_list(new_manifest_list_location)
                     .with_sequence_number(table_metadata.last_sequence_number + dsn_offset)
                     .with_summary(Summary {
                         operation: snapshot_operation,
-                        other: additional_summary.unwrap_or_default(),
+                        other: summary_fields,
                     })
                     .with_schema_id(
                         *table_metadata
@@ -290,6 +306,13 @@ impl Operation {
                     return Ok((None, Vec::new()));
                 }
 
+                // Compute summary before moving data_files and delete_files
+                let summary_fields = update_snapshot_summary(
+                    old_snapshot.map(|s| s.summary()),
+                    data_files.iter(),
+                    delete_files.iter(),
+                );
+
                 let data_files_iter = delete_files.iter().chain(data_files.iter());
 
                 let mut manifest_list_writer = if let Some(manifest_list_bytes) =
@@ -339,9 +362,13 @@ impl Operation {
                     Pin<Box<dyn Future<Output = Result<(), Error>> + Send>>,
                 > = Vec::new();
                 for (content, files, n_files) in [
-                    (Content::Data, Either::Left(new_datafile_iter), n_data_files),
                     (
-                        Content::Deletes,
+                        ManifestListContent::Data,
+                        Either::Left(new_datafile_iter),
+                        n_data_files,
+                    ),
+                    (
+                        ManifestListContent::Deletes,
                         Either::Right(new_deletefile_iter),
                         n_delete_files,
                     ),
@@ -389,14 +416,20 @@ impl Operation {
                     (_, _) => Ok(SnapshotOperation::Overwrite),
                 }?;
 
+                // Merge with any additional summary fields provided by the caller
+                let mut summary_fields = summary_fields;
+                if let Some(additional) = additional_summary {
+                    summary_fields.extend(additional);
+                }
+
                 let mut snapshot_builder = SnapshotBuilder::default();
                 snapshot_builder
                     .with_snapshot_id(snapshot_id)
                     .with_manifest_list(new_manifest_list_location)
                     .with_sequence_number(table_metadata.last_sequence_number + 1)
                     .with_summary(Summary {
                         operation: snapshot_operation,
-                        other: additional_summary.unwrap_or_default(),
+                        other: summary_fields,
                     })
                     .with_schema_id(
                         *table_metadata
@@ -495,7 +528,7 @@ impl Operation {
                         snapshot_id,
                         &manifest_schema,
                         table_metadata,
-                        Content::Data,
+                        ManifestListContent::Data,
                         branch.as_deref(),
                     )?;
 
@@ -527,7 +560,7 @@ impl Operation {
                             snapshot_id,
                             &manifest_schema,
                             table_metadata,
-                            Content::Data,
+                            ManifestListContent::Data,
                             branch.as_deref(),
                         )?;
 
@@ -550,6 +583,23 @@ impl Operation {
                     )
                     .await?;
 
+                // For Replace operation, we're replacing all data with new files
+                // Compute summary for the new state (not incremental)
+                // We set old_summary to None so totals equal the new files only
+                let data_files_iter = files.iter().filter(|f| *f.content() == Content::Data);
+                let delete_files_iter = files.iter().filter(|f| *f.content() != Content::Data);
+
+                let mut summary_fields = update_snapshot_summary(
+                    None, // No old summary - this is a full replacement
+                    data_files_iter,
+                    delete_files_iter,
+                );
+
+                // Merge with any additional summary fields
+                if let Some(additional) = additional_summary {
+                    summary_fields.extend(additional);
+                }
+
                 let mut snapshot_builder = SnapshotBuilder::default();
                 snapshot_builder
                     .with_snapshot_id(snapshot_id)
@@ -558,7 +608,7 @@ impl Operation {
                     .with_manifest_list(new_manifest_list_location)
                     .with_summary(Summary {
                         operation: iceberg_rust_spec::spec::snapshot::Operation::Overwrite,
-                        other: additional_summary.unwrap_or_default(),
+                        other: summary_fields,
                     });
                 let snapshot = snapshot_builder.build()?;
 
@@ -604,6 +654,13 @@ impl Operation {
                     return Ok((None, Vec::new()));
                 }
 
+                // Compute summary before moving data_files
+                let summary_fields = update_snapshot_summary(
+                    Some(old_snapshot.summary()),
+                    data_files.iter(),
+                    std::iter::empty::<&DataFile>(), // No separate delete files in this operation
+                );
+
                 let data_files_iter = data_files.iter();
 
                 let manifests_to_overwrite: HashSet<String> =
@@ -631,7 +688,8 @@ impl Operation {
                     )
                     .await?;
 
-                let n_splits = manifest_list_writer.n_splits(n_data_files, Content::Data);
+                let n_splits =
+                    manifest_list_writer.n_splits(n_data_files, ManifestListContent::Data);
 
                 let new_datafile_iter = data_files.into_iter().map(|data_file| {
                     ManifestEntry::builder()
@@ -649,9 +707,9 @@ impl Operation {
                     .map(|x| x.manifest_path.clone())
                     .ok_or(Error::NotFound("Selected manifest".to_owned()))?;
 
-                let data_files = files_to_overwrite.get(&selected_manifest_location);
+                let files_to_filter = files_to_overwrite.get(&selected_manifest_location);
 
-                let filter = if let Some(filter_files) = data_files {
+                let filter = if let Some(filter_files) = files_to_filter {
                     let filter_files: HashSet<String> =
                         filter_files.iter().map(ToOwned::to_owned).collect();
                     Some(move |file: &Result<ManifestEntry, Error>| {
@@ -672,7 +730,7 @@ impl Operation {
                             snapshot_id,
                             filter,
                             object_store.clone(),
-                            Content::Data,
+                            ManifestListContent::Data,
                         )
                         .await?;
                     manifest_list_writer
@@ -686,7 +744,7 @@ impl Operation {
                             n_splits,
                             filter,
                             object_store.clone(),
-                            Content::Data,
+                            ManifestListContent::Data,
                         )
                         .await?;
                     manifest_list_writer
@@ -696,14 +754,20 @@ impl Operation {
 
                 let snapshot_operation = SnapshotOperation::Overwrite;
 
+                // Merge with any additional summary fields
+                let mut summary_fields = summary_fields;
+                if let Some(additional) = additional_summary {
+                    summary_fields.extend(additional);
+                }
+
                 let mut snapshot_builder = SnapshotBuilder::default();
                 snapshot_builder
                     .with_snapshot_id(snapshot_id)
                     .with_manifest_list(new_manifest_list_location)
                     .with_sequence_number(table_metadata.last_sequence_number + 1)
                     .with_summary(Summary {
                         operation: snapshot_operation,
-                        other: additional_summary.unwrap_or_default(),
+                        other: summary_fields,
                     })
                     .with_schema_id(
                         *table_metadata
@@ -840,6 +904,98 @@ fn prefetch_manifest_list(
         })
 }
 
+/// Updates snapshot summary fields incrementally based on operation changes.
+///
+/// This function computes snapshot summary metrics by:
+/// 1. Parsing existing totals from the old snapshot summary
+/// 2. Computing deltas from the new data/delete files being added
+/// 3. Adding deltas to existing totals
+///
+/// # Arguments
+/// * `old_summary` - The summary from the parent snapshot (if it exists)
+/// * `data_files` - Iterator over new data files being added in this operation
+/// * `delete_files` - Iterator over new delete files being added in this operation
+///
+/// # Returns
+/// A HashMap with updated summary fields including:
+/// - `total-records`: Total record count across all data files
+/// - `total-data-files`: Total number of data files
+/// - `total-delete-files`: Total number of delete files
+/// - `total-file-size-bytes`: Total size of all files in bytes
+/// - `added-records`: Records added in this operation
+/// - `added-data-files`: Data files added in this operation
+/// - `added-files-size-bytes`: Size of files added in this operation
+pub fn update_snapshot_summary<'files>(
+    old_summary: Option<&Summary>,
+    data_files: impl Iterator<Item = &'files DataFile>,
+    delete_files: impl Iterator<Item = &'files DataFile>,
+) -> HashMap<String, String> {
+    // Parse existing values from old summary
+    let old_other = old_summary.map(|s| &s.other);
+
+    let parse_i64 = |key: &str| -> i64 {
+        old_other
+            .and_then(|m| m.get(key))
+            .and_then(|v| v.parse::<i64>().ok())
+            .unwrap_or(0)
+    };
+
+    let old_total_records = parse_i64("total-records");
+    let old_total_data_files = parse_i64("total-data-files");
+    let old_total_delete_files = parse_i64("total-delete-files");
+    let old_total_file_size = parse_i64("total-file-size-bytes");
+
+    // Compute deltas from new files - we need to iterate to count and sum
+    let mut added_data_files = 0i64;
+    let mut added_records = 0i64;
+    let mut added_data_files_size = 0i64;
+
+    for file in data_files {
+        added_data_files += 1;
+        if *file.content() == Content::Data {
+            added_records += file.record_count();
+        }
+        added_data_files_size += file.file_size_in_bytes();
+    }
+
+    let mut added_delete_files = 0i64;
+    let mut added_delete_files_size = 0i64;
+
+    for file in delete_files {
+        added_delete_files += 1;
+        added_delete_files_size += file.file_size_in_bytes();
+    }
+
+    let added_files_size = added_data_files_size + added_delete_files_size;
+
+    // Compute new totals
+    let total_records = old_total_records + added_records;
+    let total_data_files = old_total_data_files + added_data_files;
+    let total_delete_files = old_total_delete_files + added_delete_files;
+    let total_file_size = old_total_file_size + added_files_size;
+
+    // Build result map
+    let mut result = HashMap::new();
+    result.insert("total-records".to_string(), total_records.to_string());
+    result.insert("total-data-files".to_string(), total_data_files.to_string());
+    result.insert(
+        "total-delete-files".to_string(),
+        total_delete_files.to_string(),
+    );
+    result.insert(
+        "total-file-size-bytes".to_string(),
+        total_file_size.to_string(),
+    );
+    result.insert("added-records".to_string(), added_records.to_string());
+    result.insert("added-data-files".to_string(), added_data_files.to_string());
+    result.insert(
+        "added-files-size-bytes".to_string(),
+        added_files_size.to_string(),
+    );
+
+    result
+}
+
 pub(crate) fn new_manifest_location(
     table_metadata_location: &str,
     commit_uuid: &str,