feat: expose local task plans to routing handlers

docs/source/advanced/06-worker-routing.md

@@ -66,3 +66,21 @@ workers by hashing the file path, so each worker can serve them from an
 in-memory cache on repeat queries, see the
 [custom_worker_url_routing.rs](https://github.com/datafusion-contrib/datafusion-distributed/blob/main/examples/custom_worker_url_routing.rs)
 example.
+
+### Inspecting native scan metadata
+
+`RouteTaskEvent::task_specialized_plan` is the worker-facing plan, whose work-unit
+feeds are remote handles. `local_task_plan` provides the same task's coordinator
+view before that conversion, retaining native providers and their metadata.
+It is for inspection, not execution: do not consume its work-unit streams.
+The supplied dialer continues to send the worker-facing plan.
+
+Use `visit_local_task_plan` to inspect scans with their scan-local
+`DistributedTaskContext`. This skips unrelated branches and accounts for task
+indices/counts inside isolated UNION children. The event's stage-wide task number
+is not necessarily a scan's task index.
+
+For eager Iceberg scans, `IcebergDataSource::assigned_files` borrows the files for
+that context using the same assignment policy as delivery. It returns `None` for
+lazy scans and remote feeds. Applications can use the file descriptors in their
+own routing policy; DFD does not prescribe a cache key or worker-affinity policy.

docs/upgrade/5.0.0.md

@@ -111,3 +111,11 @@ 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. Local task plans in routing events
+
+If you construct `RouteTaskEvent` directly (for example in handler tests), provide
+its new `local_task_plan` field: the task-specialized coordinator plan before
+native work-unit feeds are converted to remote handles. Keep
+`task_specialized_plan` as the worker-facing plan. Exhaustive destructuring must
+include the new field or use `..`.

iceberg/src/data_source.rs

@@ -17,10 +17,11 @@ use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSe
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{DisplayFormatType, SortOrderPushdownResult};
 use datafusion::prelude::Expr;
-use datafusion_distributed::WorkUnitFeed;
+use datafusion_distributed::{DistributedTaskContext, WorkUnitFeed};
 use futures::{StreamExt, TryStreamExt};
 use iceberg::arrow::ArrowReaderBuilder;
 use iceberg::io::FileIO;
+use iceberg::scan::FileScanTask;
 use iceberg::spec::SnapshotRef;
 use iceberg::table::Table;
 use iceberg_datafusion::read::convert_filters_to_predicate;
@@ -29,6 +30,7 @@ use tokio::sync::oneshot;
 pub use crate::column_statistics::compute_column_stats;
 use crate::column_statistics::selected_column_statistics;
 use crate::common::{df_err, iceberg_err};
+use crate::file_partitions::FilePartitions;
 use crate::planned_files::PlannedFiles;
 use crate::{IcebergConfig, IcebergWorkUnitFeed};
 
@@ -282,6 +284,42 @@ impl IcebergDataSource {
 }
 
 impl IcebergDataSource {
+    /// Borrows the eagerly planned files assigned to one scan-local task.
+    ///
+    /// Returns `None` for lazy scans and remote feeds, and an empty iterator for
+    /// an eager task with no files. Uses the same assignment policy as delivery,
+    /// without reading metadata or consuming the feed. Greedy assignment allocates
+    /// temporary destination indices; file descriptors are borrowed, not cloned.
+    ///
+    /// Pass the context supplied by `RouteTaskEvent::visit_local_task_plan` and
+    /// the query's task context so isolated UNION branches and configuration agree
+    /// with execution. Invalid task/partition counts return an error.
+    pub async fn assigned_files(
+        &self,
+        task: DistributedTaskContext,
+        context: &TaskContext,
+    ) -> Result<Option<impl Iterator<Item = &FileScanTask>>> {
+        if task.task_count == 0 || task.task_index >= task.task_count {
+            return exec_err!("Invalid Iceberg task index/count");
+        }
+        let Some(planned) = self.planned_files() else {
+            return Ok(None);
+        };
+        let partitions_per_task = self.partitioning.partition_count();
+        let destinations = FilePartitions::new(
+            Some(planned.tasks()),
+            IcebergConfig::from_session_config(context.session_config()).greedy_file_assignment,
+            task.task_count,
+            partitions_per_task,
+        )
+        .await?;
+        Ok(Some(planned.tasks().iter().zip(destinations).filter_map(
+            move |(file, partition)| {
+                (partition / partitions_per_task == task.task_index).then_some(file)
+            },
+        )))
+    }
+
     /// Returns the [WorkUnitFeed] implementation that feeds this
     /// DataSource with [iceberg::scan::FileScanTask] messages.
     pub fn feed(&self) -> &WorkUnitFeed<IcebergWorkUnitFeed> {

iceberg/src/file_partitions.rs

@@ -2,14 +2,15 @@ use std::cmp::Reverse;
 use std::collections::BinaryHeap;
 use std::vec::IntoIter;
 
+use datafusion::common::{Result, exec_datafusion_err, exec_err};
 use iceberg::scan::FileScanTask;
 
 // Largest file first; object identity breaks ties independently of manifest order.
 type FilePriority<'a> = (u64, Reverse<&'a str>, Reverse<u64>, Reverse<usize>);
 type TaskLoad = Reverse<(u128, usize, usize)>;
 
-/// Destination partitions in file-stream order. The dispatcher validates nonzero
-/// counts and their product before constructing either routing policy.
+/// Destination partitions in file-stream order. `new` validates nonzero counts
+/// and their product before constructing either assignment policy.
 #[derive(Debug)]
 pub(crate) enum FilePartitions {
     RoundRobin {
@@ -21,12 +22,34 @@ pub(crate) enum FilePartitions {
 }
 
 impl FilePartitions {
-    /// Build destinations in the original retained-file order, without reordering files.
-    pub(crate) async fn greedy(
-        files: &[FileScanTask],
+    pub(crate) async fn new(
+        files: Option<&[FileScanTask]>,
+        greedy: bool,
         task_count: usize,
         partitions_per_task: usize,
-    ) -> Self {
+    ) -> Result<Self> {
+        if task_count
+            .checked_mul(partitions_per_task)
+            .is_none_or(|total| total == 0)
+        {
+            return exec_err!("Invalid Iceberg feed partition count");
+        }
+        if greedy {
+            let files = files.ok_or_else(|| {
+                exec_datafusion_err!("Greedy Iceberg assignment requires eager file planning")
+            })?;
+            Ok(Self::greedy(files, task_count, partitions_per_task).await)
+        } else {
+            Ok(Self::RoundRobin {
+                next_file: 0,
+                task_count,
+                partitions_per_task,
+            })
+        }
+    }
+
+    /// Build destinations in the original retained-file order, without reordering files.
+    async fn greedy(files: &[FileScanTask], task_count: usize, partitions_per_task: usize) -> Self {
         let mut order: BinaryHeap<FilePriority<'_>> = BinaryHeap::with_capacity(files.len());
         let mut partitions = Vec::with_capacity(files.len());
         for (index, file) in files.iter().enumerate() {

iceberg/src/work_unit_feed.rs

@@ -294,20 +294,13 @@ 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::new(
+                        planned.as_ref().map(|files| files.tasks()),
+                        greedy,
+                        task_count,
+                        partitions_per_task,
+                    )
+                    .await?;
                     let mut files: BoxStream<'_, Result<FileScanTask>> =
                         if let Some(planned) = &planned {
                             // The plan retains reusable descriptors; execution owns its copies.

iceberg/tests/file_planning.rs

@@ -10,13 +10,68 @@ mod tests {
     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::{DistributedExt, DistributedTaskContext, WorkUnitFeedProvider};
     use datafusion_distributed_iceberg::test_utils::{FIXTURE_URI, IcebergTestHarness};
     use datafusion_distributed_iceberg::{IcebergConfig, IcebergDataSource, IcebergWorkUnitFeed};
     use futures::TryStreamExt;
     use iceberg::spec::{DataContentType, DataFileBuilder, DataFileFormat, Literal, Struct};
     use prost::Message;
 
+    #[tokio::test]
+    async fn inspected_assignment_matches_delivered_files() -> Result<(), Box<dyn Error>> {
+        for sizes in [vec![100, 90, 10, 5, 4, 3], vec![10], vec![]] {
+            let harness = file_harness(&sizes, false).await?;
+            let scan = harness.scan().await?;
+            let source = scan
+                .downcast_ref::<DataSourceExec>()
+                .unwrap()
+                .data_source()
+                .downcast_ref::<IcebergDataSource>()
+                .unwrap();
+            let mut config = SessionConfig::new();
+            let mut iceberg = IcebergConfig::default();
+            iceberg.greedy_file_assignment = true;
+            config.options_mut().extensions.insert(iceberg);
+            let context = TaskContext::default().with_session_config(config);
+            let delivered = harness.scan_file_tasks(2).await?;
+            for task_index in 0..2 {
+                let mut expected: Vec<_> = delivered[task_index * 2..task_index * 2 + 2]
+                    .iter()
+                    .flatten()
+                    .map(|file| file.data_file_path())
+                    .collect();
+                let mut inspected: Vec<_> = source
+                    .assigned_files(
+                        DistributedTaskContext {
+                            task_index,
+                            task_count: 2,
+                        },
+                        &context,
+                    )
+                    .await?
+                    .expect("eager scan")
+                    .map(|file| file.data_file_path())
+                    .collect();
+                expected.sort();
+                inspected.sort();
+                assert_eq!(inspected, expected);
+            }
+            assert!(
+                source
+                    .assigned_files(
+                        DistributedTaskContext {
+                            task_index: 2,
+                            task_count: 2
+                        },
+                        &context,
+                    )
+                    .await
+                    .is_err()
+            );
+        }
+        Ok(())
+    }
+
     #[tokio::test]
     async fn greedy_routing_balances_bytes_across_tasks_and_local_partitions()
     -> Result<(), Box<dyn Error>> {