From 521bdd8b6aca87ed81c60556c537091576ea524a Mon Sep 17 00:00:00 2001 From: osipovartem Date: Tue, 22 Sep 2026 18:15:03 +0300 Subject: [PATCH] Read Iceberg v3 deletion vectors --- Cargo.lock | 14 ++ Cargo.toml | 1 + datafusion_iceberg/Cargo.toml | 2 + .../src/deletion_vector_filter.rs | 198 ++++++++++++++++++ datafusion_iceberg/src/lib.rs | 1 + datafusion_iceberg/src/table/mod.rs | 179 ++++++++++++++-- datafusion_iceberg/tests/position_delete.rs | 146 +++++++++++++ iceberg-rust-spec/Cargo.toml | 2 + iceberg-rust-spec/src/spec/deletion_vector.rs | 197 +++++++++++++++++ iceberg-rust-spec/src/spec/manifest.rs | 5 + iceberg-rust-spec/src/spec/mod.rs | 1 + 11 files changed, 725 insertions(+), 21 deletions(-) create mode 100644 datafusion_iceberg/src/deletion_vector_filter.rs create mode 100644 iceberg-rust-spec/src/spec/deletion_vector.rs diff --git a/Cargo.lock b/Cargo.lock index 5b969d3d..6524290c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2707,6 +2707,7 @@ dependencies = [ "async-trait", "bytes", "chrono", + "crc32fast", "dashmap", "datafusion", "datafusion-expr", @@ -2725,6 +2726,7 @@ dependencies = [ "pyo3", "regex", "reqwest", + "roaring", "rstest 0.26.1", "serde_json", "sqlx", @@ -3823,6 +3825,7 @@ dependencies = [ "apache-avro", "arrow-schema 59.2.0", "chrono", + "crc32fast", "derive-getters", "derive_builder", "fastnum", @@ -3830,6 +3833,7 @@ dependencies = [ "itertools 0.14.0", "murmur3", "ordered-float 5.3.0", + "roaring", "rstest 0.23.0", "serde", "serde_bytes", @@ -5501,6 +5505,16 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "roaring" +version = "0.10.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19e8d2cfa184d94d0726d650a9f4a1be7f9b76ac9fdb954219878dc00c1c1e7b" +dependencies = [ + "bytemuck", + "byteorder", +] + [[package]] name = "rsa" version = "0.9.10" diff --git a/Cargo.toml b/Cargo.toml index 43a989fb..691885af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ murmur3 = { version = "0.5.2" } parquet = { version = "59", features = ["async", "object_store", "variant_experimental"] } pin-project-lite = "0.2" regex = "1.11.1" +roaring = "0.10.12" serde = "^1.0" serde_derive = "^1.0" serde_json = "^1.0" diff --git a/datafusion_iceberg/Cargo.toml b/datafusion_iceberg/Cargo.toml index 3ca514ef..25cc3713 100644 --- a/datafusion_iceberg/Cargo.toml +++ b/datafusion_iceberg/Cargo.toml @@ -36,6 +36,7 @@ uuid = { workspace = true } # The version should match the Arrow version used by duckdb arrow-ipc = "58" bytes = { workspace = true } +crc32fast = "1.5" duckdb = { version = "1.10503", features = ["bundled"] } iceberg-rest-catalog = { path = "../catalogs/iceberg-rest-catalog" } iceberg-rust = { path = "../iceberg-rust", version = "0.10.0", features = ["test-utils"] } @@ -43,6 +44,7 @@ iceberg-sql-catalog = { path = "../catalogs/iceberg-sql-catalog" } parquet = { workspace = true } pyo3 = { version = "0.28", features = ["auto-initialize"] } reqwest = "0.12" +roaring = { workspace = true } rstest = "0.26" sqlx = { version = "0.8", features = ["runtime-tokio", "tls-rustls", "any", "sqlite", "postgres", "mysql"], default-features = false } tempfile = "3.27.0" diff --git a/datafusion_iceberg/src/deletion_vector_filter.rs b/datafusion_iceberg/src/deletion_vector_filter.rs new file mode 100644 index 00000000..27d26297 --- /dev/null +++ b/datafusion_iceberg/src/deletion_vector_filter.rs @@ -0,0 +1,198 @@ +use std::collections::HashMap; +use std::fmt::{self, Display}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use datafusion::arrow::array::{Array, BooleanBuilder, Int64Array, RecordBatch, StringArray}; +use datafusion::arrow::datatypes::{DataType, Schema}; +use datafusion::common::{exec_err, Result}; +use datafusion::physical_plan::PhysicalExpr; +use datafusion::scalar::ScalarValue; +use datafusion_expr::ColumnarValue; +use iceberg_rust::spec::deletion_vector::DeletionVector; + +/// Batch predicate backed by compact roaring deletion vectors. +/// +/// DataFusion's file stream emits each record batch from one file, so the Iceberg path +/// partition column is constant within a batch. Looking it up once keeps the hot loop to one +/// row-position read and one roaring membership test per row. +#[derive(Debug)] +pub(crate) struct DeletionVectorPredicate { + file_path: Arc, + row_position: Arc, + vectors: Arc>, +} + +impl DeletionVectorPredicate { + pub(crate) fn new( + file_path: Arc, + row_position: Arc, + vectors: HashMap, + ) -> Self { + Self { + file_path, + row_position, + vectors: Arc::new(vectors), + } + } +} + +impl PartialEq for DeletionVectorPredicate { + fn eq(&self, other: &Self) -> bool { + self.file_path.eq(&other.file_path) + && self.row_position.eq(&other.row_position) + && Arc::ptr_eq(&self.vectors, &other.vectors) + } +} + +impl Eq for DeletionVectorPredicate {} + +impl Hash for DeletionVectorPredicate { + fn hash(&self, state: &mut H) { + self.file_path.hash(state); + self.row_position.hash(state); + Arc::as_ptr(&self.vectors).hash(state); + } +} + +impl Display for DeletionVectorPredicate { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "iceberg_deletion_vector_keep({}, {})", + self.file_path, self.row_position + ) + } +} + +impl PhysicalExpr for DeletionVectorPredicate { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + if batch.num_rows() == 0 { + return Ok(ColumnarValue::Array(Arc::new( + datafusion::arrow::array::BooleanArray::from(Vec::::new()), + ))); + } + + let paths = self + .file_path + .evaluate(batch)? + .into_array(batch.num_rows())?; + let Some(paths) = paths.as_any().downcast_ref::() else { + return exec_err!( + "Iceberg deletion-vector file path must be Utf8, got {}", + paths.data_type() + ); + }; + if paths.is_null(0) { + return exec_err!("Iceberg deletion-vector file path cannot be null"); + } + let Some(vector) = self.vectors.get(paths.value(0)) else { + return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)))); + }; + + let positions = self + .row_position + .evaluate(batch)? + .into_array(batch.num_rows())?; + let Some(positions) = positions.as_any().downcast_ref::() else { + return exec_err!( + "Iceberg deletion-vector row position must be Int64, got {}", + positions.data_type() + ); + }; + + let mut keep = BooleanBuilder::with_capacity(batch.num_rows()); + for position in positions.iter() { + let Some(position) = position else { + return exec_err!("Iceberg deletion-vector row position cannot be null"); + }; + let position = u64::try_from(position).map_err(|_| { + datafusion::common::DataFusionError::Execution(format!( + "Iceberg deletion-vector row position cannot be negative: {position}" + )) + })?; + keep.append_value(!vector.contains(position)); + } + Ok(ColumnarValue::Array(Arc::new(keep.finish()))) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.file_path, &self.row_position] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [file_path, row_position]: [Arc; 2] = + children.try_into().map_err(|children: Vec<_>| { + datafusion::common::DataFusionError::Internal(format!( + "Iceberg deletion-vector predicate requires 2 children, got {}", + children.len() + )) + })?; + Ok(Arc::new(Self { + file_path, + row_position, + vectors: Arc::clone(&self.vectors), + })) + } + + fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Display::fmt(self, f) + } +} + +#[cfg(test)] +mod tests { + use datafusion::arrow::array::{ArrayRef, Int64Array, StringArray}; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::physical_plan::expressions::Column; + use roaring::RoaringTreemap; + + use super::*; + + #[test] + fn filters_a_batch_without_materializing_delete_rows() { + let schema = Arc::new(Schema::new(vec![ + Field::new("path", DataType::Utf8, false), + Field::new("pos", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["a.parquet"; 5])) as ArrayRef, + Arc::new(Int64Array::from_iter_values(0..5)) as ArrayRef, + ], + ) + .unwrap(); + let mut vectors = HashMap::new(); + vectors.insert( + "a.parquet".to_string(), + DeletionVector::new([1, 3].into_iter().collect::()), + ); + let predicate = DeletionVectorPredicate::new( + Arc::new(Column::new("path", 0)), + Arc::new(Column::new("pos", 1)), + vectors, + ); + + let result = predicate.evaluate(&batch).unwrap().into_array(5).unwrap(); + let result = result + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + result.iter().collect::>(), + vec![Some(true), Some(false), Some(true), Some(false), Some(true)] + ); + } +} diff --git a/datafusion_iceberg/src/lib.rs b/datafusion_iceberg/src/lib.rs index 7da70c77..1c5c676a 100644 --- a/datafusion_iceberg/src/lib.rs +++ b/datafusion_iceberg/src/lib.rs @@ -1,4 +1,5 @@ pub mod catalog; +mod deletion_vector_filter; pub mod error; pub mod materialized_view; mod parquet_data_cache; diff --git a/datafusion_iceberg/src/table/mod.rs b/datafusion_iceberg/src/table/mod.rs index aac472fd..6144d2e7 100644 --- a/datafusion_iceberg/src/table/mod.rs +++ b/datafusion_iceberg/src/table/mod.rs @@ -26,8 +26,8 @@ use iceberg_rust::spec::table_metadata::{ use itertools::Itertools; use lru::LruCache; use object_store::path::Path; -use object_store::ObjectMeta; use object_store::ObjectStoreExt; +use object_store::{ObjectMeta, ObjectStore}; use std::collections::BTreeMap; use std::thread::available_parallelism; use std::{ @@ -39,6 +39,7 @@ use std::{ use tokio::sync::mpsc::{self}; use tracing::{instrument, Instrument}; +use crate::deletion_vector_filter::DeletionVectorPredicate; use crate::row_lineage::{ RowLineageExpr, RowLineageKind, FIRST_ROW_ID_COLUMN, LAST_UPDATED_SEQUENCE_NUMBER_COLUMN, LAST_UPDATED_SEQUENCE_NUMBER_FIELD_ID, PHYSICAL_LAST_UPDATED_SEQUENCE_NUMBER_COLUMN, @@ -88,6 +89,7 @@ use datafusion::{ physical_optimizer::pruning::PruningPredicateBuilder, physical_plan::{ expressions::{BinaryExpr, Column}, + filter::FilterExec, joins::{ utils::{ColumnIndex, JoinFilter}, HashJoinExec, PartitionMode, @@ -115,7 +117,7 @@ use iceberg_rust::{ use iceberg_rust::{ spec::{ arrow::schema::PARQUET_FIELD_ID_META_KEY, - manifest::{Content, ManifestEntry, Status}, + manifest::{Content, FileFormat as IcebergFileFormat, ManifestEntry, Status}, util, values::{Struct, Value}, }, @@ -1012,6 +1014,11 @@ async fn table_scan( .cmp(&y.1.sequence_number().unwrap()) }); + let active_data_file_paths = data_files + .iter() + .map(|(_, entry)| entry.data_file().file_path().clone()) + .collect::>(); + let mut data_file_iter = data_files.into_iter().peekable(); // Gather the complete equality projection up-front, since in general the requested @@ -1289,7 +1296,10 @@ async fn table_scan( delete_files.position_deletes, &object_store_url, parquet_reader_factory, - )?; + table.object_store(), + &active_data_file_paths, + ) + .await?; } Ok::<_, DataFusionError>(Arc::new(ProjectionExec::try_new(projection_expr, plan)?) @@ -1834,12 +1844,65 @@ fn generate_partitioned_file( Ok(file) } -fn apply_position_deletes( - data_plan: Arc, +async fn apply_position_deletes( + mut data_plan: Arc, delete_files: Vec<(ManifestPath, ManifestEntry)>, object_store_url: &ObjectStoreUrl, parquet_reader_factory: Arc, + object_store: Arc, + active_data_file_paths: &HashSet, ) -> Result, DataFusionError> { + let mut parquet_delete_files = Vec::new(); + let mut deletion_vector_files = Vec::new(); + for delete_file in delete_files { + let data_file = delete_file.1.data_file(); + let pointer = ( + *data_file.content_offset(), + *data_file.content_size_in_bytes(), + data_file.referenced_data_file().clone(), + ); + match pointer { + (None, None, None) => parquet_delete_files.push(delete_file), + (Some(_), Some(_), Some(ref referenced_data_file)) => { + if data_file.file_format() != &IcebergFileFormat::Puffin { + return Err(DataFusionError::Execution(format!( + "Iceberg deletion vector {} must use PUFFIN file format", + data_file.file_path() + ))); + } + if active_data_file_paths.contains(referenced_data_file) { + deletion_vector_files.push(delete_file); + } + } + _ => { + return Err(DataFusionError::Execution(format!( + "Iceberg deletion vector {} must set referenced_data_file, content_offset, and content_size_in_bytes together", + data_file.file_path() + ))); + } + } + } + + if !deletion_vector_files.is_empty() { + let vectors = load_deletion_vectors(deletion_vector_files, object_store).await?; + let predicate = Arc::new(DeletionVectorPredicate::new( + Arc::new(Column::new_with_schema( + DATA_FILE_PATH_COLUMN, + &data_plan.schema(), + )?), + Arc::new(Column::new_with_schema( + DATA_FILE_ROW_POSITION_COLUMN, + &data_plan.schema(), + )?), + vectors, + )); + data_plan = Arc::new(FilterExec::try_new(predicate, data_plan)?); + } + + if parquet_delete_files.is_empty() { + return Ok(data_plan); + } + let delete_schema = Arc::new(ArrowSchema::new(vec![ Field::new(POSITION_DELETE_FILE_PATH_COLUMN, DataType::Utf8, false).with_metadata( HashMap::from([( @@ -1855,15 +1918,10 @@ fn apply_position_deletes( ), ])); - let files = delete_files + let files = parquet_delete_files .into_iter() .map(|(_, entry)| { let data_file = entry.data_file(); - if data_file.content_offset().is_some() || data_file.content_size_in_bytes().is_some() { - return not_impl_err!( - "Iceberg v3 deletion vectors are not supported yet; use v2 position delete files" - ); - } let mut file = PartitionedFile::new( util::strip_prefix(data_file.file_path()), u64::try_from(*data_file.file_size_in_bytes()).map_err(|_| { @@ -1877,16 +1935,12 @@ fn apply_position_deletes( // delete file should also cost one ranged read, not two or three. file.metadata_size_hint = crate::parquet_metadata_cache::metadata_size_hint(file.object_meta.size); - let sequence_number = entry - .sequence_number() - .as_ref() - .copied() - .ok_or_else(|| { - DataFusionError::Execution(format!( - "Position delete file {} has no sequence number", - data_file.file_path() - )) - })?; + let sequence_number = entry.sequence_number().as_ref().copied().ok_or_else(|| { + DataFusionError::Execution(format!( + "Position delete file {} has no sequence number", + data_file.file_path() + )) + })?; file.partition_values .push(ScalarValue::Int64(Some(sequence_number))); Ok(file) @@ -1948,6 +2002,89 @@ fn apply_position_deletes( )?)) } +async fn load_deletion_vectors( + delete_files: Vec<(ManifestPath, ManifestEntry)>, + object_store: Arc, +) -> Result, DataFusionError> { + let concurrency = available_parallelism() + .map(|parallelism| parallelism.get()) + .unwrap_or(4) + .clamp(1, 16); + + let vectors = stream::iter(delete_files) + .map(|(_, entry)| { + let object_store = Arc::clone(&object_store); + async move { + let data_file = entry.data_file(); + let offset = u64::try_from(data_file.content_offset().ok_or_else(|| { + DataFusionError::Execution("Deletion vector has no content offset".to_string()) + })?) + .map_err(|_| { + DataFusionError::Execution(format!( + "Deletion vector {} has a negative content offset", + data_file.file_path() + )) + })?; + let size = u64::try_from(data_file.content_size_in_bytes().ok_or_else(|| { + DataFusionError::Execution("Deletion vector has no content size".to_string()) + })?) + .map_err(|_| { + DataFusionError::Execution(format!( + "Deletion vector {} has a negative content size", + data_file.file_path() + )) + })?; + let end = offset.checked_add(size).ok_or_else(|| { + DataFusionError::Execution(format!( + "Deletion vector {} byte range overflows", + data_file.file_path() + )) + })?; + let path = Path::from(util::strip_prefix(data_file.file_path())); + let blob = object_store + .get_range(&path, offset..end) + .await + .map_err(|error| DataFusionError::External(Box::new(error)))?; + let vector = iceberg_rust::spec::deletion_vector::DeletionVector::decode(&blob) + .map_err(|error| DataFusionError::External(Box::new(error)))?; + let expected = u64::try_from(*data_file.record_count()).map_err(|_| { + DataFusionError::Execution(format!( + "Deletion vector {} has a negative record count", + data_file.file_path() + )) + })?; + if vector.len() != expected { + return Err(DataFusionError::Execution(format!( + "Deletion vector {} contains {} positions, expected {expected}", + data_file.file_path(), + vector.len() + ))); + } + let referenced_data_file = + data_file.referenced_data_file().clone().ok_or_else(|| { + DataFusionError::Execution( + "Deletion vector has no referenced data file".to_string(), + ) + })?; + Ok((referenced_data_file, vector)) + } + }) + .buffer_unordered(concurrency) + .try_collect::>() + .await?; + + let mut merged = HashMap::with_capacity(vectors.len()); + for (data_file, vector) in vectors { + match merged.entry(data_file) { + std::collections::hash_map::Entry::Occupied(mut entry) => *entry.get_mut() |= vector, + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(vector); + } + } + } + Ok(merged) +} + fn position_delete_sequence_filter( delete_schema: &SchemaRef, data_schema: &SchemaRef, diff --git a/datafusion_iceberg/tests/position_delete.rs b/datafusion_iceberg/tests/position_delete.rs index 95f5fe20..59274db7 100644 --- a/datafusion_iceberg/tests/position_delete.rs +++ b/datafusion_iceberg/tests/position_delete.rs @@ -23,6 +23,7 @@ use iceberg_rust::{ spec::{ manifest::{Content, DataFile, FileFormat, Status}, namespace::Namespace, + puffin::{Blob, PuffinWriter, STANDARD_BLOB_TYPE_DELETION_VECTOR_V1}, schema::Schema, types::{PrimitiveType, StructField, Type}, values::{Struct, Value}, @@ -31,6 +32,7 @@ use iceberg_rust::{ }; use iceberg_sql_catalog::SqlCatalog; use object_store::local::LocalFileSystem; +use roaring::RoaringTreemap; use tempfile::TempDir; const FILE_PATH_FIELD_ID: i32 = i32::MAX - 101; @@ -90,6 +92,59 @@ fn write_position_delete_file(path: &str, data_file_path: &str, positions: &[i64 .unwrap() } +fn encode_deletion_vector(positions: impl IntoIterator) -> Vec { + const MAGIC: [u8; 4] = [0xD1, 0xD3, 0x39, 0x64]; + let positions = positions.into_iter().collect::(); + let mut roaring = Vec::with_capacity(positions.serialized_size()); + positions.serialize_into(&mut roaring).unwrap(); + let mut body = MAGIC.to_vec(); + body.extend_from_slice(&roaring); + + let mut blob = Vec::with_capacity(4 + body.len() + 4); + blob.extend_from_slice(&u32::try_from(body.len()).unwrap().to_be_bytes()); + blob.extend_from_slice(&body); + blob.extend_from_slice(&crc32fast::hash(&body).to_be_bytes()); + blob +} + +fn write_deletion_vector_file(path: &str, data_file_path: &str, positions: &[u64]) -> DataFile { + let vector = encode_deletion_vector(positions.iter().copied()); + let mut writer = PuffinWriter::new(); + writer + .write_blob(Blob { + blob_type: STANDARD_BLOB_TYPE_DELETION_VECTOR_V1.to_string(), + fields: Vec::new(), + snapshot_id: -1, + sequence_number: -1, + compression_codec: None, + properties: HashMap::new(), + payload: &vector, + }) + .unwrap(); + let puffin = writer.finish().unwrap(); + std::fs::write(path, &puffin).unwrap(); + + DataFile::builder() + .with_content(Content::PositionDeletes) + .with_file_path(path.to_string()) + .with_file_format(FileFormat::Puffin) + .with_partition(Struct::from_iter(Vec::<(String, Option)>::new())) + .with_record_count(i64::try_from(positions.len()).unwrap()) + .with_file_size_in_bytes(i64::try_from(puffin.len()).unwrap()) + .with_column_sizes(None) + .with_value_counts(None) + .with_null_value_counts(None) + .with_nan_value_counts(None) + .with_distinct_counts(None) + .with_lower_bounds(None) + .with_upper_bounds(None) + .with_referenced_data_file(Some(data_file_path.to_string())) + .with_content_offset(Some(4)) + .with_content_size_in_bytes(Some(i64::try_from(vector.len()).unwrap())) + .build() + .unwrap() +} + #[tokio::test] async fn applies_v2_position_deletes() { let temp_dir = TempDir::new().unwrap(); @@ -264,3 +319,94 @@ async fn applies_v2_position_deletes() { &batches ); } + +#[tokio::test] +async fn applies_v3_puffin_deletion_vector() { + let temp_dir = TempDir::new().unwrap(); + let table_dir = format!("{}/test/orders_v3", temp_dir.path().display()); + let object_store = ObjectStoreBuilder::Filesystem(Arc::new(LocalFileSystem::new())); + let catalog: Arc = Arc::new( + SqlCatalog::new("sqlite://", "warehouse", object_store) + .await + .unwrap(), + ); + catalog + .create_namespace(&Namespace::try_new(&["test".to_string()]).unwrap(), None) + .await + .unwrap(); + + let schema = Schema::builder() + .with_struct_field(StructField { + id: 1, + name: "id".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + initial_default: None, + write_default: None, + }) + .build() + .unwrap(); + Table::builder() + .with_name("orders_v3") + .with_location(&table_dir) + .with_schema(schema) + .with_property(("format-version".to_string(), "3".to_string())) + .build(&["test".to_owned()], catalog.clone()) + .await + .unwrap(); + + let ctx = SessionContext::new(); + ctx.register_catalog( + "warehouse", + Arc::new(IcebergCatalog::new(catalog.clone(), None).await.unwrap()), + ); + run_query( + "INSERT INTO warehouse.test.orders_v3 VALUES (10), (20), (30), (40), (50)", + &ctx, + ) + .await; + + let identifier = Identifier::new(&["test".to_string()], "orders_v3"); + let Tabular::Table(mut table) = catalog.clone().load_tabular(&identifier).await.unwrap() else { + panic!("orders_v3 should be an Iceberg table"); + }; + let manifests = table.manifests(None, None).await.unwrap(); + let data_files = table + .datafiles(&manifests, None, (None, None)) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let data_file_path = data_files + .iter() + .find(|(_, entry)| { + entry.status() != &Status::Deleted && entry.data_file().content() == &Content::Data + }) + .unwrap() + .1 + .data_file() + .file_path() + .clone(); + + let delete_dir = format!("{table_dir}/data"); + std::fs::create_dir_all(&delete_dir).unwrap(); + let deletion_vector = write_deletion_vector_file( + &format!("{delete_dir}/deletion-vector.puffin"), + &data_file_path, + &[1, 3], + ); + table + .new_transaction(None) + .append_delete(vec![deletion_vector]) + .commit() + .await + .unwrap(); + + let batches = run_query("SELECT id FROM warehouse.test.orders_v3 ORDER BY id", &ctx).await; + assert_batches_eq!( + ["+----+", "| id |", "+----+", "| 10 |", "| 30 |", "| 50 |", "+----+",], + &batches + ); +} diff --git a/iceberg-rust-spec/Cargo.toml b/iceberg-rust-spec/Cargo.toml index cb512c87..90795a65 100644 --- a/iceberg-rust-spec/Cargo.toml +++ b/iceberg-rust-spec/Cargo.toml @@ -12,6 +12,7 @@ description = "Unofficial implementation of the Iceberg table format specificati apache-avro = { workspace = true } arrow-schema = { workspace = true } chrono = { workspace = true } +crc32fast = "1.5" derive-getters = { workspace = true } derive_builder = { workspace = true } fastnum = { workspace = true } @@ -19,6 +20,7 @@ getrandom = { workspace = true } itertools = { workspace = true } murmur3 = { workspace = true } ordered-float = { version = "5.3.0", features = ["serde"] } +roaring = { workspace = true } serde = { workspace = true } serde_bytes = "0.11.15" serde_derive = { workspace = true } diff --git a/iceberg-rust-spec/src/spec/deletion_vector.rs b/iceberg-rust-spec/src/spec/deletion_vector.rs new file mode 100644 index 00000000..aec04e29 --- /dev/null +++ b/iceberg-rust-spec/src/spec/deletion_vector.rs @@ -0,0 +1,197 @@ +//! Iceberg v3 deletion-vector encoding. + +use std::io::Read; +use std::ops::BitOrAssign; + +use roaring::{RoaringBitmap, RoaringTreemap}; + +use crate::error::Error; + +const MAGIC: [u8; 4] = [0xD1, 0xD3, 0x39, 0x64]; +const LENGTH_BYTES: usize = 4; +const MAGIC_BYTES: usize = MAGIC.len(); +const CRC_BYTES: usize = 4; +const MIN_BLOB_BYTES: usize = LENGTH_BYTES + MAGIC_BYTES + CRC_BYTES; + +/// Compact set of absolute row positions deleted from one Iceberg data file. +#[derive(Debug, Default, PartialEq)] +pub struct DeletionVector { + positions: RoaringTreemap, +} + +impl DeletionVector { + /// Creates a deletion vector from absolute row positions. + #[must_use] + pub fn new(positions: RoaringTreemap) -> Self { + Self { positions } + } + + /// Returns whether an absolute row position is deleted. + #[must_use] + pub fn contains(&self, position: u64) -> bool { + self.positions.contains(position) + } + + /// Returns the number of deleted row positions. + #[must_use] + pub fn len(&self) -> u64 { + self.positions.len() + } + + /// Returns whether the vector contains no positions. + #[must_use] + pub fn is_empty(&self) -> bool { + self.positions.is_empty() + } + + /// Decodes an Iceberg `deletion-vector-v1` Puffin blob. + /// + /// The wire layout is `[length: u32 BE][magic][portable roaring64][crc32: u32 BE]`. + /// The length covers the magic and roaring payload, and the CRC covers the same bytes. + pub fn decode(blob: &[u8]) -> Result { + if blob.len() < MIN_BLOB_BYTES { + return Err(invalid(format!( + "deletion-vector-v1 blob is {} bytes, shorter than the {MIN_BLOB_BYTES}-byte minimum", + blob.len() + ))); + } + + let body = &blob[LENGTH_BYTES..blob.len() - CRC_BYTES]; + let declared_len = u32::from_be_bytes(blob[..LENGTH_BYTES].try_into()?) as usize; + if declared_len != body.len() { + return Err(invalid(format!( + "deletion-vector-v1 length prefix is {declared_len}, expected {}", + body.len() + ))); + } + + let stored_crc = u32::from_be_bytes(blob[blob.len() - CRC_BYTES..].try_into()?); + let computed_crc = crc32fast::hash(body); + if stored_crc != computed_crc { + return Err(invalid(format!( + "deletion-vector-v1 CRC mismatch: computed {computed_crc:#010x}, stored {stored_crc:#010x}" + ))); + } + + if body[..MAGIC_BYTES] != MAGIC { + return Err(invalid(format!( + "deletion-vector-v1 magic mismatch: {:02x?}, expected {MAGIC:02x?}", + &body[..MAGIC_BYTES] + ))); + } + + let positions = decode_roaring_directory(&body[MAGIC_BYTES..])?; + Ok(Self { positions }) + } +} + +impl BitOrAssign for DeletionVector { + fn bitor_assign(&mut self, rhs: Self) { + self.positions.bitor_assign(rhs.positions); + } +} + +fn decode_roaring_directory(mut input: &[u8]) -> Result { + let bitmap_count = read_u64_le(&mut input)?; + if bitmap_count > u64::from(u32::MAX) { + return Err(invalid(format!( + "deletion-vector-v1 roaring bitmap count {bitmap_count} exceeds the 32-bit key space" + ))); + } + + let mut bitmaps = Vec::with_capacity(usize::try_from(bitmap_count)?); + let mut previous_key = None; + for _ in 0..bitmap_count { + let key = read_u32_le(&mut input)?; + if let Some(previous) = previous_key { + if key <= previous { + return Err(invalid(format!( + "deletion-vector-v1 roaring keys are not strictly ordered: {key} follows {previous}" + ))); + } + } + previous_key = Some(key); + let bitmap = RoaringBitmap::deserialize_from(&mut input).map_err(Error::from)?; + bitmaps.push((key, bitmap)); + } + if !input.is_empty() { + return Err(invalid(format!( + "deletion-vector-v1 roaring payload has {} trailing bytes", + input.len() + ))); + } + Ok(RoaringTreemap::from_bitmaps(bitmaps)) +} + +fn read_u32_le(input: &mut &[u8]) -> Result { + let mut bytes = [0; 4]; + input.read_exact(&mut bytes)?; + Ok(u32::from_le_bytes(bytes)) +} + +fn read_u64_le(input: &mut &[u8]) -> Result { + let mut bytes = [0; 8]; + input.read_exact(&mut bytes)?; + Ok(u64::from_le_bytes(bytes)) +} + +fn invalid(message: String) -> Error { + Error::InvalidFormat(message) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn encode(positions: impl IntoIterator) -> Vec { + let positions = positions.into_iter().collect::(); + let mut roaring = Vec::with_capacity(positions.serialized_size()); + positions.serialize_into(&mut roaring).unwrap(); + let mut body = MAGIC.to_vec(); + body.extend_from_slice(&roaring); + + let mut blob = Vec::with_capacity(LENGTH_BYTES + body.len() + CRC_BYTES); + blob.extend_from_slice(&u32::try_from(body.len()).unwrap().to_be_bytes()); + blob.extend_from_slice(&body); + blob.extend_from_slice(&crc32fast::hash(&body).to_be_bytes()); + blob + } + + #[test] + fn decodes_positions_across_64_bit_keys() { + let expected = [0, 5, 1 << 33, (1 << 33) + 7]; + let vector = DeletionVector::decode(&encode(expected)).unwrap(); + assert_eq!(vector.len(), expected.len() as u64); + for position in expected { + assert!(vector.contains(position)); + } + assert!(!vector.contains(6)); + } + + #[test] + fn rejects_corrupt_crc() { + let mut blob = encode([1, 2, 3]); + let last = blob.len() - 1; + blob[last] ^= 0xff; + assert!(DeletionVector::decode(&blob) + .unwrap_err() + .to_string() + .contains("CRC mismatch")); + } + + #[test] + fn rejects_trailing_roaring_bytes() { + let mut blob = encode([1]); + let crc_start = blob.len() - CRC_BYTES; + blob.insert(crc_start, 0); + let body_len = blob.len() - LENGTH_BYTES - CRC_BYTES; + blob[..LENGTH_BYTES].copy_from_slice(&u32::try_from(body_len).unwrap().to_be_bytes()); + let crc = crc32fast::hash(&blob[LENGTH_BYTES..blob.len() - CRC_BYTES]); + let crc_start = blob.len() - CRC_BYTES; + blob[crc_start..].copy_from_slice(&crc.to_be_bytes()); + assert!(DeletionVector::decode(&blob) + .unwrap_err() + .to_string() + .contains("trailing bytes")); + } +} diff --git a/iceberg-rust-spec/src/spec/manifest.rs b/iceberg-rust-spec/src/spec/manifest.rs index 9edd5559..e9791102 100644 --- a/iceberg-rust-spec/src/spec/manifest.rs +++ b/iceberg-rust-spec/src/spec/manifest.rs @@ -419,6 +419,8 @@ pub enum FileFormat { Orc = 1, /// Parquet file Parquet = 2, + /// Puffin file + Puffin = 3, } /// Serialize for PrimitiveType wit special handling for @@ -433,6 +435,7 @@ impl Serialize for FileFormat { Avro => serializer.serialize_str("AVRO"), Orc => serializer.serialize_str("ORC"), Parquet => serializer.serialize_str("PARQUET"), + Puffin => serializer.serialize_str("PUFFIN"), } } } @@ -451,6 +454,8 @@ impl<'de> Deserialize<'de> for FileFormat { Ok(FileFormat::Orc) } else if s == "PARQUET" { Ok(FileFormat::Parquet) + } else if s == "PUFFIN" { + Ok(FileFormat::Puffin) } else { Err(serde::de::Error::custom("Invalid data file format.")) } diff --git a/iceberg-rust-spec/src/spec/mod.rs b/iceberg-rust-spec/src/spec/mod.rs index 89ff5bf3..f809f0ce 100644 --- a/iceberg-rust-spec/src/spec/mod.rs +++ b/iceberg-rust-spec/src/spec/mod.rs @@ -14,6 +14,7 @@ //! serialization/deserialization and validation logic. pub mod decimal; +pub mod deletion_vector; pub mod expressions; pub mod identifier; pub mod manifest;