CTTY commented on code in PR #3474: URL: https://github.com/apache/iceberg-rust/pull/3474#discussion_r4342812373
########## crates/iceberg/src/transaction/simple_snapshot_producer.rs: ########## @@ -0,0 +1,630 @@ +// 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 futures::TryStreamExt; +use futures::stream::FuturesUnordered; +use uuid::Uuid; + +use crate::error::{Result, invalid_data}; +use crate::spec::{ + DataContentType, DataFile, FormatVersion, ManifestContentType, ManifestEntry, ManifestFile, + ManifestStatus, Operation, +}; +use crate::table::Table; +use crate::transaction::ActionCommit; +use crate::transaction::snapshot_helpers::{ + SnapshotCommitInput, current_manifests, generate_snapshot_id, manifest_list_path, + manifest_path, new_manifest_writer, validate_partition_value, write_snapshot_commit, +}; +use crate::{Error, ErrorKind}; + +/// 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 { + /// 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: cleanup is not implemented yet; terminal cleanup (stateful + // transaction RFC, milestone 3) consumes this set. + #[allow(dead_code)] + owned_artifacts: HashSet<String>, +} + +/// The manifests written for one append's added files, together with the +/// table-level inputs they were written under. +/// +/// A manifest file embeds its schema and partition spec and is encoded for a +/// specific format version, so the set is only reusable while the table it is +/// rebased onto still has the same schema and format version. The partition +/// spec is implied: every file is validated against the default spec, so the +/// spec is fixed for the lifetime of the producer by that check. +/// +/// The guard is table-level because its inputs are table-level: when they +/// change, every manifest in the set is invalid together. +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>, + schema_id: i32, + format_version: FormatVersion, +} + +impl SimpleSnapshotProducer { + /// Creates a producer. `commit_uuid` overrides the generated commit UUID. + pub(crate) fn new(commit_uuid: Option<Uuid>) -> Self { + Self { + commit_uuid: commit_uuid.unwrap_or_else(Uuid::now_v7), + snapshot_id: None, + attempt: 0, + manifest_counter: 0, + added_manifests: None, + owned_artifacts: HashSet::new(), + } + } + + /// Checks that every file is a data file written under the table's default + /// partition spec with a compatible partition value. + pub(crate) fn validate_added_data_files( + &self, + table: &Table, + data_files: &[DataFile], + ) -> Result<()> { + let metadata = table.metadata(); + for data_file in data_files { + if data_file.content_type() != DataContentType::Data { + return Err(invalid_data!( + "Only data content type is allowed for fast append" + )); + } + // 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(()) + } + + /// Fails if any of `data_files` is already live in the current snapshot. + /// + /// This must re-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 unchanged manifests are not re-read on retry. + pub(crate) async fn validate_duplicate_files( + &self, + 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(()) + } + + /// Runs one attempt: produces the resulting manifest set for the current + /// table and writes this attempt's snapshot. + pub(crate) async fn apply( + &mut self, + table: &Table, + properties: HashMap<String, String>, + data_files: &[DataFile], + ) -> Result<ActionCommit> { + // 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 data_files.is_empty() && properties.is_empty() { + return Err(Error::new( + ErrorKind::PreconditionFailed, + "No added data files or added snapshot properties found when write a manifest file", + )); + } + + let snapshot_id = self.snapshot_id(table); + + let mut manifests = current_manifests(table).await?; + if !data_files.is_empty() { + manifests.extend(self.added_manifests(table, snapshot_id, data_files).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, + added_files: data_files, + removed_files: &[], + manifests, + truncate_full_table: false, + }) + .await + } + + 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 schema and format version. + async fn added_manifests( + &mut self, + table: &Table, + snapshot_id: i64, + data_files: &[DataFile], + ) -> Result<Vec<ManifestFile>> { + let metadata = table.metadata(); + let schema_id = metadata.current_schema_id(); + let format_version = metadata.format_version(); + + if let Some(cached) = &self.added_manifests Review Comment: After some thoughts, I think even schema id comparison won't be needed. schema in manifest will be only used to bind to partition spec -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
