datafusion_distributed/worker/
impl_execute_task.rs

1use crate::ExecuteTaskRequest;
2use crate::worker::worker_service::Worker;
3use datafusion::common::exec_datafusion_err;
4use datafusion::common::{Result, exec_err};
5use datafusion::error::DataFusionError;
6use datafusion::execution::{SendableRecordBatchStream, TaskContext};
7use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
8use std::sync::Arc;
9use std::time::Duration;
10
11const WAIT_PLAN_TIMEOUT_SECS: u64 = 10;
12
13/// Builds several per-partition streams by retrieving the appropriate entry from [TaskDataEntries]
14/// based on the task key extracted from [ExecuteTaskRequest].
15///
16/// This method is async mainly for the key retrieval operation from [TaskDataEntries], but it does
17/// not start polling any stream, it just instantiates them.
18impl Worker {
19    pub async fn execute_task(
20        &self,
21        request: ExecuteTaskRequest,
22    ) -> Result<(Vec<SendableRecordBatchStream>, Arc<TaskContext>)> {
23        let entry = self
24            .task_data_entries
25            .get_with(request.task_key, async { Default::default() })
26            .await;
27
28        // Other request is responsible for writing the plan that belongs to this TaskKey, so
29        // we'll resolve immediately if it was already there, or wait until it's ready.
30        let task_data = entry
31            .read(Duration::from_secs(WAIT_PLAN_TIMEOUT_SECS))
32            .await
33            .map_err(|e| exec_datafusion_err!("Worker::execute_task timed-out while waiting for the plan to be set by the coordinator. ({e})"))?
34            .map_err(DataFusionError::Shared)?;
35        task_data.task_data_metrics.mark_execution_started_once();
36
37        let plan = task_data.plan(request.producer_head)?;
38        let task_ctx = task_data.task_ctx;
39        let partition_count = plan.properties().partitioning.partition_count();
40        let plan_name = plan.name();
41
42        // Execute all the requested partitions at once, and collect all the streams so that they
43        // can be merged into a single one at the end of this function.
44        let n_streams = request.target_partition_end - request.target_partition_start;
45        let mut streams = Vec::with_capacity(n_streams);
46        for partition in request.target_partition_start..request.target_partition_end {
47            if partition >= partition_count {
48                return exec_err!(
49                    "partition {partition} not available. The head plan {plan_name} of the stage just has {partition_count} partitions"
50                );
51            }
52
53            let stream = plan.execute(partition, Arc::clone(&task_ctx))?;
54            let stream_schema = plan.schema();
55
56            streams.push(Box::pin(RecordBatchStreamAdapter::new(stream_schema, stream)) as _);
57        }
58        Ok((streams, task_ctx))
59    }
60}
datafusion_distributed/worker/impl_execute_task.rs

datafusion_distributed/worker/
impl_execute_task.rs
