diff --git a/docs/source/advanced/06-worker-routing.md b/docs/source/advanced/06-worker-routing.md index 4b1fee80b..ad7413564 100644 --- a/docs/source/advanced/06-worker-routing.md +++ b/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. diff --git a/docs/upgrade/5.0.0.md b/docs/upgrade/5.0.0.md index 6f9684c21..0730b6373 100644 --- a/docs/upgrade/5.0.0.md +++ b/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 `..`. diff --git a/iceberg/src/data_source.rs b/iceberg/src/data_source.rs index 36a49495a..b2cd3d772 100644 --- a/iceberg/src/data_source.rs +++ b/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>> { + 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 { diff --git a/iceberg/src/file_partitions.rs b/iceberg/src/file_partitions.rs index feb8385ab..319e25a3c 100644 --- a/iceberg/src/file_partitions.rs +++ b/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, Reverse); 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 { + 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> = BinaryHeap::with_capacity(files.len()); let mut partitions = Vec::with_capacity(files.len()); for (index, file) in files.iter().enumerate() { diff --git a/iceberg/src/work_unit_feed.rs b/iceberg/src/work_unit_feed.rs index 979a65000..b5051e8b3 100644 --- a/iceberg/src/work_unit_feed.rs +++ b/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> = if let Some(planned) = &planned { // The plan retains reusable descriptors; execution owns its copies. diff --git a/iceberg/tests/file_planning.rs b/iceberg/tests/file_planning.rs index a2110b1ae..0525f385f 100644 --- a/iceberg/tests/file_planning.rs +++ b/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> { + 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::() + .unwrap() + .data_source() + .downcast_ref::() + .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> { diff --git a/iceberg/tests/local_task_routing.rs b/iceberg/tests/local_task_routing.rs new file mode 100644 index 000000000..8efffbbbd --- /dev/null +++ b/iceberg/tests/local_task_routing.rs @@ -0,0 +1,150 @@ +#[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::physical_plan::ExecutionPlan; + use datafusion_distributed::{ + DistributedExt, DistributedTaskContext, RouteTaskEvent, RouteTaskEventResponse, + RouteTaskHandler, + }; + use datafusion_distributed_iceberg::test_utils::IcebergTestHarness; + use datafusion_distributed_iceberg::{IcebergConfig, IcebergDataSource}; + + #[tokio::test] + async fn routing_can_inspect_native_files_without_consuming_them() -> 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_visits_only_assigned_union_branches_with_local_task_context() -> 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.task, + DistributedTaskContext { + task_index: 0, + task_count: 1 + } + ); + assert_eq!(record.files.as_ref().unwrap().len(), 1); + } + assert_eq!(records[0].files, records[1].files); + 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 { + task: DistributedTaskContext, + files: Option>, + } + + #[derive(Clone, Default)] + struct Routes(Arc>>); + + impl Routes { + fn take(&self) -> Vec { + std::mem::take(&mut *self.0.lock().unwrap()) + } + + async fn inspect(&self, ev: RouteTaskEvent<'_>) -> Result<()> { + // The existing routing view still describes what the worker receives. + ev.task_specialized_plan.apply(|plan| { + if let Some(source) = iceberg_source(plan) { + assert!(source.feed().inner().is_none()); + } + Ok(TreeNodeRecursion::Continue) + })?; + let mut scans = Vec::new(); + ev.visit_local_task_plan(|plan, task| { + if iceberg_source(plan).is_some() { + scans.push((Arc::clone(plan), task)); + } + Ok(TreeNodeRecursion::Continue) + })?; + for (plan, task) in scans { + let source = iceberg_source(&plan).unwrap(); + assert!(source.feed().inner().is_some()); + let files = source + .assigned_files(task, ev.task_ctx) + .await? + .map(|files| files.map(|file| file.data_file_path().to_owned()).collect()); + self.0.lock().unwrap().push(Record { task, files }); + } + Ok(()) + } + } + + #[async_trait] + impl RouteTaskHandler for Routes { + async fn handle(&self, ev: RouteTaskEvent<'_>) -> Option> { + if let Err(error) = self.inspect(ev).await { + return Some(Err(error)); + } + // Let the built-in router execute the same task after inspection. + None + } + } + + fn iceberg_source(plan: &Arc) -> Option<&IcebergDataSource> { + plan.downcast_ref::()? + .data_source() + .downcast_ref::() + } +} diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 3ec35ebff..81c75edf9 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -152,6 +152,7 @@ impl<'a> StageCoordinator<'a> { let session_config = self.task_ctx.session_config(); let TaskSpecializedPlan { + local_plan, plan, work_unit_feed_declarations, dynamic_filter_remote_producer_ids, @@ -249,6 +250,7 @@ impl<'a> StageCoordinator<'a> { metrics: self.metrics_set, worker_resolver: worker_resolver.as_ref(), task_specialized_plan: &plan, + local_task_plan: &local_plan, task_key, task_count: self.task_count, dialer: &dialer, @@ -472,22 +474,6 @@ impl<'a> StageCoordinator<'a> { let plan = Arc::clone(self.plan); let transformed = plan.transform_down_with_dt_ctx(d_ctx, |plan, d_ctx| { - if let Some(wuf) = wuf_registry.get_work_unit_feed(&plan) { - work_unit_feed_declarations.push(WorkUnitFeedDeclaration { - id: wuf.id(), - partitions: plan.properties().partitioning.partition_count(), - }); - - // WorkUnitFeeds are transitioned to remote mode during proto conversion. - // Right now, there's no other way for a WorkUnitFeed to be transitioned to - // remote mode so that it can pull WorkUnits over the WorkerChannel. - // - // Doing this roundtrip here is not super clean, but it transitions the feed with - // very little specialized code. - let plan = roundtrip_pb(plan, self.task_ctx)?; - return Ok(Transformed::yes(plan)); - }; - if let Some(ciu) = plan.downcast_ref::() { let ciu = ciu.to_task_specialized(d_ctx.task_index); return Ok(Transformed::yes(Arc::new(ciu))); @@ -512,6 +498,20 @@ impl<'a> StageCoordinator<'a> { // we are explicitly not retransforming the entire plan. if other operators cause shared state errors they will error out. Ok(Transformed::no(plan)) })?; + // Keep native providers available to routing. Only the worker-facing view + // replaces them with remote feeds; retained file metadata stays shared. + let local_plan = transformed.data; + let transformed = + Arc::clone(&local_plan).transform_down_with_dt_ctx(d_ctx, |plan, _| { + if let Some(wuf) = wuf_registry.get_work_unit_feed(&plan) { + work_unit_feed_declarations.push(WorkUnitFeedDeclaration { + id: wuf.id(), + partitions: plan.properties().partitioning.partition_count(), + }); + return Ok(Transformed::yes(roundtrip_pb(plan, self.task_ctx)?)); + } + Ok(Transformed::no(plan)) + })?; let plan = if dynamic_filtering_enabled { maybe_roundtrip_plan_to_sever_in_memory_dynamic_filter_relationships( Arc::clone(&transformed.data), @@ -527,6 +527,7 @@ impl<'a> StageCoordinator<'a> { vec![] }; Ok(TaskSpecializedPlan { + local_plan, plan, work_unit_feed_declarations, dynamic_filter_remote_producer_ids, @@ -548,6 +549,7 @@ enum RuntimeUpdate { } struct TaskSpecializedPlan { + local_plan: Arc, plan: Arc, work_unit_feed_declarations: Vec, dynamic_filter_remote_producer_ids: Vec, diff --git a/src/events/route_tasks.rs b/src/events/route_tasks.rs index 375980ef9..1f30cff8f 100644 --- a/src/events/route_tasks.rs +++ b/src/events/route_tasks.rs @@ -1,7 +1,10 @@ +use crate::DistributedTaskContext; +use crate::common::TreeNodeExt; use crate::events::common::EventHandlerChain; use crate::{TaskKey, WorkerResolver, WorkerToCoordinatorMsg}; use async_trait::async_trait; use datafusion::common::Result; +use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::execution::TaskContext; use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; use datafusion::physical_plan::ExecutionPlan; @@ -24,6 +27,11 @@ pub struct RouteTaskEvent<'a> { pub task_count: usize, /// The exact plan variant that will be sent for this task. pub task_specialized_plan: &'a Arc, + /// Task-specialized coordinator view, before native feeds become remote handles. + /// Use this view for inspecting connector metadata, not for execution. Native + /// providers may be shared with other tasks; do not consume their work streams. + /// Use [`Self::visit_local_task_plan`] to account for isolated UNION branches. + pub local_task_plan: &'a Arc, /// Establishes a coordinator-to-worker connection for a candidate URL. Users are expected /// to call `dialer.dial()` inside their [RouteTaskHandler] implementations for establishing /// a connection with the remote worker. @@ -33,6 +41,27 @@ pub struct RouteTaskEvent<'a> { pub dialer: &'a dyn CoordinatorToWorkerDialer, } +impl RouteTaskEvent<'_> { + /// Visits this task's local operators with their scan-local task index/count. + /// Skips other tasks' branches and stops at network boundaries. This preserves + /// the task context of isolated UNION children when inspecting native feeds. + pub fn visit_local_task_plan( + &self, + visitor: impl FnMut( + &Arc, + DistributedTaskContext, + ) -> Result, + ) -> Result { + self.local_task_plan.apply_with_dt_ctx( + DistributedTaskContext { + task_index: self.task_key.task_number, + task_count: self.task_count, + }, + visitor, + ) + } +} + /// A worker connection selected by a [`RouteTaskHandler`]. pub struct RouteTaskEventResponse { /// The URL of the selected worker. Once this URL is returned, the worker officially starts