Merge pull request JanKaul#251 from splitgraph/new-dsn-2

AppendSequenceGroups operation

Author: JanKaul

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

@@ -21,6 +21,7 @@ use tracing::debug;
 use iceberg_rust_spec::spec::{manifest::DataFile, schema::Schema, snapshot::SnapshotReference};
 
 use crate::table::transaction::append::append_summary;
+use crate::table::transaction::operation::SequenceGroup;
 use crate::{catalog::commit::CommitTable, error::Error, table::Table};
 
 use self::operation::Operation;
@@ -32,12 +33,13 @@ pub(crate) mod overwrite;
 pub(crate) static ADD_SCHEMA_INDEX: usize = 0;
 pub(crate) static SET_DEFAULT_SPEC_INDEX: usize = 1;
 pub(crate) static APPEND_INDEX: usize = 2;
-pub(crate) static REPLACE_INDEX: usize = 3;
-pub(crate) static OVERWRITE_INDEX: usize = 4;
-pub(crate) static UPDATE_PROPERTIES_INDEX: usize = 5;
-pub(crate) static SET_SNAPSHOT_REF_INDEX: usize = 6;
+pub(crate) static APPEND_SEQUENCE_GROUPS_INDEX: usize = 3;
+pub(crate) static REPLACE_INDEX: usize = 4;
+pub(crate) static OVERWRITE_INDEX: usize = 5;
+pub(crate) static UPDATE_PROPERTIES_INDEX: usize = 6;
+pub(crate) static SET_SNAPSHOT_REF_INDEX: usize = 7;
 
-pub(crate) static NUM_OPERATIONS: usize = 7;
+pub(crate) static NUM_OPERATIONS: usize = 8;
 
 /// A transaction that can perform multiple operations on a table atomically
 ///
