laskoviymishka commented on code in PR #2752: URL: https://github.com/apache/iceberg-rust/pull/2752#discussion_r4045170891
########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1390 @@ +// 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 written to replacement files. + /// + /// Rows the rewriter emitted for files that turned out unchanged are not + /// counted, so this always matches the row counts of `added_data_files`. + pub output_rows: u64, + /// Number of input batches that changed, including batches the rewriter + /// dropped entirely (`output: None`) even if it did not flag them. + 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. + /// + /// Files whose visible rows were all removed by delete files are NOT + /// included here: they read as zero rows and are reported in + /// `removed_data_files` with no replacement, so a commit adapter can drop + /// them together with the delete files that reference them. + 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 { + let CowRewriteFile { + old_data_file, + scan_task, + } = file; + // 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 = scan_task.schema_ref(); + let has_delete_files = !scan_task.deletes().is_empty(); + + // 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. + // + // Worst-case footprint: for a file that never changes (or whose + // first change sits at its very end) the prefix holds the entire + // decoded source file in memory. Files are processed sequentially, + // so peak usage is one file at a time, but that can still be several + // GB for a compaction-sized file. A size-capped fallback that starts + // writing the replacement once the buffer crosses a threshold is + // left for follow-up work. + let mut prefix: Vec<RecordBatch> = Vec::new(); + let mut file_changed = false; + let mut file_input_rows = 0_u64; + let mut file_output_rows = 0_u64; + 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(scan_task)])) 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; + file_input_rows += batch.num_rows() as u64; + + let rewrite = rewriter.rewrite_batch(batch)?; + // `output: None` means the batch is fully removed, which is + // itself a change. Derive the effective flag instead of + // trusting every rewriter to keep `changed` consistent with + // `output` — otherwise a `{changed: false, output: None}` + // batch would silently drop its rows while leaving the file + // marked unchanged. + let changed = rewrite.changed || rewrite.output.is_none(); + if changed { + file_changed = true; + result.stats.changed_batches += 1; + } + + if let Some(output) = rewrite.output { + file_output_rows += output.num_rows() as u64; + + if file_changed { + if writer.is_none() { + let partition_key = + source_partition_key(self.table, &old_data_file, &write_schema)?; + writer = Some( + writer::build_replacement_writer( + self.table, + write_schema.clone(), + Some(partition_key), + ) + .await?, + ); + } + let Some(writer) = writer.as_mut() else { + unreachable!("writer initialized above"); Review Comment: `unreachable!` compiles a panic into a library path, which is the one thing we don't want here — even a "can't happen" invariant should return `Err` rather than abort the caller's process. If you take the prefix-flush fix from the data-loss comment above and build the writer eagerly when the file first flips to changed, this branch disappears on its own. If it stays, I'd make it `writer.as_mut().ok_or_else(|| Error::new(ErrorKind::Unexpected, "COW writer missing after file marked changed"))?`. ########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1390 @@ +// 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. Review Comment: I'd tighten this line — "become redundant" reads like a signal that a commit adapter can drop these equality delete files, and that's not safe. An equality delete applies by sequence number, so the same file can still apply to other data files that weren't part of this rewrite; dropping it would resurrect deleted rows elsewhere. Since this module doc is the contract the overwrite/row-delta adapters will read, I'd say it explicitly: only position deletes and deletion vectors that exclusively reference a removed file should be dropped, and equality deletes must be left in place. ########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1390 @@ +// 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 written to replacement files. + /// + /// Rows the rewriter emitted for files that turned out unchanged are not + /// counted, so this always matches the row counts of `added_data_files`. + pub output_rows: u64, + /// Number of input batches that changed, including batches the rewriter + /// dropped entirely (`output: None`) even if it did not flag them. + 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. + /// + /// Files whose visible rows were all removed by delete files are NOT + /// included here: they read as zero rows and are reported in + /// `removed_data_files` with no replacement, so a commit adapter can drop + /// them together with the delete files that reference them. + 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 { + let CowRewriteFile { + old_data_file, + scan_task, + } = file; + // 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 = scan_task.schema_ref(); + let has_delete_files = !scan_task.deletes().is_empty(); + + // 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. + // + // Worst-case footprint: for a file that never changes (or whose + // first change sits at its very end) the prefix holds the entire + // decoded source file in memory. Files are processed sequentially, + // so peak usage is one file at a time, but that can still be several + // GB for a compaction-sized file. A size-capped fallback that starts + // writing the replacement once the buffer crosses a threshold is + // left for follow-up work. + let mut prefix: Vec<RecordBatch> = Vec::new(); + let mut file_changed = false; + let mut file_input_rows = 0_u64; + let mut file_output_rows = 0_u64; + 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(scan_task)])) 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; + file_input_rows += batch.num_rows() as u64; + + let rewrite = rewriter.rewrite_batch(batch)?; + // `output: None` means the batch is fully removed, which is + // itself a change. Derive the effective flag instead of + // trusting every rewriter to keep `changed` consistent with + // `output` — otherwise a `{changed: false, output: None}` + // batch would silently drop its rows while leaving the file + // marked unchanged. + let changed = rewrite.changed || rewrite.output.is_none(); + if changed { + file_changed = true; + result.stats.changed_batches += 1; + } + + if let Some(output) = rewrite.output { + file_output_rows += output.num_rows() as u64; + + if file_changed { + if writer.is_none() { + let partition_key = + source_partition_key(self.table, &old_data_file, &write_schema)?; + writer = Some( + writer::build_replacement_writer( + self.table, + write_schema.clone(), + Some(partition_key), + ) + .await?, + ); + } + let Some(writer) = writer.as_mut() else { + unreachable!("writer initialized above"); + }; + for prefix_batch in prefix.drain(..) { + writer.write(prefix_batch).await?; + } + writer.write(output).await?; + } else { + prefix.push(output); + } + } + } + + // A candidate whose visible rows were all removed by its delete + // files reads as zero rows and the loop above never runs. Treat it + // the same as a rewriter that dropped every batch — changed with + // no replacement — so the file and its delete files can be + // compacted away instead of being pinned in the table forever. + let fully_removed_by_deletes = Review Comment: The `fully_removed_by_deletes` classification is exactly what I asked for last round and reads correctly — this is the path that lets a stale data+delete pair actually get compacted away. What's missing is a test that drives it through real delete files: every current test reaches the removed-with-no-replacement state via the batch rewriter, not via the delete-file reader returning zero rows. I'd add one that writes a data file plus a position delete covering all its rows, runs a `KeepAll` rewriter, and asserts the file lands in `removed_data_files` with `added_data_files` empty. Without it, a reader change that emits one empty batch instead of zero batches would silently send the file back to `unchanged_data_files` and re-pin the pair forever. ########## crates/iceberg/src/cow_rewrite/mod.rs: ########## @@ -0,0 +1,1390 @@ +// 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 written to replacement files. + /// + /// Rows the rewriter emitted for files that turned out unchanged are not + /// counted, so this always matches the row counts of `added_data_files`. + pub output_rows: u64, + /// Number of input batches that changed, including batches the rewriter + /// dropped entirely (`output: None`) even if it did not flag them. + 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. + /// + /// Files whose visible rows were all removed by delete files are NOT + /// included here: they read as zero rows and are reported in + /// `removed_data_files` with no replacement, so a commit adapter can drop + /// them together with the delete files that reference them. + 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 { + let CowRewriteFile { + old_data_file, + scan_task, + } = file; + // 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 = scan_task.schema_ref(); + let has_delete_files = !scan_task.deletes().is_empty(); + + // 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. + // + // Worst-case footprint: for a file that never changes (or whose + // first change sits at its very end) the prefix holds the entire + // decoded source file in memory. Files are processed sequentially, + // so peak usage is one file at a time, but that can still be several + // GB for a compaction-sized file. A size-capped fallback that starts + // writing the replacement once the buffer crosses a threshold is + // left for follow-up work. + let mut prefix: Vec<RecordBatch> = Vec::new(); + let mut file_changed = false; + let mut file_input_rows = 0_u64; + let mut file_output_rows = 0_u64; + 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(scan_task)])) 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; + file_input_rows += batch.num_rows() as u64; + + let rewrite = rewriter.rewrite_batch(batch)?; + // `output: None` means the batch is fully removed, which is + // itself a change. Derive the effective flag instead of + // trusting every rewriter to keep `changed` consistent with + // `output` — otherwise a `{changed: false, output: None}` + // batch would silently drop its rows while leaving the file + // marked unchanged. + let changed = rewrite.changed || rewrite.output.is_none(); + if changed { + file_changed = true; + result.stats.changed_batches += 1; + } + + if let Some(output) = rewrite.output { Review Comment: This is the same data-loss shape I flagged last round, just from the other side. The `changed || output.is_none()` derivation landed and handles the drop-first/keep-later case, but the prefix only ever gets flushed inside this `if let Some(output)` arm — so when a file accumulates a prefix of `{output: Some, changed: false}` batches and then flips to changed via a later `{output: None, changed: true}` batch, we skip this whole block, exit the loop with `writer == None`, and the post-loop "fully deleted" branch drops the buffered prefix. The kept rows vanish and the file still moves to `removed_data_files`. Concrete trigger with `batch_size=2` and rows `[1, 3, 2, 4]`: batch one is `{Some([1,3]), changed: false}` → prefix, batch two is `{None, changed: true}` → the file flips changed but nothing flushes. Rows 1 and 3 are lost, and `output_rows` still reports 2, which also breaks its documented invariant. I'd flush the prefix and build the writer at the edge where `file_changed` first becomes true, regardless of whether that batch carried output — then `writer == None` reliably means "nothing was ever written" and the fully-deleted branch is correct. A keep-first/drop-second regression test (the mirror of `cow_rewrite_silent_drop_before_change_loses_no_rows`) would lock it. Same condition as last round: I'd want this resolved before we merge. -- 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]
