diff --git a/iceberg/src/common/scan_config.rs b/iceberg/src/common/scan_config.rs index fdd918eae..6b2800997 100644 --- a/iceberg/src/common/scan_config.rs +++ b/iceberg/src/common/scan_config.rs @@ -9,14 +9,17 @@ use datafusion::physical_plan::limit::LimitStream; use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::TryStreamExt; -use iceberg::arrow::ArrowReaderBuilder; +use iceberg::arrow::{ArrowReader, ArrowReaderBuilder}; use iceberg::io::FileIO; use iceberg::scan::FileScanTaskStream; use iceberg::spec::SnapshotRef; +use std::fmt; use std::sync::Arc; /// Execution settings shared by assigned-file and feed-backed scans. -#[derive(Debug, Clone)] +/// Reader settings are captured from the planning or decoding context. Clones +/// share the process-local reader; decoding constructs a fresh worker reader. +#[derive(Clone)] pub(crate) struct ScanConfig { pub(crate) schema: SchemaRef, pub(crate) partitioning: Partitioning, @@ -24,26 +27,45 @@ pub(crate) struct ScanConfig { pub(crate) metrics: ExecutionPlanMetricsSet, pub(crate) iceberg_file_io: FileIO, pub(crate) iceberg_runtime: iceberg::Runtime, + // Clones share the delete loader across this source's local partitions. + reader: ArrowReader, } impl ScanConfig { + pub(crate) fn new( + schema: SchemaRef, + partitioning: Partitioning, + fetch: Option, + iceberg_file_io: FileIO, + iceberg_runtime: iceberg::Runtime, + context: &TaskContext, + ) -> Self { + let config = IcebergConfig::from_config_options(context.session_config().options()); + let reader = ArrowReaderBuilder::new(iceberg_file_io.clone(), iceberg_runtime.clone()) + .with_batch_size(context.session_config().batch_size()) + .with_data_file_concurrency_limit(config.data_file_concurrency_limit) + .with_row_group_filtering_enabled(config.row_group_filtering_enabled) + .with_row_selection_enabled(config.row_selection_enabled) + .build(); + Self { + schema, + partitioning, + fetch, + metrics: Default::default(), + iceberg_file_io, + iceberg_runtime, + reader, + } + } + pub(crate) fn open( &self, tasks: FileScanTaskStream, partition: usize, - context: Arc, ) -> Result { - let config = IcebergConfig::from_task_context(&context); - - let reader = - ArrowReaderBuilder::new(self.iceberg_file_io.clone(), self.iceberg_runtime.clone()) - .with_batch_size(context.session_config().batch_size()) - .with_data_file_concurrency_limit(config.data_file_concurrency_limit) - .with_row_group_filtering_enabled(config.row_group_filtering_enabled) - .with_row_selection_enabled(config.row_selection_enabled) - .build(); - - let stream = reader + let stream = self + .reader + .clone() .read(tasks) .map(|result| result.stream()) .map_err(df_err)? @@ -60,6 +82,19 @@ impl ScanConfig { } } +impl fmt::Debug for ScanConfig { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ScanConfig") + .field("schema", &self.schema) + .field("partitioning", &self.partitioning) + .field("fetch", &self.fetch) + .field("metrics", &self.metrics) + .field("iceberg_file_io", &self.iceberg_file_io) + .field("iceberg_runtime", &self.iceberg_runtime) + .finish_non_exhaustive() + } +} + /// Snapshot summary keys defined by the Iceberg table spec: /// https://iceberg.apache.org/spec/#optional-snapshot-summary-fields /// diff --git a/iceberg/src/eager/codec.rs b/iceberg/src/eager/codec.rs index a86417b9c..9aa97478d 100644 --- a/iceberg/src/eager/codec.rs +++ b/iceberg/src/eager/codec.rs @@ -175,14 +175,14 @@ impl EagerIcebergCodec { proto.custom_file_io.as_deref(), ctx, )?; - Ok(ScanConfig { + Ok(ScanConfig::new( schema, partitioning, fetch, - metrics: Default::default(), iceberg_file_io, - iceberg_runtime: self.iceberg_runtime.clone(), - }) + self.iceberg_runtime.clone(), + ctx, + )) } fn encode_config( diff --git a/iceberg/src/eager/data_source.rs b/iceberg/src/eager/data_source.rs index 248e0b5f3..b4017fb3a 100644 --- a/iceberg/src/eager/data_source.rs +++ b/iceberg/src/eager/data_source.rs @@ -29,6 +29,7 @@ impl EagerIcebergDataSource { /// Construct a scan from already selected files grouped by local output partition. /// Preserves group order without discovering, pruning, or assigning files. /// Descriptors are charged to `context`'s memory pool; statistics are unknown. + /// Reader settings are captured from `context` and shared across local partitions. /// /// # Errors /// Returns an error if there are no partition groups or memory reservation fails. @@ -53,14 +54,14 @@ impl EagerIcebergDataSource { } Ok(Self { statistics: Arc::new(Statistics::new_unknown(&schema)), - scan: ScanConfig { + scan: ScanConfig::new( schema, - partitioning: Partitioning::UnknownPartitioning(partitions.len()), + Partitioning::UnknownPartitioning(partitions.len()), fetch, - metrics: Default::default(), - iceberg_file_io: file_io, - iceberg_runtime: runtime, - }, + file_io, + runtime, + context, + ), files: Arc::new(PlannedFiles::from_tasks(tasks, context)?), partitions: partitions.into(), }) @@ -150,7 +151,7 @@ impl DataSource for EagerIcebergDataSource { fn open( &self, partition: usize, - context: Arc, + _context: Arc, ) -> Result { let count = self.planned_partition_tasks(partition)?.len(); let source = self.clone(); @@ -159,7 +160,7 @@ impl DataSource for EagerIcebergDataSource { Ok(source.files.tasks()[source.partitions[partition][offset]].clone()) }) .boxed(); - self.scan.open(tasks, partition, context) + self.scan.open(tasks, partition) } fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!( diff --git a/iceberg/src/lazy/codec.rs b/iceberg/src/lazy/codec.rs index ff69ac2e4..779f6e716 100644 --- a/iceberg/src/lazy/codec.rs +++ b/iceberg/src/lazy/codec.rs @@ -105,14 +105,14 @@ impl PhysicalExtensionCodec for LazyIcebergCodec { )?; Ok(DataSourceExec::from_data_source(LazyIcebergDataSource { - scan: ScanConfig { + scan: ScanConfig::new( schema, partitioning, fetch, - metrics: Default::default(), iceberg_file_io, - iceberg_runtime: self.iceberg_runtime.clone(), - }, + self.iceberg_runtime.clone(), + ctx, + ), column_stats: None, table_snapshot: None, feed, diff --git a/iceberg/src/lazy/data_source.rs b/iceberg/src/lazy/data_source.rs index c7b662d9c..38686490e 100644 --- a/iceberg/src/lazy/data_source.rs +++ b/iceberg/src/lazy/data_source.rs @@ -144,7 +144,7 @@ impl DataSource for LazyIcebergDataSource { Err(err) => Err(iceberg_err(err)), }) .boxed(); - self.scan.open(feed, partition, context) + self.scan.open(feed, partition) } fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { diff --git a/iceberg/src/table_provider/scan.rs b/iceberg/src/table_provider/scan.rs index 271296ee1..aeeb2ec94 100644 --- a/iceberg/src/table_provider/scan.rs +++ b/iceberg/src/table_provider/scan.rs @@ -51,16 +51,15 @@ pub(crate) async fn plan_scan( let predicates = convert_filters_to_predicate(opts.filters); let partitioning = Partitioning::UnknownPartitioning(context.session_config().target_partitions()); - let scan = ScanConfig { - schema: output_schema, - iceberg_file_io: table.file_io().clone(), - partitioning: partitioning.clone(), - fetch: opts.fetch, - metrics: Default::default(), - iceberg_runtime: opts - .iceberg_runtime + let scan = ScanConfig::new( + output_schema, + partitioning.clone(), + opts.fetch, + table.file_io().clone(), + opts.iceberg_runtime .unwrap_or_else(iceberg::Runtime::current), - }; + &context, + ); let planned = if IcebergConfig::from_task_context(&context).eager_file_planning { if partitioning.partition_count() == 0 { return config_err!("Iceberg scan requires at least one output partition"); diff --git a/iceberg/tests/shared_reader.rs b/iceberg/tests/shared_reader.rs new file mode 100644 index 000000000..d155fba9a --- /dev/null +++ b/iceberg/tests/shared_reader.rs @@ -0,0 +1,149 @@ +#[cfg(test)] +mod tests { + use std::error::Error; + use std::fs; + use std::sync::Arc; + + use datafusion::arrow::array::{Int64Array, RecordBatch, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::datasource::source::DataSource; + use datafusion::execution::TaskContext; + use datafusion_distributed_iceberg::{EagerIcebergDataSource, iceberg}; + use futures::TryStreamExt; + use iceberg::arrow::arrow_schema_to_schema; + use iceberg::io::FileIO; + use iceberg::metadata_columns::{ + RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS, + RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS, + }; + use iceberg::scan::{FileScanTask, FileScanTaskDeleteFile}; + use iceberg::spec::{DataContentType, DataFileFormat}; + use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY}; + use tempfile::TempDir; + + #[tokio::test] + async fn partitions_share_loaded_position_deletes() -> Result<(), Box> { + let directory = TempDir::new()?; + let fixture = files_with_shared_deletes(&directory)?; + let context = Arc::new(TaskContext::default()); + let source = EagerIcebergDataSource::from_file_groups( + fixture.schema, + fixture.groups, + FileIO::new_with_fs(), + iceberg::Runtime::current(), + None, + &context, + )?; + let mut batches: Vec = + source.open(0, Arc::clone(&context))?.try_collect().await?; + + // The second partition must use the cached deletes, including its own + // row's deletion, rather than reopening the shared delete file. + fs::remove_file(fixture.delete_path)?; + batches.extend(source.open(1, context)?.try_collect::>().await?); + insta::assert_snapshot!(pretty_format_batches(&batches)?.to_string(), @" + +-------+ + | value | + +-------+ + | 2 | + | 4 | + +-------+ + "); + Ok(()) + } + + struct DeleteFixture { + schema: SchemaRef, + groups: Vec>, + delete_path: String, + } + + /// Two data files, each with its first row removed by one shared delete file. + fn files_with_shared_deletes(directory: &TempDir) -> Result> { + let schema = Arc::new(Schema::new(vec![field("value", DataType::Int64, 1)])); + let iceberg_schema = Arc::new(arrow_schema_to_schema(&schema)?); + let paths = ["first.parquet", "second.parquet"] + .map(|name| directory.path().join(name).to_str().unwrap().to_owned()); + for (path, values) in paths.iter().zip([[1, 2], [3, 4]]) { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values.to_vec()))], + )?; + write_parquet(path, &batch)?; + } + let delete_schema = Arc::new(Schema::new(vec![ + field( + RESERVED_COL_NAME_DELETE_FILE_PATH, + DataType::Utf8, + RESERVED_FIELD_ID_DELETE_FILE_PATH, + ), + field( + RESERVED_COL_NAME_DELETE_FILE_POS, + DataType::Int64, + RESERVED_FIELD_ID_DELETE_FILE_POS, + ), + ])); + let deletes = RecordBatch::try_new( + delete_schema, + vec![ + Arc::new(StringArray::from( + paths.iter().map(String::as_str).collect::>(), + )), + Arc::new(Int64Array::from(vec![0, 0])), + ], + )?; + let delete_path = directory + .path() + .join("deletes.parquet") + .to_str() + .unwrap() + .to_owned(); + write_parquet(&delete_path, &deletes)?; + let delete = FileScanTaskDeleteFile::builder() + .with_file_path(delete_path.clone()) + .with_file_size_in_bytes(fs::metadata(&delete_path)?.len()) + .with_file_type(DataContentType::PositionDeletes) + .with_file_format(DataFileFormat::Parquet) + .with_partition_spec_id(0) + .build()?; + let mut groups = Vec::new(); + for path in paths { + let size = fs::metadata(&path)?.len(); + groups.push(vec![ + FileScanTask::builder() + .with_file_size_in_bytes(size) + .with_start(0) + .with_length(size) + .with_record_count(Some(2)) + .with_data_file_path(path) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(Arc::clone(&iceberg_schema)) + .with_project_field_ids(vec![1]) + .with_deletes(vec![delete.clone()]) + .with_case_sensitive(true) + .build()?, + ]); + } + Ok(DeleteFixture { + schema, + groups, + delete_path, + }) + } + + fn field(name: &str, data_type: DataType, id: i32) -> Field { + Field::new(name, data_type, false).with_metadata( + [(PARQUET_FIELD_ID_META_KEY.to_owned(), id.to_string())] + .into_iter() + .collect(), + ) + } + + fn write_parquet(path: &str, batch: &RecordBatch) -> Result<(), Box> { + let mut writer = ArrowWriter::try_new(fs::File::create(path)?, batch.schema(), None)?; + writer.write(batch)?; + writer.close()?; + Ok(()) + } +}