From ecb9fbfc7326c066e9a8f35b2b0dbbd84efb0e23 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Tue, 15 Sep 2026 15:16:23 +0300 Subject: [PATCH 1/2] Inherit Iceberg v3 row IDs from manifests --- iceberg-rust-spec/src/spec/manifest.rs | 10 ++ iceberg-rust/src/table/mod.rs | 203 ++++++++++++++++++++++--- 2 files changed, 191 insertions(+), 22 deletions(-) diff --git a/iceberg-rust-spec/src/spec/manifest.rs b/iceberg-rust-spec/src/spec/manifest.rs index f188092f..d938721c 100644 --- a/iceberg-rust-spec/src/spec/manifest.rs +++ b/iceberg-rust-spec/src/spec/manifest.rs @@ -85,6 +85,11 @@ impl ManifestEntry { pub fn snapshot_id_mut(&mut self) -> &mut Option { &mut self.snapshot_id } + + /// Returns a mutable reference to the data file stored in this manifest entry. + pub fn data_file_mut(&mut self) -> &mut DataFile { + &mut self.data_file + } } impl ManifestEntry { @@ -646,6 +651,11 @@ impl DataFile { pub fn builder() -> DataFileBuilder { DataFileBuilder::default() } + + /// Returns a mutable reference to the first row ID assigned to this data file. + pub fn first_row_id_mut(&mut self) -> &mut Option { + &mut self.first_row_id + } } impl DataFile { diff --git a/iceberg-rust/src/table/mod.rs b/iceberg-rust/src/table/mod.rs index 97769c2d..d50164eb 100644 --- a/iceberg-rust/src/table/mod.rs +++ b/iceberg-rust/src/table/mod.rs @@ -22,8 +22,8 @@ use futures::{stream, Stream, StreamExt, TryFutureExt, TryStreamExt}; use iceberg_rust_spec::util::{self}; use iceberg_rust_spec::{ spec::{ - manifest::{Content, ManifestEntry}, - manifest_list::ManifestListEntry, + manifest::{Content, ManifestEntry, Status}, + manifest_list::{Content as ManifestListContent, ManifestListEntry}, schema::Schema, table_metadata::TableMetadata, }, @@ -327,6 +327,8 @@ async fn datafiles( let object_store = object_store.clone(); let manifest_path = file.manifest_path.clone(); let manifest_sequence_number = file.sequence_number; + let manifest_first_row_id = file.first_row_id; + let manifest_content = file.content; async move { // Manifest files are immutable by path. Key by the original // URI so equal store-relative paths cannot alias across stores. @@ -345,26 +347,37 @@ async fn datafiles( }; let bytes = Cursor::new(Vec::from(data)); - ManifestReader::new(bytes)? - .filter_map_ok(|mut x| { - let sequence_number = if let Some(sequence_number) = x.sequence_number() { - *sequence_number - } else { - *x.sequence_number_mut() = Some(manifest_sequence_number); - manifest_sequence_number - }; - - let keep = match sequence_number_range { - (Some(start), Some(end)) => { - start < sequence_number && sequence_number <= end - } - (Some(start), None) => start < sequence_number, - (None, Some(end)) => sequence_number <= end, - _ => true, - }; - keep.then(|| (manifest_path.clone(), x)) - }) - .collect::, Error>>() + let mut entries = ManifestReader::new(bytes)?.collect::, Error>>()?; + assign_first_row_ids( + &manifest_path, + manifest_content, + manifest_first_row_id, + &mut entries, + )?; + Ok::<_, Error>( + entries + .into_iter() + .filter_map(|mut x| { + let sequence_number = if let Some(sequence_number) = x.sequence_number() + { + *sequence_number + } else { + *x.sequence_number_mut() = Some(manifest_sequence_number); + manifest_sequence_number + }; + + let keep = match sequence_number_range { + (Some(start), Some(end)) => { + start < sequence_number && sequence_number <= end + } + (Some(start), None) => start < sequence_number, + (None, Some(end)) => sequence_number <= end, + _ => true, + }; + keep.then(|| (manifest_path.clone(), x)) + }) + .collect::>(), + ) } }) .collect(); @@ -378,6 +391,54 @@ async fn datafiles( .try_flatten()) } +fn assign_first_row_ids( + manifest_path: &str, + content: ManifestListContent, + manifest_first_row_id: Option, + entries: &mut [ManifestEntry], +) -> Result<(), Error> { + if content != ManifestListContent::Data { + return Ok(()); + } + + let Some(mut next_row_id) = manifest_first_row_id else { + for entry in entries { + *entry.data_file_mut().first_row_id_mut() = None; + } + return Ok(()); + }; + + if next_row_id < 0 { + return Err(Error::InvalidFormat(format!( + "Manifest {manifest_path} has a negative first_row_id: {next_row_id}" + ))); + } + + for entry in entries { + if *entry.status() == Status::Deleted { + continue; + } + let data_file = entry.data_file_mut(); + if data_file.first_row_id().is_some() { + continue; + } + let record_count = *data_file.record_count(); + if record_count < 0 { + return Err(Error::InvalidFormat(format!( + "Data file {} has a negative record count: {record_count}", + data_file.file_path() + ))); + } + *data_file.first_row_id_mut() = Some(next_row_id); + next_row_id = next_row_id.checked_add(record_count).ok_or_else(|| { + Error::InvalidFormat(format!( + "Row ID overflow while reading manifest {manifest_path}" + )) + })?; + } + Ok(()) +} + /// delete all datafiles, manifests and metadata files, does not remove table from catalog pub(crate) async fn delete_all_table_files( metadata: &TableMetadata, @@ -436,8 +497,106 @@ pub(crate) async fn delete_all_table_files( #[cfg(test)] mod tests { + use iceberg_rust_spec::spec::{ + manifest::{Content, DataFile, FileFormat, ManifestEntry, Status}, + manifest_list::Content as ManifestListContent, + table_metadata::FormatVersion, + values::{Struct, Value}, + }; use rstest::rstest; + use super::assign_first_row_ids; + + fn data_entry( + status: Status, + path: &str, + record_count: i64, + first_row_id: Option, + ) -> ManifestEntry { + let mut data_file = DataFile::builder(); + data_file + .with_content(Content::Data) + .with_file_path(path.to_owned()) + .with_file_format(FileFormat::Parquet) + .with_partition(Struct::from_iter(Vec::<(String, Option)>::new())) + .with_record_count(record_count) + .with_file_size_in_bytes(100) + .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_first_row_id(first_row_id); + ManifestEntry::builder() + .with_format_version(FormatVersion::V3) + .with_status(status) + .with_data_file(data_file.build().expect("build data file")) + .with_sequence_number(1) + .build() + .expect("build manifest entry") + } + + #[test] + fn inherits_v3_first_row_ids_before_scan_filtering() { + let mut entries = vec![ + data_entry(Status::Added, "a.parquet", 2, None), + data_entry(Status::Existing, "explicit.parquet", 1, Some(100)), + data_entry(Status::Added, "b.parquet", 3, None), + data_entry(Status::Deleted, "deleted.parquet", 4, None), + ]; + + assign_first_row_ids( + "manifest.avro", + ManifestListContent::Data, + Some(10), + &mut entries, + ) + .expect("assign first row ids"); + + assert_eq!(*entries[0].data_file().first_row_id(), Some(10)); + assert_eq!(*entries[1].data_file().first_row_id(), Some(100)); + assert_eq!(*entries[2].data_file().first_row_id(), Some(12)); + assert_eq!(*entries[3].data_file().first_row_id(), None); + } + + #[test] + fn clears_inherited_row_ids_without_manifest_lineage() { + let mut entries = vec![data_entry(Status::Added, "old.parquet", 2, Some(100))]; + + assign_first_row_ids( + "manifest.avro", + ManifestListContent::Data, + None, + &mut entries, + ) + .expect("clear stale first row ids"); + + assert_eq!(*entries[0].data_file().first_row_id(), None); + } + + #[test] + fn rejects_invalid_first_row_id_ranges() { + let mut entries = vec![data_entry(Status::Added, "a.parquet", 2, None)]; + assert!(assign_first_row_ids( + "manifest.avro", + ManifestListContent::Data, + Some(-1), + &mut entries, + ) + .is_err()); + + let mut entries = vec![data_entry(Status::Added, "a.parquet", 2, None)]; + assert!(assign_first_row_ids( + "manifest.avro", + ManifestListContent::Data, + Some(i64::MAX), + &mut entries, + ) + .is_err()); + } + // ----------------------------------------------------------------------- // Placeholders for upstream scan + planning + metadata-table tests. // Scan planning lives partly in iceberg-rust and partly in datafusion_iceberg; From 9f3542316de7a47427a6cd1c0883392525bb60ec Mon Sep 17 00:00:00 2001 From: osipovartem Date: Tue, 15 Sep 2026 15:16:39 +0300 Subject: [PATCH 2/2] Add vectorized Iceberg v3 row lineage scans --- datafusion_iceberg/src/lib.rs | 1 + datafusion_iceberg/src/row_lineage.rs | 262 +++++++++++++++ datafusion_iceberg/src/table/mod.rs | 307 ++++++++++++++++-- .../src/variant_schema_adapter.rs | 93 +++++- 4 files changed, 632 insertions(+), 31 deletions(-) create mode 100644 datafusion_iceberg/src/row_lineage.rs diff --git a/datafusion_iceberg/src/lib.rs b/datafusion_iceberg/src/lib.rs index 36ccd494..7f8360be 100644 --- a/datafusion_iceberg/src/lib.rs +++ b/datafusion_iceberg/src/lib.rs @@ -4,6 +4,7 @@ pub mod materialized_view; mod parquet_metadata_cache; pub mod planner; mod pruning_statistics; +mod row_lineage; mod statistics; pub mod table; mod variant_schema_adapter; diff --git a/datafusion_iceberg/src/row_lineage.rs b/datafusion_iceberg/src/row_lineage.rs new file mode 100644 index 00000000..b92cfa1e --- /dev/null +++ b/datafusion_iceberg/src/row_lineage.rs @@ -0,0 +1,262 @@ +use std::fmt::{self, Display}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch}; +use datafusion::arrow::compute::is_not_null; +use datafusion::arrow::compute::kernels::{numeric::add, zip::zip}; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Schema}; +use datafusion::common::{exec_err, Result}; +use datafusion::physical_plan::PhysicalExpr; +use datafusion_expr::ColumnarValue; +use iceberg_rust::spec::arrow::schema::PARQUET_FIELD_ID_META_KEY; + +pub(crate) const ROW_ID_COLUMN: &str = "_row_id"; +pub(crate) const LAST_UPDATED_SEQUENCE_NUMBER_COLUMN: &str = "_last_updated_sequence_number"; +pub(crate) const PHYSICAL_ROW_ID_COLUMN: &str = "__iceberg_physical_row_id"; +pub(crate) const PHYSICAL_LAST_UPDATED_SEQUENCE_NUMBER_COLUMN: &str = + "__iceberg_physical_last_updated_sequence_number"; +pub(crate) const FIRST_ROW_ID_COLUMN: &str = "__iceberg_first_row_id"; +pub(crate) const ROW_ID_FIELD_ID: i32 = i32::MAX - 107; +pub(crate) const LAST_UPDATED_SEQUENCE_NUMBER_FIELD_ID: i32 = i32::MAX - 108; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) enum RowLineageKind { + RowId, + LastUpdatedSequenceNumber, +} + +#[derive(Debug, Eq)] +pub(crate) struct RowLineageExpr { + kind: RowLineageKind, + physical: Arc, + first_row_id: Arc, + fallback: Arc, +} + +impl RowLineageExpr { + pub(crate) fn new( + kind: RowLineageKind, + physical: Arc, + first_row_id: Arc, + fallback: Arc, + ) -> Self { + Self { + kind, + physical, + first_row_id, + fallback, + } + } + + fn output_field(&self) -> FieldRef { + let (name, field_id) = match self.kind { + RowLineageKind::RowId => (ROW_ID_COLUMN, ROW_ID_FIELD_ID), + RowLineageKind::LastUpdatedSequenceNumber => ( + LAST_UPDATED_SEQUENCE_NUMBER_COLUMN, + LAST_UPDATED_SEQUENCE_NUMBER_FIELD_ID, + ), + }; + Arc::new( + Field::new(name, DataType::Int64, true).with_metadata( + [(PARQUET_FIELD_ID_META_KEY.to_owned(), field_id.to_string())].into(), + ), + ) + } + + fn evaluate_int64( + expr: &Arc, + batch: &RecordBatch, + name: &str, + ) -> Result { + let array = expr.evaluate(batch)?.into_array(batch.num_rows())?; + if array.data_type() != &DataType::Int64 { + return exec_err!( + "Iceberg row lineage {name} must be Int64, got {}", + array.data_type() + ); + } + Ok(array) + } +} + +impl PartialEq for RowLineageExpr { + fn eq(&self, other: &Self) -> bool { + self.kind == other.kind + && self.physical.eq(&other.physical) + && self.first_row_id.eq(&other.first_row_id) + && self.fallback.eq(&other.fallback) + } +} + +impl Hash for RowLineageExpr { + fn hash(&self, state: &mut H) { + self.kind.hash(state); + self.physical.hash(state); + self.first_row_id.hash(state); + self.fallback.hash(state); + } +} + +impl Display for RowLineageExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "iceberg_{:?}", self.kind) + } +} + +impl PhysicalExpr for RowLineageExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Int64) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let physical = Self::evaluate_int64(&self.physical, batch, "physical column")?; + let first_row_id = Self::evaluate_int64(&self.first_row_id, batch, "first_row_id")?; + let fallback_source = Self::evaluate_int64(&self.fallback, batch, "fallback")?; + + let fallback = match self.kind { + RowLineageKind::RowId => add(&first_row_id, &fallback_source)?, + RowLineageKind::LastUpdatedSequenceNumber => fallback_source, + }; + let value = zip(&is_not_null(&physical)?, &physical, &fallback)?; + let nulls: ArrayRef = Arc::new(Int64Array::new_null(batch.num_rows())); + let value = zip(&is_not_null(&first_row_id)?, &value, &nulls)?; + Ok(ColumnarValue::Array(value)) + } + + fn return_field(&self, _input_schema: &Schema) -> Result { + Ok(self.output_field()) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.physical, &self.first_row_id, &self.fallback] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [physical, first_row_id, fallback]: [Arc; 3] = + children.try_into().map_err(|children: Vec<_>| { + datafusion::common::DataFusionError::Internal(format!( + "Iceberg row lineage expression requires 3 children, got {}", + children.len() + )) + })?; + Ok(Arc::new(Self::new( + self.kind, + physical, + first_row_id, + fallback, + ))) + } + + fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Display::fmt(self, f) + } +} + +#[cfg(test)] +mod tests { + use datafusion::arrow::array::Array; + use datafusion::arrow::datatypes::SchemaRef; + use datafusion::physical_plan::expressions::Column; + + use super::*; + + fn batch( + physical: Vec>, + first_row_id: Vec>, + fallback: Vec>, + ) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("physical", DataType::Int64, true), + Field::new("first_row_id", DataType::Int64, true), + Field::new("fallback", DataType::Int64, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(physical)), + Arc::new(Int64Array::from(first_row_id)), + Arc::new(Int64Array::from(fallback)), + ], + ) + .expect("build row lineage test batch") + } + + fn values(kind: RowLineageKind, batch: &RecordBatch) -> Vec> { + let expr = RowLineageExpr::new( + kind, + Arc::new(Column::new("physical", 0)), + Arc::new(Column::new("first_row_id", 1)), + Arc::new(Column::new("fallback", 2)), + ); + let ColumnarValue::Array(result) = expr.evaluate(batch).expect("evaluate lineage") else { + panic!("row lineage expression must return an array"); + }; + let result = result + .as_any() + .downcast_ref::() + .expect("Int64 row lineage result"); + (0..result.len()) + .map(|index| (!result.is_null(index)).then(|| result.value(index))) + .collect() + } + + #[test] + fn row_id_is_vectorized_and_coalesces_physical_values() { + let input = batch( + vec![None, Some(7), Some(8), None], + vec![Some(100), Some(100), None, None], + vec![Some(0), Some(1), Some(2), Some(3)], + ); + assert_eq!( + values(RowLineageKind::RowId, &input), + vec![Some(100), Some(7), None, None] + ); + } + + #[test] + fn last_updated_sequence_is_gated_by_first_row_id() { + let input = batch( + vec![None, Some(7), Some(8), None], + vec![Some(100), Some(100), None, None], + vec![Some(4), Some(4), Some(5), Some(5)], + ); + assert_eq!( + values(RowLineageKind::LastUpdatedSequenceNumber, &input), + vec![Some(4), Some(7), None, None] + ); + } + + #[test] + fn output_fields_use_reserved_iceberg_ids() { + let schema: SchemaRef = Arc::new(Schema::empty()); + for (kind, name, id) in [ + (RowLineageKind::RowId, ROW_ID_COLUMN, ROW_ID_FIELD_ID), + ( + RowLineageKind::LastUpdatedSequenceNumber, + LAST_UPDATED_SEQUENCE_NUMBER_COLUMN, + LAST_UPDATED_SEQUENCE_NUMBER_FIELD_ID, + ), + ] { + let expr = RowLineageExpr::new( + kind, + Arc::new(Column::new("physical", 0)), + Arc::new(Column::new("first_row_id", 1)), + Arc::new(Column::new("fallback", 2)), + ); + let field = expr.return_field(schema.as_ref()).expect("lineage field"); + assert_eq!(field.name(), name); + assert_eq!( + field.metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&id.to_string()) + ); + } + } +} diff --git a/datafusion_iceberg/src/table/mod.rs b/datafusion_iceberg/src/table/mod.rs index c23b8750..8302200a 100644 --- a/datafusion_iceberg/src/table/mod.rs +++ b/datafusion_iceberg/src/table/mod.rs @@ -39,6 +39,11 @@ use std::{ use tokio::sync::mpsc::{self}; use tracing::{instrument, Instrument}; +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, + PHYSICAL_ROW_ID_COLUMN, ROW_ID_COLUMN, ROW_ID_FIELD_ID, +}; use crate::statistics::statistics_from_datafiles; use crate::variant_schema_adapter::IcebergPhysicalExprAdapterFactory; use crate::{ @@ -126,6 +131,14 @@ static POSITION_DELETE_FILE_PATH_COLUMN: &str = "file_path"; static POSITION_DELETE_POS_COLUMN: &str = "pos"; const POSITION_DELETE_FILE_PATH_FIELD_ID: i32 = i32::MAX - 101; const POSITION_DELETE_POS_FIELD_ID: i32 = i32::MAX - 102; +type PhysicalProjection = Vec<(Arc, String)>; + +fn row_lineage_field(name: &str, field_id: i32) -> Field { + Field::new(name, DataType::Int64, true).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_owned(), + field_id.to_string(), + )])) +} /// When the view tracks source arrow ids (the `overrides` map is non-empty), /// reshape each top-level field's `PARQUET:field_id` metadata so it matches @@ -213,6 +226,12 @@ pub struct DataFusionTableConfig { /// With this option, an additional "__manifest_file_path" column is added to the output of the /// TableProvider that contains the path of the manifest-file for the data-file the row originates from. enable_manifest_file_path_column: bool, + /// Expose the Iceberg v3 `_row_id` metadata column. + #[builder(default)] + enable_row_id_column: bool, + /// Expose the Iceberg v3 `_last_updated_sequence_number` metadata column. + #[builder(default)] + enable_last_updated_sequence_number_column: bool, } impl DataFusionTable { @@ -262,6 +281,23 @@ impl DataFusionTable { .with_extension_type(RowNumber), ); } + if config + .as_ref() + .map(|x| x.enable_row_id_column) + .unwrap_or_default() + { + builder.push(row_lineage_field(ROW_ID_COLUMN, ROW_ID_FIELD_ID)); + } + if config + .as_ref() + .map(|x| x.enable_last_updated_sequence_number_column) + .unwrap_or_default() + { + builder.push(row_lineage_field( + LAST_UPDATED_SEQUENCE_NUMBER_COLUMN, + LAST_UPDATED_SEQUENCE_NUMBER_FIELD_ID, + )); + } Arc::new(builder.finish()) } Tabular::View(view) => { @@ -549,6 +585,14 @@ async fn table_scan( .map(|x| x.enable_manifest_file_path_column) .unwrap_or_default(); + let enable_row_id_column = config.map(|x| x.enable_row_id_column).unwrap_or_default(); + + let enable_last_updated_sequence_number_column = config + .map(|x| x.enable_last_updated_sequence_number_column) + .unwrap_or_default(); + + let enable_row_lineage = enable_row_id_column || enable_last_updated_sequence_number_column; + let partition_fields = &snapshot_range .1 .and_then(|snapshot_id| table.metadata().partition_fields(snapshot_id).ok()) @@ -561,7 +605,14 @@ async fn table_scan( .unwrap(); // If there is a filter expression the manifests to read are pruned based on the pruning statistics available in the manifest_list file. - let physical_predicate = conjunction(filters.iter().cloned()) + // Row lineage values may be synthesized from manifest metadata, so they must not be + // evaluated by the Parquet reader against only the physically stored columns. + let parquet_filters = filters.iter().filter(|expr| { + expr.column_refs().iter().all(|column| { + column.name != ROW_ID_COLUMN && column.name != LAST_UPDATED_SEQUENCE_NUMBER_COLUMN + }) + }); + let physical_predicate = conjunction(parquet_filters.clone().cloned()) .map(|predicate| { create_physical_expr( &predicate, @@ -574,7 +625,18 @@ async fn table_scan( let mut table_partition_cols = datafusion_partition_columns(partition_fields)?; - let file_schema: SchemaRef = Arc::new((schema.fields()).try_into().unwrap()); + let mut file_schema_builder = + SchemaBuilder::from(TryInto::::try_into(schema.fields()).unwrap()); + if enable_row_id_column { + file_schema_builder.push(row_lineage_field(PHYSICAL_ROW_ID_COLUMN, ROW_ID_FIELD_ID)); + } + if enable_last_updated_sequence_number_column { + file_schema_builder.push(row_lineage_field( + PHYSICAL_LAST_UPDATED_SEQUENCE_NUMBER_COLUMN, + LAST_UPDATED_SEQUENCE_NUMBER_FIELD_ID, + )); + } + let file_schema: SchemaRef = Arc::new(file_schema_builder.finish()); // The ordering the table declares, if any, expressed on the file schema. // Files attest it individually (manifest `sort_order_id`), so the claim is @@ -590,18 +652,6 @@ async fn table_scan( .cloned() .unwrap_or((0..arrow_schema.fields().len()).collect_vec()); - let projection_expr: Vec<_> = requested_projection - .iter() - .enumerate() - .map(|(i, id)| { - let name = arrow_schema.fields[*id].name(); - ( - Arc::new(Column::new(name, i)) as Arc, - name.to_owned(), - ) - }) - .collect(); - if enable_data_file_path_column { table_partition_cols.push(Field::new(DATA_FILE_PATH_COLUMN, DataType::Utf8, false)); } @@ -686,6 +736,13 @@ async fn table_scan( pruning_predicate.prune(&PruneDataFiles::new(&schema, &arrow_schema, &data_files))?; let mut statistics = statistics_from_datafiles(&schema, &data_files); + for _ in 0..usize::from(enable_row_id_column) + + usize::from(enable_last_updated_sequence_number_column) + { + statistics + .column_statistics + .push(ColumnStatistics::new_unknown()); + } // Add placeholder statistics for partition/metadata columns // This prevents index out of bounds when projecting statistics for _ in &table_partition_cols { @@ -713,6 +770,13 @@ async fn table_scan( .map_err(DataFusionIcebergError::from)?; let mut statistics = statistics_from_datafiles(&schema, &data_files); + for _ in 0..usize::from(enable_row_id_column) + + usize::from(enable_last_updated_sequence_number_column) + { + statistics + .column_statistics + .push(ColumnStatistics::new_unknown()); + } // Add placeholder statistics for partition/metadata columns // This prevents index out of bounds when projecting statistics for _ in &table_partition_cols { @@ -799,7 +863,13 @@ async fn table_scan( .column_statistics .push(ColumnStatistics::new_unknown()); } - if has_position_deletes { + if enable_row_lineage { + table_partition_cols.push(Field::new(FIRST_ROW_ID_COLUMN, DataType::Int64, true)); + statistics + .column_statistics + .push(ColumnStatistics::new_unknown()); + } + if has_position_deletes || enable_last_updated_sequence_number_column { table_partition_cols.push(Field::new( DATA_FILE_SEQUENCE_NUMBER_COLUMN, DataType::Int64, @@ -818,7 +888,7 @@ async fn table_scan( .map(Arc::new) .collect::>(), ); - if has_position_deletes || enable_data_file_row_position_column { + if has_position_deletes || enable_data_file_row_position_column || enable_row_id_column { table_schema_builder = table_schema_builder.with_virtual_columns(vec![Arc::new( Field::new(DATA_FILE_ROW_POSITION_COLUMN, DataType::Int64, false) .with_extension_type(RowNumber), @@ -830,13 +900,15 @@ async fn table_scan( let table_schema = table_schema_builder.build(); // File schema plus partition columns: what scan orderings are expressed on. let scan_schema: SchemaRef = table_schema.table_schema().clone(); - let scan_projection = requested_projection - .iter() - .map(|index| scan_schema.index_of(arrow_schema.field(*index).name())) - .collect::, _>>()?; + let (scan_projection, projection_expr) = row_lineage_projection( + &arrow_schema, + &scan_schema, + &requested_projection, + enable_row_id_column, + enable_last_updated_sequence_number_column, + )?; // See `use_parquet_row_filter_pushdown` for the wide-scan / narrow-predicate rationale. - let filter_columns: std::collections::HashSet<_> = filters - .iter() + let filter_columns: std::collections::HashSet<_> = parquet_filters .flat_map(|f| f.column_refs().into_iter().cloned()) .collect(); let pushdown_filters = @@ -942,7 +1014,8 @@ async fn table_scan( &data_manifest.1, last_updated_ms, include_data_file_path_column, - has_position_deletes, + enable_row_lineage, + has_position_deletes || enable_last_updated_sequence_number_column, manifest_path, ) .unwrap(); @@ -999,6 +1072,7 @@ async fn table_scan( last_updated_ms, enable_data_file_path_column, false, + false, manifest_path, )?; @@ -1096,7 +1170,8 @@ async fn table_scan( &x.1, last_updated_ms, include_data_file_path_column, - has_position_deletes, + enable_row_lineage, + has_position_deletes || enable_last_updated_sequence_number_column, manifest_path, ) }) @@ -1171,7 +1246,8 @@ async fn table_scan( &entry, last_updated_ms, include_data_file_path_column, - has_position_deletes, + enable_row_lineage, + has_position_deletes || enable_last_updated_sequence_number_column, manifest_path, )?; if is_attested { @@ -1198,10 +1274,17 @@ async fn table_scan( .with_limit(limit) .build(); - let other_plan: Arc = + let mut other_plan: Arc = tracing::debug_span!("datafusion_iceberg::create_physical_plan_scan_data_files") .in_scope(|| DataSourceExec::from_data_source(file_scan_config)); + if enable_row_lineage { + other_plan = Arc::new(ProjectionExec::try_new( + projection_expr.clone(), + other_plan, + )?); + } + plans.push(other_plan); } @@ -1217,13 +1300,20 @@ async fn table_scan( .with_limit(limit) .build(); - let sorted_plan = ParquetFormat::default() + let mut sorted_plan = ParquetFormat::default() .create_physical_plan(session, file_scan_config) .instrument(tracing::debug_span!( "datafusion_iceberg::create_physical_plan_scan_sorted_data_files" )) .await?; + if enable_row_lineage { + sorted_plan = Arc::new(ProjectionExec::try_new( + projection_expr.clone(), + sorted_plan, + )?); + } + plans.push(sorted_plan); } @@ -1237,6 +1327,73 @@ async fn table_scan( } } +fn row_lineage_projection( + output_schema: &SchemaRef, + scan_schema: &SchemaRef, + requested_projection: &[usize], + enable_row_id: bool, + enable_last_updated_sequence_number: bool, +) -> Result<(Vec, PhysicalProjection), DataFusionError> { + fn projected_column( + scan_projection: &mut Vec, + scan_schema: &SchemaRef, + name: &str, + ) -> Result, DataFusionError> { + let scan_index = scan_schema.index_of(name)?; + let projected_index = match scan_projection + .iter() + .position(|index| *index == scan_index) + { + Some(index) => index, + None => { + scan_projection.push(scan_index); + scan_projection.len() - 1 + } + }; + Ok(Arc::new(Column::new(name, projected_index))) + } + + let mut scan_projection = Vec::with_capacity(requested_projection.len() + 3); + let mut output_projection = Vec::with_capacity(requested_projection.len()); + + for output_index in requested_projection { + let name = output_schema.field(*output_index).name(); + let expression: Arc = if enable_row_id && name == ROW_ID_COLUMN { + Arc::new(RowLineageExpr::new( + RowLineageKind::RowId, + projected_column(&mut scan_projection, scan_schema, PHYSICAL_ROW_ID_COLUMN)?, + projected_column(&mut scan_projection, scan_schema, FIRST_ROW_ID_COLUMN)?, + projected_column( + &mut scan_projection, + scan_schema, + DATA_FILE_ROW_POSITION_COLUMN, + )?, + )) + } else if enable_last_updated_sequence_number && name == LAST_UPDATED_SEQUENCE_NUMBER_COLUMN + { + Arc::new(RowLineageExpr::new( + RowLineageKind::LastUpdatedSequenceNumber, + projected_column( + &mut scan_projection, + scan_schema, + PHYSICAL_LAST_UPDATED_SEQUENCE_NUMBER_COLUMN, + )?, + projected_column(&mut scan_projection, scan_schema, FIRST_ROW_ID_COLUMN)?, + projected_column( + &mut scan_projection, + scan_schema, + DATA_FILE_SEQUENCE_NUMBER_COLUMN, + )?, + )) + } else { + projected_column(&mut scan_projection, scan_schema, name)? + }; + output_projection.push((expression, name.to_owned())); + } + + Ok((scan_projection, output_projection)) +} + /// Maps a table sort order onto a DataFusion ordering over `file_schema`. /// /// Only the leading run of identity-transformed fields is claimed: rows @@ -1438,6 +1595,7 @@ fn generate_partitioned_file( manifest: &ManifestEntry, last_updated_ms: i64, enable_data_file_path: bool, + include_first_row_id: bool, include_sequence_number: bool, manifest_file_path: Option, ) -> Result { @@ -1463,6 +1621,10 @@ fn generate_partitioned_file( partition_values.push(ScalarValue::Utf8(Some(manifest_file_path))); } + if include_first_row_id { + partition_values.push(ScalarValue::Int64(*manifest.data_file().first_row_id())); + } + if include_sequence_number { let sequence_number = manifest .sequence_number() @@ -2120,6 +2282,97 @@ mod tests { ); } + #[tokio::test] + async fn test_v3_row_lineage_columns_are_synthesized_by_scan() { + let object_store = ObjectStoreBuilder::memory(); + let catalog: Arc = Arc::new( + SqlCatalog::new("sqlite://", "test", object_store) + .await + .expect("create catalog"), + ); + let schema = Schema::builder() + .with_struct_field(StructField { + id: 1, + name: "id".to_owned(), + required: true, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + initial_default: None, + write_default: None, + }) + .build() + .expect("build schema"); + let table = Table::builder() + .with_name("lineage_numbers") + .with_location("memory:///test/lineage_numbers") + .with_schema(schema) + .with_property(("format-version".to_owned(), "3".to_owned())) + .build(&["test".to_owned()], catalog) + .await + .expect("create v3 table"); + let writable = Arc::new(DataFusionTable::from(table)); + let ctx = SessionContext::new(); + ctx.register_table("lineage_numbers", writable.clone()) + .expect("register writable table"); + ctx.sql("INSERT INTO lineage_numbers (id) VALUES (10), (20), (30)") + .await + .expect("plan v3 insert") + .collect() + .await + .expect("write v3 rows"); + + let updated_table = { + let tabular = writable.tabular.read().expect("read updated table"); + let Tabular::Table(table) = &*tabular else { + panic!("expected an Iceberg table"); + }; + table.clone() + }; + let config = super::DataFusionTableConfigBuilder::default() + .enable_data_file_path_column(false) + .enable_data_file_row_position_column(false) + .enable_manifest_file_path_column(false) + .enable_row_id_column(true) + .enable_last_updated_sequence_number_column(true) + .build() + .expect("build row lineage scan config"); + let lineage = Arc::new(DataFusionTable::new_with_config( + Tabular::Table(updated_table), + None, + None, + None, + Some(config), + )); + ctx.deregister_table("lineage_numbers") + .expect("deregister writable table"); + ctx.register_table("lineage_numbers", lineage) + .expect("register lineage table"); + + let batches = ctx + .sql( + "SELECT id, _row_id, _last_updated_sequence_number \ + FROM lineage_numbers WHERE _row_id >= 1 ORDER BY _row_id", + ) + .await + .expect("plan lineage scan") + .collect() + .await + .expect("scan row lineage"); + let batch = batches.first().expect("lineage result batch"); + let values = |column: usize| { + batch + .column(column) + .as_any() + .downcast_ref::() + .expect("Int64 result") + .iter() + .collect::>() + }; + assert_eq!(values(0), vec![Some(20), Some(30)]); + assert_eq!(values(1), vec![Some(1), Some(2)]); + assert_eq!(values(2), vec![Some(1), Some(1)]); + } + #[tokio::test] pub async fn test_datafusion_table_insert() { let object_store = ObjectStoreBuilder::memory(); diff --git a/datafusion_iceberg/src/variant_schema_adapter.rs b/datafusion_iceberg/src/variant_schema_adapter.rs index 69ec03b2..7f659c0e 100644 --- a/datafusion_iceberg/src/variant_schema_adapter.rs +++ b/datafusion_iceberg/src/variant_schema_adapter.rs @@ -14,8 +14,14 @@ use datafusion::physical_expr_adapter::{ use datafusion::physical_plan::expressions::Column; use datafusion::physical_plan::PhysicalExpr; use datafusion_expr::ColumnarValue; +use iceberg_rust::spec::arrow::schema::PARQUET_FIELD_ID_META_KEY; use parquet_variant_compute::{unshred_variant, VariantArray}; +use crate::row_lineage::{ + LAST_UPDATED_SEQUENCE_NUMBER_FIELD_ID, PHYSICAL_LAST_UPDATED_SEQUENCE_NUMBER_COLUMN, + PHYSICAL_ROW_ID_COLUMN, ROW_ID_FIELD_ID, +}; + const PARQUET_VARIANT_EXTENSION_NAME: &str = "arrow.parquet.variant"; #[derive(Debug)] @@ -70,14 +76,55 @@ impl IcebergPhysicalExprAdapter { Arc::new(logical_field.clone()), ))) } + + fn row_lineage_field_id(name: &str) -> Option { + match name { + PHYSICAL_ROW_ID_COLUMN => Some(ROW_ID_FIELD_ID), + PHYSICAL_LAST_UPDATED_SEQUENCE_NUMBER_COLUMN => { + Some(LAST_UPDATED_SEQUENCE_NUMBER_FIELD_ID) + } + _ => None, + } + } + + fn rewrite_row_lineage_column( + &self, + _column: &Column, + field_id: i32, + ) -> Result>> { + let physical = self + .physical_file_schema + .fields() + .iter() + .enumerate() + .find(|(_, field)| { + field + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .and_then(|id| id.parse::().ok()) + == Some(field_id) + }); + let Some((index, field)) = physical else { + return Ok(None); + }; + if field.data_type() != &DataType::Int64 { + return internal_err!( + "Iceberg row lineage field id {field_id} must be Int64, got {}", + field.data_type() + ); + } + Ok(Some(Arc::new(Column::new(field.name(), index)))) + } } impl PhysicalExprAdapter for IcebergPhysicalExprAdapter { fn rewrite(&self, expr: Arc) -> Result> { - let contains_variant = collect_columns(&expr) - .iter() - .any(|column| self.is_logical_variant_column(column.name())); - if !contains_variant { + let columns = collect_columns(&expr); + let requires_iceberg_rewrite = columns.iter().any(|column| { + self.is_logical_variant_column(column.name()) + || Self::row_lineage_field_id(column.name()).is_some() + }); + if !requires_iceberg_rewrite { return self.default.rewrite(expr); } @@ -86,6 +133,13 @@ impl PhysicalExprAdapter for IcebergPhysicalExprAdapter { return Ok(Transformed::no(expr)); }; + if let Some(field_id) = Self::row_lineage_field_id(column.name()) { + if let Some(physical) = self.rewrite_row_lineage_column(column, field_id)? { + return Ok(Transformed::yes(physical)); + } + return self.default.rewrite(expr).map(Transformed::yes); + } + if self.is_logical_variant_column(column.name()) { return self.rewrite_variant_column(column).map(Transformed::yes); } @@ -258,4 +312,35 @@ mod tests { assert_eq!(format!("{:?}", variant.try_value(0)?), "BooleanTrue"); Ok(()) } + + #[test] + fn maps_physical_row_lineage_by_reserved_field_id() -> Result<()> { + let logical_schema = Arc::new(Schema::new(vec![Arc::new( + Field::new(PHYSICAL_ROW_ID_COLUMN, DataType::Int64, true).with_metadata( + [( + PARQUET_FIELD_ID_META_KEY.to_owned(), + ROW_ID_FIELD_ID.to_string(), + )] + .into(), + ), + )])); + let physical_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("writer_specific_row_id", DataType::Int64, true).with_metadata( + [( + PARQUET_FIELD_ID_META_KEY.to_owned(), + ROW_ID_FIELD_ID.to_string(), + )] + .into(), + ), + ])); + let adapter = IcebergPhysicalExprAdapterFactory.create(logical_schema, physical_schema)?; + let rewritten = adapter.rewrite(Arc::new(Column::new(PHYSICAL_ROW_ID_COLUMN, 0)))?; + let column = rewritten + .downcast_ref::() + .expect("physical lineage column"); + assert_eq!(column.name(), "writer_specific_row_id"); + assert_eq!(column.index(), 1); + Ok(()) + } }