laskoviymishka commented on code in PR #2752: URL: https://github.com/apache/iceberg-rust/pull/2752#discussion_r4024142305
########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1261 @@ +// 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. + +//! Copy-on-write rewrite primitives. +//! +//! This module plans candidate data files, reads their visible rows, applies a +//! caller-provided batch rewriter, and writes replacement data files. It returns +//! old and new file sets that can be committed by an overwrite-style transaction +//! action. +//! +//! The primitive does not parse SQL and does not commit metadata by itself. +//! Rewriters must emit batches compatible with the schema rows were read in +//! (the planned snapshot's schema) and must preserve each source file's +//! partition values; this primitive does not repartition rewritten rows. +//! +//! The result carries data files only. A commit adapter consuming these file +//! lists must also account for delete files that reference removed files — +//! for example deletion vectors whose referenced data file is being removed, +//! and position deletes scoped to it; equality deletes remain valid but +//! become redundant once their target rows are rewritten. +//! +//! ```rust,no_run +//! # use std::sync::Arc; +//! # use arrow_array::RecordBatch; +//! # use iceberg::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder}; +//! # use iceberg::table::Table; +//! # use iceberg::Result; +//! struct KeepAll; +//! +//! impl CowBatchRewriter for KeepAll { +//! fn rewrite_batch(&self, batch: RecordBatch) -> Result<CowBatchRewrite> { +//! Ok(CowBatchRewrite { +//! output: Some(batch), +//! changed: false, +//! }) +//! } +//! } +//! +//! # async fn example(table: &Table) -> Result<()> { +//! let result = CowRewriteBuilder::new(table) +//! .with_rewriter(Arc::new(KeepAll)) +//! .rewrite() +//! .await?; +//! +//! assert!(!result.has_changes()); +//! # Ok(()) +//! # } +//! ``` + +mod plan; +mod rewriter; +pub(crate) mod writer; + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use futures::TryStreamExt; +pub use plan::CowRewriteFile; +pub use rewriter::{CowBatchRewrite, CowBatchRewriter}; + +use crate::expr::Predicate; +use crate::scan::FileScanTaskStream; +use crate::spec::{DataFile, PartitionKey}; +use crate::table::Table; +use crate::{Error, ErrorKind, Result}; + +/// Counters produced by a copy-on-write rewrite. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CowRewriteStats { + /// Number of candidate files selected by planning. + pub candidate_files: usize, + /// Number of old files that have replacement output or are fully removed. + pub rewritten_files: usize, + /// Number of candidate files that did not change after row rewriting. + pub unchanged_files: usize, + /// Visible input row count read from candidate files. + pub input_rows: u64, + /// Output row count emitted by the batch rewriter. + pub output_rows: u64, + /// Number of input batches where the rewriter reported changes. + pub changed_batches: u64, +} + +/// Result of a copy-on-write rewrite operation. +#[derive(Debug, Default)] +pub struct CowRewriteResult { + /// Old data files that should be removed by the commit action. + pub removed_data_files: Vec<DataFile>, + /// New data files that should be added by the commit action. + pub added_data_files: Vec<DataFile>, + /// Candidate files that were read and left unchanged. + /// + /// This includes files whose visible rows were all removed by delete + /// files: with no surviving rows the rewriter never runs, so the file is + /// kept as-is rather than dropped. + pub unchanged_data_files: Vec<DataFile>, + /// Rewrite counters. + pub stats: CowRewriteStats, +} + +impl CowRewriteResult { + /// Returns true if the rewrite produced any table changes. + pub fn has_changes(&self) -> bool { + !self.removed_data_files.is_empty() || !self.added_data_files.is_empty() + } +} + +/// Builder for orchestrating copy-on-write data file rewrites. +pub struct CowRewriteBuilder<'a> { + table: &'a Table, + predicate: Predicate, + snapshot_id: Option<i64>, + batch_size: Option<usize>, + case_sensitive: bool, + rewriter: Option<Arc<dyn CowBatchRewriter>>, +} + +impl<'a> CowRewriteBuilder<'a> { + /// Creates a copy-on-write rewrite builder for `table`. + pub fn new(table: &'a Table) -> Self { + Self { + table, + predicate: Predicate::AlwaysTrue, + snapshot_id: None, + batch_size: None, + case_sensitive: true, + rewriter: None, + } + } + + /// Sets the row predicate used to plan candidate files. + pub fn with_predicate(mut self, predicate: Predicate) -> Self { + self.predicate = predicate; + self + } + + /// Sets the snapshot id used to plan candidate files. + pub fn with_snapshot_id(mut self, snapshot_id: i64) -> Self { + self.snapshot_id = Some(snapshot_id); + self + } + + /// Sets the Arrow reader batch size. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = Some(batch_size); + self + } + + /// Sets the case sensitivity used to bind the planning predicate. + pub fn with_case_sensitive(mut self, case_sensitive: bool) -> Self { + self.case_sensitive = case_sensitive; + self + } + + /// Sets the record batch rewriter. + pub fn with_rewriter(mut self, rewriter: Arc<dyn CowBatchRewriter>) -> Self { + self.rewriter = Some(rewriter); + self + } + + /// Plans, reads, rewrites, and writes replacement data files. + pub async fn rewrite(self) -> Result<CowRewriteResult> { + let rewriter = self.rewriter.ok_or_else(|| { + Error::new( + ErrorKind::PreconditionFailed, + "COW rewrite requires a batch rewriter", + ) + })?; + let files = plan::plan_cow_rewrite_files( + self.table, + Some(self.predicate), + self.snapshot_id, + self.case_sensitive, + ) + .await?; + + let mut result = CowRewriteResult { + stats: CowRewriteStats { + candidate_files: files.len(), + ..CowRewriteStats::default() + }, + ..CowRewriteResult::default() + }; + + for file in files { + // Schema the rows are read in (the planned snapshot's schema). The + // replacement files must be written with this schema so that batches + // remain compatible when the table's current schema has evolved past + // the snapshot the source files belong to. + let write_schema = file.scan_task.schema_ref(); + + // Batches produced before the first changed batch. They are buffered + // rather than written immediately because the primitive must not + // emit a replacement file for a source file that turns out to be + // unchanged. Once a changed batch is observed the buffered prefix is + // flushed to the writer and all subsequent batches stream straight + // through, so the in-memory footprint is bounded by the rows that Review Comment: The comment says the footprint is bounded by the rows that precede the first change — but for an unchanged file that bound is the whole file. Every batch gets pushed to `prefix` and nothing drains until we drop it in the `else` branch, so a no-op `KeepAll` over a compaction-sized file buffers the entire decoded Arrow file in memory. A DELETE whose rows sit near the end of each file approaches the same bound. Since we process files sequentially the peak is one file at a time, but that's still potentially several GB. At minimum I'd fix the comment to state the real worst case. Longer term a `with_max_prefix_bytes` escape hatch that falls back to just writing the replacement once the buffer crosses a threshold would cap it — happy to leave that for a follow-up as long as the bound is documented here. ########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1261 @@ +// 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. + +//! Copy-on-write rewrite primitives. +//! +//! This module plans candidate data files, reads their visible rows, applies a +//! caller-provided batch rewriter, and writes replacement data files. It returns +//! old and new file sets that can be committed by an overwrite-style transaction +//! action. +//! +//! The primitive does not parse SQL and does not commit metadata by itself. +//! Rewriters must emit batches compatible with the schema rows were read in +//! (the planned snapshot's schema) and must preserve each source file's +//! partition values; this primitive does not repartition rewritten rows. +//! +//! The result carries data files only. A commit adapter consuming these file +//! lists must also account for delete files that reference removed files — +//! for example deletion vectors whose referenced data file is being removed, +//! and position deletes scoped to it; equality deletes remain valid but +//! become redundant once their target rows are rewritten. +//! +//! ```rust,no_run +//! # use std::sync::Arc; +//! # use arrow_array::RecordBatch; +//! # use iceberg::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder}; +//! # use iceberg::table::Table; +//! # use iceberg::Result; +//! struct KeepAll; +//! +//! impl CowBatchRewriter for KeepAll { +//! fn rewrite_batch(&self, batch: RecordBatch) -> Result<CowBatchRewrite> { +//! Ok(CowBatchRewrite { +//! output: Some(batch), +//! changed: false, +//! }) +//! } +//! } +//! +//! # async fn example(table: &Table) -> Result<()> { +//! let result = CowRewriteBuilder::new(table) +//! .with_rewriter(Arc::new(KeepAll)) +//! .rewrite() +//! .await?; +//! +//! assert!(!result.has_changes()); +//! # Ok(()) +//! # } +//! ``` + +mod plan; +mod rewriter; +pub(crate) mod writer; + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use futures::TryStreamExt; +pub use plan::CowRewriteFile; +pub use rewriter::{CowBatchRewrite, CowBatchRewriter}; + +use crate::expr::Predicate; +use crate::scan::FileScanTaskStream; +use crate::spec::{DataFile, PartitionKey}; +use crate::table::Table; +use crate::{Error, ErrorKind, Result}; + +/// Counters produced by a copy-on-write rewrite. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CowRewriteStats { + /// Number of candidate files selected by planning. + pub candidate_files: usize, + /// Number of old files that have replacement output or are fully removed. + pub rewritten_files: usize, + /// Number of candidate files that did not change after row rewriting. + pub unchanged_files: usize, + /// Visible input row count read from candidate files. + pub input_rows: u64, + /// Output row count emitted by the batch rewriter. + pub output_rows: u64, + /// Number of input batches where the rewriter reported changes. + pub changed_batches: u64, +} + +/// Result of a copy-on-write rewrite operation. +#[derive(Debug, Default)] +pub struct CowRewriteResult { + /// Old data files that should be removed by the commit action. + pub removed_data_files: Vec<DataFile>, + /// New data files that should be added by the commit action. + pub added_data_files: Vec<DataFile>, + /// Candidate files that were read and left unchanged. + /// + /// This includes files whose visible rows were all removed by delete + /// files: with no surviving rows the rewriter never runs, so the file is + /// kept as-is rather than dropped. + pub unchanged_data_files: Vec<DataFile>, + /// Rewrite counters. + pub stats: CowRewriteStats, +} + +impl CowRewriteResult { + /// Returns true if the rewrite produced any table changes. + pub fn has_changes(&self) -> bool { + !self.removed_data_files.is_empty() || !self.added_data_files.is_empty() + } +} + +/// Builder for orchestrating copy-on-write data file rewrites. +pub struct CowRewriteBuilder<'a> { + table: &'a Table, + predicate: Predicate, + snapshot_id: Option<i64>, + batch_size: Option<usize>, + case_sensitive: bool, + rewriter: Option<Arc<dyn CowBatchRewriter>>, +} + +impl<'a> CowRewriteBuilder<'a> { + /// Creates a copy-on-write rewrite builder for `table`. + pub fn new(table: &'a Table) -> Self { + Self { + table, + predicate: Predicate::AlwaysTrue, + snapshot_id: None, + batch_size: None, + case_sensitive: true, + rewriter: None, + } + } + + /// Sets the row predicate used to plan candidate files. + pub fn with_predicate(mut self, predicate: Predicate) -> Self { + self.predicate = predicate; + self + } + + /// Sets the snapshot id used to plan candidate files. + pub fn with_snapshot_id(mut self, snapshot_id: i64) -> Self { + self.snapshot_id = Some(snapshot_id); + self + } + + /// Sets the Arrow reader batch size. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = Some(batch_size); + self + } + + /// Sets the case sensitivity used to bind the planning predicate. + pub fn with_case_sensitive(mut self, case_sensitive: bool) -> Self { + self.case_sensitive = case_sensitive; + self + } + + /// Sets the record batch rewriter. + pub fn with_rewriter(mut self, rewriter: Arc<dyn CowBatchRewriter>) -> Self { + self.rewriter = Some(rewriter); + self + } + + /// Plans, reads, rewrites, and writes replacement data files. + pub async fn rewrite(self) -> Result<CowRewriteResult> { + let rewriter = self.rewriter.ok_or_else(|| { + Error::new( + ErrorKind::PreconditionFailed, + "COW rewrite requires a batch rewriter", + ) + })?; + let files = plan::plan_cow_rewrite_files( + self.table, + Some(self.predicate), + self.snapshot_id, + self.case_sensitive, + ) + .await?; + + let mut result = CowRewriteResult { + stats: CowRewriteStats { + candidate_files: files.len(), + ..CowRewriteStats::default() + }, + ..CowRewriteResult::default() + }; + + for file in files { + // Schema the rows are read in (the planned snapshot's schema). The + // replacement files must be written with this schema so that batches + // remain compatible when the table's current schema has evolved past + // the snapshot the source files belong to. + let write_schema = file.scan_task.schema_ref(); + + // Batches produced before the first changed batch. They are buffered + // rather than written immediately because the primitive must not + // emit a replacement file for a source file that turns out to be + // unchanged. Once a changed batch is observed the buffered prefix is + // flushed to the writer and all subsequent batches stream straight + // through, so the in-memory footprint is bounded by the rows that + // precede the first change instead of the entire source file. + let mut prefix: Vec<RecordBatch> = Vec::new(); + let mut file_changed = false; + let mut writer: Option<Box<dyn crate::writer::IcebergWriter>> = None; + + // Planning already cleared the row predicate (see + // `ManifestEntryContext::into_cow_rewrite_file`), so this task + // reads every row of the source file. + let tasks = Box::pin(futures::stream::iter(vec![Ok(file.scan_task.clone())])) + as FileScanTaskStream; + + // Each candidate file gets its own reader so the per-file prefix + // and lazy-writer semantics stay intact; the delete-file cache is + // therefore also per file, and equality deletes shared by several + // candidates are fetched once per file. + let mut reader_builder = self.table.reader_builder(); + if let Some(batch_size) = self.batch_size { + reader_builder = reader_builder.with_batch_size(batch_size); + } + + let mut batches = reader_builder.build().read(tasks)?.stream(); + while let Some(batch) = batches.try_next().await? { + result.stats.input_rows += batch.num_rows() as u64; + + let rewrite = rewriter.rewrite_batch(batch)?; + if rewrite.changed { Review Comment: I think there's a data-loss path here. `file_changed` only tracks `rewrite.changed`, and we only ever look at `rewrite.output` through the `if let Some(output)` below — so a batch that comes back `{changed: false, output: None}` silently drops its rows and leaves `file_changed` untouched. The trap is that a *later* batch for the same file can still flip `file_changed` to true. Say batch 1 returns `{changed: false, output: None}` (rows 1–2 meant to be dropped) and batch 2 returns `{changed: true, output: Some([3,4])}`: we write a replacement holding only [3,4], move the original [1,2,3,4] to removed, and rows 1–2 vanish with no error. The docstring on `CowBatchRewrite` says to set `changed` whenever output differs including `None`, but correctness shouldn't hinge on every rewriter getting that exactly right. I'd make the orchestrator self-enforcing — derive the effective flag and use it everywhere we currently read `rewrite.changed`: ```rust let changed = rewrite.changed || rewrite.output.is_none(); ``` Alternatively, reject `{changed: false, output: None}` up front with a `PreconditionFailed`. Either's fine, but I'd want one of them before this lands. Worth a regression test on exactly this case too — nothing in the suite exercises it today, so it won't surface on its own. wdyt? ########## crates/iceberg/src/cow_rewrite/rewriter.rs: ########## @@ -0,0 +1,43 @@ +// 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 arrow_array::RecordBatch; + +use crate::Result; + +/// Result of rewriting a single record batch. +pub struct CowBatchRewrite { Review Comment: Could we derive `Debug` on `CowBatchRewrite`? It's a public return type but not printable, so callers can't `dbg!` it or fold it into error context, and it shows up as a gap in `public-api.txt`. Both fields are already `Debug`, so it's a one-line derive. ########## crates/iceberg/src/cow_rewrite/writer.rs: ########## @@ -0,0 +1,353 @@ +// 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::str::FromStr; + +use uuid::Uuid; + +use crate::Result; +#[cfg(test)] +use crate::spec::DataFile; +use crate::spec::{DataFileFormat, PartitionKey, SchemaRef}; +use crate::table::Table; +use crate::writer::IcebergWriterBuilder; +use crate::writer::base_writer::data_file_writer::DataFileWriterBuilder; +use crate::writer::file_writer::ParquetWriterBuilder; +use crate::writer::file_writer::location_generator::{ + DefaultFileNameGenerator, DefaultLocationGenerator, +}; +use crate::writer::file_writer::rolling_writer::RollingFileWriterBuilder; + +/// Builds a boxed replacement-data-file writer. +/// +/// `write_schema` is the schema the input batches are encoded in. It must match +/// the schema the rows were read in (the planned snapshot's schema), not the +/// table's possibly-evolved current schema; otherwise the parquet writer will +/// reject batches that lack columns added after the source files were written. +/// +/// Building the writer is cheap: no physical file is opened until the first +/// batch is written, so it is safe to construct one optimistically and only +/// write to it once a source file is known to have changed. +pub(crate) async fn build_replacement_writer( + table: &Table, + write_schema: SchemaRef, + partition_key: Option<PartitionKey>, +) -> Result<Box<dyn crate::writer::IcebergWriter>> { + let location_generator = DefaultLocationGenerator::new(table.metadata())?; Review Comment: We always construct a `DefaultLocationGenerator` here, which ignores `write.object-storage.enabled`. On a table using the object-storage layout the replacement files land in the flat `write.data.path` layout instead of the hash-entropy dirs, and if `write.object-storage.path` differs from `write.data.path` they can end up under a different prefix entirely. Not corruption, but it makes COW output an inconsistent outlier and breaks the S3 prefix-sharding those tables opt into. I'd branch the same way the other write paths do — read `write.object-storage.enabled` and pick `ObjectStorageLocationGenerator` or `DefaultLocationGenerator` accordingly. ########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1261 @@ +// 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. + +//! Copy-on-write rewrite primitives. +//! +//! This module plans candidate data files, reads their visible rows, applies a +//! caller-provided batch rewriter, and writes replacement data files. It returns +//! old and new file sets that can be committed by an overwrite-style transaction +//! action. +//! +//! The primitive does not parse SQL and does not commit metadata by itself. +//! Rewriters must emit batches compatible with the schema rows were read in +//! (the planned snapshot's schema) and must preserve each source file's +//! partition values; this primitive does not repartition rewritten rows. +//! +//! The result carries data files only. A commit adapter consuming these file +//! lists must also account for delete files that reference removed files — +//! for example deletion vectors whose referenced data file is being removed, +//! and position deletes scoped to it; equality deletes remain valid but +//! become redundant once their target rows are rewritten. +//! +//! ```rust,no_run +//! # use std::sync::Arc; +//! # use arrow_array::RecordBatch; +//! # use iceberg::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder}; +//! # use iceberg::table::Table; +//! # use iceberg::Result; +//! struct KeepAll; +//! +//! impl CowBatchRewriter for KeepAll { +//! fn rewrite_batch(&self, batch: RecordBatch) -> Result<CowBatchRewrite> { +//! Ok(CowBatchRewrite { +//! output: Some(batch), +//! changed: false, +//! }) +//! } +//! } +//! +//! # async fn example(table: &Table) -> Result<()> { +//! let result = CowRewriteBuilder::new(table) +//! .with_rewriter(Arc::new(KeepAll)) +//! .rewrite() +//! .await?; +//! +//! assert!(!result.has_changes()); +//! # Ok(()) +//! # } +//! ``` + +mod plan; +mod rewriter; +pub(crate) mod writer; + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use futures::TryStreamExt; +pub use plan::CowRewriteFile; +pub use rewriter::{CowBatchRewrite, CowBatchRewriter}; + +use crate::expr::Predicate; +use crate::scan::FileScanTaskStream; +use crate::spec::{DataFile, PartitionKey}; +use crate::table::Table; +use crate::{Error, ErrorKind, Result}; + +/// Counters produced by a copy-on-write rewrite. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CowRewriteStats { + /// Number of candidate files selected by planning. + pub candidate_files: usize, + /// Number of old files that have replacement output or are fully removed. + pub rewritten_files: usize, + /// Number of candidate files that did not change after row rewriting. + pub unchanged_files: usize, + /// Visible input row count read from candidate files. + pub input_rows: u64, + /// Output row count emitted by the batch rewriter. + pub output_rows: u64, + /// Number of input batches where the rewriter reported changes. + pub changed_batches: u64, +} + +/// Result of a copy-on-write rewrite operation. +#[derive(Debug, Default)] +pub struct CowRewriteResult { + /// Old data files that should be removed by the commit action. + pub removed_data_files: Vec<DataFile>, + /// New data files that should be added by the commit action. + pub added_data_files: Vec<DataFile>, + /// Candidate files that were read and left unchanged. + /// + /// This includes files whose visible rows were all removed by delete + /// files: with no surviving rows the rewriter never runs, so the file is + /// kept as-is rather than dropped. + pub unchanged_data_files: Vec<DataFile>, + /// Rewrite counters. + pub stats: CowRewriteStats, +} + +impl CowRewriteResult { + /// Returns true if the rewrite produced any table changes. + pub fn has_changes(&self) -> bool { + !self.removed_data_files.is_empty() || !self.added_data_files.is_empty() + } +} + +/// Builder for orchestrating copy-on-write data file rewrites. +pub struct CowRewriteBuilder<'a> { + table: &'a Table, + predicate: Predicate, + snapshot_id: Option<i64>, + batch_size: Option<usize>, + case_sensitive: bool, + rewriter: Option<Arc<dyn CowBatchRewriter>>, +} + +impl<'a> CowRewriteBuilder<'a> { + /// Creates a copy-on-write rewrite builder for `table`. + pub fn new(table: &'a Table) -> Self { + Self { + table, + predicate: Predicate::AlwaysTrue, + snapshot_id: None, + batch_size: None, + case_sensitive: true, + rewriter: None, + } + } + + /// Sets the row predicate used to plan candidate files. + pub fn with_predicate(mut self, predicate: Predicate) -> Self { + self.predicate = predicate; + self + } + + /// Sets the snapshot id used to plan candidate files. + pub fn with_snapshot_id(mut self, snapshot_id: i64) -> Self { + self.snapshot_id = Some(snapshot_id); + self + } + + /// Sets the Arrow reader batch size. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = Some(batch_size); + self + } + + /// Sets the case sensitivity used to bind the planning predicate. + pub fn with_case_sensitive(mut self, case_sensitive: bool) -> Self { + self.case_sensitive = case_sensitive; + self + } + + /// Sets the record batch rewriter. + pub fn with_rewriter(mut self, rewriter: Arc<dyn CowBatchRewriter>) -> Self { + self.rewriter = Some(rewriter); + self + } + + /// Plans, reads, rewrites, and writes replacement data files. + pub async fn rewrite(self) -> Result<CowRewriteResult> { + let rewriter = self.rewriter.ok_or_else(|| { + Error::new( + ErrorKind::PreconditionFailed, + "COW rewrite requires a batch rewriter", + ) + })?; + let files = plan::plan_cow_rewrite_files( + self.table, + Some(self.predicate), + self.snapshot_id, + self.case_sensitive, + ) + .await?; + + let mut result = CowRewriteResult { + stats: CowRewriteStats { + candidate_files: files.len(), + ..CowRewriteStats::default() + }, + ..CowRewriteResult::default() + }; + + for file in files { + // Schema the rows are read in (the planned snapshot's schema). The + // replacement files must be written with this schema so that batches + // remain compatible when the table's current schema has evolved past + // the snapshot the source files belong to. + let write_schema = file.scan_task.schema_ref(); + + // Batches produced before the first changed batch. They are buffered + // rather than written immediately because the primitive must not + // emit a replacement file for a source file that turns out to be + // unchanged. Once a changed batch is observed the buffered prefix is + // flushed to the writer and all subsequent batches stream straight + // through, so the in-memory footprint is bounded by the rows that + // precede the first change instead of the entire source file. + let mut prefix: Vec<RecordBatch> = Vec::new(); + let mut file_changed = false; + let mut writer: Option<Box<dyn crate::writer::IcebergWriter>> = None; + + // Planning already cleared the row predicate (see + // `ManifestEntryContext::into_cow_rewrite_file`), so this task + // reads every row of the source file. + let tasks = Box::pin(futures::stream::iter(vec![Ok(file.scan_task.clone())])) + as FileScanTaskStream; + + // Each candidate file gets its own reader so the per-file prefix + // and lazy-writer semantics stay intact; the delete-file cache is + // therefore also per file, and equality deletes shared by several + // candidates are fetched once per file. + let mut reader_builder = self.table.reader_builder(); + if let Some(batch_size) = self.batch_size { + reader_builder = reader_builder.with_batch_size(batch_size); + } + + let mut batches = reader_builder.build().read(tasks)?.stream(); + while let Some(batch) = batches.try_next().await? { + result.stats.input_rows += batch.num_rows() as u64; + + let rewrite = rewriter.rewrite_batch(batch)?; + if rewrite.changed { + file_changed = true; + result.stats.changed_batches += 1; + } + + if let Some(output) = rewrite.output { + result.stats.output_rows += output.num_rows() as u64; Review Comment: `output_rows` gets bumped for every `Some(output)` batch, including prefix batches on files that end up unchanged — so a `KeepAll` over a 3-row file reports `output_rows == 3` while `rewritten_files == 0` and nothing is written. `cow_rewrite_keep_all_produces_no_changes` actually asserts this, so it's currently baked in as intended. The problem is a caller can't read this as "rows persisted to replacement files" — anyone cross-checking against `added_data_files` row counts sees a phantom mismatch. I'd either only accumulate it when `file_changed`, or rename to something like `rewriter_emitted_rows` and add a separate `written_rows`. wdyt? ########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1261 @@ +// 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. + +//! Copy-on-write rewrite primitives. +//! +//! This module plans candidate data files, reads their visible rows, applies a +//! caller-provided batch rewriter, and writes replacement data files. It returns +//! old and new file sets that can be committed by an overwrite-style transaction +//! action. +//! +//! The primitive does not parse SQL and does not commit metadata by itself. +//! Rewriters must emit batches compatible with the schema rows were read in +//! (the planned snapshot's schema) and must preserve each source file's +//! partition values; this primitive does not repartition rewritten rows. +//! +//! The result carries data files only. A commit adapter consuming these file +//! lists must also account for delete files that reference removed files — +//! for example deletion vectors whose referenced data file is being removed, +//! and position deletes scoped to it; equality deletes remain valid but +//! become redundant once their target rows are rewritten. +//! +//! ```rust,no_run +//! # use std::sync::Arc; +//! # use arrow_array::RecordBatch; +//! # use iceberg::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder}; +//! # use iceberg::table::Table; +//! # use iceberg::Result; +//! struct KeepAll; +//! +//! impl CowBatchRewriter for KeepAll { +//! fn rewrite_batch(&self, batch: RecordBatch) -> Result<CowBatchRewrite> { +//! Ok(CowBatchRewrite { +//! output: Some(batch), +//! changed: false, +//! }) +//! } +//! } +//! +//! # async fn example(table: &Table) -> Result<()> { +//! let result = CowRewriteBuilder::new(table) +//! .with_rewriter(Arc::new(KeepAll)) +//! .rewrite() +//! .await?; +//! +//! assert!(!result.has_changes()); +//! # Ok(()) +//! # } +//! ``` + +mod plan; +mod rewriter; +pub(crate) mod writer; + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use futures::TryStreamExt; +pub use plan::CowRewriteFile; +pub use rewriter::{CowBatchRewrite, CowBatchRewriter}; + +use crate::expr::Predicate; +use crate::scan::FileScanTaskStream; +use crate::spec::{DataFile, PartitionKey}; +use crate::table::Table; +use crate::{Error, ErrorKind, Result}; + +/// Counters produced by a copy-on-write rewrite. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CowRewriteStats { + /// Number of candidate files selected by planning. + pub candidate_files: usize, + /// Number of old files that have replacement output or are fully removed. + pub rewritten_files: usize, + /// Number of candidate files that did not change after row rewriting. + pub unchanged_files: usize, + /// Visible input row count read from candidate files. + pub input_rows: u64, + /// Output row count emitted by the batch rewriter. + pub output_rows: u64, + /// Number of input batches where the rewriter reported changes. + pub changed_batches: u64, +} + +/// Result of a copy-on-write rewrite operation. +#[derive(Debug, Default)] +pub struct CowRewriteResult { + /// Old data files that should be removed by the commit action. + pub removed_data_files: Vec<DataFile>, + /// New data files that should be added by the commit action. + pub added_data_files: Vec<DataFile>, + /// Candidate files that were read and left unchanged. + /// + /// This includes files whose visible rows were all removed by delete + /// files: with no surviving rows the rewriter never runs, so the file is + /// kept as-is rather than dropped. + pub unchanged_data_files: Vec<DataFile>, + /// Rewrite counters. + pub stats: CowRewriteStats, +} + +impl CowRewriteResult { + /// Returns true if the rewrite produced any table changes. + pub fn has_changes(&self) -> bool { + !self.removed_data_files.is_empty() || !self.added_data_files.is_empty() + } +} + +/// Builder for orchestrating copy-on-write data file rewrites. +pub struct CowRewriteBuilder<'a> { + table: &'a Table, + predicate: Predicate, + snapshot_id: Option<i64>, + batch_size: Option<usize>, + case_sensitive: bool, + rewriter: Option<Arc<dyn CowBatchRewriter>>, +} + +impl<'a> CowRewriteBuilder<'a> { + /// Creates a copy-on-write rewrite builder for `table`. + pub fn new(table: &'a Table) -> Self { + Self { + table, + predicate: Predicate::AlwaysTrue, + snapshot_id: None, + batch_size: None, + case_sensitive: true, + rewriter: None, + } + } + + /// Sets the row predicate used to plan candidate files. + pub fn with_predicate(mut self, predicate: Predicate) -> Self { + self.predicate = predicate; + self + } + + /// Sets the snapshot id used to plan candidate files. + pub fn with_snapshot_id(mut self, snapshot_id: i64) -> Self { + self.snapshot_id = Some(snapshot_id); + self + } + + /// Sets the Arrow reader batch size. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = Some(batch_size); + self + } + + /// Sets the case sensitivity used to bind the planning predicate. + pub fn with_case_sensitive(mut self, case_sensitive: bool) -> Self { + self.case_sensitive = case_sensitive; + self + } + + /// Sets the record batch rewriter. + pub fn with_rewriter(mut self, rewriter: Arc<dyn CowBatchRewriter>) -> Self { + self.rewriter = Some(rewriter); + self + } + + /// Plans, reads, rewrites, and writes replacement data files. + pub async fn rewrite(self) -> Result<CowRewriteResult> { + let rewriter = self.rewriter.ok_or_else(|| { + Error::new( + ErrorKind::PreconditionFailed, + "COW rewrite requires a batch rewriter", + ) + })?; + let files = plan::plan_cow_rewrite_files( + self.table, + Some(self.predicate), + self.snapshot_id, + self.case_sensitive, + ) + .await?; + + let mut result = CowRewriteResult { + stats: CowRewriteStats { + candidate_files: files.len(), + ..CowRewriteStats::default() + }, + ..CowRewriteResult::default() + }; + + for file in files { + // Schema the rows are read in (the planned snapshot's schema). The + // replacement files must be written with this schema so that batches + // remain compatible when the table's current schema has evolved past + // the snapshot the source files belong to. + let write_schema = file.scan_task.schema_ref(); + + // Batches produced before the first changed batch. They are buffered + // rather than written immediately because the primitive must not + // emit a replacement file for a source file that turns out to be + // unchanged. Once a changed batch is observed the buffered prefix is + // flushed to the writer and all subsequent batches stream straight + // through, so the in-memory footprint is bounded by the rows that + // precede the first change instead of the entire source file. + let mut prefix: Vec<RecordBatch> = Vec::new(); + let mut file_changed = false; + let mut writer: Option<Box<dyn crate::writer::IcebergWriter>> = None; + + // Planning already cleared the row predicate (see + // `ManifestEntryContext::into_cow_rewrite_file`), so this task + // reads every row of the source file. + let tasks = Box::pin(futures::stream::iter(vec![Ok(file.scan_task.clone())])) + as FileScanTaskStream; + + // Each candidate file gets its own reader so the per-file prefix + // and lazy-writer semantics stay intact; the delete-file cache is + // therefore also per file, and equality deletes shared by several + // candidates are fetched once per file. + let mut reader_builder = self.table.reader_builder(); + if let Some(batch_size) = self.batch_size { + reader_builder = reader_builder.with_batch_size(batch_size); + } + + let mut batches = reader_builder.build().read(tasks)?.stream(); + while let Some(batch) = batches.try_next().await? { + result.stats.input_rows += batch.num_rows() as u64; + + let rewrite = rewriter.rewrite_batch(batch)?; + if rewrite.changed { + file_changed = true; + result.stats.changed_batches += 1; + } + + if let Some(output) = rewrite.output { + result.stats.output_rows += output.num_rows() as u64; + + if file_changed { + if writer.is_none() { + let partition_key = source_partition_key( + self.table, + &file.old_data_file, + &write_schema, + )?; + writer = Some( + writer::build_replacement_writer( + self.table, + write_schema.clone(), + Some(partition_key), + ) + .await?, + ); + } + let writer = writer.as_mut().expect("writer just built"); Review Comment: This `.expect("writer just built")` is sound today given the `is_none()` check right above, but a future reorder would turn it into a runtime panic mid-stream rather than a compile error. I'd either use `unreachable!` to signal it's an invariant, or restructure as `if let Some(w) = writer.as_mut()` after the init block so the compiler enforces it. ########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1261 @@ +// 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. + +//! Copy-on-write rewrite primitives. +//! +//! This module plans candidate data files, reads their visible rows, applies a +//! caller-provided batch rewriter, and writes replacement data files. It returns +//! old and new file sets that can be committed by an overwrite-style transaction +//! action. +//! +//! The primitive does not parse SQL and does not commit metadata by itself. +//! Rewriters must emit batches compatible with the schema rows were read in +//! (the planned snapshot's schema) and must preserve each source file's +//! partition values; this primitive does not repartition rewritten rows. +//! +//! The result carries data files only. A commit adapter consuming these file +//! lists must also account for delete files that reference removed files — +//! for example deletion vectors whose referenced data file is being removed, +//! and position deletes scoped to it; equality deletes remain valid but +//! become redundant once their target rows are rewritten. +//! +//! ```rust,no_run +//! # use std::sync::Arc; +//! # use arrow_array::RecordBatch; +//! # use iceberg::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder}; +//! # use iceberg::table::Table; +//! # use iceberg::Result; +//! struct KeepAll; +//! +//! impl CowBatchRewriter for KeepAll { +//! fn rewrite_batch(&self, batch: RecordBatch) -> Result<CowBatchRewrite> { +//! Ok(CowBatchRewrite { +//! output: Some(batch), +//! changed: false, +//! }) +//! } +//! } +//! +//! # async fn example(table: &Table) -> Result<()> { +//! let result = CowRewriteBuilder::new(table) +//! .with_rewriter(Arc::new(KeepAll)) +//! .rewrite() +//! .await?; +//! +//! assert!(!result.has_changes()); +//! # Ok(()) +//! # } +//! ``` + +mod plan; +mod rewriter; +pub(crate) mod writer; + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use futures::TryStreamExt; +pub use plan::CowRewriteFile; +pub use rewriter::{CowBatchRewrite, CowBatchRewriter}; + +use crate::expr::Predicate; +use crate::scan::FileScanTaskStream; +use crate::spec::{DataFile, PartitionKey}; +use crate::table::Table; +use crate::{Error, ErrorKind, Result}; + +/// Counters produced by a copy-on-write rewrite. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CowRewriteStats { + /// Number of candidate files selected by planning. + pub candidate_files: usize, + /// Number of old files that have replacement output or are fully removed. + pub rewritten_files: usize, + /// Number of candidate files that did not change after row rewriting. + pub unchanged_files: usize, + /// Visible input row count read from candidate files. + pub input_rows: u64, + /// Output row count emitted by the batch rewriter. + pub output_rows: u64, + /// Number of input batches where the rewriter reported changes. + pub changed_batches: u64, +} + +/// Result of a copy-on-write rewrite operation. +#[derive(Debug, Default)] +pub struct CowRewriteResult { + /// Old data files that should be removed by the commit action. + pub removed_data_files: Vec<DataFile>, + /// New data files that should be added by the commit action. + pub added_data_files: Vec<DataFile>, + /// Candidate files that were read and left unchanged. + /// + /// This includes files whose visible rows were all removed by delete + /// files: with no surviving rows the rewriter never runs, so the file is + /// kept as-is rather than dropped. + pub unchanged_data_files: Vec<DataFile>, + /// Rewrite counters. + pub stats: CowRewriteStats, +} + +impl CowRewriteResult { + /// Returns true if the rewrite produced any table changes. + pub fn has_changes(&self) -> bool { + !self.removed_data_files.is_empty() || !self.added_data_files.is_empty() + } +} + +/// Builder for orchestrating copy-on-write data file rewrites. +pub struct CowRewriteBuilder<'a> { + table: &'a Table, + predicate: Predicate, + snapshot_id: Option<i64>, + batch_size: Option<usize>, + case_sensitive: bool, + rewriter: Option<Arc<dyn CowBatchRewriter>>, +} + +impl<'a> CowRewriteBuilder<'a> { + /// Creates a copy-on-write rewrite builder for `table`. + pub fn new(table: &'a Table) -> Self { + Self { + table, + predicate: Predicate::AlwaysTrue, + snapshot_id: None, + batch_size: None, + case_sensitive: true, + rewriter: None, + } + } + + /// Sets the row predicate used to plan candidate files. + pub fn with_predicate(mut self, predicate: Predicate) -> Self { + self.predicate = predicate; + self + } + + /// Sets the snapshot id used to plan candidate files. + pub fn with_snapshot_id(mut self, snapshot_id: i64) -> Self { + self.snapshot_id = Some(snapshot_id); + self + } + + /// Sets the Arrow reader batch size. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = Some(batch_size); + self + } + + /// Sets the case sensitivity used to bind the planning predicate. + pub fn with_case_sensitive(mut self, case_sensitive: bool) -> Self { + self.case_sensitive = case_sensitive; + self + } + + /// Sets the record batch rewriter. + pub fn with_rewriter(mut self, rewriter: Arc<dyn CowBatchRewriter>) -> Self { + self.rewriter = Some(rewriter); + self + } + + /// Plans, reads, rewrites, and writes replacement data files. + pub async fn rewrite(self) -> Result<CowRewriteResult> { + let rewriter = self.rewriter.ok_or_else(|| { + Error::new( + ErrorKind::PreconditionFailed, + "COW rewrite requires a batch rewriter", + ) + })?; + let files = plan::plan_cow_rewrite_files( + self.table, + Some(self.predicate), + self.snapshot_id, + self.case_sensitive, + ) + .await?; + + let mut result = CowRewriteResult { + stats: CowRewriteStats { + candidate_files: files.len(), + ..CowRewriteStats::default() + }, + ..CowRewriteResult::default() + }; + + for file in files { + // Schema the rows are read in (the planned snapshot's schema). The + // replacement files must be written with this schema so that batches + // remain compatible when the table's current schema has evolved past + // the snapshot the source files belong to. + let write_schema = file.scan_task.schema_ref(); + + // Batches produced before the first changed batch. They are buffered + // rather than written immediately because the primitive must not + // emit a replacement file for a source file that turns out to be + // unchanged. Once a changed batch is observed the buffered prefix is + // flushed to the writer and all subsequent batches stream straight + // through, so the in-memory footprint is bounded by the rows that + // precede the first change instead of the entire source file. + let mut prefix: Vec<RecordBatch> = Vec::new(); + let mut file_changed = false; + let mut writer: Option<Box<dyn crate::writer::IcebergWriter>> = None; + + // Planning already cleared the row predicate (see + // `ManifestEntryContext::into_cow_rewrite_file`), so this task + // reads every row of the source file. + let tasks = Box::pin(futures::stream::iter(vec![Ok(file.scan_task.clone())])) Review Comment: Small one, non-blocker — `file.scan_task.clone()` here and `file.old_data_file.clone()` at the end of the loop are both avoidable since `file` isn't used again after the if/else. Destructuring at the top (`let CowRewriteFile { old_data_file, scan_task } = file;`) lets you move both instead, and `FileScanTask` / `DataFile` aren't cheap to clone (paths, projection vecs, stats maps) when it's once per candidate. Just while we're here. ########## crates/iceberg/src/cow_rewrite/rewriter.rs: ########## @@ -0,0 +1,43 @@ +// 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 arrow_array::RecordBatch; + +use crate::Result; + +/// Result of rewriting a single record batch. +pub struct CowBatchRewrite { + /// Rewritten output batch, or `None` when the input batch is fully removed. + /// + /// Output batches must use the same schema as their input batch — the + /// planned snapshot's schema, which may be older than the table's current + /// schema. Rewriters must also preserve each source file's partition + /// values: this primitive writes replacements into the source file's + /// partition and does not repartition rows. + pub output: Option<RecordBatch>, + /// Whether the rewrite changed the input batch contents. + /// + /// Set this to `true` whenever `output` differs from the input batch, + /// including filtered rows, updated values, reordered rows, or `None`. + pub changed: bool, +} + +/// Rewrites record batches for copy-on-write operations. +pub trait CowBatchRewriter: Send + Sync { + /// Rewrites a record batch and reports whether it changed. + fn rewrite_batch(&self, batch: RecordBatch) -> Result<CowBatchRewrite>; Review Comment: `rewrite_batch` is synchronous, so a rewriter that needs async I/O — catalog enrichment, cross-table dedup, an audit lookup — has to block the runtime thread or shell out to `spawn_blocking`. That rules out a fair chunk of what people will want to plug in here. Making it `async` later is a breaking change once this is stabilized, so I'd rather decide now. If object safety is the reason it's sync (I see `cow_batch_rewriter_is_object_safe`), that's a fair constraint — but then I'd document in the trait doc that blocking work isn't supported, so nobody learns it the hard way. wdyt? ########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1261 @@ +// 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. + +//! Copy-on-write rewrite primitives. +//! +//! This module plans candidate data files, reads their visible rows, applies a +//! caller-provided batch rewriter, and writes replacement data files. It returns +//! old and new file sets that can be committed by an overwrite-style transaction +//! action. +//! +//! The primitive does not parse SQL and does not commit metadata by itself. +//! Rewriters must emit batches compatible with the schema rows were read in +//! (the planned snapshot's schema) and must preserve each source file's +//! partition values; this primitive does not repartition rewritten rows. +//! +//! The result carries data files only. A commit adapter consuming these file +//! lists must also account for delete files that reference removed files — +//! for example deletion vectors whose referenced data file is being removed, +//! and position deletes scoped to it; equality deletes remain valid but +//! become redundant once their target rows are rewritten. +//! +//! ```rust,no_run +//! # use std::sync::Arc; +//! # use arrow_array::RecordBatch; +//! # use iceberg::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder}; +//! # use iceberg::table::Table; +//! # use iceberg::Result; +//! struct KeepAll; +//! +//! impl CowBatchRewriter for KeepAll { +//! fn rewrite_batch(&self, batch: RecordBatch) -> Result<CowBatchRewrite> { +//! Ok(CowBatchRewrite { +//! output: Some(batch), +//! changed: false, +//! }) +//! } +//! } +//! +//! # async fn example(table: &Table) -> Result<()> { +//! let result = CowRewriteBuilder::new(table) +//! .with_rewriter(Arc::new(KeepAll)) +//! .rewrite() +//! .await?; +//! +//! assert!(!result.has_changes()); +//! # Ok(()) +//! # } +//! ``` + +mod plan; +mod rewriter; +pub(crate) mod writer; + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use futures::TryStreamExt; +pub use plan::CowRewriteFile; +pub use rewriter::{CowBatchRewrite, CowBatchRewriter}; + +use crate::expr::Predicate; +use crate::scan::FileScanTaskStream; +use crate::spec::{DataFile, PartitionKey}; +use crate::table::Table; +use crate::{Error, ErrorKind, Result}; + +/// Counters produced by a copy-on-write rewrite. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CowRewriteStats { + /// Number of candidate files selected by planning. + pub candidate_files: usize, + /// Number of old files that have replacement output or are fully removed. + pub rewritten_files: usize, + /// Number of candidate files that did not change after row rewriting. + pub unchanged_files: usize, + /// Visible input row count read from candidate files. + pub input_rows: u64, + /// Output row count emitted by the batch rewriter. + pub output_rows: u64, + /// Number of input batches where the rewriter reported changes. + pub changed_batches: u64, +} + +/// Result of a copy-on-write rewrite operation. +#[derive(Debug, Default)] +pub struct CowRewriteResult { + /// Old data files that should be removed by the commit action. + pub removed_data_files: Vec<DataFile>, + /// New data files that should be added by the commit action. + pub added_data_files: Vec<DataFile>, + /// Candidate files that were read and left unchanged. + /// + /// This includes files whose visible rows were all removed by delete + /// files: with no surviving rows the rewriter never runs, so the file is + /// kept as-is rather than dropped. + pub unchanged_data_files: Vec<DataFile>, + /// Rewrite counters. + pub stats: CowRewriteStats, +} + +impl CowRewriteResult { + /// Returns true if the rewrite produced any table changes. + pub fn has_changes(&self) -> bool { + !self.removed_data_files.is_empty() || !self.added_data_files.is_empty() + } +} + +/// Builder for orchestrating copy-on-write data file rewrites. +pub struct CowRewriteBuilder<'a> { + table: &'a Table, + predicate: Predicate, + snapshot_id: Option<i64>, + batch_size: Option<usize>, + case_sensitive: bool, + rewriter: Option<Arc<dyn CowBatchRewriter>>, +} + +impl<'a> CowRewriteBuilder<'a> { + /// Creates a copy-on-write rewrite builder for `table`. + pub fn new(table: &'a Table) -> Self { + Self { + table, + predicate: Predicate::AlwaysTrue, + snapshot_id: None, + batch_size: None, + case_sensitive: true, + rewriter: None, + } + } + + /// Sets the row predicate used to plan candidate files. + pub fn with_predicate(mut self, predicate: Predicate) -> Self { + self.predicate = predicate; + self + } + + /// Sets the snapshot id used to plan candidate files. + pub fn with_snapshot_id(mut self, snapshot_id: i64) -> Self { + self.snapshot_id = Some(snapshot_id); + self + } + + /// Sets the Arrow reader batch size. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = Some(batch_size); + self + } + + /// Sets the case sensitivity used to bind the planning predicate. + pub fn with_case_sensitive(mut self, case_sensitive: bool) -> Self { + self.case_sensitive = case_sensitive; + self + } + + /// Sets the record batch rewriter. + pub fn with_rewriter(mut self, rewriter: Arc<dyn CowBatchRewriter>) -> Self { + self.rewriter = Some(rewriter); + self + } + + /// Plans, reads, rewrites, and writes replacement data files. + pub async fn rewrite(self) -> Result<CowRewriteResult> { + let rewriter = self.rewriter.ok_or_else(|| { + Error::new( + ErrorKind::PreconditionFailed, + "COW rewrite requires a batch rewriter", + ) + })?; + let files = plan::plan_cow_rewrite_files( + self.table, + Some(self.predicate), + self.snapshot_id, + self.case_sensitive, + ) + .await?; + + let mut result = CowRewriteResult { + stats: CowRewriteStats { + candidate_files: files.len(), + ..CowRewriteStats::default() + }, + ..CowRewriteResult::default() + }; + + for file in files { + // Schema the rows are read in (the planned snapshot's schema). The + // replacement files must be written with this schema so that batches + // remain compatible when the table's current schema has evolved past + // the snapshot the source files belong to. + let write_schema = file.scan_task.schema_ref(); + + // Batches produced before the first changed batch. They are buffered + // rather than written immediately because the primitive must not + // emit a replacement file for a source file that turns out to be + // unchanged. Once a changed batch is observed the buffered prefix is + // flushed to the writer and all subsequent batches stream straight + // through, so the in-memory footprint is bounded by the rows that + // precede the first change instead of the entire source file. + let mut prefix: Vec<RecordBatch> = Vec::new(); + let mut file_changed = false; + let mut writer: Option<Box<dyn crate::writer::IcebergWriter>> = None; + + // Planning already cleared the row predicate (see + // `ManifestEntryContext::into_cow_rewrite_file`), so this task + // reads every row of the source file. + let tasks = Box::pin(futures::stream::iter(vec![Ok(file.scan_task.clone())])) + as FileScanTaskStream; + + // Each candidate file gets its own reader so the per-file prefix + // and lazy-writer semantics stay intact; the delete-file cache is + // therefore also per file, and equality deletes shared by several + // candidates are fetched once per file. + let mut reader_builder = self.table.reader_builder(); + if let Some(batch_size) = self.batch_size { + reader_builder = reader_builder.with_batch_size(batch_size); + } + + let mut batches = reader_builder.build().read(tasks)?.stream(); + while let Some(batch) = batches.try_next().await? { + result.stats.input_rows += batch.num_rows() as u64; + + let rewrite = rewriter.rewrite_batch(batch)?; + if rewrite.changed { + file_changed = true; + result.stats.changed_batches += 1; + } + + if let Some(output) = rewrite.output { + result.stats.output_rows += output.num_rows() as u64; + + if file_changed { + if writer.is_none() { + let partition_key = source_partition_key( + self.table, + &file.old_data_file, + &write_schema, + )?; + writer = Some( + writer::build_replacement_writer( + self.table, + write_schema.clone(), + Some(partition_key), + ) + .await?, + ); + } + let writer = writer.as_mut().expect("writer just built"); + for prefix_batch in prefix.drain(..) { + writer.write(prefix_batch).await?; + } + writer.write(output).await?; + } else { + prefix.push(output); + } + } + } + + if file_changed { + result.stats.rewritten_files += 1; + result.removed_data_files.push(file.old_data_file.clone()); + + if let Some(mut writer) = writer { + let added_data_files = writer.close().await?; + result.added_data_files.extend(added_data_files); + } + // If `writer` is `None`, the source file was fully deleted + // (every batch dropped to `output: None`), so no replacement + // file is written. + } else { + result.stats.unchanged_files += 1; + result.unchanged_data_files.push(file.old_data_file); + // `prefix` is dropped here; no replacement file was written. + } + } + + Ok(result) + } +} + +fn source_partition_key( + table: &Table, + data_file: &DataFile, + schema: &crate::spec::SchemaRef, +) -> Result<PartitionKey> { + let spec = table + .metadata() + .partition_spec_by_id(data_file.partition_spec_id) + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Missing partition spec {} for COW rewrite source file", + data_file.partition_spec_id + ), + ) + })? + .as_ref() + .clone(); + spec.partition_type(schema).map_err(|err| { Review Comment: This calls `partition_type` purely for the early-error side effect and drops the `StructType`, then `PartitionKey::new` below likely computes it again. It reads like dead code — a future reader could easily delete it and quietly remove the validation. If `PartitionKey::new` already validates, I'd drop this pre-call; if it doesn't, a one-line comment on why we're binding here would save the next person the double-take. ########## crates/iceberg/src/cow_rewrite/plan.rs: ########## @@ -0,0 +1,158 @@ +// 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 futures::TryStreamExt; + +use crate::scan::FileScanTask; +use crate::spec::DataFile; + +/// A data file selected for COW rewrite. +/// +/// The planning entry points that produce candidates are crate-internal and +/// reached through [`crate::cow_rewrite::CowRewriteBuilder`]; the type itself +/// is public so follow-up commit-adapter work (overwrite and row-delta +/// actions) can consume planned candidates directly. +#[derive(Debug, Clone)] +pub struct CowRewriteFile { Review Comment: The doc says this type is public so follow-up commit-adapter work can consume planned candidates directly, but both fields are `pub(crate)` with only read accessors — so an adapter in another crate can read a `CowRewriteFile` but can't construct one, and can't reach `plan_cow_rewrite_files` either. If those adapters are meant to live outside this crate, the intended contract isn't actually exposed yet. Since this is the public surface future work builds on, I'd rather settle now whether it's read-only-via-getters or genuinely constructable externally. wdyt? ########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1261 @@ +// 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. + +//! Copy-on-write rewrite primitives. +//! +//! This module plans candidate data files, reads their visible rows, applies a +//! caller-provided batch rewriter, and writes replacement data files. It returns +//! old and new file sets that can be committed by an overwrite-style transaction +//! action. +//! +//! The primitive does not parse SQL and does not commit metadata by itself. +//! Rewriters must emit batches compatible with the schema rows were read in +//! (the planned snapshot's schema) and must preserve each source file's +//! partition values; this primitive does not repartition rewritten rows. +//! +//! The result carries data files only. A commit adapter consuming these file +//! lists must also account for delete files that reference removed files — +//! for example deletion vectors whose referenced data file is being removed, +//! and position deletes scoped to it; equality deletes remain valid but +//! become redundant once their target rows are rewritten. +//! +//! ```rust,no_run +//! # use std::sync::Arc; +//! # use arrow_array::RecordBatch; +//! # use iceberg::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder}; +//! # use iceberg::table::Table; +//! # use iceberg::Result; +//! struct KeepAll; +//! +//! impl CowBatchRewriter for KeepAll { +//! fn rewrite_batch(&self, batch: RecordBatch) -> Result<CowBatchRewrite> { +//! Ok(CowBatchRewrite { +//! output: Some(batch), +//! changed: false, +//! }) +//! } +//! } +//! +//! # async fn example(table: &Table) -> Result<()> { +//! let result = CowRewriteBuilder::new(table) +//! .with_rewriter(Arc::new(KeepAll)) +//! .rewrite() +//! .await?; +//! +//! assert!(!result.has_changes()); +//! # Ok(()) +//! # } +//! ``` + +mod plan; +mod rewriter; +pub(crate) mod writer; + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use futures::TryStreamExt; +pub use plan::CowRewriteFile; +pub use rewriter::{CowBatchRewrite, CowBatchRewriter}; + +use crate::expr::Predicate; +use crate::scan::FileScanTaskStream; +use crate::spec::{DataFile, PartitionKey}; +use crate::table::Table; +use crate::{Error, ErrorKind, Result}; + +/// Counters produced by a copy-on-write rewrite. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct CowRewriteStats { + /// Number of candidate files selected by planning. + pub candidate_files: usize, + /// Number of old files that have replacement output or are fully removed. + pub rewritten_files: usize, + /// Number of candidate files that did not change after row rewriting. + pub unchanged_files: usize, + /// Visible input row count read from candidate files. + pub input_rows: u64, + /// Output row count emitted by the batch rewriter. + pub output_rows: u64, + /// Number of input batches where the rewriter reported changes. + pub changed_batches: u64, +} + +/// Result of a copy-on-write rewrite operation. +#[derive(Debug, Default)] +pub struct CowRewriteResult { + /// Old data files that should be removed by the commit action. + pub removed_data_files: Vec<DataFile>, + /// New data files that should be added by the commit action. + pub added_data_files: Vec<DataFile>, + /// Candidate files that were read and left unchanged. + /// + /// This includes files whose visible rows were all removed by delete Review Comment: This one I'd push back on. When a file's visible rows are all removed by delete files the reader yields zero batches, the loop never runs, and the file lands in `unchanged_data_files` — carrying its delete files with it, forever. The Java side (RewriteDataFiles) drops the file in that case: the data-file-plus-delete-file pair gets compacted away rather than pinned in the table. As written, any DELETE/UPDATE adapter built on this primitive can never compact an all-deleted file, so snapshot size and delete-file read I/O only grow. I'd treat a scan task that has delete files and produces zero rows as changed-with-no-replacement (the same terminal state as "rewriter dropped every batch"). If we do want to defer it, I'd at least frame it in the struct doc as a known gap the commit adapter has to handle, rather than as correct behavior. wdyt? -- 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]
