This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-3474-8119f796c461283f559d7b7354963ee3e44dac04 in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git
commit 7ae9f4a44a4f8afae1ba6e7caa1456e4cde790e6 Author: Shawn Chang <[email protected]> AuthorDate: Tue Oct 27 21:36:38 2026 +0000 feat(transaction): Add simple snapshot producer, and wire fast append (#3474) * Add simple snapshot producer, and wire fast append * Leave some todo * port changes from 3363 to this PR * delete snapshot * Validate seq num and snapshot id in manifest cache * address comments * update comments * address v1 manifest seq num issue, remove AddedManifests * cleanup * clean up * add transaction/snapshot mod, move some validations to fast append --- crates/iceberg/src/io/object_cache.rs | 106 +++- crates/iceberg/src/spec/manifest/reader.rs | 29 + crates/iceberg/src/transaction/append.rs | 343 +++++++--- crates/iceberg/src/transaction/mod.rs | 25 +- crates/iceberg/src/transaction/snapshot.rs | 741 ---------------------- crates/iceberg/src/transaction/snapshot/mod.rs | 667 +++++++++++++++++++ crates/iceberg/src/transaction/snapshot/simple.rs | 476 ++++++++++++++ 7 files changed, 1543 insertions(+), 844 deletions(-) diff --git a/crates/iceberg/src/io/object_cache.rs b/crates/iceberg/src/io/object_cache.rs index dfc6c124c..874c9d348 100644 --- a/crates/iceberg/src/io/object_cache.rs +++ b/crates/iceberg/src/io/object_cache.rs @@ -42,11 +42,33 @@ pub(crate) enum CachedObjectKey { // shares that entry. The version can still shift under a live cache across // an in-place upgrade (`Table::with_metadata`), which is why the key omits it. ManifestList(String), - // The manifest-level `first_row_id` is part of the key because the parsed - // manifest inherits it onto its entries: the same physical manifest can be - // referenced with different offsets across snapshots and branches, so it - // cannot be shared under the path alone. - Manifest((String, Option<u64>)), + // Keyed by the path plus every manifest-list value that `ManifestReader` + // inherits onto the parsed entries. The same physical manifest can be + // referenced under different values: `first_row_id` differs across + // snapshots and branches, and `sequence_number` differs when a transaction + // retry reuses a manifest it wrote under an earlier attempt's manifest + // list. A key missing any of these would serve one reference's inherited + // values for another. + Manifest(ManifestKey), +} + +#[derive(Clone, Debug, Hash, Eq, PartialEq)] +pub(crate) struct ManifestKey { + path: String, + added_snapshot_id: i64, + sequence_number: i64, + first_row_id: Option<u64>, +} + +impl ManifestKey { + fn new(manifest_file: &ManifestFile) -> Self { + Self { + path: manifest_file.manifest_path.clone(), + added_snapshot_id: manifest_file.added_snapshot_id, + sequence_number: manifest_file.sequence_number, + first_row_id: manifest_file.first_row_id, + } + } } /// Caches metadata objects deserialized from immutable files @@ -114,10 +136,7 @@ impl ObjectCache { .map(Arc::new); } - let key = CachedObjectKey::Manifest(( - manifest_file.manifest_path.clone(), - manifest_file.first_row_id, - )); + let key = CachedObjectKey::Manifest(ManifestKey::new(manifest_file)); let cache_entry = self .cache @@ -662,4 +681,73 @@ mod tests { Some(2000) ); } + + /// A transaction retry reuses the manifest it wrote during an earlier + /// attempt, but the retry's manifest list assigns it a newer sequence + /// number. Entries inherit that number on read, so a cache that already + /// holds the earlier attempt's parse must not serve it for the retry. + #[tokio::test] + async fn test_get_manifest_keys_on_sequence_number() { + use crate::spec::{NestedField, PartitionSpec, PrimitiveType, Schema, SchemaRef, Type}; + + let io = FileIO::new_with_memory(); + let path = "memory:///metadata/retried_manifest.avro"; + + let schema: SchemaRef = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + ]) + .build() + .unwrap() + .into(); + let partition_spec = PartitionSpec::builder(schema.clone()) + .with_spec_id(0) + .build() + .unwrap(); + + // A v2 manifest: entries carry no sequence number of their own and + // `first_row_id` stays `None`, so only the sequence number distinguishes + // the two references below. + let mut writer = ManifestWriterBuilder::new( + io.new_output(path).unwrap(), + Some(1), + schema, + partition_spec, + ) + .build_v2_data(); + writer + .add_entry( + ManifestEntry::builder() + .status(ManifestStatus::Added) + .data_file( + DataFileBuilder::default() + .content(DataContentType::Data) + .file_path(path.to_string()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(100) + .record_count(1) + .build() + .unwrap(), + ) + .build(), + ) + .unwrap(); + let manifest_file = writer.write_manifest_file().await.unwrap(); + assert_eq!(manifest_file.first_row_id, None); + + // First attempt's manifest list assigns 35; after a rebase the retry's + // manifest list assigns 36 to the same physical manifest. + let mut first_attempt = manifest_file.clone(); + first_attempt.sequence_number = 35; + let mut retry = manifest_file; + retry.sequence_number = 36; + + let object_cache = ObjectCache::new(io, None); + + let first = object_cache.get_manifest(&first_attempt).await.unwrap(); + let second = object_cache.get_manifest(&retry).await.unwrap(); + + assert_eq!(first.entries()[0].sequence_number(), Some(35)); + assert_eq!(second.entries()[0].sequence_number(), Some(36)); + } } diff --git a/crates/iceberg/src/spec/manifest/reader.rs b/crates/iceberg/src/spec/manifest/reader.rs index a056d207c..14348ca49 100644 --- a/crates/iceberg/src/spec/manifest/reader.rs +++ b/crates/iceberg/src/spec/manifest/reader.rs @@ -178,6 +178,35 @@ mod tests { ); } + /// Per spec, inheritance does not apply to v1: a v1 manifest's entries read + /// with sequence number 0 even when the manifest list entry referencing them + /// carries a nonzero one. + #[tokio::test] + async fn test_read_v1_manifest_defaults_sequence_number_to_zero() { + let schema = test_schema(); + let partition_spec = PartitionSpec::builder(schema.clone()) + .with_spec_id(0) + .build() + .unwrap(); + + let io = FileIO::new_with_memory(); + let path = "memory:///table/metadata/v1.avro"; + let mut writer = ManifestWriterBuilder::new( + io.new_output(path).unwrap(), + Some(1), + schema.clone(), + partition_spec, + ) + .build_v1(); + writer.add_entry(test_entry()).unwrap(); + let mut manifest_file = writer.write_manifest_file().await.unwrap(); + manifest_file.sequence_number = 36; + + let manifest = ManifestReader::new(io).read(&manifest_file).await.unwrap(); + assert_eq!(manifest.entries()[0].sequence_number(), Some(0)); + assert_eq!(manifest.entries()[0].file_sequence_number, Some(0)); + } + #[tokio::test] async fn test_read_encrypted_manifest_roundtrip() { let schema = test_schema(); diff --git a/crates/iceberg/src/transaction/append.rs b/crates/iceberg/src/transaction/append.rs index 9c438759b..08ac900ad 100644 --- a/crates/iceberg/src/transaction/append.rs +++ b/crates/iceberg/src/transaction/append.rs @@ -20,23 +20,26 @@ use std::collections::{HashMap, HashSet}; use async_trait::async_trait; use uuid::Uuid; -use crate::error::Result; -use crate::spec::{DataFile, ManifestEntry, ManifestFile, Operation}; +use crate::error::{Result, invalid_data}; +use crate::spec::{DataContentType, DataFile}; use crate::table::Table; use crate::transaction::action::CommitStatus; -use crate::transaction::snapshot::{ - DefaultManifestProcess, SnapshotProduceOperation, SnapshotProducer, -}; +use crate::transaction::snapshot::simple::{AppendIntent, SimpleSnapshotProducer}; +use crate::transaction::snapshot::validate_files_not_live; use crate::transaction::{ActionCommit, TransactionAction}; +use crate::{Error, ErrorKind}; /// FastAppendAction is a transaction action for fast append data files to the table. #[derive(Clone)] pub struct FastAppendAction { check_duplicate: bool, - // below are properties used to create SnapshotProducer when commit + /// Overrides the commit UUID of the producer created in `new_state`. commit_uuid: Option<Uuid>, snapshot_properties: HashMap<String, String>, + /// Free of duplicate paths; see [`Self::add_data_files`]. added_data_files: Vec<DataFile>, + /// Paths already in `added_data_files`. + added_paths: HashSet<String>, } impl FastAppendAction { @@ -46,6 +49,7 @@ impl FastAppendAction { commit_uuid: None, snapshot_properties: HashMap::default(), added_data_files: vec![], + added_paths: HashSet::new(), } } @@ -56,8 +60,16 @@ impl FastAppendAction { } /// Add data files to the snapshot. + /// + /// A path added more than once keeps its first occurrence, so a manifest + /// never references the same file twice. This is independent of the + /// `check_duplicate`-gated check against files already in the table. pub fn add_data_files(mut self, data_files: impl IntoIterator<Item = DataFile>) -> Self { - self.added_data_files.extend(data_files); + for data_file in data_files { + if self.added_paths.insert(data_file.file_path.clone()) { + self.added_data_files.push(data_file); + } + } self } @@ -73,94 +85,59 @@ impl FastAppendAction { self } - /// Collapse files sharing a path to their first occurrence, so a single - /// manifest never references the same file twice. Always runs (unlike the - /// `check_duplicate`-gated cross-snapshot check) since it is in-memory only. - fn dedupe_added_files(&self) -> Vec<DataFile> { - let mut seen = HashSet::with_capacity(self.added_data_files.len()); - self.added_data_files + /// Checks the parts of this action that do not depend on the table. + fn validate(&self) -> Result<()> { + // TODO: Allowing snapshot property setup with no added data files is a workaround. + // We should clean it up after all necessary actions are supported. + // For details, please refer to https://github.com/apache/iceberg-rust/issues/1548 + if self.added_data_files.is_empty() && self.snapshot_properties.is_empty() { + return Err(Error::new( + ErrorKind::PreconditionFailed, + "No added data files or added snapshot properties found when write a manifest file", + )); + } + if let Some(data_file) = self + .added_data_files .iter() - .filter(|data_file| seen.insert(data_file.file_path.as_str())) - .cloned() - .collect() + .find(|f| f.content_type() != DataContentType::Data) + { + return Err(invalid_data!( + "Only data content type is allowed for fast append, got {:?} for {}", + data_file.content_type(), + data_file.file_path() + )); + } + Ok(()) } } #[async_trait] impl TransactionAction for FastAppendAction { - // TODO: replace with a persistent `SimpleSnapshotProducer` once snapshot - // production is migrated (stateful transaction RFC, milestone 2). - type State = (); - - fn new_state(&self) -> Self::State {} - - async fn commit(&self, _state: &mut (), table: &Table) -> Result<ActionCommit> { - let snapshot_producer = SnapshotProducer::new( - table, - self.commit_uuid.unwrap_or_else(Uuid::now_v7), - self.snapshot_properties.clone(), - self.dedupe_added_files(), - vec![], - ); - - // validate added files - snapshot_producer.validate_added_data_files()?; - - // Checks duplicate files - if self.check_duplicate { - snapshot_producer.validate_duplicate_files().await?; - } - - snapshot_producer - .commit(FastAppendOperation, DefaultManifestProcess) - .await - } - - async fn cleanup(self, _state: (), _table: &Table, _status: CommitStatus) { - // Stateless for now: generated manifests are not yet tracked for - // terminal cleanup (stateful transaction RFC, milestone 2/3). + type State = SimpleSnapshotProducer; + + fn new_state(&self) -> SimpleSnapshotProducer { + SimpleSnapshotProducer::new(AppendIntent { + commit_uuid: self.commit_uuid, + snapshot_properties: self.snapshot_properties.clone(), + data_files: self.added_data_files.clone(), + }) } -} - -struct FastAppendOperation; -impl SnapshotProduceOperation for FastAppendOperation { - fn operation(&self) -> Operation { - Operation::Append - } - - async fn delete_entries( + async fn commit( &self, - _snapshot_produce: &SnapshotProducer<'_>, - ) -> Result<Vec<ManifestEntry>> { - Ok(vec![]) + state: &mut SimpleSnapshotProducer, + table: &Table, + ) -> Result<ActionCommit> { + self.validate()?; + if self.check_duplicate { + validate_files_not_live(table, &self.added_data_files).await?; + } + state.apply(table).await } - async fn existing_manifest( - &self, - snapshot_produce: &SnapshotProducer<'_>, - ) -> Result<Vec<ManifestFile>> { - let Some(snapshot) = snapshot_produce.table.metadata().current_snapshot() else { - return Ok(vec![]); - }; - - let manifest_list = snapshot_produce - .table - .manifest_list_reader(snapshot) - .load() - .await?; - - Ok(manifest_list - .entries() - .iter() - .filter(|entry| { - // Keep delete-only manifests too: they record which files were removed and - // must persist across snapshots until `expire_snapshots` cleans them up. - // Dropping them lets the removed files reappear as live data (see #2148). - entry.has_added_files() || entry.has_existing_files() || entry.has_deleted_files() - }) - .cloned() - .collect()) + async fn cleanup(self, _state: SimpleSnapshotProducer, _table: &Table, _status: CommitStatus) { + // TODO(#3417): remove the producer's uncommitted artifacts once terminal + // cleanup is wired up. } } @@ -169,11 +146,13 @@ mod tests { use std::collections::HashMap; use std::fs; use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; use minijinja::{AutoEscape, Environment, Value, context}; use tempfile::TempDir; use uuid::Uuid; + use crate::catalog::MockCatalog; use crate::encryption::kms::MemoryKeyManagementClient; use crate::encryption::{SensitiveBytes, StandardKeyMetadata}; use crate::io::FileIO; @@ -184,9 +163,11 @@ mod tests { }; use crate::table::Table; use crate::test_utils::{commit_with_fresh_state, make_encrypted_table, test_runtime}; - use crate::transaction::Transaction; - use crate::transaction::tests::make_v2_minimal_table; - use crate::{TableIdent, TableRequirement, TableUpdate}; + use crate::transaction::tests::{ + make_v2_minimal_table, make_v3_minimal_table, with_retry_properties, + }; + use crate::transaction::{ApplyTransactionAction, FastAppendAction, Transaction}; + use crate::{ErrorKind, TableCommit, TableIdent, TableRequirement, TableUpdate}; fn render_template(template: &str, ctx: Value) -> String { let mut env = Environment::new(); @@ -493,6 +474,30 @@ mod tests { ); } + #[tokio::test] + async fn test_fast_append_rejects_delete_files() { + let table = make_v2_minimal_table(); + let delete_file = DataFileBuilder::default() + .content(DataContentType::EqualityDeletes) + .file_path("test/1-deletes.parquet".to_string()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(100) + .record_count(1) + .partition_spec_id(table.metadata().default_partition_spec_id()) + .partition(Struct::from_iter([Some(Literal::long(300))])) + .equality_ids(Some(vec![1])) + .build() + .unwrap(); + + let tx = Transaction::new(&table); + let action = tx.fast_append().add_data_files(vec![delete_file]); + let err = commit_with_fresh_state(action, &table) + .await + .err() + .expect("delete files must be rejected"); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + #[tokio::test] async fn test_empty_data_append_action() { let table = make_v2_minimal_table(); @@ -879,12 +884,12 @@ mod tests { .unwrap() }; - // Same path three times: the manifest keeps a single entry, the first one. - let action = tx.fast_append().add_data_files(vec![ - make_file(100, 10), - make_file(200, 20), - make_file(300, 30), - ]); + // Same path three times, across two `add_data_files` calls: the manifest + // keeps a single entry, the first one. + let action = tx + .fast_append() + .add_data_files(vec![make_file(100, 10), make_file(200, 20)]) + .add_data_files(vec![make_file(300, 30)]); let mut action_commit = commit_with_fresh_state(action, &table).await.unwrap(); let files = committed_data_files(&table, &action_commit.take_updates()).await; assert_eq!(1, files.len()); @@ -995,4 +1000,156 @@ mod tests { ); assert_eq!(data_file, *manifest.entries()[0].data_file()); } + + /// Two writers append concurrently. Writer A starts from S0, but writer B + /// commits S1 first, so A's first attempt fails the `RefSnapshotIdMatch` + /// requirement and A rebases onto S1 to commit S2. + /// + /// The producer state survives the retry: A's manifest is written once and + /// reused, while parent, sequence number, manifest list and summary are + /// recomputed against S1 (see #3418). + #[tokio::test] + async fn test_fast_append_retry_rebases_onto_concurrent_commit() { + let (s1, s2) = rebase_onto_concurrent_commit(make_v2_minimal_table()).await; + assert_eq!(s1.first_row_id(), None); + assert_eq!(s2.first_row_id(), None); + } + + /// On a v3 table the retry must also move the row range: S2's first row id + /// is S1's `next_row_id`, not S0's. + #[tokio::test] + async fn test_fast_append_retry_rebases_row_range_onto_concurrent_commit() { + let (s1, s2) = rebase_onto_concurrent_commit(make_v3_minimal_table()).await; + + // B's single-row file takes row 0 in S1, so A's file starts at row 1. + assert_eq!(s1.row_range(), Some((0, 1))); + assert_eq!(s2.row_range(), Some((1, 1))); + } + + /// Runs the concurrent-append rebase scenario on `s0` and returns the + /// snapshots committed by writer B (S1) and writer A (S2). + async fn rebase_onto_concurrent_commit(s0: Table) -> (SnapshotRef, SnapshotRef) { + // `TableCommit::apply` derives the next metadata location from the + // current one, so it needs a well-formed name. + let s0 = with_retry_properties(s0, "3").with_metadata_location( + "memory:///table/metadata/00000-9c12d441-03fe-4693-9a96-a0705ddf69c1.metadata.json" + .to_string(), + ); + let data_file = |path: &str| { + DataFileBuilder::default() + .content(DataContentType::Data) + .file_path(path.to_string()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(100) + .record_count(1) + .partition_spec_id(s0.metadata().default_partition_spec_id()) + .partition(Struct::from_iter([Some(Literal::long(300))])) + .build() + .unwrap() + }; + + // Writer B commits S1 on top of S0. + let mut b_commit = commit_with_fresh_state( + FastAppendAction::new().add_data_files(vec![data_file("test/b.parquet")]), + &s0, + ) + .await + .unwrap(); + let s1 = TableCommit::builder() + .ident(s0.identifier().clone()) + .updates(b_commit.take_updates()) + .requirements(b_commit.take_requirements()) + .build() + .apply(s0.clone()) + .unwrap(); + let s1_snapshot = s1.metadata().current_snapshot().unwrap().clone(); + + // The catalog serves S0 to A's first attempt and S1 afterwards; every + // update is validated against S1, so the first attempt's + // `RefSnapshotIdMatch { snapshot_id: None }` fails with a real conflict. + let mut mock_catalog = MockCatalog::new(); + let loads = AtomicU32::new(0); + let (load_s0, load_s1) = (s0.clone(), s1.clone()); + mock_catalog + .expect_load_table() + .times(2) + .returning_st(move |_| { + let table = if loads.fetch_add(1, Ordering::SeqCst) == 0 { + load_s0.clone() + } else { + load_s1.clone() + }; + Box::pin(async move { Ok(table) }) + }); + let apply_s1 = s1.clone(); + mock_catalog + .expect_update_table() + .times(2) + .returning_st(move |commit| { + let s1 = apply_s1.clone(); + Box::pin(async move { commit.apply(s1) }) + }); + + // Writer A starts from S0. + let a_uuid = Uuid::now_v7(); + let tx = Transaction::new(&s0); + let tx = tx + .fast_append() + .set_commit_uuid(a_uuid) + .add_data_files(vec![data_file("test/a.parquet")]) + .apply(tx) + .unwrap(); + let s2 = tx.commit(&mock_catalog).await.unwrap(); + let s2_snapshot = s2.metadata().current_snapshot().unwrap().clone(); + + // S2 is rebased onto S1, not S0. + assert_eq!( + s2_snapshot.parent_snapshot_id(), + Some(s1_snapshot.snapshot_id()) + ); + assert_eq!( + s2_snapshot.sequence_number(), + s1_snapshot.sequence_number() + 1 + ); + // The committed manifest list is A's second attempt. + assert!( + s2_snapshot + .manifest_list() + .ends_with(&format!("-1-{a_uuid}.avro")), + "unexpected manifest list path {}", + s2_snapshot.manifest_list() + ); + + // S2 carries B's manifest forward and adds A's. A's manifest is the one + // written during the first attempt (counter 0), reused after rebase. + let manifest_list = s2.manifest_list_reader(&s2_snapshot).load().await.unwrap(); + assert_eq!(manifest_list.entries().len(), 2); + let b_manifest = manifest_list + .entries() + .iter() + .find(|m| m.added_snapshot_id == s1_snapshot.snapshot_id()) + .expect("S2 should carry B's manifest forward"); + let a_manifest = manifest_list + .entries() + .iter() + .find(|m| m.added_snapshot_id == s2_snapshot.snapshot_id()) + .expect("S2 should add A's manifest"); + assert!( + a_manifest + .manifest_path + .ends_with(&format!("{a_uuid}-m0.avro")), + "unexpected manifest path {}", + a_manifest.manifest_path + ); + assert_eq!(a_manifest.sequence_number, s2_snapshot.sequence_number()); + assert_eq!(b_manifest.sequence_number, s1_snapshot.sequence_number()); + + // Summary totals are computed against S1, not S0. + let summary = &s2_snapshot.summary().additional_properties; + assert_eq!(summary.get("added-data-files").unwrap(), "1"); + assert_eq!(summary.get("total-data-files").unwrap(), "2"); + assert_eq!(summary.get("total-records").unwrap(), "2"); + + (s1_snapshot, s2_snapshot) + } } diff --git a/crates/iceberg/src/transaction/mod.rs b/crates/iceberg/src/transaction/mod.rs index 57520447d..cd31ca0a2 100644 --- a/crates/iceberg/src/transaction/mod.rs +++ b/crates/iceberg/src/transaction/mod.rs @@ -332,6 +332,26 @@ mod tests { .unwrap() } + pub(crate) fn make_v3_minimal_table() -> Table { + let file = File::open(format!( + "{}/testdata/table_metadata/{}", + env!("CARGO_MANIFEST_DIR"), + "TableMetadataV3ValidMinimal.json" + )) + .unwrap(); + let reader = BufReader::new(file); + let resp = serde_json::from_reader::<_, TableMetadata>(reader).unwrap(); + + Table::builder() + .metadata(resp) + .metadata_location("s3://bucket/test/location/metadata/v1.json") + .identifier(TableIdent::from_strs(["ns1", "test1"]).unwrap()) + .file_io(FileIO::new_with_memory()) + .runtime(test_runtime()) + .build() + .unwrap() + } + async fn make_v2_minimal_readonly_table() -> Table { let table = make_v2_minimal_table(); StaticTable::from_metadata( @@ -379,8 +399,11 @@ mod tests { /// Helper function to create a test table with retry properties pub(super) fn setup_test_table(num_retries: &str) -> Table { - let table = make_v2_table(); + with_retry_properties(make_v2_table(), num_retries) + } + /// Sets fast commit retry properties on `table`. + pub(super) fn with_retry_properties(table: Table, num_retries: &str) -> Table { // Set retry properties let mut props = HashMap::new(); props.insert("commit.retry.min-wait-ms".to_string(), "10".to_string()); diff --git a/crates/iceberg/src/transaction/snapshot.rs b/crates/iceberg/src/transaction/snapshot.rs deleted file mode 100644 index 7e4d6fb72..000000000 --- a/crates/iceberg/src/transaction/snapshot.rs +++ /dev/null @@ -1,741 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::{HashMap, HashSet}; -use std::future::Future; -use std::sync::atomic::{AtomicU64, Ordering}; - -use futures::TryStreamExt; -use futures::stream::FuturesUnordered; -use uuid::Uuid; - -use crate::error::{Result, invalid_data}; -use crate::spec::{ - DataFile, DataFileFormat, FormatVersion, MAIN_BRANCH, ManifestContentType, ManifestEntry, - ManifestFile, ManifestListWriter, ManifestWriter, ManifestWriterBuilder, Operation, - PartitionSpecRef, SchemaRef, Snapshot, SnapshotReference, SnapshotRetention, - SnapshotSummaryCollector, Struct, StructType, Summary, TableProperties, set_iceberg_version, - update_snapshot_summaries, -}; -use crate::table::Table; -use crate::transaction::ActionCommit; -use crate::{Error, ErrorKind, TableRequirement, TableUpdate}; - -/// A trait that defines how different table operations produce new snapshots. -/// -/// `SnapshotProduceOperation` is used by [`SnapshotProducer`] to customize snapshot creation -/// based on the type of operation being performed (e.g., `Append`, `Overwrite`, `Delete`, etc.). -/// Each operation type implements this trait to specify: -/// - Which operation type to record in the snapshot summary -/// - Which existing manifest files should be included in the new snapshot -/// - Which manifest entries should be marked as deleted -/// -/// # When it accomplishes -/// -/// This trait is used during the snapshot creation process in [`SnapshotProducer::commit()`]: -/// -/// 1. **Operation Type Recording**: The `operation()` method determines which operation type -/// (e.g., `Operation::Append`, `Operation::Overwrite`) is recorded in the snapshot summary. -/// This metadata helps track what kind of change was made to the table. -/// -/// 2. **Manifest File Selection**: The `existing_manifest()` method determines which existing -/// manifest files from the current snapshot should be carried forward to the new snapshot. -/// For example: -/// - An `Append` operation typically includes all existing manifests plus new ones -/// - An `Overwrite` operation might exclude manifests for partitions being overwritten -/// -/// 3. **Delete Entry Processing**: The `delete_entries()` method is intended for future delete -/// operations to specify which manifest entries should be marked as deleted. -pub(crate) trait SnapshotProduceOperation: Send + Sync { - /// Returns the operation type that will be recorded in the snapshot summary. - /// - /// This determines what kind of operation is being performed (e.g., `Append`, `Overwrite`), - /// which is stored in the snapshot metadata for tracking and auditing purposes. - fn operation(&self) -> Operation; - - /// Returns manifest entries that should be marked as deleted in the new snapshot. - #[allow(unused)] - fn delete_entries( - &self, - snapshot_produce: &SnapshotProducer, - ) -> impl Future<Output = Result<Vec<ManifestEntry>>> + Send; - - /// Returns existing manifest files that should be included in the new snapshot. - /// - /// This method determines which manifest files from the current snapshot should be - /// carried forward to the new snapshot. The selection depends on the operation type: - /// - /// - **Append operations**: Typically include all existing manifests - /// - **Overwrite operations**: May exclude manifests for partitions being overwritten - /// - **Delete operations**: May exclude manifests for partitions being deleted - fn existing_manifest( - &self, - snapshot_produce: &SnapshotProducer<'_>, - ) -> impl Future<Output = Result<Vec<ManifestFile>>> + Send; - - /// Returns whether the snapshot summary resets the table totals. Defaults to every - /// overwrite; an operation that overwrites only some files returns `false`. - fn truncate_full_table(&self) -> bool { - self.operation() == Operation::Overwrite - } -} - -pub(crate) struct DefaultManifestProcess; - -impl ManifestProcess for DefaultManifestProcess { - fn process_manifests( - &self, - _snapshot_produce: &SnapshotProducer<'_>, - manifests: Vec<ManifestFile>, - ) -> Vec<ManifestFile> { - manifests - } -} - -pub(crate) trait ManifestProcess: Send + Sync { - fn process_manifests( - &self, - snapshot_produce: &SnapshotProducer<'_>, - manifests: Vec<ManifestFile>, - ) -> Vec<ManifestFile>; -} - -pub(crate) struct SnapshotProducer<'a> { - pub(crate) table: &'a Table, - snapshot_id: i64, - commit_uuid: Uuid, - snapshot_properties: HashMap<String, String>, - added_data_files: Vec<DataFile>, - /// Must match the data files marked Deleted by this commit, including their metrics. - deleted_data_files: Vec<DataFile>, - // Numbers this commit's manifest files. Atomic so `new_manifest_writer` can take `&self`. - manifest_counter: AtomicU64, -} - -impl<'a> SnapshotProducer<'a> { - pub(crate) fn new( - table: &'a Table, - commit_uuid: Uuid, - snapshot_properties: HashMap<String, String>, - added_data_files: Vec<DataFile>, - deleted_data_files: Vec<DataFile>, - ) -> Self { - Self { - table, - snapshot_id: Self::generate_unique_snapshot_id(table), - commit_uuid, - snapshot_properties, - added_data_files, - deleted_data_files, - manifest_counter: AtomicU64::new(0), - } - } - - pub(crate) fn validate_added_data_files(&self) -> Result<()> { - for data_file in &self.added_data_files { - if data_file.content_type() != crate::spec::DataContentType::Data { - return Err(invalid_data!( - "Only data content type is allowed for fast append" - )); - } - // Check if the data file partition spec id matches the table default partition spec id. - if self.table.metadata().default_partition_spec_id() != data_file.partition_spec_id { - return Err(invalid_data!( - "Data file partition spec id does not match table default partition spec id" - )); - } - Self::validate_partition_value( - data_file.partition(), - self.table.metadata().default_partition_type(), - )?; - } - - Ok(()) - } - - pub(crate) async fn validate_duplicate_files(&self) -> Result<()> { - let Some(current_snapshot) = self.table.metadata().current_snapshot() else { - return Ok(()); - }; - - let new_files: HashSet<&str> = self - .added_data_files - .iter() - .map(|df| df.file_path.as_str()) - .collect(); - - let runtime = self.table.runtime(); - let manifest_list = self - .table - .manifest_list_reader(current_snapshot) - .load() - .await?; - - let new_files_ref = &new_files; - let referenced_files: Vec<String> = manifest_list - .consume_entries() - .into_iter() - .map(|entry| { - let reader = self.table.manifest_reader(); - runtime.io().spawn(async move { reader.read(&entry).await }) - }) - .collect::<FuturesUnordered<_>>() - .try_fold(Vec::new(), |mut acc, manifest| async move { - acc.extend( - manifest? - .entries() - .iter() - .filter(|e| new_files_ref.contains(e.file_path()) && e.is_alive()) - .map(|e| e.file_path().to_string()), - ); - Ok(acc) - }) - .await?; - - if !referenced_files.is_empty() { - return Err(invalid_data!( - "Cannot add files that are already referenced by table, files: {}", - referenced_files.join(", ") - )); - } - - Ok(()) - } - - fn generate_unique_snapshot_id(table: &Table) -> i64 { - let generate_random_id = || -> i64 { - let (lhs, rhs) = Uuid::new_v4().as_u64_pair(); - let snapshot_id = (lhs ^ rhs) as i64; - if snapshot_id < 0 { - -snapshot_id - } else { - snapshot_id - } - }; - let mut snapshot_id = generate_random_id(); - - while table - .metadata() - .snapshots() - .any(|s| s.snapshot_id() == snapshot_id) - { - snapshot_id = generate_random_id(); - } - snapshot_id - } - - /// Returns a writer for the next manifest file of this commit, written under the given - /// schema and partition spec. - pub(crate) fn new_manifest_writer( - &self, - content: ManifestContentType, - schema: SchemaRef, - partition_spec: PartitionSpecRef, - ) -> Result<ManifestWriter> { - let new_manifest_path = format!( - "{}/{}-m{}.{}", - self.table.metadata().metadata_location()?, - self.commit_uuid, - self.manifest_counter.fetch_add(1, Ordering::Relaxed), - DataFileFormat::Avro - ); - let output_file = self.table.file_io().new_output(new_manifest_path)?; - let partition_spec = partition_spec.as_ref().clone(); - - let builder = if let Some(em) = self.table.encryption_manager() { - ManifestWriterBuilder::new_from_encrypted( - em.encrypt(output_file), - Some(self.snapshot_id), - schema, - partition_spec, - )? - } else { - ManifestWriterBuilder::new(output_file, Some(self.snapshot_id), schema, partition_spec) - }; - - match self.table.metadata().format_version() { - FormatVersion::V1 => Ok(builder.build_v1()), - FormatVersion::V2 => match content { - ManifestContentType::Data => Ok(builder.build_v2_data()), - ManifestContentType::Deletes => Ok(builder.build_v2_deletes()), - }, - FormatVersion::V3 => match content { - ManifestContentType::Data => Ok(builder.build_v3_data()), - ManifestContentType::Deletes => Ok(builder.build_v3_deletes()), - }, - } - } - - // Check if the partition value is compatible with the partition type. - fn validate_partition_value( - partition_value: &Struct, - partition_type: &StructType, - ) -> Result<()> { - if partition_value.fields().len() != partition_type.fields().len() { - return Err(invalid_data!( - "Partition value is not compatible with partition type" - )); - } - - for (value, field) in partition_value.fields().iter().zip(partition_type.fields()) { - let field = field.field_type.as_primitive_type().ok_or_else(|| { - Error::new( - ErrorKind::Unexpected, - "Partition field should only be primitive type.", - ) - })?; - if let Some(value) = value - && !field.compatible(&value.as_primitive_literal().unwrap()) - { - return Err(invalid_data!( - "Partition value is not compatible partition type" - )); - } - } - Ok(()) - } - - // Write manifest file for added data files and return the ManifestFile for ManifestList. - async fn write_added_manifest(&mut self) -> Result<ManifestFile> { - let added_data_files = std::mem::take(&mut self.added_data_files); - if added_data_files.is_empty() { - return Err(Error::new( - ErrorKind::PreconditionFailed, - "No added data files found when write an added manifest file", - )); - } - - let snapshot_id = self.snapshot_id; - let format_version = self.table.metadata().format_version(); - let manifest_entries = added_data_files.into_iter().map(|data_file| { - let builder = ManifestEntry::builder() - .status(crate::spec::ManifestStatus::Added) - .data_file(data_file); - if format_version == FormatVersion::V1 { - builder.snapshot_id(snapshot_id).build() - } else { - // For format version > 1, we set the snapshot id at the inherited time to avoid rewrite the manifest file when - // commit failed. - builder.build() - } - }); - let mut writer = self.new_manifest_writer( - ManifestContentType::Data, - self.table.metadata().current_schema().clone(), - self.table.metadata().default_partition_spec().clone(), - )?; - for entry in manifest_entries { - writer.add_entry(entry)?; - } - writer.write_manifest_file().await - } - - /// Creates new manifests for data files added or removed, - /// and collects all of the manifests to be included in the new snapshot as [ManifestFile] entries. - async fn produce_manifests<OP: SnapshotProduceOperation, MP: ManifestProcess>( - &mut self, - snapshot_produce_operation: &OP, - manifest_process: &MP, - ) -> Result<Vec<ManifestFile>> { - // Assert current snapshot producer contains new content to add to new snapshot. - // - // TODO: Allowing snapshot property setup with no added data files is a workaround. - // We should clean it up after all necessary actions are supported. - // For details, please refer to https://github.com/apache/iceberg-rust/issues/1548 - if self.added_data_files.is_empty() - && self.deleted_data_files.is_empty() - && self.snapshot_properties.is_empty() - { - return Err(Error::new( - ErrorKind::PreconditionFailed, - "No added or deleted data files or added snapshot properties found when write a manifest file", - )); - } - - let existing_manifests = snapshot_produce_operation.existing_manifest(self).await?; - let mut manifest_files = existing_manifests; - - // Process added entries. - if !self.added_data_files.is_empty() { - let added_manifest = self.write_added_manifest().await?; - manifest_files.push(added_manifest); - } - - // # TODO - // Support process delete entries. - - debug_assert_eq!( - manifest_files - .iter() - .filter(|m| m.added_snapshot_id == self.snapshot_id - && m.content == ManifestContentType::Data) - .filter_map(|m| m.deleted_files_count) - .map(u64::from) - .sum::<u64>(), - self.deleted_data_files.len() as u64, - "deleted files must match this commit's manifest entries" - ); - - let manifest_files = manifest_process.process_manifests(self, manifest_files); - Ok(manifest_files) - } - - // Returns a `Summary` of the current snapshot - fn summary<OP: SnapshotProduceOperation>( - &self, - snapshot_produce_operation: &OP, - ) -> Result<Summary> { - let mut summary_collector = SnapshotSummaryCollector::default(); - let table_metadata = self.table.metadata_ref(); - - let partition_summary_limit = if let Some(limit) = table_metadata - .properties() - .get(TableProperties::PROPERTY_WRITE_PARTITION_SUMMARY_LIMIT) - { - if let Ok(limit) = limit.parse::<u64>() { - limit - } else { - TableProperties::PROPERTY_WRITE_PARTITION_SUMMARY_LIMIT_DEFAULT - } - } else { - TableProperties::PROPERTY_WRITE_PARTITION_SUMMARY_LIMIT_DEFAULT - }; - - summary_collector.set_partition_summary_limit(partition_summary_limit); - - for data_file in &self.added_data_files { - summary_collector.add_file( - data_file, - table_metadata.current_schema().clone(), - table_metadata.default_partition_spec().clone(), - ); - } - - for data_file in &self.deleted_data_files { - let partition_spec = table_metadata - .partition_spec_by_id(data_file.partition_spec_id) - .ok_or_else(|| { - invalid_data!("Unknown partition spec id {}", data_file.partition_spec_id) - })?; - summary_collector.remove_file( - data_file, - table_metadata.current_schema().clone(), - partition_spec.clone(), - ); - } - - let previous_snapshot = table_metadata.current_snapshot(); - - // User-supplied snapshot properties are applied first, then the computed - // metrics overwrite any colliding keys. This matches iceberg-java - // (`SnapshotProducer.summary`), where computed `added-*`/`total-*` values - // are written after user properties so a user cannot shadow them with a - // bad (or merely wrong) value that would corrupt the snapshot summary. - let mut additional_properties = self.snapshot_properties.clone(); - additional_properties.extend(summary_collector.build()); - - let summary = Summary { - operation: snapshot_produce_operation.operation(), - additional_properties, - }; - - let mut summary = update_snapshot_summaries( - summary, - previous_snapshot.map(|s| s.summary()), - snapshot_produce_operation.truncate_full_table(), - )?; - set_iceberg_version(&mut summary.additional_properties); - Ok(summary) - } - - fn generate_manifest_list_file_path(&self, attempt: i64) -> Result<String> { - Ok(format!( - "{}/snap-{}-{}-{}.{}", - self.table.metadata().metadata_location()?, - self.snapshot_id, - attempt, - self.commit_uuid, - DataFileFormat::Avro - )) - } - - /// Finished building the action and return the [`ActionCommit`] to the transaction. - pub(crate) async fn commit<OP: SnapshotProduceOperation, MP: ManifestProcess>( - mut self, - snapshot_produce_operation: OP, - process: MP, - ) -> Result<ActionCommit> { - let manifest_list_path = self.generate_manifest_list_file_path(0)?; - let next_seq_num = self.table.metadata().next_sequence_number(); - let first_row_id = self.table.metadata().next_row_id(); - - let raw_output = self - .table - .file_io() - .new_output(manifest_list_path.clone())?; - - let (writer, encrypted_output) = match self.table.encryption_manager() { - Some(em) => { - let encrypted_output = em.encrypt(raw_output); - ( - encrypted_output.writer().await?, - Some((em.clone(), encrypted_output)), - ) - } - None => (raw_output.writer().await?, None), - }; - - let parent_snapshot_id = self.table.metadata().current_snapshot_id(); - let mut manifest_list_writer = match self.table.metadata().format_version() { - FormatVersion::V1 => { - ManifestListWriter::v1(writer, self.snapshot_id, parent_snapshot_id) - } - FormatVersion::V2 => { - ManifestListWriter::v2(writer, self.snapshot_id, parent_snapshot_id, next_seq_num) - } - FormatVersion::V3 => ManifestListWriter::v3( - writer, - self.snapshot_id, - parent_snapshot_id, - next_seq_num, - Some(first_row_id), - ), - }; - - // Calling self.summary() before self.produce_manifests() is important because self.added_data_files - // will be set to an empty vec after self.produce_manifests() returns, resulting in an empty summary - // being generated. - let summary = self.summary(&snapshot_produce_operation).map_err(|err| { - Error::new(ErrorKind::Unexpected, "Failed to create snapshot summary.").with_source(err) - })?; - - let new_manifests = self - .produce_manifests(&snapshot_produce_operation, &process) - .await?; - - manifest_list_writer.add_manifests(new_manifests.into_iter())?; - let writer_next_row_id = manifest_list_writer.next_row_id(); - let file_metadata = manifest_list_writer.close().await?; - let encryption_key_id = match encrypted_output { - Some((em, encrypted_output)) => Some( - em.encrypt_manifest_list_key_metadata( - &encrypted_output.key_metadata_with_saved_file_metadata(&file_metadata), - ) - .await?, - ), - None => None, - }; - - let commit_ts = chrono::Utc::now().timestamp_millis(); - let new_snapshot = Snapshot::builder() - .with_manifest_list(manifest_list_path) - .with_snapshot_id(self.snapshot_id) - .with_parent_snapshot_id(self.table.metadata().current_snapshot_id()) - .with_sequence_number(next_seq_num) - .with_summary(summary) - .with_schema_id(self.table.metadata().current_schema_id()) - .with_encryption_key_id(encryption_key_id) - .with_timestamp_ms(commit_ts); - - let new_snapshot = if let Some(writer_next_row_id) = writer_next_row_id { - let assigned_rows = writer_next_row_id - self.table.metadata().next_row_id(); - new_snapshot - .with_row_range(first_row_id, assigned_rows) - .build() - } else { - new_snapshot.build() - }; - - let encryption_key_updates: Vec<TableUpdate> = self - .table - .encryption_manager() - .map(|em| { - em.with_encryption_keys(|keys| { - keys.values() - .filter(|k| self.table.metadata().encryption_key(k.key_id()).is_none()) - .map(|k| TableUpdate::AddEncryptionKey { - encryption_key: k.clone(), - }) - .collect() - }) - }) - .unwrap_or_default(); - - let updates = [encryption_key_updates, vec![ - TableUpdate::AddSnapshot { - snapshot: new_snapshot, - }, - TableUpdate::SetSnapshotRef { - ref_name: MAIN_BRANCH.to_string(), - reference: SnapshotReference::new( - self.snapshot_id, - SnapshotRetention::branch(None, None, None), - ), - }, - ]] - .concat(); - - let requirements = vec![ - TableRequirement::UuidMatch { - uuid: self.table.metadata().uuid(), - }, - TableRequirement::RefSnapshotIdMatch { - r#ref: MAIN_BRANCH.to_string(), - snapshot_id: self.table.metadata().current_snapshot_id(), - }, - ]; - - Ok(ActionCommit::new(updates, requirements)) - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::sync::Arc; - - use uuid::Uuid; - - use super::*; - use crate::spec::{ - DataContentType, DataFileBuilder, Literal, Manifest, NestedField, PartitionSpec, - PrimitiveType, Schema, Transform, Type, - }; - use crate::transaction::tests::make_v2_minimal_table; - - struct DeleteOperation; - - impl SnapshotProduceOperation for DeleteOperation { - fn operation(&self) -> Operation { - Operation::Delete - } - - async fn delete_entries(&self, _: &SnapshotProducer<'_>) -> Result<Vec<ManifestEntry>> { - Ok(vec![]) - } - - async fn existing_manifest(&self, _: &SnapshotProducer<'_>) -> Result<Vec<ManifestFile>> { - Ok(vec![]) - } - } - - #[test] - fn test_summary_counts_deleted_data_files() { - let table = make_v2_minimal_table(); - let deleted = DataFileBuilder::default() - .content(DataContentType::Data) - .file_path("test/deleted.parquet".to_string()) - .file_format(DataFileFormat::Parquet) - .file_size_in_bytes(100) - .record_count(2) - .partition_spec_id(table.metadata().default_partition_spec_id()) - .partition(Struct::from_iter([Some(Literal::long(300))])) - .build() - .unwrap(); - - let producer = SnapshotProducer::new(&table, Uuid::now_v7(), HashMap::new(), vec![], vec![ - deleted, - ]); - let summary = producer.summary(&DeleteOperation).unwrap(); - let properties = &summary.additional_properties; - - assert_eq!(properties.get("deleted-data-files").unwrap(), "1"); - assert_eq!(properties.get("deleted-records").unwrap(), "2"); - assert_eq!(properties.get("removed-files-size").unwrap(), "100"); - } - - #[cfg(debug_assertions)] - #[tokio::test] - #[should_panic(expected = "deleted files must match this commit's manifest entries")] - async fn test_commit_rejects_unrecorded_deletions() { - let table = make_v2_minimal_table(); - let deleted = DataFileBuilder::default() - .content(DataContentType::Data) - .file_path("test/deleted.parquet".to_string()) - .file_format(DataFileFormat::Parquet) - .file_size_in_bytes(100) - .record_count(2) - .partition_spec_id(table.metadata().default_partition_spec_id()) - .partition(Struct::from_iter([Some(Literal::long(300))])) - .build() - .unwrap(); - SnapshotProducer::new(&table, Uuid::now_v7(), HashMap::new(), vec![], vec![ - deleted, - ]) - .commit(DeleteOperation, DefaultManifestProcess) - .await - .unwrap(); - } - - #[tokio::test] - async fn test_new_manifest_writer_uses_the_given_schema_and_spec() { - let table = make_v2_minimal_table(); - let schema = Arc::new( - Schema::builder() - .with_schema_id(7) - .with_fields(vec![ - NestedField::required(1, "a", Type::Primitive(PrimitiveType::Long)).into(), - ]) - .build() - .unwrap(), - ); - let partition_spec = Arc::new( - PartitionSpec::builder(schema.clone()) - .with_spec_id(3) - .add_partition_field("a", "a", Transform::Identity) - .unwrap() - .build() - .unwrap(), - ); - - let producer = - SnapshotProducer::new(&table, Uuid::now_v7(), HashMap::new(), vec![], vec![]); - let first = producer - .new_manifest_writer( - ManifestContentType::Data, - schema.clone(), - partition_spec.clone(), - ) - .unwrap() - .write_manifest_file() - .await - .unwrap(); - let second = producer - .new_manifest_writer(ManifestContentType::Data, schema.clone(), partition_spec) - .unwrap() - .write_manifest_file() - .await - .unwrap(); - - assert_ne!(first.manifest_path, second.manifest_path); - - let bytes = table - .file_io() - .new_input(&first.manifest_path) - .unwrap() - .read() - .await - .unwrap(); - let metadata = Manifest::parse_avro(&bytes).unwrap().into_parts().1; - - // Not the table's defaults, which are schema 0 and spec 0. - assert_eq!(metadata.schema_id(), 7); - assert_eq!(metadata.schema().as_ref(), schema.as_ref()); - assert_eq!(metadata.partition_spec().spec_id(), 3); - assert_eq!(first.partition_spec_id, 3); - } -} diff --git a/crates/iceberg/src/transaction/snapshot/mod.rs b/crates/iceberg/src/transaction/snapshot/mod.rs new file mode 100644 index 000000000..36616d012 --- /dev/null +++ b/crates/iceberg/src/transaction/snapshot/mod.rs @@ -0,0 +1,667 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Snapshot production for transaction actions. +//! +//! Producers (see [`simple`]) own retry-persistent state: identity, reusable +//! artifacts, and the action's intent. The functions in this module are the +//! stateless building blocks they share: they translate explicit inputs into +//! Iceberg metadata and files, never select operation policy, and never cache +//! across calls. + +pub(crate) mod simple; + +use std::collections::{HashMap, HashSet}; + +use futures::TryStreamExt; +use futures::stream::FuturesUnordered; +use uuid::Uuid; + +use crate::error::{Result, invalid_data}; +use crate::spec::{ + DataFile, DataFileFormat, FormatVersion, MAIN_BRANCH, ManifestContentType, ManifestFile, + ManifestListWriter, ManifestWriter, ManifestWriterBuilder, Operation, PartitionSpecRef, + SchemaRef, Snapshot, SnapshotReference, SnapshotRetention, SnapshotSummaryCollector, Struct, + StructType, Summary, TableProperties, set_iceberg_version, update_snapshot_summaries, +}; +use crate::table::Table; +use crate::transaction::ActionCommit; +use crate::{Error, ErrorKind, TableRequirement, TableUpdate}; + +/// Generates a positive snapshot id not already present in `table`. +pub(crate) fn generate_snapshot_id(table: &Table) -> i64 { + loop { + let (lhs, rhs) = Uuid::new_v4().as_u64_pair(); + let snapshot_id = (lhs ^ rhs) as i64 & i64::MAX; + if table.metadata().snapshot_by_id(snapshot_id).is_none() { + return snapshot_id; + } + } +} + +/// Returns the metadata path for a producer-owned manifest. +/// +/// `counter` must be unique per `commit_uuid` for the lifetime of the producer, +/// so that no attempt overwrites a manifest written by an earlier one. +pub(crate) fn manifest_path(table: &Table, commit_uuid: Uuid, counter: u64) -> Result<String> { + Ok(format!( + "{}/{}-m{}.{}", + table.metadata().metadata_location()?, + commit_uuid, + counter, + DataFileFormat::Avro + )) +} + +/// Returns the metadata path for one attempt's manifest list. +/// +/// Each attempt gets a distinct `attempt` number so the path is write-once. +pub(crate) fn manifest_list_path( + table: &Table, + snapshot_id: i64, + attempt: u64, + commit_uuid: Uuid, +) -> Result<String> { + Ok(format!( + "{}/snap-{}-{}-{}.{}", + table.metadata().metadata_location()?, + snapshot_id, + attempt, + commit_uuid, + DataFileFormat::Avro + )) +} + +/// Creates a manifest writer at `path`, honoring the table's format version and +/// encryption configuration, under the given schema and partition spec. +pub(crate) fn new_manifest_writer( + table: &Table, + path: String, + snapshot_id: i64, + content: ManifestContentType, + schema: SchemaRef, + partition_spec: PartitionSpecRef, +) -> Result<ManifestWriter> { + let output_file = table.file_io().new_output(path)?; + let partition_spec = partition_spec.as_ref().clone(); + + let builder = if let Some(em) = table.encryption_manager() { + ManifestWriterBuilder::new_from_encrypted( + em.encrypt(output_file), + Some(snapshot_id), + schema, + partition_spec, + )? + } else { + ManifestWriterBuilder::new(output_file, Some(snapshot_id), schema, partition_spec) + }; + + match table.metadata().format_version() { + FormatVersion::V1 => Ok(builder.build_v1()), + FormatVersion::V2 => match content { + ManifestContentType::Data => Ok(builder.build_v2_data()), + ManifestContentType::Deletes => Ok(builder.build_v2_deletes()), + }, + FormatVersion::V3 => match content { + ManifestContentType::Data => Ok(builder.build_v3_data()), + ManifestContentType::Deletes => Ok(builder.build_v3_deletes()), + }, + } +} + +/// Checks that a partition value is compatible with the partition type. +pub(crate) fn validate_partition_value( + partition_value: &Struct, + partition_type: &StructType, +) -> Result<()> { + if partition_value.fields().len() != partition_type.fields().len() { + return Err(invalid_data!( + "Partition value is not compatible with partition type" + )); + } + + for (value, field) in partition_value.fields().iter().zip(partition_type.fields()) { + let field = field.field_type.as_primitive_type().ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + "Partition field should only be primitive type.", + ) + })?; + if let Some(value) = value + && !field.compatible(&value.as_primitive_literal().unwrap()) + { + return Err(invalid_data!( + "Partition value is not compatible partition type" + )); + } + } + Ok(()) +} + +/// Returns the current snapshot's manifests that a new snapshot carries forward. +/// +/// Delete-only manifests are kept too: they record which files were removed and +/// must persist across snapshots until `expire_snapshots` cleans them up. +/// Dropping them lets the removed files reappear as live data (see #2148). +pub(crate) async fn current_manifests(table: &Table) -> Result<Vec<ManifestFile>> { + let Some(snapshot) = table.metadata().current_snapshot() else { + return Ok(vec![]); + }; + + let manifest_list = table.manifest_list_reader(snapshot).load().await?; + + Ok(manifest_list + .consume_entries() + .into_iter() + .filter(|entry| { + entry.has_added_files() || entry.has_existing_files() || entry.has_deleted_files() + }) + .collect()) +} + +/// Fails if any of `data_files` is already live in the table's current snapshot. +/// +/// Must run against each attempt's base, since a concurrent commit may have +/// added one of these files. Manifest reads go through the table's object +/// cache, so manifests unchanged by a rebase are not re-read. +pub(crate) async fn validate_files_not_live(table: &Table, data_files: &[DataFile]) -> Result<()> { + let Some(current_snapshot) = table.metadata().current_snapshot() else { + return Ok(()); + }; + + let new_files: HashSet<&str> = data_files.iter().map(|df| df.file_path()).collect(); + if new_files.is_empty() { + return Ok(()); + } + + let manifest_list = table.manifest_list_reader(current_snapshot).load().await?; + let runtime = table.runtime(); + let object_cache = table.object_cache(); + + let new_files_ref = &new_files; + let referenced_files: Vec<String> = manifest_list + .consume_entries() + .into_iter() + .map(|entry| { + let object_cache = object_cache.clone(); + runtime + .io() + .spawn(async move { object_cache.get_manifest(&entry).await }) + }) + .collect::<FuturesUnordered<_>>() + .try_fold(Vec::new(), |mut acc, manifest| async move { + acc.extend( + manifest? + .entries() + .iter() + .filter(|e| new_files_ref.contains(e.file_path()) && e.is_alive()) + .map(|e| e.file_path().to_string()), + ); + Ok(acc) + }) + .await?; + + if !referenced_files.is_empty() { + return Err(invalid_data!( + "Cannot add files that are already referenced by table, files: {}", + referenced_files.join(", ") + )); + } + + Ok(()) +} + +/// Builds a snapshot summary from an explicit set of file changes. +/// +/// User-supplied `properties` are applied first, then the computed metrics +/// overwrite any colliding keys. This matches iceberg-java +/// (`SnapshotProducer.summary`), where computed `added-*`/`total-*` values are +/// written after user properties so a user cannot shadow them with a bad (or +/// merely wrong) value that would corrupt the snapshot summary. +fn build_summary( + table: &Table, + operation: Operation, + mut properties: HashMap<String, String>, + added_files: &[DataFile], + removed_files: &[DataFile], + truncate_full_table: bool, +) -> Result<Summary> { + let metadata = table.metadata(); + let mut collector = SnapshotSummaryCollector::default(); + + collector.set_partition_summary_limit( + metadata + .properties() + .get(TableProperties::PROPERTY_WRITE_PARTITION_SUMMARY_LIMIT) + .and_then(|limit| limit.parse().ok()) + .unwrap_or(TableProperties::PROPERTY_WRITE_PARTITION_SUMMARY_LIMIT_DEFAULT), + ); + + for data_file in added_files { + collector.add_file( + data_file, + metadata.current_schema().clone(), + metadata.default_partition_spec().clone(), + ); + } + + for data_file in removed_files { + let partition_spec = metadata + .partition_spec_by_id(data_file.partition_spec_id) + .ok_or_else(|| { + invalid_data!("Unknown partition spec id {}", data_file.partition_spec_id) + })?; + collector.remove_file( + data_file, + metadata.current_schema().clone(), + partition_spec.clone(), + ); + } + + properties.extend(collector.build()); + + let mut summary = update_snapshot_summaries( + Summary { + operation, + additional_properties: properties, + }, + metadata.current_snapshot().map(|s| s.summary()), + truncate_full_table, + )?; + // Recorded last so it also overrides a user-supplied `iceberg-version`. + set_iceberg_version(&mut summary.additional_properties); + Ok(summary) +} + +/// Everything one attempt needs to turn a resulting manifest set into a snapshot. +/// +/// The producer resolves identity (`snapshot_id`, `manifest_list_path`) and the +/// manifest set; the attempt-sensitive values (parent, sequence number, row-id +/// allocation) are read from `table` by [`write_snapshot_commit`]. +pub(crate) struct SnapshotCommitInput<'a> { + pub snapshot_id: i64, + pub manifest_list_path: String, + pub operation: Operation, + pub properties: HashMap<String, String>, + pub added_files: &'a [DataFile], + /// Must match the data files marked Deleted by this commit, including their metrics. + pub removed_files: &'a [DataFile], + pub manifests: Vec<ManifestFile>, + /// Whether the summary resets table totals (a full-table overwrite). + pub truncate_full_table: bool, +} + +/// Writes the attempt's manifest list, assembles the snapshot, and returns the +/// resulting updates and requirements as an [`ActionCommit`]. +pub(crate) async fn write_snapshot_commit( + table: &Table, + input: SnapshotCommitInput<'_>, +) -> Result<ActionCommit> { + let SnapshotCommitInput { + snapshot_id, + manifest_list_path, + operation, + properties, + added_files, + removed_files, + manifests, + truncate_full_table, + } = input; + + debug_assert_eq!( + manifests + .iter() + .filter(|m| m.added_snapshot_id == snapshot_id && m.content == ManifestContentType::Data) + .filter_map(|m| m.deleted_files_count) + .map(u64::from) + .sum::<u64>(), + removed_files.len() as u64, + "deleted files must match this commit's manifest entries" + ); + + let metadata = table.metadata(); + let parent_snapshot_id = metadata.current_snapshot_id(); + let next_seq_num = metadata.next_sequence_number(); + let first_row_id = metadata.next_row_id(); + + let summary = build_summary( + table, + operation, + properties, + added_files, + removed_files, + truncate_full_table, + ) + .map_err(|err| { + Error::new(ErrorKind::Unexpected, "Failed to create snapshot summary.").with_source(err) + })?; + + let raw_output = table.file_io().new_output(manifest_list_path.clone())?; + let (writer, encrypted_output) = match table.encryption_manager() { + Some(em) => { + let encrypted_output = em.encrypt(raw_output); + ( + encrypted_output.writer().await?, + Some((em.clone(), encrypted_output)), + ) + } + None => (raw_output.writer().await?, None), + }; + + let mut manifest_list_writer = match metadata.format_version() { + FormatVersion::V1 => ManifestListWriter::v1(writer, snapshot_id, parent_snapshot_id), + FormatVersion::V2 => { + ManifestListWriter::v2(writer, snapshot_id, parent_snapshot_id, next_seq_num) + } + FormatVersion::V3 => ManifestListWriter::v3( + writer, + snapshot_id, + parent_snapshot_id, + next_seq_num, + Some(first_row_id), + ), + }; + + manifest_list_writer.add_manifests(manifests.into_iter())?; + let writer_next_row_id = manifest_list_writer.next_row_id(); + let file_metadata = manifest_list_writer.close().await?; + + let encryption_key_id = match encrypted_output { + Some((em, encrypted_output)) => Some( + em.encrypt_manifest_list_key_metadata( + &encrypted_output.key_metadata_with_saved_file_metadata(&file_metadata), + ) + .await?, + ), + None => None, + }; + + let snapshot = Snapshot::builder() + .with_manifest_list(manifest_list_path) + .with_snapshot_id(snapshot_id) + .with_parent_snapshot_id(parent_snapshot_id) + .with_sequence_number(next_seq_num) + .with_summary(summary) + .with_schema_id(metadata.current_schema_id()) + .with_encryption_key_id(encryption_key_id) + .with_timestamp_ms(chrono::Utc::now().timestamp_millis()); + + let snapshot = match writer_next_row_id { + Some(next_row_id) => snapshot + .with_row_range(first_row_id, next_row_id - first_row_id) + .build(), + None => snapshot.build(), + }; + + let mut updates: Vec<TableUpdate> = table + .encryption_manager() + .map(|em| { + em.with_encryption_keys(|keys| { + keys.values() + .filter(|k| metadata.encryption_key(k.key_id()).is_none()) + .map(|k| TableUpdate::AddEncryptionKey { + encryption_key: k.clone(), + }) + .collect() + }) + }) + .unwrap_or_default(); + + updates.extend([ + TableUpdate::AddSnapshot { snapshot }, + TableUpdate::SetSnapshotRef { + ref_name: MAIN_BRANCH.to_string(), + reference: SnapshotReference::new( + snapshot_id, + SnapshotRetention::branch(None, None, None), + ), + }, + ]); + + let requirements = vec![ + TableRequirement::UuidMatch { + uuid: metadata.uuid(), + }, + TableRequirement::RefSnapshotIdMatch { + r#ref: MAIN_BRANCH.to_string(), + snapshot_id: parent_snapshot_id, + }, + ]; + + Ok(ActionCommit::new(updates, requirements)) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use uuid::Uuid; + + use super::simple::{AppendIntent, SimpleSnapshotProducer}; + use super::*; + use crate::spec::{ + DataContentType, DataFileBuilder, Literal, Manifest, NestedField, PartitionSpec, + PrimitiveType, Schema, SnapshotRef, Transform, Type, + }; + use crate::transaction::tests::make_v2_minimal_table; + use crate::{ErrorKind, TableUpdate}; + + fn data_file(table: &Table, path: &str, record_count: u64) -> DataFile { + DataFileBuilder::default() + .content(DataContentType::Data) + .file_path(path.to_string()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(100) + .record_count(record_count) + .partition_spec_id(table.metadata().default_partition_spec_id()) + .partition(Struct::from_iter([Some(Literal::long(300))])) + .build() + .unwrap() + } + + fn deleted_file(table: &Table) -> DataFile { + data_file(table, "test/deleted.parquet", 2) + } + + /// Commits `data_files` to `table` through a fresh producer, returning the + /// table with that snapshot as current. + async fn append_to(table: &Table, data_files: Vec<DataFile>) -> Table { + let mut producer = SimpleSnapshotProducer::new(AppendIntent { + commit_uuid: None, + snapshot_properties: HashMap::new(), + data_files, + }); + let snapshot = producer + .apply(table) + .await + .unwrap() + .take_updates() + .into_iter() + .find_map(|u| match u { + TableUpdate::AddSnapshot { snapshot } => Some(SnapshotRef::new(snapshot)), + _ => None, + }) + .expect("append should add a snapshot"); + let metadata = table + .metadata() + .clone() + .into_builder(None) + .add_snapshot(snapshot.as_ref().clone()) + .unwrap() + .set_ref( + MAIN_BRANCH, + SnapshotReference::new( + snapshot.snapshot_id(), + SnapshotRetention::branch(None, None, None), + ), + ) + .unwrap() + .build() + .unwrap() + .metadata; + table.clone().with_metadata(Arc::new(metadata)) + } + + #[tokio::test] + async fn test_validate_files_not_live_sees_files_added_by_rebase() { + // No current snapshot: nothing is live, so nothing is a duplicate. + let table = make_v2_minimal_table(); + let files = vec![data_file(&table, "test/dup.parquet", 1)]; + validate_files_not_live(&table, &files).await.unwrap(); + + // A concurrent writer commits the same file; validating against that + // base must detect it. + let rebased = append_to(&table, files.clone()).await; + let err = validate_files_not_live(&rebased, &files) + .await + .expect_err("file committed concurrently must be detected"); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.message().contains("test/dup.parquet")); + + // A different file is not flagged. + validate_files_not_live(&rebased, &[data_file(&table, "test/other.parquet", 1)]) + .await + .unwrap(); + } + + #[test] + fn test_metadata_paths_carry_identity_and_attempt() { + let table = make_v2_minimal_table(); + let uuid = Uuid::nil(); + + assert!( + manifest_path(&table, uuid, 3) + .unwrap() + .ends_with("00000000-0000-0000-0000-000000000000-m3.avro") + ); + assert!( + manifest_list_path(&table, 7, 2, uuid) + .unwrap() + .ends_with("snap-7-2-00000000-0000-0000-0000-000000000000.avro") + ); + } + + #[test] + fn test_summary_counts_deleted_data_files() { + let table = make_v2_minimal_table(); + let deleted = deleted_file(&table); + + let summary = build_summary( + &table, + Operation::Delete, + HashMap::new(), + &[], + &[deleted], + false, + ) + .unwrap(); + let properties = &summary.additional_properties; + + assert_eq!(properties.get("deleted-data-files").unwrap(), "1"); + assert_eq!(properties.get("deleted-records").unwrap(), "2"); + assert_eq!(properties.get("removed-files-size").unwrap(), "100"); + } + + #[cfg(debug_assertions)] + #[tokio::test] + #[should_panic(expected = "deleted files must match this commit's manifest entries")] + async fn test_commit_rejects_unrecorded_deletions() { + let table = make_v2_minimal_table(); + let deleted = deleted_file(&table); + let snapshot_id = generate_snapshot_id(&table); + + write_snapshot_commit(&table, SnapshotCommitInput { + snapshot_id, + manifest_list_path: manifest_list_path(&table, snapshot_id, 0, Uuid::now_v7()).unwrap(), + operation: Operation::Delete, + properties: HashMap::new(), + added_files: &[], + removed_files: &[deleted], + manifests: vec![], + truncate_full_table: false, + }) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_new_manifest_writer_uses_the_given_schema_and_spec() { + let table = make_v2_minimal_table(); + let schema = Arc::new( + Schema::builder() + .with_schema_id(7) + .with_fields(vec![ + NestedField::required(1, "a", Type::Primitive(PrimitiveType::Long)).into(), + ]) + .build() + .unwrap(), + ); + let partition_spec = Arc::new( + PartitionSpec::builder(schema.clone()) + .with_spec_id(3) + .add_partition_field("a", "a", Transform::Identity) + .unwrap() + .build() + .unwrap(), + ); + + let commit_uuid = Uuid::now_v7(); + let snapshot_id = generate_snapshot_id(&table); + let first = new_manifest_writer( + &table, + manifest_path(&table, commit_uuid, 0).unwrap(), + snapshot_id, + ManifestContentType::Data, + schema.clone(), + partition_spec.clone(), + ) + .unwrap() + .write_manifest_file() + .await + .unwrap(); + let second = new_manifest_writer( + &table, + manifest_path(&table, commit_uuid, 1).unwrap(), + snapshot_id, + ManifestContentType::Data, + schema.clone(), + partition_spec, + ) + .unwrap() + .write_manifest_file() + .await + .unwrap(); + + assert_ne!(first.manifest_path, second.manifest_path); + + let bytes = table + .file_io() + .new_input(&first.manifest_path) + .unwrap() + .read() + .await + .unwrap(); + let metadata = Manifest::parse_avro(&bytes).unwrap().into_parts().1; + + // Not the table's defaults, which are schema 0 and spec 0. + assert_eq!(metadata.schema_id(), 7); + assert_eq!(metadata.schema().as_ref(), schema.as_ref()); + assert_eq!(metadata.partition_spec().spec_id(), 3); + assert_eq!(first.partition_spec_id, 3); + } +} diff --git a/crates/iceberg/src/transaction/snapshot/simple.rs b/crates/iceberg/src/transaction/snapshot/simple.rs new file mode 100644 index 000000000..7aaa50539 --- /dev/null +++ b/crates/iceberg/src/transaction/snapshot/simple.rs @@ -0,0 +1,476 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Persistent producer for append-shaped snapshot actions. + +use std::collections::{HashMap, HashSet}; + +use uuid::Uuid; + +use super::{ + SnapshotCommitInput, current_manifests, generate_snapshot_id, manifest_list_path, + manifest_path, new_manifest_writer, validate_partition_value, write_snapshot_commit, +}; +use crate::error::{Result, invalid_data}; +use crate::spec::{ + DataFile, FormatVersion, ManifestContentType, ManifestEntry, ManifestFile, ManifestStatus, + Operation, +}; +use crate::table::Table; +use crate::transaction::ActionCommit; + +/// What an append-shaped action asks the producer to do. +/// +/// Fixed for one logical execution; the producer derives everything else from +/// it and the table of the current attempt. +pub(crate) struct AppendIntent { + /// Overrides the generated commit UUID. + pub commit_uuid: Option<Uuid>, + pub snapshot_properties: HashMap<String, String>, + /// Data files to add. The action validates them before each attempt. + pub data_files: Vec<DataFile>, +} + +/// Retry-persistent state for an append-shaped action: added data files go +/// into new manifests, and the current snapshot's manifests are carried +/// forward untouched (no filtering, no merging). +/// +/// One producer lives for one logical execution of an action. It is driven +/// once per attempt via [`SimpleSnapshotProducer::apply`]: logical identity +/// (commit UUID, snapshot id) is fixed on first use, the manifests written for +/// the added files are reused across attempts while their semantic inputs are +/// unchanged, and every attempt writes a fresh manifest list under a distinct +/// path. +pub(crate) struct SimpleSnapshotProducer { + intent: AppendIntent, + /// Fixed when the producer is created, so every attempt shares one UUID. + commit_uuid: Uuid, + /// Resolved against the table on the first attempt, then stable. + snapshot_id: Option<i64>, + /// Incremented per attempt; distinguishes manifest-list paths. + attempt: u64, + /// Numbers this producer's manifests. Never reset, so reused manifests keep + /// their paths and new ones never collide with earlier attempts. + manifest_counter: u64, + /// Manifests written for the added files, reused across attempts. + added_manifests: Option<AddedManifests>, + /// Metadata files written by this producer (manifests and every attempt's + /// manifest list). + // TODO(#3417): cleanup is not implemented yet; terminal cleanup consumes + // this set. + #[allow(dead_code)] + owned_artifacts: HashSet<String>, +} + +/// The manifests written for one append's added files. +/// +/// Reusable while the table's format version is unchanged. Per spec, sequence +/// number inheritance does not apply to v1 manifests, so one written before a +/// concurrent upgrade cannot carry the v2 sequence number a retry assigns and +/// must be rewritten. The schema is not a guard: the manifest's embedded schema +/// describes the files it was written for, which do not change. +struct AddedManifests { + /// One manifest today. Several once added files are grouped by partition + /// spec (#3500) or rolled over at a target size (#3499). + manifests: Vec<ManifestFile>, + format_version: FormatVersion, +} + +impl SimpleSnapshotProducer { + pub(crate) fn new(intent: AppendIntent) -> Self { + Self { + commit_uuid: intent.commit_uuid.unwrap_or_else(Uuid::now_v7), + intent, + snapshot_id: None, + attempt: 0, + manifest_counter: 0, + added_manifests: None, + owned_artifacts: HashSet::new(), + } + } + + /// Produces the resulting manifest set for the current + /// table and writes this attempt's snapshot. + pub(crate) async fn apply(&mut self, table: &Table) -> Result<ActionCommit> { + let snapshot_id = self.snapshot_id(table); + + let mut manifests = current_manifests(table).await?; + if !self.intent.data_files.is_empty() { + manifests.extend(self.added_manifests(table, snapshot_id).await?); + } + + let manifest_list_path = self.next_manifest_list_path(table, snapshot_id)?; + + write_snapshot_commit(table, SnapshotCommitInput { + snapshot_id, + manifest_list_path, + operation: Operation::Append, + properties: self.intent.snapshot_properties.clone(), + added_files: &self.intent.data_files, + removed_files: &[], + manifests, + truncate_full_table: false, + }) + .await + } + + /// Checks that every file is written under the table's default partition + /// spec with a compatible partition value. + /// + /// Runs against the table the manifests are written for; a later rebase + /// does not invalidate manifests already written under a spec the table + /// still has. + fn validate_added_data_files(&self, table: &Table) -> Result<()> { + let metadata = table.metadata(); + for data_file in &self.intent.data_files { + // TODO(#3500): accept any spec known to the table and group the + // files by spec when writing manifests, instead of requiring the + // default spec. + if metadata.default_partition_spec_id() != data_file.partition_spec_id { + return Err(invalid_data!( + "Data file partition spec id does not match table default partition spec id" + )); + } + validate_partition_value(data_file.partition(), metadata.default_partition_type())?; + } + Ok(()) + } + + fn snapshot_id(&mut self, table: &Table) -> i64 { + *self + .snapshot_id + .get_or_insert_with(|| generate_snapshot_id(table)) + } + + fn next_manifest_list_path(&mut self, table: &Table, snapshot_id: i64) -> Result<String> { + let path = manifest_list_path(table, snapshot_id, self.attempt, self.commit_uuid)?; + self.attempt += 1; + self.owned_artifacts.insert(path.clone()); + Ok(path) + } + + /// Returns the manifests for the added files, writing them only if no + /// reusable set exists for the current format version. + async fn added_manifests( + &mut self, + table: &Table, + snapshot_id: i64, + ) -> Result<Vec<ManifestFile>> { + let format_version = table.metadata().format_version(); + + if let Some(cached) = &self.added_manifests + && cached.format_version == format_version + { + return Ok(cached.manifests.clone()); + } + + self.validate_added_data_files(table)?; + + // Allocate every path before writing, so each is recorded as owned + // even if the write fails part-way: the path is write-once and a + // partial file is an orphan to clean up like any other. + // + // TODO(#3500): group the files by partition spec and write one + // manifest per group. + // TODO(#3499): roll a group over at the target manifest size, and + // write the resulting manifests concurrently. + let path = manifest_path(table, self.commit_uuid, self.manifest_counter)?; + self.manifest_counter += 1; + self.owned_artifacts.insert(path.clone()); + + let manifests = vec![ + write_added_manifest( + table, + path, + snapshot_id, + table.metadata().default_partition_spec_id(), + &self.intent.data_files, + ) + .await?, + ]; + + self.added_manifests = Some(AddedManifests { + manifests: manifests.clone(), + format_version, + }); + Ok(manifests) + } +} + +/// Writes one data manifest at `path` for `data_files`, all of which must +/// belong to `partition_spec_id`, and returns its manifest-list entry. +async fn write_added_manifest( + table: &Table, + path: String, + snapshot_id: i64, + partition_spec_id: i32, + data_files: &[DataFile], +) -> Result<ManifestFile> { + let metadata = table.metadata(); + let partition_spec = metadata + .partition_spec_by_id(partition_spec_id) + .ok_or_else(|| invalid_data!("Unknown partition spec id {partition_spec_id}"))?; + + let mut writer = new_manifest_writer( + table, + path, + snapshot_id, + ManifestContentType::Data, + metadata.current_schema().clone(), + partition_spec.clone(), + )?; + let format_version = metadata.format_version(); + for data_file in data_files { + let builder = ManifestEntry::builder() + .status(ManifestStatus::Added) + .data_file(data_file.clone()); + let entry = if format_version == FormatVersion::V1 { + builder.snapshot_id(snapshot_id).build() + } else { + // For format version > 1, the snapshot id is inherited from the + // manifest list, so a rebase does not require rewriting the manifest. + builder.build() + }; + writer.add_entry(entry)?; + } + writer.write_manifest_file().await +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use uuid::Uuid; + + use super::{AppendIntent, SimpleSnapshotProducer}; + use crate::TableUpdate; + use crate::spec::{ + DataContentType, DataFile, DataFileBuilder, DataFileFormat, FormatVersion, Literal, + NestedField, PrimitiveType, SnapshotRef, Struct, Type, + }; + use crate::table::Table; + use crate::transaction::tests::{make_v1_table, make_v2_minimal_table}; + + fn data_file(table: &Table, path: &str) -> DataFile { + DataFileBuilder::default() + .content(DataContentType::Data) + .file_path(path.to_string()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(100) + .record_count(1) + .partition_spec_id(table.metadata().default_partition_spec_id()) + .partition(Struct::from_iter([Some(Literal::long(300))])) + .build() + .unwrap() + } + + fn producer(data_files: Vec<DataFile>) -> SimpleSnapshotProducer { + producer_with_uuid(None, data_files) + } + + fn producer_with_uuid( + commit_uuid: Option<Uuid>, + data_files: Vec<DataFile>, + ) -> SimpleSnapshotProducer { + SimpleSnapshotProducer::new(AppendIntent { + commit_uuid, + snapshot_properties: HashMap::new(), + data_files, + }) + } + + fn added_snapshot(commit: &mut crate::transaction::ActionCommit) -> SnapshotRef { + commit + .take_updates() + .into_iter() + .find_map(|u| match u { + TableUpdate::AddSnapshot { snapshot } => Some(SnapshotRef::new(snapshot)), + _ => None, + }) + .expect("append should add a snapshot") + } + + /// The minimal table has no current snapshot, so `current_manifests` is + /// empty and the only manifest in the result is the one we wrote. + #[tokio::test] + async fn test_retry_reuses_identity_and_added_manifest() { + let table = make_v2_minimal_table(); + let requested_uuid = Uuid::now_v7(); + let mut producer = producer_with_uuid(Some(requested_uuid), vec![data_file( + &table, + "test/1.parquet", + )]); + + let mut first = producer.apply(&table).await.unwrap(); + let mut second = producer.apply(&table).await.unwrap(); + + let first_snapshot = added_snapshot(&mut first); + let second_snapshot = added_snapshot(&mut second); + + // Identity is stable across attempts. + assert_eq!(producer.commit_uuid, requested_uuid); + assert_eq!(first_snapshot.snapshot_id(), second_snapshot.snapshot_id()); + assert_eq!(producer.snapshot_id, Some(first_snapshot.snapshot_id())); + + // Each attempt writes its own manifest list under a distinct path. + assert_ne!( + first_snapshot.manifest_list(), + second_snapshot.manifest_list() + ); + assert!(first_snapshot.manifest_list().contains("-0-")); + assert!(second_snapshot.manifest_list().contains("-1-")); + + // Both manifest lists point at the same, single added manifest. + let first_list = table + .manifest_list_reader(&first_snapshot) + .load() + .await + .unwrap(); + let second_list = table + .manifest_list_reader(&second_snapshot) + .load() + .await + .unwrap(); + assert_eq!(first_list.entries().len(), 1); + assert_eq!(second_list.entries().len(), 1); + assert_eq!( + first_list.entries()[0].manifest_path, + second_list.entries()[0].manifest_path + ); + assert_eq!(producer.manifest_counter, 1); + + // One manifest plus two manifest lists are owned. + assert_eq!(producer.owned_artifacts.len(), 3); + } + + /// The manifest's embedded schema still describes the same files after a + /// concurrent schema change, so the manifest stays valid (as in Java). + #[tokio::test] + async fn test_added_manifest_is_reused_across_schema_change() { + let table = make_v2_minimal_table(); + let mut producer = producer(vec![data_file(&table, "test/1.parquet")]); + + let mut first = producer.apply(&table).await.unwrap(); + let first_snapshot = added_snapshot(&mut first); + + // Rebase onto a table whose current schema changed. + let new_schema = table + .metadata() + .current_schema() + .as_ref() + .clone() + .into_builder() + .with_fields([ + NestedField::optional(100, "added", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build() + .unwrap(); + let metadata = table + .metadata() + .clone() + .into_builder(None) + .add_current_schema(new_schema) + .unwrap() + .build() + .unwrap() + .metadata; + let rebased = table.clone().with_metadata(Arc::new(metadata)); + + let mut second = producer.apply(&rebased).await.unwrap(); + let second_snapshot = added_snapshot(&mut second); + + let first_list = table + .manifest_list_reader(&first_snapshot) + .load() + .await + .unwrap(); + let second_list = rebased + .manifest_list_reader(&second_snapshot) + .load() + .await + .unwrap(); + assert_eq!( + first_list.entries()[0].manifest_path, + second_list.entries()[0].manifest_path + ); + assert_eq!(producer.manifest_counter, 1); + } + + /// Per spec, a v1 manifest's entries always read with sequence number 0, so + /// a v1 manifest written before a concurrent format upgrade cannot carry the + /// sequence number the retry's v2 manifest list assigns. The upgrade must + /// invalidate the cached manifest. + #[tokio::test] + async fn test_added_manifest_is_rewritten_when_format_version_changes() { + let table = make_v1_table(); + let mut producer = producer(vec![data_file(&table, "test/1.parquet")]); + + let mut first = producer.apply(&table).await.unwrap(); + let first_snapshot = added_snapshot(&mut first); + + // A concurrent commit upgraded the table; schema and spec are unchanged. + let metadata = table + .metadata() + .clone() + .into_builder(None) + .upgrade_format_version(FormatVersion::V2) + .unwrap() + .build() + .unwrap() + .metadata; + let upgraded = table.clone().with_metadata(Arc::new(metadata)); + assert_eq!( + upgraded.metadata().current_schema_id(), + table.metadata().current_schema_id() + ); + + let mut second = producer.apply(&upgraded).await.unwrap(); + let second_snapshot = added_snapshot(&mut second); + + let first_list = table + .manifest_list_reader(&first_snapshot) + .load() + .await + .unwrap(); + let second_list = upgraded + .manifest_list_reader(&second_snapshot) + .load() + .await + .unwrap(); + let first_added = first_list + .entries() + .iter() + .find(|m| m.added_snapshot_id == first_snapshot.snapshot_id()) + .unwrap(); + let second_added = second_list + .entries() + .iter() + .find(|m| m.added_snapshot_id == second_snapshot.snapshot_id()) + .unwrap(); + assert_ne!(first_added.manifest_path, second_added.manifest_path); + assert_eq!(producer.manifest_counter, 2); + + // The rewritten v2 manifest's entries inherit the v2 sequence number. + let manifest = upgraded.manifest_reader().read(second_added).await.unwrap(); + assert_eq!( + manifest.entries()[0].sequence_number(), + Some(second_snapshot.sequence_number()) + ); + assert_ne!(second_snapshot.sequence_number(), 0); + } +}
