From 6d6f37a4a91f82512d3f4b7eb69a66ba0c242148 Mon Sep 17 00:00:00 2001 From: kid Date: Fri, 17 Jul 2026 14:45:09 +0800 Subject: [PATCH 1/3] feat: add core COW rewrite primitive Add a core copy-on-write rewrite primitive that plans candidate data files, applies a caller-provided RecordBatch rewriter, writes replacement data files, and returns removed/added/unchanged file lists for future overwrite-style commit actions. - Plan candidates through the existing scan path so manifest pruning and delete-file application stay consistent with normal reads; the row predicate is cleared only for the full-file rewrite input. - Stream rewritten batches to the replacement writer: buffer only the prefix before the first changed batch, then open the writer lazily; unchanged files open no writer, fully-deleted files produce no replacement. - Bind the replacement writer and partition key to the planned snapshot's schema so rewrites survive schema evolution. - Build the replacement parquet writer via ParquetWriterBuilder::from_table_properties so replacement files honor the table's write.parquet.* properties, carry the table's encryption manager so encrypted tables are not downgraded to plaintext, and refuse to run when write.format.default is not parquet. - Keep the planner internal; CowRewriteBuilder is the only public entry and CowRewriteFile fields sit behind accessors. Cover delete-style, update-style, no-op, full-file delete, no-match delete, delete-file planning, writer edge cases, schema evolution, multi-source replacement path uniqueness, and partition-value preservation (including the null partition) end-to-end. Part of #2269. Co-Authored-By: Claude --- crates/iceberg/public-api.txt | 52 + crates/iceberg/src/cow_rewrite/mod.rs | 1261 ++++++++++++++++++++ crates/iceberg/src/cow_rewrite/plan.rs | 158 +++ crates/iceberg/src/cow_rewrite/rewriter.rs | 43 + crates/iceberg/src/cow_rewrite/writer.rs | 353 ++++++ crates/iceberg/src/lib.rs | 1 + crates/iceberg/src/scan/context.rs | 17 + crates/iceberg/src/scan/mod.rs | 54 +- crates/iceberg/src/scan/task.rs | 10 + 9 files changed, 1935 insertions(+), 14 deletions(-) create mode 100644 crates/iceberg/src/cow_rewrite/mod.rs create mode 100644 crates/iceberg/src/cow_rewrite/plan.rs create mode 100644 crates/iceberg/src/cow_rewrite/rewriter.rs create mode 100644 crates/iceberg/src/cow_rewrite/writer.rs diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index e2e78073e9..15b4e75481 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -175,6 +175,58 @@ impl serde_core::ser::Serialize for iceberg::compression::CompressionCodec pub fn iceberg::compression::CompressionCodec::serialize(&self, serializer: S) -> core::result::Result<::Ok, ::Error> impl<'de> serde_core::de::Deserialize<'de> for iceberg::compression::CompressionCodec pub fn iceberg::compression::CompressionCodec::deserialize>(deserializer: D) -> core::result::Result::Error> +pub mod iceberg::cow_rewrite +pub struct iceberg::cow_rewrite::CowBatchRewrite +pub iceberg::cow_rewrite::CowBatchRewrite::changed: bool +pub iceberg::cow_rewrite::CowBatchRewrite::output: core::option::Option +pub struct iceberg::cow_rewrite::CowRewriteBuilder<'a> +impl<'a> iceberg::cow_rewrite::CowRewriteBuilder<'a> +pub fn iceberg::cow_rewrite::CowRewriteBuilder<'a>::new(table: &'a iceberg::table::Table) -> Self +pub async fn iceberg::cow_rewrite::CowRewriteBuilder<'a>::rewrite(self) -> iceberg::Result +pub fn iceberg::cow_rewrite::CowRewriteBuilder<'a>::with_batch_size(self, batch_size: usize) -> Self +pub fn iceberg::cow_rewrite::CowRewriteBuilder<'a>::with_case_sensitive(self, case_sensitive: bool) -> Self +pub fn iceberg::cow_rewrite::CowRewriteBuilder<'a>::with_predicate(self, predicate: iceberg::expr::Predicate) -> Self +pub fn iceberg::cow_rewrite::CowRewriteBuilder<'a>::with_rewriter(self, rewriter: alloc::sync::Arc) -> Self +pub fn iceberg::cow_rewrite::CowRewriteBuilder<'a>::with_snapshot_id(self, snapshot_id: i64) -> Self +pub struct iceberg::cow_rewrite::CowRewriteFile +impl iceberg::cow_rewrite::CowRewriteFile +pub fn iceberg::cow_rewrite::CowRewriteFile::old_data_file(&self) -> &iceberg::spec::DataFile +pub fn iceberg::cow_rewrite::CowRewriteFile::scan_task(&self) -> &iceberg::scan::FileScanTask +impl core::clone::Clone for iceberg::cow_rewrite::CowRewriteFile +pub fn iceberg::cow_rewrite::CowRewriteFile::clone(&self) -> iceberg::cow_rewrite::CowRewriteFile +impl core::fmt::Debug for iceberg::cow_rewrite::CowRewriteFile +pub fn iceberg::cow_rewrite::CowRewriteFile::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct iceberg::cow_rewrite::CowRewriteResult +pub iceberg::cow_rewrite::CowRewriteResult::added_data_files: alloc::vec::Vec +pub iceberg::cow_rewrite::CowRewriteResult::removed_data_files: alloc::vec::Vec +pub iceberg::cow_rewrite::CowRewriteResult::stats: iceberg::cow_rewrite::CowRewriteStats +pub iceberg::cow_rewrite::CowRewriteResult::unchanged_data_files: alloc::vec::Vec +impl iceberg::cow_rewrite::CowRewriteResult +pub fn iceberg::cow_rewrite::CowRewriteResult::has_changes(&self) -> bool +impl core::default::Default for iceberg::cow_rewrite::CowRewriteResult +pub fn iceberg::cow_rewrite::CowRewriteResult::default() -> iceberg::cow_rewrite::CowRewriteResult +impl core::fmt::Debug for iceberg::cow_rewrite::CowRewriteResult +pub fn iceberg::cow_rewrite::CowRewriteResult::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct iceberg::cow_rewrite::CowRewriteStats +pub iceberg::cow_rewrite::CowRewriteStats::candidate_files: usize +pub iceberg::cow_rewrite::CowRewriteStats::changed_batches: u64 +pub iceberg::cow_rewrite::CowRewriteStats::input_rows: u64 +pub iceberg::cow_rewrite::CowRewriteStats::output_rows: u64 +pub iceberg::cow_rewrite::CowRewriteStats::rewritten_files: usize +pub iceberg::cow_rewrite::CowRewriteStats::unchanged_files: usize +impl core::clone::Clone for iceberg::cow_rewrite::CowRewriteStats +pub fn iceberg::cow_rewrite::CowRewriteStats::clone(&self) -> iceberg::cow_rewrite::CowRewriteStats +impl core::cmp::Eq for iceberg::cow_rewrite::CowRewriteStats +impl core::cmp::PartialEq for iceberg::cow_rewrite::CowRewriteStats +pub fn iceberg::cow_rewrite::CowRewriteStats::eq(&self, other: &iceberg::cow_rewrite::CowRewriteStats) -> bool +impl core::default::Default for iceberg::cow_rewrite::CowRewriteStats +pub fn iceberg::cow_rewrite::CowRewriteStats::default() -> iceberg::cow_rewrite::CowRewriteStats +impl core::fmt::Debug for iceberg::cow_rewrite::CowRewriteStats +pub fn iceberg::cow_rewrite::CowRewriteStats::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::Copy for iceberg::cow_rewrite::CowRewriteStats +impl core::marker::StructuralPartialEq for iceberg::cow_rewrite::CowRewriteStats +pub trait iceberg::cow_rewrite::CowBatchRewriter: core::marker::Send + core::marker::Sync +pub fn iceberg::cow_rewrite::CowBatchRewriter::rewrite_batch(&self, batch: arrow_array::record_batch::RecordBatch) -> iceberg::Result pub mod iceberg::encryption pub mod iceberg::encryption::kms pub struct iceberg::encryption::kms::GeneratedKey diff --git a/crates/iceberg/src/cow_rewrite/mod.rs b/crates/iceberg/src/cow_rewrite/mod.rs new file mode 100644 index 0000000000..940be14cdd --- /dev/null +++ b/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 { +//! 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, + /// New data files that should be added by the commit action. + pub added_data_files: Vec, + /// 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, + /// 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, + batch_size: Option, + case_sensitive: bool, + rewriter: Option>, +} + +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) -> Self { + self.rewriter = Some(rewriter); + self + } + + /// Plans, reads, rewrites, and writes replacement data files. + pub async fn rewrite(self) -> Result { + 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 = Vec::new(); + let mut file_changed = false; + let mut writer: Option> = 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 { + 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| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Cannot bind partition spec {} to the planned snapshot schema for COW rewrite", + data_file.partition_spec_id + ), + ) + .with_source(err) + })?; + + Ok(PartitionKey::new( + spec, + schema.clone(), + data_file.partition().clone(), + )) +} + +#[cfg(test)] +mod tests { + use std::collections::{HashMap, HashSet}; + use std::sync::Arc; + + use arrow_array::{Array, ArrayRef, BooleanArray, Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema as ArrowSchema}; + use futures::TryStreamExt; + use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use tempfile::TempDir; + + use crate::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder}; + use crate::io::LocalFsStorageFactory; + use crate::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; + use crate::scan::{FileScanTask, FileScanTaskStream}; + use crate::spec::{ + DataFile, Literal, NestedField, PrimitiveType, Schema, Struct, TableProperties, Transform, + Type, + }; + use crate::table::Table; + use crate::transaction::{AddColumn, ApplyTransactionAction, Transaction}; + use crate::{Catalog, CatalogBuilder, Error, ErrorKind, NamespaceIdent, Result, TableCreation}; + + struct KeepAll; + + impl CowBatchRewriter for KeepAll { + fn rewrite_batch(&self, batch: RecordBatch) -> Result { + Ok(CowBatchRewrite { + output: Some(batch), + changed: false, + }) + } + } + + struct DeleteEvenIds; + + impl CowBatchRewriter for DeleteEvenIds { + fn rewrite_batch(&self, batch: RecordBatch) -> Result { + let ids = batch + .column_by_name("id") + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "missing id column"))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "id must be Int32"))?; + + let keep = + BooleanArray::from_iter((0..ids.len()).map(|row| Some(ids.value(row) % 2 != 0))); + let filtered = arrow_select::filter::filter_record_batch(&batch, &keep) + .map_err(|err| Error::new(ErrorKind::Unexpected, err.to_string()))?; + + Ok(CowBatchRewrite { + changed: filtered.num_rows() != batch.num_rows(), + output: (filtered.num_rows() > 0).then_some(filtered), + }) + } + } + + struct IncrementValueForEvenIds; + + impl CowBatchRewriter for IncrementValueForEvenIds { + fn rewrite_batch(&self, batch: RecordBatch) -> Result { + let ids = batch + .column_by_name("id") + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "missing id column"))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "id must be Int32"))?; + let values = batch + .column_by_name("value") + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "missing value column"))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "value must be Int32"))?; + + let mut changed = false; + let updated_values = Int32Array::from_iter((0..values.len()).map(|row| { + let value = values.value(row); + if ids.value(row) % 2 == 0 { + changed = true; + Some(value + 10) + } else { + Some(value) + } + })); + let output = RecordBatch::try_new(batch.schema(), vec![ + batch.column(0).clone(), + Arc::new(updated_values), + ]) + .map_err(|err| Error::new(ErrorKind::Unexpected, err.to_string()))?; + + Ok(CowBatchRewrite { + output: Some(output), + changed, + }) + } + } + + struct CowRewriteFixture { + _temp_dir: TempDir, + table: Table, + } + + async fn test_table_with_ids(ids: Vec) -> Result { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("cow_rewrite_fixture".to_string()) + .schema(schema) + .build(), + ) + .await?; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + let batch = RecordBatch::try_new(arrow_schema, vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef + ])?; + let data_files = super::writer::write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + None, + futures::stream::iter(vec![Ok(batch)]), + ) + .await?; + + let tx = Transaction::new(&table); + let tx = tx.fast_append().add_data_files(data_files).apply(tx)?; + let table = tx.commit(&catalog).await?; + + Ok(CowRewriteFixture { + _temp_dir: temp_dir, + table, + }) + } + + async fn test_table_with_id_batches(batches: Vec>) -> Result { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("cow_rewrite_fixture".to_string()) + .schema(schema) + .properties(HashMap::from([( + TableProperties::PROPERTY_WRITE_TARGET_FILE_SIZE_BYTES.to_string(), + "1".to_string(), + )])) + .build(), + ) + .await?; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + let input = batches.into_iter().map(|ids| { + Ok(RecordBatch::try_new(arrow_schema.clone(), vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + ])?) + }); + let data_files = super::writer::write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + None, + futures::stream::iter(input), + ) + .await?; + + let tx = Transaction::new(&table); + let tx = tx.fast_append().add_data_files(data_files).apply(tx)?; + let table = tx.commit(&catalog).await?; + + Ok(CowRewriteFixture { + _temp_dir: temp_dir, + table, + }) + } + + async fn test_table_with_id_value_rows(rows: Vec<(i32, i32)>) -> Result { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "value", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("cow_rewrite_fixture".to_string()) + .schema(schema) + .build(), + ) + .await?; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + Field::new("value", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "2".to_string(), + )])), + ])); + let ids = rows.iter().map(|(id, _)| *id).collect::>(); + let values = rows.iter().map(|(_, value)| *value).collect::>(); + let batch = RecordBatch::try_new(arrow_schema, vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(Int32Array::from(values)) as ArrayRef, + ])?; + let data_files = super::writer::write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + None, + futures::stream::iter(vec![Ok(batch)]), + ) + .await?; + + let tx = Transaction::new(&table); + let tx = tx.fast_append().add_data_files(data_files).apply(tx)?; + let table = tx.commit(&catalog).await?; + + Ok(CowRewriteFixture { + _temp_dir: temp_dir, + table, + }) + } + + async fn read_ids(table: &Table, files: &[DataFile]) -> Result> { + let schema = table.metadata().current_schema().clone(); + let project_field_ids = schema + .as_struct() + .fields() + .iter() + .map(|field| field.id) + .collect::>(); + let tasks = files + .iter() + .map(|data_file| { + FileScanTask::builder() + .with_file_size_in_bytes(data_file.file_size_in_bytes()) + .with_start(0) + .with_length(data_file.file_size_in_bytes()) + .with_record_count(Some(data_file.record_count())) + .with_data_file_path(data_file.file_path().to_string()) + .with_data_file_format(data_file.file_format()) + .with_schema(schema.clone()) + .with_project_field_ids(project_field_ids.clone()) + .with_case_sensitive(true) + .build() + }) + .collect::>(); + let task_stream = Box::pin(futures::stream::iter(tasks)) as FileScanTaskStream; + let batches = table + .reader_builder() + .build() + .read(task_stream)? + .stream() + .try_collect::>() + .await?; + + let mut ids = Vec::new(); + for batch in batches { + let column = batch + .column_by_name("id") + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "missing id column"))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "id must be Int32"))?; + ids.extend((0..column.len()).map(|row| column.value(row))); + } + ids.sort_unstable(); + Ok(ids) + } + + async fn read_ids_and_values(table: &Table, files: &[DataFile]) -> Result> { + let schema = table.metadata().current_schema().clone(); + let project_field_ids = schema + .as_struct() + .fields() + .iter() + .map(|field| field.id) + .collect::>(); + let tasks = files + .iter() + .map(|data_file| { + FileScanTask::builder() + .with_file_size_in_bytes(data_file.file_size_in_bytes()) + .with_start(0) + .with_length(data_file.file_size_in_bytes()) + .with_record_count(Some(data_file.record_count())) + .with_data_file_path(data_file.file_path().to_string()) + .with_data_file_format(data_file.file_format()) + .with_schema(schema.clone()) + .with_project_field_ids(project_field_ids.clone()) + .with_case_sensitive(true) + .build() + }) + .collect::>(); + let task_stream = Box::pin(futures::stream::iter(tasks)) as FileScanTaskStream; + let batches = table + .reader_builder() + .build() + .read(task_stream)? + .stream() + .try_collect::>() + .await?; + + let mut rows = Vec::new(); + for batch in batches { + let ids = batch + .column_by_name("id") + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "missing id column"))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "id must be Int32"))?; + let values = batch + .column_by_name("value") + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "missing value column"))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "value must be Int32"))?; + rows.extend((0..ids.len()).map(|row| (ids.value(row), values.value(row)))); + } + rows.sort_unstable_by_key(|(id, _)| *id); + Ok(rows) + } + + #[tokio::test] + async fn cow_rewrite_keep_all_produces_no_changes() -> Result<()> { + let fixture = test_table_with_ids(vec![1, 2, 3]).await?; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_rewriter(Arc::new(KeepAll)) + .rewrite() + .await?; + + assert!(!result.has_changes()); + assert_eq!(result.removed_data_files.len(), 0); + assert_eq!(result.added_data_files.len(), 0); + assert_eq!(result.unchanged_data_files.len(), 1); + assert_eq!(result.stats.candidate_files, 1); + assert_eq!(result.stats.unchanged_files, 1); + assert_eq!(result.stats.input_rows, 3); + assert_eq!(result.stats.output_rows, 3); + + Ok(()) + } + + #[tokio::test] + async fn cow_rewrite_requires_rewriter() -> Result<()> { + let fixture = test_table_with_ids(vec![1]).await?; + + let err = CowRewriteBuilder::new(&fixture.table) + .rewrite() + .await + .expect_err("missing rewriter should fail"); + + assert_eq!(err.kind(), ErrorKind::PreconditionFailed); + + Ok(()) + } + + #[tokio::test] + async fn cow_rewrite_delete_rows_removes_old_file_and_adds_replacement() -> Result<()> { + let fixture = test_table_with_ids(vec![1, 2, 3, 4]).await?; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_rewriter(Arc::new(DeleteEvenIds)) + .rewrite() + .await?; + + assert!(result.has_changes()); + assert_eq!(result.removed_data_files.len(), 1); + assert_eq!(result.added_data_files.len(), 1); + assert_eq!(result.stats.input_rows, 4); + assert_eq!(result.stats.output_rows, 2); + + let ids = read_ids(&fixture.table, &result.added_data_files).await?; + assert_eq!(ids, vec![1, 3]); + + Ok(()) + } + + #[tokio::test] + async fn cow_rewrite_update_rows_rewrites_file_with_updated_values() -> Result<()> { + let fixture = + test_table_with_id_value_rows(vec![(1, 10), (2, 20), (3, 30), (4, 40)]).await?; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_rewriter(Arc::new(IncrementValueForEvenIds)) + .rewrite() + .await?; + + assert_eq!(result.removed_data_files.len(), 1); + assert_eq!(result.added_data_files.len(), 1); + assert_eq!(result.stats.input_rows, 4); + assert_eq!(result.stats.output_rows, 4); + + let rows = read_ids_and_values(&fixture.table, &result.added_data_files).await?; + assert_eq!(rows, vec![(1, 10), (2, 30), (3, 30), (4, 50)]); + + Ok(()) + } + + #[tokio::test] + async fn cow_rewrite_full_file_delete_removes_old_file_without_replacement() -> Result<()> { + let fixture = test_table_with_ids(vec![2, 4]).await?; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_rewriter(Arc::new(DeleteEvenIds)) + .rewrite() + .await?; + + assert!(result.has_changes()); + assert_eq!(result.removed_data_files.len(), 1); + assert_eq!(result.added_data_files.len(), 0); + assert_eq!(result.stats.input_rows, 2); + assert_eq!(result.stats.output_rows, 0); + + Ok(()) + } + + #[tokio::test] + async fn cow_rewrite_delete_no_matching_rows_keeps_old_file() -> Result<()> { + let fixture = test_table_with_ids(vec![1, 3]).await?; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_rewriter(Arc::new(DeleteEvenIds)) + .rewrite() + .await?; + + assert!(!result.has_changes()); + assert_eq!(result.removed_data_files.len(), 0); + assert_eq!(result.added_data_files.len(), 0); + assert_eq!(result.unchanged_data_files.len(), 1); + assert_eq!(result.stats.input_rows, 2); + assert_eq!(result.stats.output_rows, 2); + + Ok(()) + } + + #[tokio::test] + async fn cow_rewrite_uses_unique_replacement_paths_for_multiple_source_files() -> Result<()> { + let fixture = test_table_with_id_batches(vec![vec![1, 2], vec![3, 4]]).await?; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_rewriter(Arc::new(DeleteEvenIds)) + .rewrite() + .await?; + + assert_eq!(result.stats.candidate_files, 2); + assert_eq!(result.removed_data_files.len(), 2); + assert_eq!(result.added_data_files.len(), 2); + + let added_paths = result + .added_data_files + .iter() + .map(|file| file.file_path().to_string()) + .collect::>(); + let removed_paths = result + .removed_data_files + .iter() + .map(|file| file.file_path().to_string()) + .collect::>(); + + assert_eq!(added_paths.len(), result.added_data_files.len()); + assert!(added_paths.is_disjoint(&removed_paths)); + + let ids = read_ids(&fixture.table, &result.added_data_files).await?; + assert_eq!(ids, vec![1, 3]); + + Ok(()) + } + + #[test] + fn cow_batch_rewriter_is_object_safe() { + let _rewriter: Arc = Arc::new(KeepAll); + } + + #[test] + fn cow_rewrite_result_reports_no_changes() { + let result = crate::cow_rewrite::CowRewriteResult { + removed_data_files: vec![], + added_data_files: vec![], + unchanged_data_files: vec![], + stats: crate::cow_rewrite::CowRewriteStats { + candidate_files: 0, + rewritten_files: 0, + unchanged_files: 0, + input_rows: 0, + output_rows: 0, + changed_batches: 0, + }, + }; + + assert!(!result.has_changes()); + assert_eq!(result.stats.candidate_files, 0); + } + + /// `DeleteIfModThree` keeps rows whose `id` is not divisible by 3. + struct DeleteIfModThree; + + impl CowBatchRewriter for DeleteIfModThree { + fn rewrite_batch(&self, batch: RecordBatch) -> Result { + let ids = batch + .column_by_name("id") + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "missing id column"))? + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "id must be Int32"))?; + + let keep = + BooleanArray::from_iter((0..ids.len()).map(|row| Some(ids.value(row) % 3 != 0))); + let filtered = arrow_select::filter::filter_record_batch(&batch, &keep) + .map_err(|err| Error::new(ErrorKind::Unexpected, err.to_string()))?; + + Ok(CowBatchRewrite { + changed: filtered.num_rows() != batch.num_rows(), + output: (filtered.num_rows() > 0).then_some(filtered), + }) + } + } + + /// After adding an optional column (`value`) to the schema, a COW rewrite + /// of the pre-existing data must not fail. The replacement writer must use + /// the schema the batches were read in (the snapshot schema), not the + /// table's evolved current schema, otherwise the parquet writer rejects + /// the input batch whose column set does not include `value`. + #[tokio::test] + async fn cow_rewrite_after_optional_column_add() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog: Arc = Arc::new( + MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?, + ); + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("evolved".to_string()) + .schema(schema) + .build(), + ) + .await?; + + // Write the original data file with the {id}-only schema. + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(Int32Array::from(vec![ + 1, 2, 3, 4, + ])) as ArrayRef])?; + let data_files = super::writer::write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + None, + futures::stream::iter(vec![Ok(batch)]), + ) + .await?; + let tx = Transaction::new(&table); + let tx = tx.fast_append().add_data_files(data_files).apply(tx)?; + let table = tx.commit(&*catalog).await?; + + // Evolve the schema: add an optional `value` column. + let tx = Transaction::new(&table); + let tx = tx + .update_schema() + .add_column(AddColumn::optional( + "value", + Type::Primitive(PrimitiveType::Int), + )) + .apply(tx)?; + let table = tx.commit(&*catalog).await?; + + // The current schema now has {id, value} but the current snapshot's + // schema is still the original {id} schema. + let current_snapshot = table.metadata().current_snapshot().unwrap(); + assert_ne!( + table.metadata().current_schema_id(), + current_snapshot.schema_id().unwrap() + ); + + // Before the fix this fails when the parquet writer rejects the {id}-only + // batch while configured with the {id, value} current schema. + let result = CowRewriteBuilder::new(&table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_rewriter(Arc::new(DeleteIfModThree)) + .rewrite() + .await?; + + // id 3 is divisible by 3 and is dropped; remaining ids are 1, 2, 4. + assert!(result.has_changes()); + assert_eq!(result.removed_data_files.len(), 1); + assert_eq!(result.added_data_files.len(), 1); + + let ids = read_ids(&table, &result.added_data_files).await?; + assert_eq!(ids, vec![1, 2, 4]); + + Ok(()) + } + + /// Builds a two-file table whose files have disjoint id ranges so predicate + /// metrics pruning can exclude one file entirely. + async fn two_file_table_disjoint_ids() -> Result { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("disjoint".to_string()) + .schema(schema) + .properties(HashMap::from([( + TableProperties::PROPERTY_WRITE_TARGET_FILE_SIZE_BYTES.to_string(), + "1".to_string(), + )])) + .build(), + ) + .await?; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + // First batch ids 1..=4, second batch ids 100..=103. Force a split via + // tiny target file size so each batch lands in its own data file. + let input = vec![vec![1, 2, 3, 4], vec![100, 101, 102, 103]] + .into_iter() + .map(|ids| { + Ok(RecordBatch::try_new(arrow_schema.clone(), vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + ])?) + }); + let data_files = super::writer::write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + None, + futures::stream::iter(input), + ) + .await?; + + let tx = Transaction::new(&table); + let tx = tx.fast_append().add_data_files(data_files).apply(tx)?; + let table = tx.commit(&catalog).await?; + + Ok(CowRewriteFixture { + _temp_dir: temp_dir, + table, + }) + } + + /// Predicate-based planning must exclude files whose metrics cannot match + /// the predicate. With `id > 50`, the file holding ids 1..=4 must never be + /// planned as a candidate. + #[tokio::test] + async fn cow_rewrite_predicate_prunes_non_matching_files() -> Result<()> { + let fixture = two_file_table_disjoint_ids().await?; + + let predicate = crate::expr::Reference::new("id").greater_than(crate::spec::Datum::int(50)); + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(predicate) + .with_rewriter(Arc::new(KeepAll)) + .rewrite() + .await?; + + // Only the high-id file should be a candidate. + assert_eq!(result.stats.candidate_files, 1); + // KeepAll does not change anything, so the single candidate is unchanged. + assert_eq!(result.removed_data_files.len(), 0); + assert_eq!(result.added_data_files.len(), 0); + assert_eq!(result.unchanged_data_files.len(), 1); + + // The surviving candidate must be the high-id file (ids 100..=103). + let ids = read_ids(&fixture.table, &result.unchanged_data_files).await?; + assert_eq!(ids, vec![100, 101, 102, 103]); + + Ok(()) + } + + /// Rewrites every batch without altering rows, so every candidate file is + /// replaced. Useful for exercising the replacement-writing path itself. + struct TouchEveryBatch; + + impl CowBatchRewriter for TouchEveryBatch { + fn rewrite_batch(&self, batch: RecordBatch) -> Result { + Ok(CowBatchRewrite { + output: Some(batch), + changed: true, + }) + } + } + + /// An end-to-end rewrite on an identity-partitioned table must preserve + /// each source file's partition: the replacement file carries the same + /// partition values and partition spec id, and lands under the same + /// partition directory — including the null partition. + #[tokio::test] + async fn cow_rewrite_preserves_partitions_end_to_end() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "value", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let partition_spec = crate::spec::PartitionSpec::builder(Arc::new(schema.clone())) + .add_partition_field("value", "value", Transform::Identity)? + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("partitioned".to_string()) + .schema(schema) + .partition_spec(partition_spec.into_unbound()) + .build(), + ) + .await?; + + // Two source files: ids 1, 2 in the `value=1` partition and ids 3, 4 + // in the null partition. + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + Field::new("value", DataType::Int32, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "2".to_string(), + )])), + ])); + let current_schema = table.metadata().current_schema().clone(); + let default_spec = table.metadata().default_partition_spec().as_ref().clone(); + + let batch_one = RecordBatch::try_new(arrow_schema.clone(), vec![ + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + Arc::new(Int32Array::from(vec![1, 1])), + ])?; + let one_partition = crate::spec::PartitionKey::new( + default_spec.clone(), + current_schema.clone(), + Struct::from_iter([Some(Literal::int(1))]), + ); + let files_one = super::writer::write_replacement_batches( + &table, + current_schema.clone(), + Some(one_partition), + futures::stream::iter(vec![Ok(batch_one)]), + ) + .await?; + + let batch_null = RecordBatch::try_new(arrow_schema.clone(), vec![ + Arc::new(Int32Array::from(vec![3, 4])) as ArrayRef, + Arc::new(Int32Array::from(vec![None::, None])), + ])?; + let null_partition = crate::spec::PartitionKey::new( + default_spec.clone(), + current_schema.clone(), + Struct::from_iter([None::]), + ); + let files_null = super::writer::write_replacement_batches( + &table, + current_schema.clone(), + Some(null_partition), + futures::stream::iter(vec![Ok(batch_null)]), + ) + .await?; + + let tx = Transaction::new(&table); + let tx = tx + .fast_append() + .add_data_files([files_one, files_null].concat()) + .apply(tx)?; + let table = tx.commit(&catalog).await?; + + let result = CowRewriteBuilder::new(&table) + .with_rewriter(Arc::new(TouchEveryBatch)) + .rewrite() + .await?; + + assert_eq!(result.stats.candidate_files, 2); + assert_eq!(result.stats.rewritten_files, 2); + assert_eq!(result.removed_data_files.len(), 2); + assert_eq!(result.added_data_files.len(), 2); + assert!(result.unchanged_data_files.is_empty()); + + // Partition values survive the rewrite unchanged. + let one_struct = Struct::from_iter([Some(Literal::int(1))]); + let null_struct = Struct::from_iter([None::]); + let mut removed_partitions = result + .removed_data_files + .iter() + .map(|file| file.partition().clone()) + .collect::>(); + let mut added_partitions = result + .added_data_files + .iter() + .map(|file| file.partition().clone()) + .collect::>(); + removed_partitions.sort_by_key(|partition| partition.is_null_at_index(0)); + added_partitions.sort_by_key(|partition| partition.is_null_at_index(0)); + assert_eq!(removed_partitions, added_partitions); + assert_eq!(added_partitions[0], one_struct); + assert_eq!(added_partitions[1], null_struct); + + for file in &result.added_data_files { + assert_eq!(file.partition_spec_id, 0); + if file.partition() == &one_struct { + assert!(file.file_path().contains("value=1"), "{}", file.file_path()); + } else { + assert!( + file.file_path().contains("value=null"), + "{}", + file.file_path() + ); + } + } + + // Replacement content matches the source rows. + let ids = read_ids(&table, &result.added_data_files).await?; + assert_eq!(ids, vec![1, 2, 3, 4]); + + Ok(()) + } +} diff --git a/crates/iceberg/src/cow_rewrite/plan.rs b/crates/iceberg/src/cow_rewrite/plan.rs new file mode 100644 index 0000000000..fbaa11e217 --- /dev/null +++ b/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 { + /// Original data file from the manifest entry. + pub(crate) old_data_file: DataFile, + /// Full-read task for visible rows in the original file. + pub(crate) scan_task: FileScanTask, +} + +impl CowRewriteFile { + /// Original data file from the manifest entry. + pub fn old_data_file(&self) -> &DataFile { + &self.old_data_file + } + + /// Full-read task for visible rows in the original file. + pub fn scan_task(&self) -> &FileScanTask { + &self.scan_task + } +} + +/// Plan data files that may need copy-on-write rewrite. +pub(crate) async fn plan_cow_rewrite_files( + table: &crate::table::Table, + predicate: Option, + snapshot_id: Option, + case_sensitive: bool, +) -> crate::Result> { + let mut scan_builder = table + .scan() + .select_all() + .with_case_sensitive(case_sensitive); + + if let Some(predicate) = predicate { + scan_builder = scan_builder.with_filter(predicate); + } + + if let Some(snapshot_id) = snapshot_id { + scan_builder = scan_builder.snapshot_id(snapshot_id); + } + + scan_builder + .build()? + .plan_cow_rewrite_files() + .await? + .try_collect() + .await +} + +#[cfg(test)] +mod tests { + use futures::TryStreamExt; + + use crate::Result; + use crate::expr::Predicate; + use crate::scan::tests::TableTestFixture; + + #[tokio::test] + async fn cow_planner_returns_old_file_and_full_read_task() -> Result<()> { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + + let mut files = + super::plan_cow_rewrite_files(&fixture.table, Some(Predicate::AlwaysTrue), None, true) + .await?; + + assert_eq!(files.len(), 2); + + files.sort_by_key(|file| file.old_data_file.file_path().to_string()); + assert_eq!( + files[0].old_data_file.file_path(), + format!("{}/1.parquet", fixture.table_location) + ); + assert_eq!( + files[1].old_data_file.file_path(), + format!("{}/3.parquet", fixture.table_location) + ); + + for file in files { + assert_eq!( + file.old_data_file.file_path(), + file.scan_task.data_file_path() + ); + assert!(file.scan_task.predicate().is_none()); + assert_eq!( + Some(file.old_data_file.record_count()), + file.scan_task.record_count() + ); + assert_eq!(0, file.scan_task.start()); + assert_eq!( + file.old_data_file.file_size_in_bytes(), + file.scan_task.length() + ); + } + + Ok(()) + } + + #[tokio::test] + async fn cow_planner_preserves_delete_files() -> Result<()> { + let mut fixture = TableTestFixture::new(); + fixture.setup_deadlock_manifests().await; + + let scan = fixture + .table + .scan() + .select_all() + .with_concurrency_limit(1) + .build()?; + + let files = tokio::time::timeout(std::time::Duration::from_secs(5), async { + scan.plan_cow_rewrite_files() + .await? + .try_collect::>() + .await + }) + .await + .expect("COW planning should not deadlock")?; + + assert_eq!(files.len(), 10); + for file in files { + assert_eq!(file.scan_task.deletes().len(), 1); + assert_eq!( + file.scan_task.deletes()[0].file_path, + format!("{}/del.parquet", fixture.table_location) + ); + } + + Ok(()) + } +} diff --git a/crates/iceberg/src/cow_rewrite/rewriter.rs b/crates/iceberg/src/cow_rewrite/rewriter.rs new file mode 100644 index 0000000000..04d9175e6d --- /dev/null +++ b/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, + /// 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; +} diff --git a/crates/iceberg/src/cow_rewrite/writer.rs b/crates/iceberg/src/cow_rewrite/writer.rs new file mode 100644 index 0000000000..a1bfcd129a --- /dev/null +++ b/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, +) -> Result> { + let location_generator = DefaultLocationGenerator::new(table.metadata())?; + let file_name_generator = DefaultFileNameGenerator::new( + format!("cow-rewrite-{}", Uuid::now_v7()), + None, + DataFileFormat::Parquet, + ); + let table_props = table.metadata().table_properties(); + // The replacement writer only produces parquet today; refuse to run on a + // table configured for another format instead of silently writing parquet + // against `write.format.default`. + let write_format_default = table_props.write_format_default()?; + if DataFileFormat::from_str(&write_format_default)? != DataFileFormat::Parquet { + return Err(crate::Error::new( + crate::ErrorKind::FeatureUnsupported, + format!("File format {write_format_default} is not supported for COW rewrite yet!"), + )); + } + let mut parquet_builder = + ParquetWriterBuilder::from_table_properties(&table_props, write_schema)?; + // Replacement files must not silently downgrade an encrypted table to + // plaintext: source rows were read through decryption, so re-encrypt the + // rewritten rows when the table carries an encryption manager. + if let Some(encryption_manager) = table.encryption_manager() { + parquet_builder = parquet_builder.with_encryption_manager(encryption_manager.clone()); + } + let rolling_builder = RollingFileWriterBuilder::new( + parquet_builder, + table_props.write_target_file_size_bytes()?, + table.file_io().clone(), + location_generator, + file_name_generator, + ); + let data_writer_builder = DataFileWriterBuilder::new(rolling_builder); + Ok(Box::new(data_writer_builder.build(partition_key).await?)) +} + +/// Writes replacement record batches as Iceberg data files. +/// +/// This is a convenience wrapper around [`build_replacement_writer`] that +/// consumes a stream end-to-end. Empty input or all-zero-row batches produce no +/// data files. +#[cfg(test)] +pub(crate) async fn write_replacement_batches( + table: &Table, + write_schema: SchemaRef, + partition_key: Option, + mut batches: S, +) -> Result> +where + S: futures::Stream> + Unpin, +{ + use futures::TryStreamExt as _; + + let mut writer = build_replacement_writer(table, write_schema, partition_key).await?; + + let mut wrote_rows = false; + while let Some(batch) = batches.try_next().await? { + if batch.num_rows() == 0 { + continue; + } + wrote_rows = true; + writer.write(batch).await?; + } + + if !wrote_rows { + return Ok(vec![]); + } + + writer.close().await +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use arrow_array::{ArrayRef, Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema as ArrowSchema}; + use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + use tempfile::TempDir; + + use super::write_replacement_batches; + use crate::io::LocalFsStorageFactory; + use crate::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; + use crate::spec::{ + DataContentType, NestedField, PartitionKey, PrimitiveType, Schema, Struct, Transform, Type, + }; + use crate::{Catalog, CatalogBuilder, NamespaceIdent, Result, TableCreation}; + + #[tokio::test] + async fn cow_replacement_writer_outputs_data_files() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("cow_writer".to_string()) + .schema(schema) + .build(), + ) + .await?; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(Int32Array::from(vec![ + 1, 2, 3, + ])) as ArrayRef])?; + + let data_files = write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + None, + futures::stream::iter(vec![Ok(batch)]), + ) + .await?; + + assert_eq!(data_files.len(), 1); + assert_eq!(data_files[0].record_count(), 3); + assert_eq!(data_files[0].content_type(), DataContentType::Data); + + Ok(()) + } + + #[tokio::test] + async fn cow_replacement_writer_skips_empty_input() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("cow_empty_writer".to_string()) + .schema(schema) + .build(), + ) + .await?; + + let data_files = write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + None, + futures::stream::iter(vec![]), + ) + .await?; + + assert!(data_files.is_empty()); + + Ok(()) + } + + #[tokio::test] + async fn cow_replacement_writer_skips_zero_row_batches() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("cow_zero_row_writer".to_string()) + .schema(schema) + .build(), + ) + .await?; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + let batch = + RecordBatch::try_new(arrow_schema, vec![ + Arc::new(Int32Array::from(Vec::::new())) as ArrayRef, + ])?; + + let data_files = write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + None, + futures::stream::iter(vec![Ok(batch)]), + ) + .await?; + + assert!(data_files.is_empty()); + + Ok(()) + } + + #[tokio::test] + async fn cow_replacement_writer_preserves_partition_key() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let partition_spec = crate::spec::PartitionSpec::builder(Arc::new(schema.clone())) + .add_partition_field("id", "id", Transform::Identity)? + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("cow_partitioned_writer".to_string()) + .schema(schema) + .partition_spec(partition_spec.clone().into_unbound()) + .build(), + ) + .await?; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(Int32Array::from(vec![ + 1, 1, 1, + ])) as ArrayRef])?; + let partition_key = PartitionKey::new( + table.metadata().default_partition_spec().as_ref().clone(), + table.metadata().current_schema().clone(), + Struct::from_iter([Some(crate::spec::Literal::int(1))]), + ); + + let data_files = write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + Some(partition_key), + futures::stream::iter(vec![Ok(batch)]), + ) + .await?; + + assert_eq!(data_files.len(), 1); + assert_eq!(data_files[0].partition_spec_id, 0); + assert_eq!( + data_files[0].partition(), + &Struct::from_iter([Some(crate::spec::Literal::int(1))]) + ); + + Ok(()) + } +} diff --git a/crates/iceberg/src/lib.rs b/crates/iceberg/src/lib.rs index 94641b665d..9350b3d478 100644 --- a/crates/iceberg/src/lib.rs +++ b/crates/iceberg/src/lib.rs @@ -83,6 +83,7 @@ pub mod table; mod avro; pub mod cache; pub mod compression; +pub mod cow_rewrite; pub mod io; pub mod spec; diff --git a/crates/iceberg/src/scan/context.rs b/crates/iceberg/src/scan/context.rs index 4635325fbf..d17835afa5 100644 --- a/crates/iceberg/src/scan/context.rs +++ b/crates/iceberg/src/scan/context.rs @@ -172,6 +172,23 @@ impl ManifestEntryContext { .with_sort_order(self.sort_order) .build() } + + /// Consume this `ManifestEntryContext`, returning a COW rewrite file candidate. + /// + /// The scan task keeps everything planned for a normal read (projection, + /// delete files, partition metadata, ...) except the row predicate: the + /// candidate was already selected by that predicate during planning, and + /// the rewrite must read the whole file. + pub(crate) async fn into_cow_rewrite_file(self) -> Result { + let old_data_file = self.manifest_entry.data_file().clone(); + let mut scan_task = self.into_file_scan_task().await?; + scan_task.clear_predicate(); + + Ok(crate::cow_rewrite::CowRewriteFile { + old_data_file, + scan_task, + }) + } } /// PlanContext wraps a [`SnapshotRef`] alongside all the other diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index 7b39189bfb..64cd081816 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -24,6 +24,7 @@ use context::*; mod task; use std::collections::HashMap; +use std::future::Future; use std::sync::Arc; use arrow_array::RecordBatch; @@ -397,6 +398,25 @@ pub struct TableScan { impl TableScan { /// Returns a stream of [`FileScanTask`]s. pub async fn plan_files(&self) -> Result { + self.plan_data_files(|ctx| ctx.into_file_scan_task()).await + } + + pub(crate) async fn plan_cow_rewrite_files( + &self, + ) -> Result>> { + self.plan_data_files(|ctx| ctx.into_cow_rewrite_file()) + .await + } + + async fn plan_data_files( + &self, + build_result: F, + ) -> Result>> + where + T: Send + 'static, + F: Fn(ManifestEntryContext) -> Fut + Copy + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { let Some(plan_context) = self.plan_context.as_ref() else { return Ok(Box::pin(futures::stream::empty())); }; @@ -410,8 +430,8 @@ impl TableScan { let (manifest_entry_delete_ctx_tx, manifest_entry_delete_ctx_rx) = channel(concurrency_limit_manifest_files); - // used to stream the results back to the caller - let (file_scan_task_tx, file_scan_task_rx) = channel(concurrency_limit_manifest_entries); + // used to stream the planned data file results back to the caller + let (planned_file_tx, planned_file_rx) = channel(concurrency_limit_manifest_entries); let (delete_file_idx, delete_file_tx) = DeleteFileIndex::new(self.runtime.clone()); @@ -427,9 +447,9 @@ impl TableScan { manifest_entry_delete_ctx_tx, )?; - let mut channel_for_manifest_error = file_scan_task_tx.clone(); - let mut channel_for_data_manifest_entry_error = file_scan_task_tx.clone(); - let mut channel_for_delete_manifest_entry_error = file_scan_task_tx.clone(); + let mut channel_for_manifest_error = planned_file_tx.clone(); + let mut channel_for_data_manifest_entry_error = planned_file_tx.clone(); + let mut channel_for_delete_manifest_entry_error = planned_file_tx.clone(); let rt = self.runtime.clone(); @@ -486,7 +506,7 @@ impl TableScan { let rt_inner = rt.clone(); rt.cpu().spawn(async move { let result = manifest_entry_data_ctx_rx - .map(|me_ctx| Ok((me_ctx, file_scan_task_tx.clone()))) + .map(|me_ctx| Ok((me_ctx, planned_file_tx.clone()))) .try_for_each_concurrent( concurrency_limit_manifest_entries, |(manifest_entry_context, tx)| { @@ -498,6 +518,7 @@ impl TableScan { Self::process_data_manifest_entry( manifest_entry_context, tx, + build_result, ) .await }) @@ -513,7 +534,7 @@ impl TableScan { }); } - Ok(file_scan_task_rx.boxed()) + Ok(planned_file_rx.boxed()) } /// Returns an [`ArrowRecordBatchStream`]. @@ -545,10 +566,15 @@ impl TableScan { self.plan_context.as_ref().map(|x| &x.snapshot) } - async fn process_data_manifest_entry( + async fn process_data_manifest_entry( manifest_entry_context: ManifestEntryContext, - mut file_scan_task_tx: Sender>, - ) -> Result<()> { + mut planned_file_tx: Sender>, + build_result: F, + ) -> Result<()> + where + F: Fn(ManifestEntryContext) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { // skip processing this manifest entry if it has been marked as deleted if !manifest_entry_context.manifest_entry.is_alive() { return Ok(()); @@ -593,10 +619,10 @@ impl TableScan { } // congratulations! the manifest entry has made its way through the - // entire plan without getting filtered out. Create a corresponding - // FileScanTask and push it to the result stream - file_scan_task_tx - .send(Ok(manifest_entry_context.into_file_scan_task().await?)) + // entire plan without getting filtered out. Build the planned file before + // sending so delete-file lookup preserves the original scan timing. + planned_file_tx + .send(Ok(build_result(manifest_entry_context).await?)) .await?; Ok(()) diff --git a/crates/iceberg/src/scan/task.rs b/crates/iceberg/src/scan/task.rs index 96a322e3d4..97ca7e5181 100644 --- a/crates/iceberg/src/scan/task.rs +++ b/crates/iceberg/src/scan/task.rs @@ -213,6 +213,16 @@ impl FileScanTask { self.predicate.as_ref() } + /// Clears the row predicate of this file scan task. + /// + /// The COW rewrite path uses this after candidate selection: candidates are + /// chosen with the predicate during planning, but each chosen file must then + /// be read in full so the rewrite sees every surviving row. Delete-file + /// application is unaffected. + pub(crate) fn clear_predicate(&mut self) { + self.predicate = None; + } + /// Returns the delete files that may need to be applied to the data file. pub fn deletes(&self) -> &[FileScanTaskDeleteFile] { &self.deletes From 56d96828442ae916693e033552d1d79411f96ae4 Mon Sep 17 00:00:00 2001 From: kid Date: Wed, 16 Sep 2026 17:17:00 +0800 Subject: [PATCH 2/3] fix: address COW rewrite review comments - derive the effective change flag as `changed || output.is_none()` so a rewriter returning {changed: false, output: None} cannot silently drop rows while leaving the file marked unchanged; add regression tests - treat a candidate whose visible rows are all removed by delete files as removed with no replacement, so it can be compacted away with its delete files instead of landing in unchanged_data_files forever - honor write.object-storage.enabled when building replacement writers (add the missing table property) so object-storage-layout tables keep hash-entropy paths - count output_rows only for rows actually written to replacement files - derive Debug for CowBatchRewrite - document the sync rewriter contract, the read-only CowRewriteFile surface, and the prefix buffer's whole-file worst case - drop avoidable FileScanTask/DataFile clones and replace the writer expect() with a let-else invariant --- crates/iceberg/public-api.txt | 5 + crates/iceberg/src/cow_rewrite/mod.rs | 173 +++++++++++++++++--- crates/iceberg/src/cow_rewrite/plan.rs | 11 +- crates/iceberg/src/cow_rewrite/rewriter.rs | 13 ++ crates/iceberg/src/cow_rewrite/writer.rs | 108 +++++++++++- crates/iceberg/src/spec/table_properties.rs | 12 ++ 6 files changed, 289 insertions(+), 33 deletions(-) diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 15b4e75481..977ff5214a 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -179,6 +179,8 @@ pub mod iceberg::cow_rewrite pub struct iceberg::cow_rewrite::CowBatchRewrite pub iceberg::cow_rewrite::CowBatchRewrite::changed: bool pub iceberg::cow_rewrite::CowBatchRewrite::output: core::option::Option +impl core::fmt::Debug for iceberg::cow_rewrite::CowBatchRewrite +pub fn iceberg::cow_rewrite::CowBatchRewrite::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::cow_rewrite::CowRewriteBuilder<'a> impl<'a> iceberg::cow_rewrite::CowRewriteBuilder<'a> pub fn iceberg::cow_rewrite::CowRewriteBuilder<'a>::new(table: &'a iceberg::table::Table) -> Self @@ -2920,6 +2922,8 @@ pub const iceberg::spec::TableProperties<'_>::PROPERTY_UUID: &'static str pub const iceberg::spec::TableProperties<'_>::PROPERTY_WRITE_DATA_LOCATION: &'static str pub const iceberg::spec::TableProperties<'_>::PROPERTY_WRITE_FOLDER_STORAGE_LOCATION: &'static str pub const iceberg::spec::TableProperties<'_>::PROPERTY_WRITE_METADATA_PATH: &'static str +pub const iceberg::spec::TableProperties<'_>::PROPERTY_WRITE_OBJECT_STORAGE_ENABLED: &'static str +pub const iceberg::spec::TableProperties<'_>::PROPERTY_WRITE_OBJECT_STORAGE_ENABLED_DEFAULT: bool pub const iceberg::spec::TableProperties<'_>::PROPERTY_WRITE_OBJECT_STORAGE_LOCATION: &'static str pub const iceberg::spec::TableProperties<'_>::PROPERTY_WRITE_OBJECT_STORAGE_PARTITIONED_PATHS: &'static str pub const iceberg::spec::TableProperties<'_>::PROPERTY_WRITE_OBJECT_STORAGE_PARTITIONED_PATHS_DEFAULT: bool @@ -2957,6 +2961,7 @@ pub fn iceberg::spec::TableProperties<'properties>::write_datafusion_fanout_enab pub fn iceberg::spec::TableProperties<'properties>::write_folder_storage_location(&self) -> iceberg::Result> pub fn iceberg::spec::TableProperties<'properties>::write_format_default(&self) -> iceberg::Result pub fn iceberg::spec::TableProperties<'properties>::write_metadata_path(&self) -> iceberg::Result> +pub fn iceberg::spec::TableProperties<'properties>::write_object_storage_enabled(&self) -> iceberg::Result pub fn iceberg::spec::TableProperties<'properties>::write_object_storage_location(&self) -> iceberg::Result> pub fn iceberg::spec::TableProperties<'properties>::write_object_storage_partitioned_paths(&self) -> iceberg::Result pub fn iceberg::spec::TableProperties<'properties>::write_target_file_size_bytes(&self) -> iceberg::Result diff --git a/crates/iceberg/src/cow_rewrite/mod.rs b/crates/iceberg/src/cow_rewrite/mod.rs index 940be14cdd..ed31f2112c 100644 --- a/crates/iceberg/src/cow_rewrite/mod.rs +++ b/crates/iceberg/src/cow_rewrite/mod.rs @@ -89,9 +89,13 @@ pub struct CowRewriteStats { pub unchanged_files: usize, /// Visible input row count read from candidate files. pub input_rows: u64, - /// Output row count emitted by the batch rewriter. + /// 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 where the rewriter reported changes. + /// 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, } @@ -104,9 +108,10 @@ pub struct CowRewriteResult { pub added_data_files: Vec, /// 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. + /// 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, /// Rewrite counters. pub stats: CowRewriteStats, @@ -197,28 +202,41 @@ impl<'a> CowRewriteBuilder<'a> { }; 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 = file.scan_task.schema_ref(); + 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, so the in-memory footprint is bounded by the rows that - // precede the first change instead of the entire source file. + // 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 = 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> = 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; + 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 @@ -232,23 +250,28 @@ impl<'a> CowRewriteBuilder<'a> { 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)?; - if rewrite.changed { + // `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 { - result.stats.output_rows += output.num_rows() as u64; + file_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, - )?; + let partition_key = + source_partition_key(self.table, &old_data_file, &write_schema)?; writer = Some( writer::build_replacement_writer( self.table, @@ -258,7 +281,9 @@ impl<'a> CowRewriteBuilder<'a> { .await?, ); } - let writer = writer.as_mut().expect("writer just built"); + let Some(writer) = writer.as_mut() else { + unreachable!("writer initialized above"); + }; for prefix_batch in prefix.drain(..) { writer.write(prefix_batch).await?; } @@ -269,9 +294,18 @@ impl<'a> CowRewriteBuilder<'a> { } } - if file_changed { + // 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 = + !file_changed && file_input_rows == 0 && has_delete_files; + + if file_changed || fully_removed_by_deletes { result.stats.rewritten_files += 1; - result.removed_data_files.push(file.old_data_file.clone()); + result.stats.output_rows += file_output_rows; + result.removed_data_files.push(old_data_file); if let Some(mut writer) = writer { let added_data_files = writer.close().await?; @@ -282,7 +316,7 @@ impl<'a> CowRewriteBuilder<'a> { // file is written. } else { result.stats.unchanged_files += 1; - result.unchanged_data_files.push(file.old_data_file); + result.unchanged_data_files.push(old_data_file); // `prefix` is dropped here; no replacement file was written. } } @@ -310,6 +344,10 @@ fn source_partition_key( })? .as_ref() .clone(); + // `PartitionKey::new` does not bind the spec to the schema, so validate the + // binding here: a spec that is incompatible with the planned snapshot + // schema should fail with a clear error instead of producing a bad + // partition path when the writer later calls `PartitionKey::to_path`. spec.partition_type(schema).map_err(|err| { Error::new( ErrorKind::DataInvalid, @@ -726,7 +764,9 @@ mod tests { assert_eq!(result.stats.candidate_files, 1); assert_eq!(result.stats.unchanged_files, 1); assert_eq!(result.stats.input_rows, 3); - assert_eq!(result.stats.output_rows, 3); + // Nothing was written, so rows emitted for the unchanged file are not + // counted as output. + assert_eq!(result.stats.output_rows, 0); Ok(()) } @@ -823,8 +863,97 @@ mod tests { assert_eq!(result.added_data_files.len(), 0); assert_eq!(result.unchanged_data_files.len(), 1); assert_eq!(result.stats.input_rows, 2); + assert_eq!(result.stats.output_rows, 0); + + Ok(()) + } + + /// Drops every batch while reporting `changed: false`, violating the + /// documented contract. The orchestrator must derive the change from + /// `output: None` itself, otherwise the file would be kept as "unchanged" + /// while its rows were meant to be dropped. + struct SilentFullDrop; + + impl CowBatchRewriter for SilentFullDrop { + fn rewrite_batch(&self, _batch: RecordBatch) -> Result { + Ok(CowBatchRewrite { + output: None, + changed: false, + }) + } + } + + #[tokio::test] + async fn cow_rewrite_none_output_implies_change() -> Result<()> { + let fixture = test_table_with_ids(vec![1, 2]).await?; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_rewriter(Arc::new(SilentFullDrop)) + .rewrite() + .await?; + + assert!(result.has_changes()); + assert_eq!(result.removed_data_files.len(), 1); + assert_eq!(result.added_data_files.len(), 0); + assert!(result.unchanged_data_files.is_empty()); + assert_eq!(result.stats.rewritten_files, 1); + assert_eq!(result.stats.input_rows, 2); + assert_eq!(result.stats.output_rows, 0); + assert_eq!(result.stats.changed_batches, 1); + + Ok(()) + } + + /// Returns `output: None` with `changed: false` for the first batch, then + /// flags the second batch as changed. The replacement must contain exactly + /// the second batch's rows — the silently dropped first batch must not + /// resurrect, and the original file must be removed. + struct DropFirstBatchSilently { + batches_seen: std::sync::atomic::AtomicUsize, + } + + impl CowBatchRewriter for DropFirstBatchSilently { + fn rewrite_batch(&self, batch: RecordBatch) -> Result { + let seen = self + .batches_seen + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if seen == 0 { + Ok(CowBatchRewrite { + output: None, + changed: false, + }) + } else { + Ok(CowBatchRewrite { + output: Some(batch), + changed: true, + }) + } + } + } + + #[tokio::test] + async fn cow_rewrite_silent_drop_before_change_loses_no_rows() -> Result<()> { + let fixture = test_table_with_ids(vec![1, 2, 3, 4]).await?; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_batch_size(2) + .with_rewriter(Arc::new(DropFirstBatchSilently { + batches_seen: std::sync::atomic::AtomicUsize::new(0), + })) + .rewrite() + .await?; + + assert!(result.has_changes()); + assert_eq!(result.removed_data_files.len(), 1); + assert_eq!(result.added_data_files.len(), 1); + assert_eq!(result.stats.input_rows, 4); assert_eq!(result.stats.output_rows, 2); + let ids = read_ids(&fixture.table, &result.added_data_files).await?; + assert_eq!(ids, vec![3, 4]); + Ok(()) } diff --git a/crates/iceberg/src/cow_rewrite/plan.rs b/crates/iceberg/src/cow_rewrite/plan.rs index fbaa11e217..1f2fa6f82e 100644 --- a/crates/iceberg/src/cow_rewrite/plan.rs +++ b/crates/iceberg/src/cow_rewrite/plan.rs @@ -22,10 +22,13 @@ 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. +/// This is a read-only output type: the planning entry points that produce +/// candidates are crate-internal and reached through +/// [`crate::cow_rewrite::CowRewriteBuilder`], and the fields are exposed +/// through accessors only. Follow-up commit-adapter work (overwrite and +/// row-delta actions) consumes planned candidates through those accessors; +/// construction from outside the crate is intentionally not exposed yet and +/// can be added when an adapter actually needs it. #[derive(Debug, Clone)] pub struct CowRewriteFile { /// Original data file from the manifest entry. diff --git a/crates/iceberg/src/cow_rewrite/rewriter.rs b/crates/iceberg/src/cow_rewrite/rewriter.rs index 04d9175e6d..4ed5e3402c 100644 --- a/crates/iceberg/src/cow_rewrite/rewriter.rs +++ b/crates/iceberg/src/cow_rewrite/rewriter.rs @@ -20,6 +20,7 @@ use arrow_array::RecordBatch; use crate::Result; /// Result of rewriting a single record batch. +#[derive(Debug)] pub struct CowBatchRewrite { /// Rewritten output batch, or `None` when the input batch is fully removed. /// @@ -37,7 +38,19 @@ pub struct CowBatchRewrite { } /// Rewrites record batches for copy-on-write operations. +/// +/// `rewrite_batch` is synchronous: it runs on the async runtime thread that +/// drives the read/write pipeline, so implementations must not perform +/// blocking work — async I/O such as catalog enrichment or cross-table +/// lookups is not supported in this contract. The method stays sync (rather +/// than returning a boxed future) so the trait remains object safe for +/// `Arc`. pub trait CowBatchRewriter: Send + Sync { /// Rewrites a record batch and reports whether it changed. + /// + /// `output: None` means the batch is fully removed, which is itself a + /// change: the orchestrator treats it as changed regardless of the + /// `changed` flag, so a rewriter cannot accidentally keep dropped rows + /// alive by reporting `changed: false` alongside a `None` output. fn rewrite_batch(&self, batch: RecordBatch) -> Result; } diff --git a/crates/iceberg/src/cow_rewrite/writer.rs b/crates/iceberg/src/cow_rewrite/writer.rs index a1bfcd129a..ce235e889b 100644 --- a/crates/iceberg/src/cow_rewrite/writer.rs +++ b/crates/iceberg/src/cow_rewrite/writer.rs @@ -28,7 +28,8 @@ 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, + DefaultFileNameGenerator, DefaultLocationGenerator, LocationGenerator, + ObjectStorageLocationGenerator, }; use crate::writer::file_writer::rolling_writer::RollingFileWriterBuilder; @@ -42,17 +43,15 @@ use crate::writer::file_writer::rolling_writer::RollingFileWriterBuilder; /// 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. +/// +/// The location layout follows `write.object-storage.enabled`: tables that opt +/// into the object-storage layout get hash-entropy paths from +/// [`ObjectStorageLocationGenerator`], everything else gets the default layout. pub(crate) async fn build_replacement_writer( table: &Table, write_schema: SchemaRef, partition_key: Option, ) -> Result> { - let location_generator = DefaultLocationGenerator::new(table.metadata())?; - let file_name_generator = DefaultFileNameGenerator::new( - format!("cow-rewrite-{}", Uuid::now_v7()), - None, - DataFileFormat::Parquet, - ); let table_props = table.metadata().table_properties(); // The replacement writer only produces parquet today; refuse to run on a // table configured for another format instead of silently writing parquet @@ -72,6 +71,28 @@ pub(crate) async fn build_replacement_writer( if let Some(encryption_manager) = table.encryption_manager() { parquet_builder = parquet_builder.with_encryption_manager(encryption_manager.clone()); } + + if table_props.write_object_storage_enabled()? { + let location_generator = ObjectStorageLocationGenerator::new(table.metadata())?; + build_data_file_writer(table, parquet_builder, location_generator, partition_key).await + } else { + let location_generator = DefaultLocationGenerator::new(table.metadata())?; + build_data_file_writer(table, parquet_builder, location_generator, partition_key).await + } +} + +async fn build_data_file_writer( + table: &Table, + parquet_builder: ParquetWriterBuilder, + location_generator: L, + partition_key: Option, +) -> Result> { + let file_name_generator = DefaultFileNameGenerator::new( + format!("cow-rewrite-{}", Uuid::now_v7()), + None, + DataFileFormat::Parquet, + ); + let table_props = table.metadata().table_properties(); let rolling_builder = RollingFileWriterBuilder::new( parquet_builder, table_props.write_target_file_size_bytes()?, @@ -285,6 +306,79 @@ mod tests { Ok(()) } + /// With `write.object-storage.enabled` the replacement file must land in + /// the hash-entropy object storage layout instead of the flat default one. + #[tokio::test] + async fn cow_replacement_writer_honors_object_storage_layout() -> Result<()> { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display()); + let catalog = MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(LocalFsStorageFactory)) + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]), + ) + .await?; + let namespace = NamespaceIdent::new("ns".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("cow_object_storage_writer".to_string()) + .schema(schema) + .properties(HashMap::from([( + crate::spec::TableProperties::PROPERTY_WRITE_OBJECT_STORAGE_ENABLED + .to_string(), + "true".to_string(), + )])) + .build(), + ) + .await?; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(Int32Array::from(vec![ + 1, 2, 3, + ])) as ArrayRef])?; + + let data_files = write_replacement_batches( + &table, + table.metadata().current_schema().clone(), + None, + futures::stream::iter(vec![Ok(batch)]), + ) + .await?; + + assert_eq!(data_files.len(), 1); + // Object storage layout: `{data_location}/{entropy dirs}/{file name}`, + // where the entropy dirs are binary strings of 4, 4, 4, and 8 chars. + let path = data_files[0].file_path(); + let after_data = path + .split("/data/") + .nth(1) + .unwrap_or_else(|| panic!("expected /data/ in {path}")); + let segments = after_data.split('/').collect::>(); + assert_eq!(segments.len(), 5, "{path}"); + for (segment, len) in segments[..4].iter().zip([4, 4, 4, 8]) { + assert_eq!(segment.len(), len, "{path}"); + assert!(segment.chars().all(|c| c == '0' || c == '1'), "{path}"); + } + assert!(segments[4].starts_with("cow-rewrite-"), "{path}"); + + Ok(()) + } + #[tokio::test] async fn cow_replacement_writer_preserves_partition_key() -> Result<()> { let temp_dir = TempDir::new().unwrap(); diff --git a/crates/iceberg/src/spec/table_properties.rs b/crates/iceberg/src/spec/table_properties.rs index 09bf71c196..f418d3a1de 100644 --- a/crates/iceberg/src/spec/table_properties.rs +++ b/crates/iceberg/src/spec/table_properties.rs @@ -316,6 +316,14 @@ pub struct TableProperties { getter )] write_object_storage_location: Option, + /// Whether new data files use the object storage location layout, which + /// injects hash entropy into file paths to spread object-store prefixes. + #[property( + key = Self::PROPERTY_WRITE_OBJECT_STORAGE_ENABLED, + default = Self::PROPERTY_WRITE_OBJECT_STORAGE_ENABLED_DEFAULT, + getter + )] + write_object_storage_enabled: bool, /// Whether partition values are included in object storage paths. #[property( key = Self::PROPERTY_WRITE_OBJECT_STORAGE_PARTITIONED_PATHS, @@ -530,6 +538,10 @@ impl TableProperties<'_> { pub const PROPERTY_WRITE_FOLDER_STORAGE_LOCATION: &'static str = "write.folder-storage.path"; /// Property key for deprecated object storage path, kept as a fallback for compatibility. pub const PROPERTY_WRITE_OBJECT_STORAGE_LOCATION: &'static str = "write.object-storage.path"; + /// Property key for enabling the object storage location layout for new data files. + pub const PROPERTY_WRITE_OBJECT_STORAGE_ENABLED: &'static str = "write.object-storage.enabled"; + /// Default value for [`TableProperties::PROPERTY_WRITE_OBJECT_STORAGE_ENABLED`] + pub const PROPERTY_WRITE_OBJECT_STORAGE_ENABLED_DEFAULT: bool = false; /// Property key for controlling whether partition values are included in object storage paths. pub const PROPERTY_WRITE_OBJECT_STORAGE_PARTITIONED_PATHS: &'static str = "write.object-storage.partitioned-paths"; From 91146da3a7ff5993a40f8559bb5ed218000c7fa7 Mon Sep 17 00:00:00 2001 From: kid Date: Thu, 24 Sep 2026 09:27:34 +0800 Subject: [PATCH 3/3] fix(cow_rewrite): flush kept prefix after a dropped batch A dropped batch could mark a file changed while skipping the Some(output) branch that initialized the writer and flushed kept rows. Open the writer whenever a changed file has buffered or current output, then flush the prefix independently of the current batch output. Cover the keep-first/drop-second data-loss case and full coverage by a real position-delete file; clarify delete-file and schema contracts. --- crates/iceberg/src/cow_rewrite/mod.rs | 121 +++++++++++++++++++------- 1 file changed, 88 insertions(+), 33 deletions(-) diff --git a/crates/iceberg/src/cow_rewrite/mod.rs b/crates/iceberg/src/cow_rewrite/mod.rs index ed31f2112c..20d2f21734 100644 --- a/crates/iceberg/src/cow_rewrite/mod.rs +++ b/crates/iceberg/src/cow_rewrite/mod.rs @@ -28,10 +28,10 @@ //! 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. +//! lists may remove position delete files and deletion vectors only when they +//! exclusively reference removed data files. Equality delete files must be +//! retained while they can still apply to other live data files by sequence +//! number. //! //! ```rust,no_run //! # use std::sync::Arc; @@ -110,8 +110,10 @@ pub struct CowRewriteResult { /// /// 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. + /// `removed_data_files` with no replacement. A commit adapter may remove + /// position deletes and deletion vectors that exclusively reference these + /// removed files, but must retain equality deletes that can still apply to + /// other live files. pub unchanged_data_files: Vec, /// Rewrite counters. pub stats: CowRewriteStats, @@ -209,7 +211,9 @@ impl<'a> CowRewriteBuilder<'a> { // 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. + // the snapshot the source files belong to. This preserves the older + // schema in replacement files rather than promoting them to the + // table's current schema during this rewrite. let write_schema = scan_task.schema_ref(); let has_delete_files = !scan_task.deletes().is_empty(); @@ -265,32 +269,36 @@ impl<'a> CowRewriteBuilder<'a> { result.stats.changed_batches += 1; } - if let Some(output) = rewrite.output { - file_output_rows += output.num_rows() as u64; + // A dropped batch can be the first change. Flush any kept + // prefix even when this batch has no output; defer opening the + // writer only if there are no rows to preserve yet. + if file_changed + && writer.is_none() + && (!prefix.is_empty() || rewrite.output.is_some()) + { + 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?, + ); + } - 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?; - } + if let Some(writer) = writer.as_mut() { + for prefix_batch in prefix.drain(..) { + writer.write(prefix_batch).await?; + } + if let Some(output) = rewrite.output { + file_output_rows += output.num_rows() as u64; writer.write(output).await?; - } else { - prefix.push(output); } + } else if let Some(output) = rewrite.output { + file_output_rows += output.num_rows() as u64; + prefix.push(output); } } @@ -311,9 +319,8 @@ impl<'a> CowRewriteBuilder<'a> { 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. + // If `writer` is `None`, no visible rows remain after delete + // files and batch rewriting, so no replacement is written. } else { result.stats.unchanged_files += 1; result.unchanged_data_files.push(old_data_file); @@ -380,6 +387,7 @@ mod tests { use crate::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder}; use crate::io::LocalFsStorageFactory; use crate::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; + use crate::scan::tests::TableTestFixture; use crate::scan::{FileScanTask, FileScanTaskStream}; use crate::spec::{ DataFile, Literal, NestedField, PrimitiveType, Schema, Struct, TableProperties, Transform, @@ -848,6 +856,29 @@ mod tests { Ok(()) } + #[tokio::test] + async fn cow_rewrite_removes_file_fully_covered_by_position_deletes() -> Result<()> { + let mut fixture = TableTestFixture::new(); + let positions = (0..300).collect::>(); + fixture.setup_multi_row_group_manifest(&positions).await; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_rewriter(Arc::new(KeepAll)) + .rewrite() + .await?; + + assert_eq!(result.stats.candidate_files, 1); + assert_eq!(result.stats.rewritten_files, 1); + assert_eq!(result.stats.input_rows, 0); + assert_eq!(result.stats.output_rows, 0); + assert_eq!(result.removed_data_files.len(), 1); + assert!(result.added_data_files.is_empty()); + assert!(result.unchanged_data_files.is_empty()); + + Ok(()) + } + #[tokio::test] async fn cow_rewrite_delete_no_matching_rows_keeps_old_file() -> Result<()> { let fixture = test_table_with_ids(vec![1, 3]).await?; @@ -957,6 +988,30 @@ mod tests { Ok(()) } + #[tokio::test] + async fn cow_rewrite_kept_prefix_survives_later_full_batch_delete() -> Result<()> { + let fixture = test_table_with_ids(vec![1, 3, 2, 4]).await?; + + let result = CowRewriteBuilder::new(&fixture.table) + .with_predicate(crate::expr::Predicate::AlwaysTrue) + .with_batch_size(2) + .with_rewriter(Arc::new(DeleteEvenIds)) + .rewrite() + .await?; + + assert_eq!(result.removed_data_files.len(), 1); + assert_eq!(result.stats.input_rows, 4); + assert_eq!(result.stats.output_rows, 2); + assert_eq!(result.added_data_files.len(), 1); + assert!(result.unchanged_data_files.is_empty()); + assert_eq!( + read_ids(&fixture.table, &result.added_data_files).await?, + vec![1, 3] + ); + + Ok(()) + } + #[tokio::test] async fn cow_rewrite_uses_unique_replacement_paths_for_multiple_source_files() -> Result<()> { let fixture = test_table_with_id_batches(vec![vec![1, 2], vec![3, 4]]).await?;