diff --git a/docs/upgrade/5.0.0.md b/docs/upgrade/5.0.0.md index 6f9684c21..34a1c37d0 100644 --- a/docs/upgrade/5.0.0.md +++ b/docs/upgrade/5.0.0.md @@ -111,3 +111,14 @@ Some(Ok(DesiredTaskCountEventResponse::desired(0.25))) or `StorageFactory` to this crate, update those dependencies to the same revision. Alternatively, use `datafusion_distributed_iceberg::iceberg` for the Iceberg types you pass to its public APIs. + +## 5. Iceberg scan inputs + +`IcebergDataSource::feed()` now returns `Option<&WorkUnitFeed>`: +only lazy scans have a feed. Use `input()` to match `ScanInput`, or +`assigned_files()` to inspect eager files by local partition. Custom work-unit +feed registrations should use `and_then(IcebergDataSource::feed)` instead of `map`. + +Eager task plans now contain assigned file descriptors. Upgrade coordinators and +workers together before enabling eager planning; older workers do not decode +this new input. Lazy feed plans retain their existing representation. diff --git a/iceberg/README.md b/iceberg/README.md index 9b1a0442d..3b7e7cf2c 100644 --- a/iceberg/README.md +++ b/iceberg/README.md @@ -35,22 +35,26 @@ while building each physical scan, then deliver those same tasks to the reader a execution. Retained metadata is charged to the query memory pool, without a file-count limit. Eager planning uses metadata I/O up front. Scan statistics and desired distributed task counts use selected file work instead of snapshot totals; -reader settings and file assignment are unchanged. Both modes use the same dispatcher -and partition channels. Eager execution explicitly copies retained file descriptors -into owned work units; lazy execution moves newly discovered descriptors into them. +one Iceberg reader handles both inputs. ```text -Planning Execution -Eager: discover → retain files ──→ copy descriptors ─┐ -Lazy: retain scan definition ───→ discover files ──┤ - ▼ - shared task-first dispatcher - ↓ - partition channels - ↓ - local or remote readers +Eager: discover → assign files → task scan with local file groups → reader +Lazy: retain scan definition → discover and dispatch via feed → reader ``` +Eager scans use `ScanInput::AssignedFiles`; lazy scans use +`ScanInput::WorkUnitFeed`. Once DFD chooses the distributed task count, the +scale-up handler creates task-specific eager scans with concrete file groups. +The existing routing event can inspect `IcebergDataSource::assigned_files()` +in its `task_specialized_plan`, without accessing a coordinator-side feed. + +The codec sends only the assigned groups for that task, preserving local +partition membership. This increases plan payload size and shifts descriptor +transport from execution-time feed messages to worker setup. Greedy assignment +runs synchronously in the existing scale-up hook; this draft does not claim a +performance improvement. Reader execution clones descriptors as needed, while +coordinator-side assignments share the retained descriptors and memory reservation. + Configure the session before creating the `SessionState`: ```rust diff --git a/iceberg/src/codec.rs b/iceberg/src/codec.rs index 173e261b6..f374a2faa 100644 --- a/iceberg/src/codec.rs +++ b/iceberg/src/codec.rs @@ -1,7 +1,8 @@ use std::sync::Arc; +use crate::{AssignedFiles, ScanInput}; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::{Result, internal_err}; +use datafusion::common::{Result, exec_datafusion_err, internal_err}; use datafusion::datasource::source::DataSourceExec; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; @@ -69,11 +70,6 @@ impl PhysicalExtensionCodec for IcebergCodec { .schema .ok_or_else(|| proto_error("IcebergDataSource is missing its schema"))?; let schema = SchemaRef::new((&schema).try_into()?); - let feed = proto - .feed - .ok_or_else(|| proto_error("IcebergDataSource is missing its work-unit feed"))?; - let feed = - WorkUnitFeed::::from_proto(WorkUnitFeedProto { id: feed.id })?; let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self); let partitioning = parse_protobuf_partitioning( proto.partitioning.as_ref(), @@ -82,6 +78,45 @@ impl PhysicalExtensionCodec for IcebergCodec { proto_converter, )? .ok_or_else(|| proto_error("IcebergDataSource is missing its partitioning"))?; + let input = match (proto.feed, proto.assigned_files) { + (Some(feed), None) => { + ScanInput::WorkUnitFeed(Box::new(WorkUnitFeed::::from_proto( + WorkUnitFeedProto { id: feed.id }, + )?)) + } + (None, Some(files)) => { + if files.partitions.len() != partitioning.partition_count() + || files.partitions.is_empty() + { + return internal_err!("Iceberg assigned files do not match scan partitioning"); + } + let groups = files + .partitions + .into_iter() + .map(|group| { + group + .tasks + .into_iter() + .map(|bytes| { + rmp_serde::from_slice(&bytes).map_err(|error| { + exec_datafusion_err!("Invalid assigned Iceberg file: {error}") + }) + }) + .collect::>>() + }) + .collect::>>()?; + ScanInput::AssignedFiles(AssignedFiles::from_partitions( + groups, + ctx, + files.may_filter_rows, + )?) + } + _ => { + return internal_err!( + "Iceberg scan requires exactly one input: assigned files or feed" + ); + } + }; let fetch = proto .fetch .map(usize::try_from) @@ -100,7 +135,7 @@ impl PhysicalExtensionCodec for IcebergCodec { table_snapshot: None, iceberg_file_io, iceberg_runtime: self.iceberg_runtime.clone(), - feed, + input, })) } @@ -120,12 +155,44 @@ impl PhysicalExtensionCodec for IcebergCodec { return internal_err!("expected DataSourceExec wrapping IcebergDataSource"); }; - let feed = source.feed.to_proto(); + let (feed, assigned_files) = match source.input() { + ScanInput::WorkUnitFeed(feed) => ( + Some(pb::WorkUnitFeed { + id: feed.to_proto().id, + }), + None, + ), + ScanInput::AssignedFiles(files) => { + let partitions = (0..files.partition_count()) + .map(|partition| { + let tasks = files + .partition(partition)? + .map(|file| { + rmp_serde::to_vec_named(file).map_err(|error| { + exec_datafusion_err!( + "Cannot encode assigned Iceberg file: {error}" + ) + }) + }) + .collect::>>()?; + Ok(pb::FileGroup { tasks }) + }) + .collect::>>()?; + ( + None, + Some(pb::AssignedFiles { + partitions, + may_filter_rows: files.planned.may_filter_rows(), + }), + ) + } + }; let proto = pb::IcebergDataSource { schema: Some(datafusion_proto::protobuf::Schema::try_from( source.schema.as_ref(), )?), - feed: Some(pb::WorkUnitFeed { id: feed.id }), + feed, + assigned_files, partitioning: Some(serialize_partitioning( &source.partitioning, self, @@ -168,7 +235,10 @@ mod tests { decoded.partitioning.to_string() ); assert_eq!(source.fetch, decoded.fetch); - assert_eq!(source.feed.to_proto(), decoded.feed.to_proto()); + assert_eq!( + source.feed().unwrap().to_proto(), + decoded.feed().unwrap().to_proto() + ); assert_eq!( source.iceberg_file_io.config().props(), decoded.iceberg_file_io.config().props() diff --git a/iceberg/src/data_source.rs b/iceberg/src/data_source.rs index 36a49495a..935b182b8 100644 --- a/iceberg/src/data_source.rs +++ b/iceberg/src/data_source.rs @@ -30,7 +30,8 @@ pub use crate::column_statistics::compute_column_stats; use crate::column_statistics::selected_column_statistics; use crate::common::{df_err, iceberg_err}; use crate::planned_files::PlannedFiles; -use crate::{IcebergConfig, IcebergWorkUnitFeed}; +use crate::{AssignedFiles, IcebergConfig, IcebergWorkUnitFeed, ScanInput}; +use iceberg::scan::FileScanTaskStream; /// Snapshot summary keys defined by the Iceberg table spec: /// https://iceberg.apache.org/spec/#optional-snapshot-summary-fields @@ -40,93 +41,8 @@ use crate::{IcebergConfig, IcebergWorkUnitFeed}; const TOTAL_RECORDS: &str = "total-records"; const TOTAL_FILE_SIZE: &str = "total-files-size"; -/// Consumes a stream of [iceberg::scan::FileScanTask]s per partition and reads the underlying -/// files into an Arrow stream. -/// -/// By default [IcebergWorkUnitFeed] discovers file tasks progressively during execution. -/// Opt-in planning retains those tasks during physical planning and feeds the same tasks to -/// this reader at execution. This works in both single-node and distributed execution: -/// -/// ## Single Node -/// -/// [iceberg::scan::FileScanTask] are streamed in-memory, with as many parallel streams as -/// partitions this [IcebergDataSource] exposes: -/// -/// ```text -/// ┌────────────────────────────────────────────┐ -/// │ IcebergDataSource │ -/// │ │ -/// │┌──────────────────────────────────────────┐│ -/// ││ IcebergWorkUnitFeed ││ -/// ││┌────────────┐┌────────────┐┌────────────┐││ -/// │││ Feed 0 ││ Feed 1 ││ Feed 2 │││ -/// ││└──────┬─────┘└──────┬─────┘└──────┬─────┘││ -/// │└───────┼─────────────┼─────────────┼──────┘│ -/// │ .─────▼─────. .─────▼─────. .─────▼─────. │ -/// │ (FileScanTask (FileScanTask (FileScanTask )│ -/// │ .───────────. `─────┬─────' .───────────. │ -/// │ (FileScanTask ) │ (FileScanTask )│ -/// │ `─────┬─────' │ .───────────. │ -/// │ │ │ (FileScanTask )│ -/// │ │ │ `─────┬─────' │ -/// │ │ │ │ │ -/// │ ┌──────▼─────┐┌──────▼─────┐┌──────▼─────┐ │ -/// │ │Partition 0 ││Partition 1 ││Partition 2 │ │ -/// │ │ArrowReader ││ArrowReader ││ArrowReader │ │ -/// │ └──────┬─────┘└──────┬─────┘└──────┬─────┘ │ -/// │ │ │ │ │ -/// │ .─────▼─────. │ .─────▼─────. │ -/// │ ( RecordBatch ).─────▼─────.( RecordBatch )│ -/// │ `─────┬─────'( RecordBatch ).───────────. │ -/// │ │ `─────┬─────'( RecordBatch )│ -/// │ │ │ `───────────' │ -/// └────────┼─────────────┼─────────────┼───────┘ -/// ▼ ▼ ▼ -/// ``` -/// -/// ## Distributed -/// -/// [iceberg::scan::FileScanTask] are streamed over the network, with as many parallel streams as -/// partitions * distributed tasks: -/// -/// ```text -/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ -/// Coordinating Context │ -/// │ -/// ┌────────────────────────────────────────────────────────────────────────────────────────┐│ -/// ││ IcebergWorkUnitFeed │ -/// │┌─────────────┐┌─────────────┐┌────────────┐┌────────────┐┌─────────────┐┌─────────────┐││ -/// │││ Feed 0 ││ Feed 1 ││ Feed 2 ││ Feed 3 ││ Feed 4 ││ Feed 5 ││ -/// │└──────┬──────┘└─────┬───────┘└────┬───────┘└───────┬────┘└───────┬─────┘└──────┬──────┘││ -/// └└───────┼─────────────┼─────────────┼────────────────┼─────────────┼─────────────┼───────┴ -/// .─────▼─────. .─────▼─────. .─────▼─────. .─────▼─────. .─────▼─────. .─────▼─────. -/// (FileScanTask (FileScanTask (FileScanTask ) (FileScanTask (FileScanTask (FileScanTask ) -/// .───────────. `─────┬─────' .───────────. `─────┬─────' .───────────. `─────┬─────' -/// (FileScanTask ) │ (FileScanTask ) │ (FileScanTask ) │ -/// `─────┬─────' │ .───────────. │ `───────────' │ -/// │ │ (FileScanTask ) │ │ │ -/// Worker 0 │ │ `─────┬─────' │ │ │ Worker 1 -/// ┌ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ┐┌ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ┐ -/// ┌───────┼─────────────┼─────────────┼───────┐┌───────┼─────────────┼─────────────┼───────┐ -/// │ │ │ IcebergD│taSource │ ││ │ IcebergD│taSource │ │ │ -/// │ │ │ │ ││ │ │ │ │ -/// │ │┌──────▼─────┐┌──────▼─────┐┌──────▼─────┐ ││┌──────▼─────┐┌──────▼─────┐┌──────▼─────┐ │ │ -/// ││Partition 0 ││Partition 1 ││Partition 2 │ │││Partition 0 ││Partition 1 ││Partition 2 │ │ -/// │ ││ArrowReader ││ArrowReader ││ArrowReader │ │││ArrowReader ││ArrowReader ││ArrowReader │ │ │ -/// │└──────┬─────┘└──────┬─────┘└──────┬─────┘ ││└──────┬─────┘└──────┬─────┘└──────┬─────┘ │ -/// │ │ │ │ │ ││ │ │ │ │ │ -/// │ .─────▼─────. │ .─────▼─────. ││ │ ▼ ▼ │ -/// │ │( RecordBatch ).─────▼─────.( RecordBatch )││ .─────▼─────. .───────────. .───────────. │ │ -/// │ `─────┬─────'( RecordBatch ).───────────. ││( RecordBatch ( RecordBatch ) RecordBatch )│ -/// │ │ │ `─────┬─────'( RecordBatch )││ `─────┬─────' `───────────' `─────┬─────' │ │ -/// │ │ │ `───────────' ││ │ ( RecordBatch ) │ │ -/// │ │ │ │ │ ││ │ `─────┬─────' │ │ │ -/// └───────┼─────────────┼─────────────┼───────┘└───────┼─────────────┼─────────────┼───────┘ -/// │ ▼ ▼ ▼ ││ ▼ ▼ ▼ │ -/// ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ -/// ``` -/// -/// This distributed mechanism is transparent to this [DataSource]. +/// Reads Iceberg file tasks through one reader. Eager scans own an immutable +/// assignment in [`ScanInput::AssignedFiles`]; lazy scans use a work-unit feed. #[derive(Debug, Clone)] pub struct IcebergDataSource { pub(crate) schema: SchemaRef, @@ -137,7 +53,7 @@ pub struct IcebergDataSource { pub(crate) table_snapshot: Option, pub(crate) iceberg_file_io: FileIO, pub(crate) iceberg_runtime: iceberg::Runtime, - pub(crate) feed: WorkUnitFeed, + pub(crate) input: ScanInput, } /// Optional fields for building an [IcebergDataSource]. @@ -184,15 +100,14 @@ impl IcebergDataSource { iceberg_runtime: opts .iceberg_runtime .unwrap_or_else(iceberg::Runtime::current), - feed: WorkUnitFeed::new(IcebergWorkUnitFeed { + input: ScanInput::WorkUnitFeed(Box::new(WorkUnitFeed::new(IcebergWorkUnitFeed { iceberg_table: table, snapshot_id: opts.snapshot_id, projection, predicates, partitioning, - planned: None, sync_manager: Default::default(), - }), + }))), table_snapshot, column_stats: None, } @@ -207,28 +122,44 @@ impl IcebergDataSource { ); } if config.eager_file_planning - && let Some(feed) = self.feed.inner_mut() + && let Some(feed) = self.feed().and_then(WorkUnitFeed::inner) { - feed.prepare_files(&self.iceberg_runtime, context).await?; + let scan = feed.table_scan()?; + let may_filter_rows = feed.predicates.is_some(); + let (_cancel_on_drop, canceled) = oneshot::channel::<()>(); + let task = self.iceberg_runtime.io().spawn(async move { + tokio::select! { + _ = canceled => exec_err!("Iceberg file planning canceled"), + result = PlannedFiles::collect(scan, &context, may_filter_rows) => result, + } + }); + let planned = Arc::new(task.await.map_err(df_err)??); + let mut assignments = AssignedFiles::assign( + planned, + 1, + self.partitioning.partition_count(), + config.greedy_file_assignment, + )?; + self.input = ScanInput::AssignedFiles(assignments.remove(0)); } Ok(self) } /// Returns locally retained file work, if available; an empty scan is still Some. - fn planned_files(&self) -> Option<&PlannedFiles> { - self.feed.inner()?.planned.as_deref() + fn planned_files(&self) -> Option<&Arc> { + self.assigned_files().map(|files| &files.planned) } pub(crate) fn planned_file_count(&self) -> Option { - self.planned_files().map(|planned| planned.tasks().len()) + self.assigned_files().map(|files| files.iter().count()) } /// Whole-file I/O used for scan parallelism, independent of projected column sizes. pub(crate) fn file_scan_bytes(&self) -> Option { - self.feed.inner()?; - if let Some(planned) = self.planned_files() { - return Some(planned.bytes()); + if let Some(files) = self.assigned_files() { + return Some(files.bytes()); } + self.feed()?.inner()?; match &self.table_snapshot { Some(snapshot) => snapshot .summary() @@ -259,7 +190,7 @@ impl IcebergDataSource { .collect(), None => fields.to_vec(), }; - let planned = self.feed.inner().and_then(|feed| feed.planned.clone()); + let planned = self.planned_files().cloned(); let snapshot = self.table_snapshot.clone(); let (_cancel_on_drop, canceled) = oneshot::channel::<()>(); let task = self.iceberg_runtime.io().spawn(async move { @@ -282,10 +213,25 @@ impl IcebergDataSource { } impl IcebergDataSource { - /// Returns the [WorkUnitFeed] implementation that feeds this - /// DataSource with [iceberg::scan::FileScanTask] messages. - pub fn feed(&self) -> &WorkUnitFeed { - &self.feed + /// Input work used by this scan, available to routing before worker placement. + pub fn input(&self) -> &ScanInput { + &self.input + } + + /// Files owned by this scan, grouped by local partition. Lazy scans return None. + pub fn assigned_files(&self) -> Option<&AssignedFiles> { + match &self.input { + ScanInput::AssignedFiles(files) => Some(files), + ScanInput::WorkUnitFeed(_) => None, + } + } + + /// Work-unit feed for lazy scans. Eager scans deliver their assigned files in the plan. + pub fn feed(&self) -> Option<&WorkUnitFeed> { + match &self.input { + ScanInput::AssignedFiles(_) => None, + ScanInput::WorkUnitFeed(feed) => Some(feed), + } } } @@ -312,14 +258,21 @@ impl DataSource for IcebergDataSource { .with_row_selection_enabled(config.row_selection_enabled) .build(); - let feed = self - .feed - .feed(partition, context)? - .map(|work_unit_or_err| match work_unit_or_err { - Ok(work_unit) => work_unit.into_task().map_err(iceberg_err), - Err(err) => Err(iceberg_err(err)), - }) - .boxed(); + let feed: FileScanTaskStream = match &self.input { + ScanInput::AssignedFiles(files) => { + // Clone only descriptors owned by this partition as the reader polls. + let files = files.clone(); + let count = files.partition(partition)?.len(); + // Retain the assignment without borrowing the source in a 'static stream. + futures::stream::iter(0..count) + .map(move |offset| Ok(files.task(partition, offset).clone())) + .boxed() + } + ScanInput::WorkUnitFeed(feed) => feed + .feed(partition, context)? + .map(|unit| unit.and_then(|unit| unit.into_task()).map_err(iceberg_err)) + .boxed(), + }; let stream = reader .read(feed) @@ -339,20 +292,19 @@ impl DataSource for IcebergDataSource { fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "format=iceberg")?; - let Some(feed) = self.feed.inner() else { - return Ok(()); - }; - if let Some(projection) = &feed.projection { - write!(f, ", projection=[{}]", projection.join(", "))?; - } - if let Some(predicate) = &feed.predicates { - write!(f, ", predicate={predicate}")?; + if let Some(feed) = self.feed().and_then(WorkUnitFeed::inner) { + if let Some(projection) = &feed.projection { + write!(f, ", projection=[{}]", projection.join(", "))?; + } + if let Some(predicate) = &feed.predicates { + write!(f, ", predicate={predicate}")?; + } } - if let Some(planned) = self.planned_files() { + if let Some(planned) = self.assigned_files() { write!( f, ", planned_files={}, planned_bytes={}", - planned.tasks().len(), + planned.len(), planned.bytes() )?; } @@ -371,20 +323,20 @@ impl DataSource for IcebergDataSource { } fn partition_statistics(&self, _partition: Option) -> Result> { - if self.feed.inner().is_none() { + if self.feed().is_some_and(|feed| feed.inner().is_none()) { return Ok(Arc::new(Statistics::new_unknown(&self.schema))); } - let planned = self.planned_files(); + let planned = self.assigned_files(); let mut stats = if let Some(planned) = planned { let mut stats = Statistics::new_unknown(&self.schema); stats.num_rows = match planned.rows() { - Some(0) if planned.tasks().is_empty() => Precision::Exact(0), - Some(rows) if !planned.may_filter_rows() => Precision::Exact(rows), + Some(0) if planned.is_empty() => Precision::Exact(0), + Some(rows) if !planned.planned.may_filter_rows() => Precision::Exact(rows), Some(rows) => Precision::Inexact(rows), None => Precision::Absent, }; // File bytes describe scan IO, not the exact size of the filtered output. - stats.total_byte_size = if planned.may_filter_rows() { + stats.total_byte_size = if planned.planned.may_filter_rows() { Precision::Inexact(planned.bytes()) } else { Precision::Exact(planned.bytes()) diff --git a/iceberg/src/distributed_scale_up_handler.rs b/iceberg/src/distributed_scale_up_handler.rs new file mode 100644 index 000000000..65212398e --- /dev/null +++ b/iceberg/src/distributed_scale_up_handler.rs @@ -0,0 +1,43 @@ +use std::sync::Arc; + +use datafusion::common::Result; +use datafusion::datasource::source::DataSourceExec; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_distributed::{ + DistributedLeafExec, ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, +}; + +use crate::{AssignedFiles, IcebergConfig, IcebergDataSource, ScanInput}; + +/// Materialize eager assignments after the planner has chosen the final task count. +pub(crate) fn iceberg_scale_up( + ev: ScaleUpLeafNodeEvent, +) -> Option> { + let exec = ev.plan.downcast_ref::()?; + let source = exec.data_source().downcast_ref::()?; + let files = source.assigned_files()?; + if ev.task_count <= 1 { + return None; + } + Some((|| { + let assignments = AssignedFiles::assign( + Arc::clone(&files.planned), + ev.task_count, + files.partition_count(), + IcebergConfig::from_session_config(ev.session_config).greedy_file_assignment, + )?; + let variants = assignments + .into_iter() + .map(|files| { + let mut source = source.clone(); + source.input = ScanInput::AssignedFiles(files); + // Whole-scan column aggregates cannot describe a task's subset. + source.column_stats = None; + source.metrics = Default::default(); + DataSourceExec::from_data_source(source) as Arc + }) + .collect::>>(); + let plan = DistributedLeafExec::try_new(Arc::clone(ev.plan), variants)?; + Ok(ScaleUpLeafNodeEventResponse::new(Arc::new(plan))) + })()) +} diff --git a/iceberg/src/file_partitions.rs b/iceberg/src/file_partitions.rs index feb8385ab..6ce4efb70 100644 --- a/iceberg/src/file_partitions.rs +++ b/iceberg/src/file_partitions.rs @@ -22,7 +22,7 @@ pub(crate) enum FilePartitions { impl FilePartitions { /// Build destinations in the original retained-file order, without reordering files. - pub(crate) async fn greedy( + pub(crate) fn greedy( files: &[FileScanTask], task_count: usize, partitions_per_task: usize, @@ -37,19 +37,12 @@ impl FilePartitions { Reverse(index), )); partitions.push(0); - if index % 256 == 0 { - tokio::task::yield_now().await; - } } // Prefer fewer files, then task index, when byte loads tie. let mut loads: BinaryHeap = BinaryHeap::with_capacity(task_count); for task in 0..task_count { loads.push(Reverse((0_u128, 0_usize, task))); - if task % 256 == 0 { - tokio::task::yield_now().await; - } } - let mut assigned = 0; while let Some((bytes, _, _, Reverse(index))) = order.pop() { { let mut least_loaded = loads.peek_mut().expect("validated nonzero task count"); @@ -58,10 +51,6 @@ impl FilePartitions { *load += u128::from(bytes); *count += 1; } - assigned += 1; - if assigned % 256 == 0 { - tokio::task::yield_now().await; - } } Self::Greedy(partitions.into_iter()) } diff --git a/iceberg/src/iceberg_ext.rs b/iceberg/src/iceberg_ext.rs index e3391c8bd..0385da967 100644 --- a/iceberg/src/iceberg_ext.rs +++ b/iceberg/src/iceberg_ext.rs @@ -10,6 +10,7 @@ use iceberg_storage_opendal::OpenDalResolvingStorageFactory; use crate::codec::IcebergCodec; use crate::distributed_desired_task_count_handler::iceberg_desired_task_count; +use crate::distributed_scale_up_handler::iceberg_scale_up; use crate::{IcebergConfig, IcebergDataSource, IcebergTableProviderFactory}; /// Configuration required to register the Iceberg SQL integration. @@ -122,10 +123,11 @@ fn set_iceberg_integration(state: &mut SessionState, options: IcebergIntegration )), ); state.set_distributed_desired_task_count_handler(iceberg_desired_task_count); + state.set_distributed_scale_up_leaf_node_handler(iceberg_scale_up); state.set_distributed_work_unit_feed(|plan: &DataSourceExec| { plan.data_source() .downcast_ref::() - .map(IcebergDataSource::feed) + .and_then(IcebergDataSource::feed) }); state.set_distributed_user_codec(codec); } @@ -145,10 +147,11 @@ impl IcebergExt for SessionStateBuilder { )), ); self.set_distributed_desired_task_count_handler(iceberg_desired_task_count); + self.set_distributed_scale_up_leaf_node_handler(iceberg_scale_up); self.set_distributed_work_unit_feed(|plan: &DataSourceExec| { plan.data_source() .downcast_ref::() - .map(IcebergDataSource::feed) + .and_then(IcebergDataSource::feed) }); self.set_distributed_user_codec(codec); } diff --git a/iceberg/src/lib.rs b/iceberg/src/lib.rs index d597ca3c5..80db00d7a 100644 --- a/iceberg/src/lib.rs +++ b/iceberg/src/lib.rs @@ -3,18 +3,20 @@ //! Read-only Apache Iceberg integration for DataFusion Distributed. //! //! This crate ports the read path from Apache Iceberg Rust's DataFusion -//! integration. It deliberately contains no distributed execution adaptation -//! and no production Iceberg write or commit implementation. +//! integration, with distributed scan assignment and work delivery. It contains +//! no production Iceberg write or commit implementation. mod column_statistics; mod common; mod config; mod data_source; mod distributed_desired_task_count_handler; +mod distributed_scale_up_handler; mod file_partitions; mod iceberg_ext; mod planned_files; mod proto; +mod scan_input; mod table_provider; mod work_unit_feed; @@ -29,6 +31,7 @@ pub use data_source::IcebergDataSource; pub use distributed_desired_task_count_handler::iceberg_desired_task_count; pub use iceberg_ext::IcebergExt; pub use iceberg_ext::IcebergIntegrationOptions; +pub use scan_input::{AssignedFiles, ScanInput}; pub use table_provider::IcebergCatalogTableProvider; pub use table_provider::IcebergStaticTableProvider; pub use table_provider::IcebergTableProviderFactory; diff --git a/iceberg/src/planned_files.rs b/iceberg/src/planned_files.rs index f8889ece2..2d7087501 100644 --- a/iceberg/src/planned_files.rs +++ b/iceberg/src/planned_files.rs @@ -37,33 +37,52 @@ impl PlannedFiles { self.may_filter_rows } + fn new(context: &TaskContext, may_filter_rows: bool) -> Self { + Self { + tasks: Vec::new(), + bytes: 0, + rows: Some(0), + may_filter_rows, + _reservation: MemoryConsumer::new("Iceberg file planning") + .register(context.memory_pool()), + } + } + + fn push(&mut self, task: FileScanTask) -> Result<()> { + // Include owned paths/delete references and allow room for Vec capacity. + self._reservation + .try_grow(retained_size(&task).saturating_mul(2))?; + self.bytes = self.bytes.saturating_add(task.length() as usize); + self.rows = self.rows.and_then(|rows| { + task.record_count() + .and_then(|n| rows.checked_add(n as usize)) + }); + self.may_filter_rows |= !task.deletes().is_empty(); + self.tasks.push(task); + Ok(()) + } + + pub(crate) fn from_tasks( + tasks: Vec, + context: &TaskContext, + may_filter_rows: bool, + ) -> Result { + let mut planned = Self::new(context, may_filter_rows); + for task in tasks { + planned.push(task)?; + } + Ok(planned) + } + pub(crate) async fn collect( scan: TableScan, context: &TaskContext, may_filter_rows: bool, ) -> Result { - let reservation = - MemoryConsumer::new("Iceberg file planning").register(context.memory_pool()); - let mut planned = Self { - tasks: Vec::new(), - bytes: 0, - rows: Some(0), - may_filter_rows, - _reservation: reservation, - }; + let mut planned = Self::new(context, may_filter_rows); let mut stream = scan.plan_files().await.map_err(df_err)?; while let Some(task) = stream.try_next().await.map_err(df_err)? { - // Account for owned paths and delete references, and allow room for Vec capacity. - planned - ._reservation - .try_grow(retained_size(&task).saturating_mul(2))?; - planned.bytes = planned.bytes.saturating_add(task.length() as usize); - planned.rows = planned.rows.and_then(|rows| { - task.record_count() - .and_then(|count| rows.checked_add(count as usize)) - }); - planned.may_filter_rows |= !task.deletes().is_empty(); - planned.tasks.push(task); + planned.push(task)?; } Ok(planned) } diff --git a/iceberg/src/proto/generated/iceberg.rs b/iceberg/src/proto/generated/iceberg.rs index f3fed4fd5..a3937852b 100644 --- a/iceberg/src/proto/generated/iceberg.rs +++ b/iceberg/src/proto/generated/iceberg.rs @@ -1,6 +1,8 @@ // This file is @generated by prost-build. #[derive(Clone, PartialEq, ::prost::Message)] pub struct IcebergDataSource { + #[prost(message, optional, tag = "6")] + pub assigned_files: ::core::option::Option, #[prost(message, optional, tag = "1")] pub schema: ::core::option::Option<::datafusion_proto::protobuf::Schema>, #[prost(message, optional, tag = "2")] @@ -18,3 +20,16 @@ pub struct WorkUnitFeed { #[prost(bytes = "vec", tag = "1")] pub id: ::prost::alloc::vec::Vec, } + +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AssignedFiles { + #[prost(message, repeated, tag = "1")] + pub partitions: ::prost::alloc::vec::Vec, + #[prost(bool, tag = "2")] + pub may_filter_rows: bool, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct FileGroup { + #[prost(bytes = "vec", repeated, tag = "1")] + pub tasks: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec>, +} diff --git a/iceberg/src/proto/iceberg.proto b/iceberg/src/proto/iceberg.proto index 1c7f4b626..cad000b90 100644 --- a/iceberg/src/proto/iceberg.proto +++ b/iceberg/src/proto/iceberg.proto @@ -7,6 +7,7 @@ message IcebergDataSource { DataFusionPartitioning partitioning = 3; optional uint64 fetch = 4; map storage_properties = 5; + AssignedFiles assigned_files = 6; // Exactly one of feed or assigned_files must be set. } message WorkUnitFeed { @@ -16,3 +17,12 @@ message WorkUnitFeed { // The generated Rust fields use DataFusion's canonical protobuf types. message DataFusionSchema {} message DataFusionPartitioning {} + +message AssignedFiles { + repeated FileGroup partitions = 1; + bool may_filter_rows = 2; +} + +message FileGroup { + repeated bytes tasks = 1; +} diff --git a/iceberg/src/scan_input.rs b/iceberg/src/scan_input.rs new file mode 100644 index 000000000..8cc33f20e --- /dev/null +++ b/iceberg/src/scan_input.rs @@ -0,0 +1,140 @@ +use std::sync::Arc; + +use datafusion::common::{Result, exec_err}; +use datafusion::execution::TaskContext; +use datafusion_distributed::WorkUnitFeed; +use iceberg::scan::FileScanTask; + +use crate::IcebergWorkUnitFeed; +use crate::file_partitions::FilePartitions; +use crate::planned_files::PlannedFiles; + +/// How a scan obtains file tasks. Assignment policy is resolved before execution +/// for retained files; lazy scans receive their work from the coordinator's feed. +#[derive(Debug, Clone)] +#[non_exhaustive] +pub enum ScanInput { + /// Immutable files assigned to this scan, grouped by local output partition. + AssignedFiles(AssignedFiles), + /// File tasks discovered and delivered during execution. + WorkUnitFeed(Box>), +} + +/// A concrete assignment, shared by routing inspection, serialization and reading. +/// File descriptors stay shared on the coordinator; only this scan's groups are +/// serialized when sending a task to a worker. +#[derive(Debug, Clone)] +pub struct AssignedFiles { + pub(crate) planned: Arc, + partitions: Arc<[Vec]>, +} + +impl AssignedFiles { + /// Number of local output partitions, including empty partitions. + pub fn partition_count(&self) -> usize { + self.partitions.len() + } + + /// The files owned by a local output partition, in delivery order. + pub fn partition( + &self, + partition: usize, + ) -> Result> { + let Some(indices) = self.partitions.get(partition) else { + return exec_err!("Invalid Iceberg scan partition {partition}"); + }; + Ok(indices.iter().map(|&index| &self.planned.tasks()[index])) + } + + /// All files owned by this scan. This does not promise row ordering. + pub fn iter(&self) -> impl Iterator { + self.partitions + .iter() + .flatten() + .map(|&index| &self.planned.tasks()[index]) + } + + pub(crate) fn assign( + planned: Arc, + task_count: usize, + partitions_per_task: usize, + greedy: bool, + ) -> Result> { + if task_count + .checked_mul(partitions_per_task) + .is_none_or(|n| n == 0) + { + return exec_err!("Invalid Iceberg assignment partition count"); + } + let destinations = if greedy { + FilePartitions::greedy(planned.tasks(), task_count, partitions_per_task) + } else { + FilePartitions::RoundRobin { + next_file: 0, + task_count, + partitions_per_task, + } + }; + let mut groups = vec![vec![Vec::new(); partitions_per_task]; task_count]; + for (index, partition) in (0..planned.tasks().len()).zip(destinations) { + groups[partition / partitions_per_task][partition % partitions_per_task].push(index); + } + Ok(groups + .into_iter() + .map(|partitions| Self { + planned: Arc::clone(&planned), + partitions: partitions.into(), + }) + .collect()) + } + + pub(crate) fn from_partitions( + groups: Vec>, + context: &TaskContext, + may_filter_rows: bool, + ) -> Result { + let mut tasks = Vec::new(); + let mut partitions = Vec::new(); + for group in groups { + let start = tasks.len(); + tasks.extend(group); + partitions.push((start..tasks.len()).collect()); + } + Ok(Self { + planned: Arc::new(PlannedFiles::from_tasks(tasks, context, may_filter_rows)?), + partitions: partitions.into(), + }) + } + + pub(crate) fn task(&self, partition: usize, offset: usize) -> &FileScanTask { + &self.planned.tasks()[self.partitions[partition][offset]] + } + + /// Number of files owned by this scan. + pub fn len(&self) -> usize { + self.partitions.iter().map(Vec::len).sum() + } + + /// Whether this scan has no assigned files. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub(crate) fn bytes(&self) -> usize { + if self.len() == self.planned.tasks().len() { + return self.planned.bytes(); + } + self.iter().fold(0usize, |total, file| { + total.saturating_add(file.length() as usize) + }) + } + + pub(crate) fn rows(&self) -> Option { + if self.len() == self.planned.tasks().len() { + return self.planned.rows(); + } + self.iter().try_fold(0usize, |total, file| { + total.checked_add(usize::try_from(file.record_count()?).ok()?) + }) + } +} diff --git a/iceberg/src/test_utils/harness.rs b/iceberg/src/test_utils/harness.rs index 70ef6701c..6f0e0c2ce 100644 --- a/iceberg/src/test_utils/harness.rs +++ b/iceberg/src/test_utils/harness.rs @@ -101,9 +101,25 @@ impl IcebergTestHarness { .data_source() .downcast_ref::() .ok_or_else(|| internal_datafusion_err!("Expected an Iceberg data source"))?; + if let Some(files) = source.assigned_files() { + let assignments = crate::AssignedFiles::assign( + Arc::clone(&files.planned), + task_count, + files.partition_count(), + crate::IcebergConfig::from_session_config(&self.ctx.copied_config()) + .greedy_file_assignment, + )?; + let mut groups = Vec::new(); + for files in assignments { + for partition in 0..files.partition_count() { + groups.push(files.partition(partition)?.cloned().collect()); + } + } + return Ok(groups); + } let feed = source .feed() - .inner() + .and_then(|feed| feed.inner()) .ok_or_else(|| internal_datafusion_err!("Expected a local Iceberg feed"))?; let partitions = scan .properties() diff --git a/iceberg/src/work_unit_feed.rs b/iceberg/src/work_unit_feed.rs index 979a65000..c26cffa9a 100644 --- a/iceberg/src/work_unit_feed.rs +++ b/iceberg/src/work_unit_feed.rs @@ -11,19 +11,16 @@ use datafusion::execution::TaskContext; use datafusion::physical_expr::Partitioning; use datafusion_distributed::{DistributedWorkUnitFeedContext, WorkUnitFeedProvider}; use futures::StreamExt; -use futures::stream::{self, BoxStream}; +use futures::stream::BoxStream; use iceberg::expr::Predicate; use iceberg::scan::{FileScanTask, TableScan}; use prost::encoding::{DecodeContext, WireType}; use prost::{DecodeError, Message}; use tokio::sync::mpsc::UnboundedReceiver; -use tokio::sync::oneshot; use tokio_stream::wrappers::UnboundedReceiverStream; -use crate::IcebergConfig; use crate::common::df_err; use crate::file_partitions::FilePartitions; -use crate::planned_files::PlannedFiles; #[derive(Clone, PartialEq, prost::Message)] struct SerializedFileScanTask { @@ -114,7 +111,7 @@ impl Message for FileScanTaskWorkUnit { /// Work unit feed implementation that yields [FileScanTask] messages at execution time. /// -/// At execution time, one dispatcher reads either retained planned files or a live +/// At execution time, one dispatcher reads a live /// Iceberg scan, and places each [FileScanTask] in P * T output channels, where: /// - P is the number of partition in each distributed task. /// - T is the number of distributed tasks @@ -188,7 +185,6 @@ pub struct IcebergWorkUnitFeed { /// Ideally, both Range partitioning and hash partitioning should be supported. pub(crate) partitioning: Partitioning, /// Query-local tasks discovered before execution, when enabled. - pub(crate) planned: Option>, /// Container for the lazily initialized dispatcher shared by both discovery modes. /// It will start as soon as the first [IcebergWorkUnitFeed::feed] is called. pub(crate) sync_manager: OnceLock>>, @@ -202,14 +198,13 @@ impl Clone for IcebergWorkUnitFeed { projection: self.projection.clone(), predicates: self.predicates.clone(), partitioning: self.partitioning.clone(), - planned: self.planned.clone(), sync_manager: Default::default(), } } } impl IcebergWorkUnitFeed { - fn table_scan(&self) -> Result { + pub(crate) fn table_scan(&self) -> Result { let builder = match self.snapshot_id { Some(id) => self.iceberg_table.scan().snapshot_id(id), None => self.iceberg_table.scan(), @@ -223,26 +218,6 @@ impl IcebergWorkUnitFeed { } builder.build().map_err(df_err) } - - pub(crate) async fn prepare_files( - &mut self, - runtime: &iceberg::Runtime, - context: Arc, - ) -> Result<()> { - let scan = self.table_scan()?; - let may_filter_rows = self.predicates.is_some(); - // Keep manifest I/O off an embedder's IO-disabled query runtime. Cancel the - // spawned planner if the query abandons physical planning. - let (_cancel_on_drop, canceled) = oneshot::channel::<()>(); - let task = runtime.io().spawn(async move { - tokio::select! { - _ = canceled => exec_err!("Iceberg file planning canceled"), - result = PlannedFiles::collect(scan, &context, may_filter_rows) => result, - } - }); - self.planned = Some(Arc::new(task.await.map_err(df_err)??)); - Ok(()) - } } type TakeableVec = Vec>>; @@ -274,8 +249,6 @@ impl WorkUnitFeedProvider for IcebergWorkUnitFeed { // Both discovery modes share one dispatcher, initialized by the first feed call. let sync_manager_or_err = self.sync_manager.get_or_init(|| { - let planned = self.planned.clone(); - let greedy = IcebergConfig::from_task_context(&ctx).greedy_file_assignment; let table_scan = self.table_scan()?; // Fanout the FileScanTask stream across P * T output channels where: @@ -294,34 +267,16 @@ impl WorkUnitFeedProvider for IcebergWorkUnitFeed { // the return streams of the `feed()` method. let task = SpawnedTask::spawn(async move { let result: Result<()> = async { - let mut destinations = if greedy { - let files = planned.as_ref().ok_or_else(|| { - exec_datafusion_err!( - "Greedy Iceberg assignment requires eager file planning" - ) - })?; - FilePartitions::greedy(files.tasks(), task_count, partitions_per_task).await - } else { - FilePartitions::RoundRobin { - next_file: 0, - task_count, - partitions_per_task, - } + let mut destinations = FilePartitions::RoundRobin { + next_file: 0, + task_count, + partitions_per_task, }; - let mut files: BoxStream<'_, Result> = - if let Some(planned) = &planned { - // The plan retains reusable descriptors; execution owns its copies. - stream::iter(planned.tasks().iter().cloned()) - .map(Ok) - .boxed() - } else { - table_scan - .plan_files() - .await - .map_err(df_err)? - .map(|task| task.map_err(df_err)) - .boxed() - }; + let mut files = table_scan + .plan_files() + .await + .map_err(df_err)? + .map(|task| task.map_err(df_err)); let mut dispatched = 0; while let Some(task) = files.next().await { diff --git a/iceberg/tests/assigned_file_routing.rs b/iceberg/tests/assigned_file_routing.rs new file mode 100644 index 000000000..6a77faf2d --- /dev/null +++ b/iceberg/tests/assigned_file_routing.rs @@ -0,0 +1,191 @@ +#[cfg(all(test, feature = "integration"))] +mod tests { + use std::sync::{Arc, Mutex}; + + use async_trait::async_trait; + use datafusion::common::Result; + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion::datasource::source::DataSourceExec; + use datafusion::execution::TaskContext; + use datafusion::physical_plan::{ExecutionPlan, collect}; + use datafusion_distributed::DistributedCodec; + use datafusion_distributed::{ + DistributedExt, RouteTaskEvent, RouteTaskEventResponse, RouteTaskHandler, + }; + use datafusion_distributed_iceberg::test_utils::IcebergTestHarness; + use datafusion_distributed_iceberg::{IcebergConfig, IcebergDataSource}; + use datafusion_proto::physical_plan::AsExecutionPlan; + use datafusion_proto::protobuf::PhysicalPlanNode; + + #[tokio::test] + async fn routing_inspects_assigned_files_in_the_worker_facing_plan() -> Result<()> { + let sql = + "SELECT pickup_date, COUNT(*) FROM taxi GROUP BY pickup_date ORDER BY pickup_date"; + let (_, expected) = IcebergTestHarness::new().await?.query(sql).await?; + for (eager, greedy) in [(false, false), (true, false), (true, true)] { + let routes = Routes::default(); + let harness = distributed(&routes, eager, greedy).await?; + for _ in 0..2 { + let (_, actual) = harness.query(sql).await?; + assert_eq!(actual, expected); + let records = routes.take(); + assert_eq!(records.len(), 3); + if eager { + let mut files: Vec<_> = + records.into_iter().flat_map(|r| r.files.unwrap()).collect(); + assert_eq!(files.len(), 7); + files.sort(); + files.dedup(); + assert_eq!(files.len(), 7); + } else { + assert!(records.iter().all(|r| r.files.is_none())); + } + } + } + Ok(()) + } + + #[tokio::test] + async fn routing_sees_only_the_assigned_union_branch() -> Result<()> { + let routes = Routes::default(); + let harness = distributed(&routes, true, true).await?; + let sql = "WITH a AS ( + SELECT 'a' AS kind, pickup_date FROM taxi + WHERE pickup_date = DATE '2024-01-08' AND trip_distance > 1 + ), b AS ( + SELECT 'b' AS kind, pickup_date FROM taxi + WHERE pickup_date = DATE '2024-01-08' AND trip_distance > 2 + ) SELECT kind, pickup_date, COUNT(*) FROM ( + SELECT * FROM a UNION ALL SELECT * FROM b + ) GROUP BY kind, pickup_date ORDER BY kind, pickup_date"; + let (_, expected) = IcebergTestHarness::new().await?.query(sql).await?; + let (_, actual) = harness.query(sql).await?; + assert_eq!(actual, expected); + let records = routes.take(); + assert_eq!(records.len(), 2); + for record in &records { + assert_eq!(record.files.as_ref().unwrap().len(), 1); + } + assert_eq!(records[0].files, records[1].files); + Ok(()) + } + + #[tokio::test] + async fn assigned_scan_and_decoded_scan_can_be_read_again() -> Result<()> { + let harness = distributed(&Routes::default(), true, true).await?; + let scan = harness.scan().await?; + let decoded = harness.roundtrip_plan(Arc::clone(&scan))?; + let mut expected = None; + for plan in [scan, decoded] { + let source = iceberg_source(&plan).unwrap(); + let files = source.assigned_files().unwrap(); + assert!(files.partition(files.partition_count()).is_err()); + for _ in 0..2 { + let batches = collect(Arc::clone(&plan), Arc::new(TaskContext::default())).await?; + let rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert!(rows > 0); + assert_eq!(rows, *expected.get_or_insert(rows)); + } + } + Ok(()) + } + + async fn distributed(routes: &Routes, eager: bool, greedy: bool) -> Result { + IcebergTestHarness::builder() + .with_workers(3) + .configure_session(|state| { + let mut config = IcebergConfig::default(); + config.eager_file_planning = eager; + config.greedy_file_assignment = greedy; + state + .with_distributed_option_extension(config) + .with_distributed_route_task_handler(routes.clone()) + .with_distributed_file_scan_config_bytes_per_partition(1) + })? + .build() + .await + } + + struct Record { + files: Option>, + } + + #[derive(Clone, Default)] + struct Routes(Arc>>); + + impl Routes { + fn take(&self) -> Vec { + std::mem::take(&mut *self.0.lock().unwrap()) + } + + fn inspect(&self, ev: RouteTaskEvent<'_>) -> Result { + let mut found = false; + ev.task_specialized_plan.apply(|plan| { + if let Some(source) = iceberg_source(plan) { + if let Some(files) = source.assigned_files() { + let codec = + DistributedCodec::new_combined_with_user(ev.task_ctx.session_config()); + let proto = + PhysicalPlanNode::try_from_physical_plan(Arc::clone(plan), &codec)?; + let decoded = proto.try_into_physical_plan(ev.task_ctx, &codec)?; + let restored = iceberg_source(&decoded).unwrap().assigned_files().unwrap(); + assert_eq!(files.partition_count(), restored.partition_count()); + for partition in 0..files.partition_count() { + let expected: Vec<_> = files + .partition(partition)? + .map(|file| (file.data_file_path(), file.length())) + .collect(); + let actual: Vec<_> = restored + .partition(partition)? + .map(|file| (file.data_file_path(), file.length())) + .collect(); + assert_eq!(expected, actual); + } + } + let files = source.assigned_files().map(|files| { + assert!(source.feed().is_none()); + files + .iter() + .map(|file| file.data_file_path().to_owned()) + .collect() + }); + if files.is_none() { + assert!(source.feed().unwrap().inner().is_none()); + } + self.0.lock().unwrap().push(Record { files }); + found = true; + } + Ok(TreeNodeRecursion::Continue) + })?; + Ok(found) + } + } + + #[async_trait] + impl RouteTaskHandler for Routes { + async fn handle(&self, ev: RouteTaskEvent<'_>) -> Option> { + match self.inspect(ev) { + Err(error) => return Some(Err(error)), + Ok(false) => return None, + Ok(true) => {} + } + let urls = match ev.worker_resolver.get_urls() { + Ok(urls) if !urls.is_empty() => urls, + Ok(_) => return None, + Err(error) => return Some(Err(error)), + }; + // Exercise the unchanged dialer with the inspected worker-facing scan. + Some( + ev.dialer + .dial(urls[ev.task_key.task_number % urls.len()].clone()) + .await, + ) + } + } + + fn iceberg_source(plan: &Arc) -> Option<&IcebergDataSource> { + plan.downcast_ref::()? + .data_source() + .downcast_ref::() + } +} diff --git a/iceberg/tests/file_planning.rs b/iceberg/tests/file_planning.rs index a2110b1ae..01a80a467 100644 --- a/iceberg/tests/file_planning.rs +++ b/iceberg/tests/file_planning.rs @@ -9,7 +9,6 @@ mod tests { use datafusion::execution::memory_pool::GreedyMemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_plan::statistics::{StatisticsArgs, StatisticsContext}; - use datafusion::prelude::SessionConfig; use datafusion_distributed::{DistributedExt, WorkUnitFeedProvider}; use datafusion_distributed_iceberg::test_utils::{FIXTURE_URI, IcebergTestHarness}; use datafusion_distributed_iceberg::{IcebergConfig, IcebergDataSource, IcebergWorkUnitFeed}; @@ -183,55 +182,29 @@ mod tests { } #[tokio::test] - async fn both_modes_claim_feeds_once_and_clones_get_fresh_execution() -> Result<()> { - for (eager, greedy) in [(false, false), (true, false), (true, true)] { - let harness = IcebergTestHarness::builder() - .configure_session(|mut state| { - let mut config = IcebergConfig::default(); - config.eager_file_planning = eager; - config.greedy_file_assignment = greedy; - state - .config() - .get_or_insert_default() - .options_mut() - .extensions - .insert(config); - Ok(state) - })? - .build() - .await?; - let scan = harness.scan().await?; - let partitions = scan.properties().partitioning.partition_count(); - let provider = scan - .downcast_ref::() - .unwrap() - .data_source() - .downcast_ref::() - .unwrap() - .feed() - .inner() - .unwrap(); - let first = read_feed(provider, partitions, greedy).await?; - assert_eq!(first.len(), 7); - assert!(provider.feed(0, Arc::new(TaskContext::default())).is_err()); - assert_eq!( - read_feed(&provider.clone(), partitions, greedy).await?, - first - ); - } + async fn lazy_feeds_are_claimed_once_and_clones_get_fresh_execution() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let scan = harness.scan().await?; + let partitions = scan.properties().partitioning.partition_count(); + let provider = scan + .downcast_ref::() + .unwrap() + .data_source() + .downcast_ref::() + .unwrap() + .feed() + .unwrap() + .inner() + .unwrap(); + let first = read_feed(provider, partitions).await?; + assert_eq!(first.len(), 7); + assert!(provider.feed(0, Arc::new(TaskContext::default())).is_err()); + assert_eq!(read_feed(&provider.clone(), partitions).await?, first); Ok(()) } - async fn read_feed( - provider: &IcebergWorkUnitFeed, - partitions: usize, - greedy: bool, - ) -> Result>> { - let mut config = SessionConfig::new(); - let mut iceberg = IcebergConfig::default(); - iceberg.greedy_file_assignment = greedy; - config.options_mut().extensions.insert(iceberg); - let context = Arc::new(TaskContext::default().with_session_config(config)); + async fn read_feed(provider: &IcebergWorkUnitFeed, partitions: usize) -> Result>> { + let context = Arc::new(TaskContext::default()); let mut files = vec![]; for partition in 0..partitions { let units: Vec<_> = provider diff --git a/iceberg/tests/work_distribution.rs b/iceberg/tests/work_distribution.rs index 924f01a5f..260783920 100644 --- a/iceberg/tests/work_distribution.rs +++ b/iceberg/tests/work_distribution.rs @@ -69,11 +69,31 @@ mod tests { GROUP BY pickup_date ORDER BY pickup_date" ); let (plan, results) = harness.query_with_metrics(&sql).await?; - let scan = plan + // Assigned inputs render one scan variant per task; feeds render one scan + // with all task metrics. Assert the same per-task distribution in either form. + let mut rows: Vec<_> = plan .lines() - .find(|line| line.contains("DataSourceExec: format=iceberg")) - .expect("query must execute an Iceberg scan"); - assert!(scan.contains(expected_rows), "{plan}"); + .filter(|line| line.contains("DataSourceExec: format=iceberg")) + .flat_map(|line| { + line.split_once("output_rows={") + .unwrap() + .1 + .split_once('}') + .unwrap() + .0 + .split(", ") + }) + .collect(); + let mut expected: Vec<_> = expected_rows + .strip_prefix("output_rows={") + .unwrap() + .strip_suffix('}') + .unwrap() + .split(", ") + .collect(); + rows.sort_unstable(); + expected.sort_unstable(); + assert_eq!(rows, expected, "{plan}"); let (_, local_results) = IcebergTestHarness::new().await?.query(&sql).await?; assert_eq!(results, local_results); Ok(())