@@ -118,6 +120,9 @@ impl<'table> TableTransaction<'table> {
     ///     .await?;
     /// ```
     pub fn append_data(mut self, files: Vec<DataFile>) -> Self {
+        if self.operations[APPEND_SEQUENCE_GROUPS_INDEX].is_some() {
+            panic!("Cannot use append and append_sequence_group in the same transaction");
+        }
         let summary = append_summary(&files);
 
         if let Some(ref mut operation) = self.operations[APPEND_INDEX] {
@@ -157,6 +162,9 @@ impl<'table> TableTransaction<'table> {
     ///     .await?;
     /// ```
     pub fn append_delete(mut self, files: Vec<DataFile>) -> Self {
+        if self.operations[APPEND_SEQUENCE_GROUPS_INDEX].is_some() {
+            panic!("Cannot use append and append_sequence_group in the same transaction");
+        }
         if let Some(ref mut operation) = self.operations[APPEND_INDEX] {
             if let Operation::Append {
                 delete_files: old, ..
@@ -174,6 +182,39 @@ impl<'table> TableTransaction<'table> {
         }
         self
     }
+
+    /// Appends a group of data and delete files to the table
+    ///
+    pub fn append_sequence_group(
+        mut self,
+        data_files: Vec<DataFile>,
+        delete_files: Vec<DataFile>,
+    ) -> Self {
+        if self.operations[APPEND_INDEX].is_some() {
+            panic!("Cannot use append and append_sequence_group in the same transaction");
+        }
+        if let Some(ref mut operation) = self.operations[APPEND_SEQUENCE_GROUPS_INDEX] {
+            if let Operation::AppendSequenceGroups {
+                sequence_groups: old,
+                ..
+            } = operation
+            {
+                old.push(SequenceGroup {
+                    delete_files,
+                    data_files,
+                });
+            }
+        } else {
+            self.operations[APPEND_SEQUENCE_GROUPS_INDEX] = Some(Operation::AppendSequenceGroups {
+                branch: self.branch.clone(),
+                sequence_groups: vec![SequenceGroup {
+                    delete_files,
+                    data_files,
+                }],
+            });
+        }
+        self
+    }
     /// Overwrites specific data files in the table with new ones
     ///
     /// This operation replaces specified existing data files with new ones, rather than

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

@@ -28,6 +28,7 @@ use tracing::{debug, instrument};
 
 use crate::table::manifest::ManifestWriter;
 use crate::table::manifest_list::ManifestListWriter;
+use crate::table::transaction::append::append_summary;
 use crate::{
     catalog::commit::{TableRequirement, TableUpdate},
     error::Error,
@@ -39,6 +40,15 @@ use super::append::split_datafiles;
 /// The target number of datafiles per manifest is dynamic, but we don't want to go below this number.
 static MIN_DATAFILES_PER_MANIFEST: usize = 4;
 
+#[derive(Debug, Clone)]
+/// Group of writes sharing a Data Sequence Number
+pub struct SequenceGroup {
+    /// Delete files. These apply to insert files from previous Sequence Groups
+    pub delete_files: Vec<DataFile>,
+    /// Insert files
+    pub data_files: Vec<DataFile>,
+}
+
 #[derive(Debug)]
 ///Table operations
 pub enum Operation {
@@ -61,6 +71,11 @@ pub enum Operation {
         delete_files: Vec<DataFile>,
         additional_summary: Option<HashMap<String, String>>,
     },
+    /// Append new change groups to the table
+    AppendSequenceGroups {
+        branch: Option<String>,
+        sequence_groups: Vec<SequenceGroup>,
+    },
     // /// Quickly append new files to the table
     // NewFastAppend {
     //     paths: Vec<String>,
@@ -99,6 +114,151 @@ impl Operation {
         object_store: Arc<dyn ObjectStore>,
     ) -> Result<(Option<TableRequirement>, Vec<TableUpdate>), Error> {
         match self {
+            Operation::AppendSequenceGroups {
+                branch,
+                sequence_groups,
+            } => {
+                let old_snapshot = table_metadata.current_snapshot(branch.as_deref())?;
+
+                let manifest_list_schema = match table_metadata.format_version {
+                    FormatVersion::V1 => manifest_list_schema_v1(),
+                    FormatVersion::V2 => manifest_list_schema_v2(),
+                };
+
+                let n_data_files = sequence_groups.iter().map(|d| d.data_files.len()).sum();
+                let n_delete_files = sequence_groups.iter().map(|d| d.delete_files.len()).sum();
+                if n_data_files + n_delete_files == 0 {
+                    return Ok((None, vec![]));
+                };
+
+                let all_files: Vec<DataFile> = sequence_groups
+                    .iter()
+                    .flat_map(|d| d.delete_files.iter().chain(d.data_files.iter()))
+                    .cloned()
+                    .collect();
+                let additional_summary = append_summary(&all_files);
+
+                let mut manifest_list_writer = if let Some(manifest_list_bytes) =
+                    prefetch_manifest_list(old_snapshot, &object_store)
+                {
+                    let bytes = manifest_list_bytes.await??;
+                    ManifestListWriter::from_existing(
+                        &bytes,
+                        all_files.iter(),
+                        manifest_list_schema,
+                        table_metadata,
+                        branch.as_deref(),
+                    )?
+                } else {
+                    ManifestListWriter::new(
+                        all_files.iter(),
+                        manifest_list_schema,
+                        table_metadata,
+                        branch.as_deref(),
+                    )?
+                };
+
+                let snapshot_id = generate_snapshot_id();
+
+                let mut dsn_offset: i64 = 0;
+                for SequenceGroup {
+                    data_files,
+                    delete_files,
+                } in sequence_groups.into_iter()
+                {
+                    dsn_offset += 1;
+                    let n_data_files_in_group = data_files.len();
+                    let n_delete_files_in_group = delete_files.len();
+
+                    let new_datafile_iter = data_files.into_iter().map(|data_file| {
+                        ManifestEntry::builder()
+                            .with_format_version(table_metadata.format_version)
+                            .with_status(Status::Added)
+                            .with_data_file(data_file)
+                            .with_sequence_number(table_metadata.last_sequence_number + dsn_offset)
+                            .build()
+                            .map_err(Error::from)
+                    });
+
+                    let new_deletefile_iter = delete_files.into_iter().map(|data_file| {
+                        ManifestEntry::builder()
+                            .with_format_version(table_metadata.format_version)
+                            .with_status(Status::Added)
+                            .with_data_file(data_file)
+                            .with_sequence_number(table_metadata.last_sequence_number + dsn_offset)
+                            .build()
+                            .map_err(Error::from)
+                    });
+
+                    // Write manifest files
+                    // Split manifest file if limit is exceeded
+                    for (content, files, n_files) in [
+                        (
+                            Content::Data,
+                            Either::Left(new_datafile_iter),
+                            n_data_files_in_group,
+                        ),
+                        (
+                            Content::Deletes,
+                            Either::Right(new_deletefile_iter),
+                            n_delete_files_in_group,
+                        ),
+                    ] {
+                        if n_files != 0 {
+                            manifest_list_writer
+                                .append(files, snapshot_id, object_store.clone(), content)
+                                .await?;
+                        }
+                    }
+                }
+
+                let new_manifest_list_location = manifest_list_writer
+                    .finish(snapshot_id, object_store)
+                    .await?;
+
+                let snapshot_operation = match (n_data_files, n_delete_files) {
+                    (0, 0) => return Ok((None, Vec::new())),
+                    (_, 0) => Ok::<_, Error>(SnapshotOperation::Append),
+                    (0, _) => Ok(SnapshotOperation::Delete),
+                    (_, _) => Ok(SnapshotOperation::Overwrite),
+                }?;
+
+                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(),
+                    })
+                    .with_schema_id(
+                        *table_metadata
+                            .current_schema(branch.as_deref())?
+                            .schema_id(),
+                    );
+                if let Some(snapshot) = old_snapshot {
+                    snapshot_builder.with_parent_snapshot_id(*snapshot.snapshot_id());
+                }
+                let snapshot = snapshot_builder.build()?;
+
+                Ok((
+                    old_snapshot.map(|x| TableRequirement::AssertRefSnapshotId {
+                        r#ref: branch.clone().unwrap_or("main".to_owned()),
+                        snapshot_id: *x.snapshot_id(),
+                    }),
+                    vec![
+                        TableUpdate::AddSnapshot { snapshot },
+                        TableUpdate::SetSnapshotRef {
+                            ref_name: branch.unwrap_or("main".to_owned()),
+                            snapshot_reference: SnapshotReference {
+                                snapshot_id,
+                                retention: SnapshotRetention::default(),
+                            },
+                        },
+                    ],
+                ))
+            }
             Operation::Append {
                 branch,
                 data_files,