datafusion_execution/
task.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::{
19    config::SessionConfig, memory_pool::MemoryPool, registry::FunctionRegistry,
20    runtime_env::RuntimeEnv,
21};
22use datafusion_common::{Result, internal_datafusion_err, plan_datafusion_err};
23use datafusion_expr::planner::ExprPlanner;
24use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF};
25use std::collections::HashSet;
26use std::{collections::HashMap, sync::Arc};
27
28/// Task Execution Context
29///
30/// A [`TaskContext`] contains the state required during a single query's
31/// execution. Please see the documentation on [`SessionContext`] for more
32/// information.
33///
34/// # Relationship with [`ExecutionProps`]
35///
36/// [`TaskContext`] is intentionally distinct from [`ExecutionProps`].
37/// [`ExecutionProps`] is state used while optimizing a logical
38/// plan and constructing a physical plan.
39///
40/// [`TaskContext`] is the runtime context passed to physical operators when
41/// executing a physical plan. It carries runtime services and session state
42/// needed at that stage, such as [`RuntimeEnv`], memory-pool access, session
43/// configuration, and function lookup.
44///
45/// Keeping these structures separate avoids threading execution/runtime state
46/// through planning APIs, and avoids making execution depend on planner-only
47/// scratch state.
48///
49/// [`SessionContext`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html
50/// [`ExecutionProps`]: datafusion_expr::execution_props::ExecutionProps
51#[derive(Debug)]
52pub struct TaskContext {
53    /// Session Id
54    session_id: String,
55    /// Optional task identity
56    task_id: Option<String>,
57    /// Session configuration
58    session_config: SessionConfig,
59    /// Scalar functions associated with this task context
60    scalar_functions: HashMap<String, Arc<ScalarUDF>>,
61    /// Higher order functions associated with this task context
62    higher_order_functions: HashMap<String, Arc<HigherOrderUDF>>,
63    /// Aggregate functions associated with this task context
64    aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
65    /// Window functions associated with this task context
66    window_functions: HashMap<String, Arc<WindowUDF>>,
67    /// Runtime environment associated with this task context
68    runtime: Arc<RuntimeEnv>,
69}
70
71impl Default for TaskContext {
72    fn default() -> Self {
73        let runtime = Arc::new(RuntimeEnv::default());
74
75        // Create a default task context, mostly useful for testing
76        Self {
77            session_id: "DEFAULT".to_string(),
78            task_id: None,
79            session_config: SessionConfig::new(),
80            scalar_functions: HashMap::new(),
81            higher_order_functions: HashMap::new(),
82            aggregate_functions: HashMap::new(),
83            window_functions: HashMap::new(),
84            runtime,
85        }
86    }
87}
88
89impl TaskContext {
90    /// Create a new [`TaskContext`] instance.
91    ///
92    /// Most users will use [`SessionContext::task_ctx`] to create [`TaskContext`]s
93    ///
94    /// [`SessionContext::task_ctx`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html#method.task_ctx
95    #[expect(clippy::too_many_arguments)]
96    pub fn new(
97        task_id: Option<String>,
98        session_id: String,
99        session_config: SessionConfig,
100        scalar_functions: HashMap<String, Arc<ScalarUDF>>,
101        higher_order_functions: HashMap<String, Arc<HigherOrderUDF>>,
102        aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
103        window_functions: HashMap<String, Arc<WindowUDF>>,
104        runtime: Arc<RuntimeEnv>,
105    ) -> Self {
106        Self {
107            task_id,
108            session_id,
109            session_config,
110            scalar_functions,
111            higher_order_functions,
112            aggregate_functions,
113            window_functions,
114            runtime,
115        }
116    }
117
118    /// Return the SessionConfig associated with this [TaskContext]
119    pub fn session_config(&self) -> &SessionConfig {
120        &self.session_config
121    }
122
123    /// Return the `session_id` of this [TaskContext]
124    pub fn session_id(&self) -> String {
125        self.session_id.clone()
126    }
127
128    /// Return the `task_id` of this [TaskContext]
129    pub fn task_id(&self) -> Option<String> {
130        self.task_id.clone()
131    }
132
133    /// Return the [`MemoryPool`] associated with this [TaskContext]
134    pub fn memory_pool(&self) -> &Arc<dyn MemoryPool> {
135        &self.runtime.memory_pool
136    }
137
138    /// Return the [RuntimeEnv] associated with this [TaskContext]
139    pub fn runtime_env(&self) -> Arc<RuntimeEnv> {
140        Arc::clone(&self.runtime)
141    }
142
143    pub fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>> {
144        &self.scalar_functions
145    }
146
147    pub fn higher_order_functions(&self) -> &HashMap<String, Arc<HigherOrderUDF>> {
148        &self.higher_order_functions
149    }
150
151    pub fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
152        &self.aggregate_functions
153    }
154
155    pub fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>> {
156        &self.window_functions
157    }
158
159    /// Update the [`SessionConfig`]
160    pub fn with_session_config(mut self, session_config: SessionConfig) -> Self {
161        self.session_config = session_config;
162        self
163    }
164
165    /// Update the [`RuntimeEnv`]
166    pub fn with_runtime(mut self, runtime: Arc<RuntimeEnv>) -> Self {
167        self.runtime = runtime;
168        self
169    }
170
171    /// Update the `task_id`
172    pub fn with_task_id(mut self, task_id: String) -> Self {
173        self.task_id = Some(task_id);
174        self
175    }
176}
177
178impl FunctionRegistry for TaskContext {
179    fn udfs(&self) -> HashSet<String> {
180        self.scalar_functions.keys().cloned().collect()
181    }
182
183    fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>> {
184        let result = self.scalar_functions.get(name);
185
186        result.cloned().ok_or_else(|| {
187            plan_datafusion_err!("There is no UDF named \"{name}\" in the TaskContext")
188        })
189    }
190
191    fn higher_order_function(&self, name: &str) -> Result<Arc<HigherOrderUDF>> {
192        let result = self.higher_order_functions.get(name);
193
194        result.cloned().ok_or_else(|| {
195            plan_datafusion_err!(
196                "There is no higher-order function named \"{name}\" in the TaskContext"
197            )
198        })
199    }
200
201    fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {
202        let result = self.aggregate_functions.get(name);
203
204        result.cloned().ok_or_else(|| {
205            plan_datafusion_err!("There is no UDAF named \"{name}\" in the TaskContext")
206        })
207    }
208
209    fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>> {
210        let result = self.window_functions.get(name);
211
212        result.cloned().ok_or_else(|| {
213            internal_datafusion_err!(
214                "There is no UDWF named \"{name}\" in the TaskContext"
215            )
216        })
217    }
218    fn register_udaf(
219        &mut self,
220        udaf: Arc<AggregateUDF>,
221    ) -> Result<Option<Arc<AggregateUDF>>> {
222        udaf.aliases().iter().for_each(|alias| {
223            self.aggregate_functions
224                .insert(alias.clone(), Arc::clone(&udaf));
225        });
226        Ok(self.aggregate_functions.insert(udaf.name().into(), udaf))
227    }
228    fn register_udwf(&mut self, udwf: Arc<WindowUDF>) -> Result<Option<Arc<WindowUDF>>> {
229        udwf.aliases().iter().for_each(|alias| {
230            self.window_functions
231                .insert(alias.clone(), Arc::clone(&udwf));
232        });
233        Ok(self.window_functions.insert(udwf.name().into(), udwf))
234    }
235    fn register_udf(&mut self, udf: Arc<ScalarUDF>) -> Result<Option<Arc<ScalarUDF>>> {
236        udf.aliases().iter().for_each(|alias| {
237            self.scalar_functions
238                .insert(alias.clone(), Arc::clone(&udf));
239        });
240        Ok(self.scalar_functions.insert(udf.name().into(), udf))
241    }
242
243    fn register_higher_order_function(
244        &mut self,
245        function: Arc<HigherOrderUDF>,
246    ) -> Result<Option<Arc<HigherOrderUDF>>> {
247        function.aliases().iter().for_each(|alias| {
248            self.higher_order_functions
249                .insert(alias.clone(), Arc::clone(&function));
250        });
251        Ok(self
252            .higher_order_functions
253            .insert(function.name().into(), function))
254    }
255
256    fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
257        vec![]
258    }
259
260    fn higher_order_function_names(&self) -> HashSet<String> {
261        self.higher_order_functions.keys().cloned().collect()
262    }
263
264    fn udafs(&self) -> HashSet<String> {
265        self.aggregate_functions.keys().cloned().collect()
266    }
267
268    fn udwfs(&self) -> HashSet<String> {
269        self.window_functions.keys().cloned().collect()
270    }
271}
272
273/// Produce the [`TaskContext`].
274pub trait TaskContextProvider {
275    fn task_ctx(&self) -> Arc<TaskContext>;
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281    use datafusion_common::{
282        config::{ConfigExtension, ConfigOptions, Extensions},
283        extensions_options,
284    };
285
286    extensions_options! {
287        struct TestExtension {
288            value: usize, default = 42
289            option_value: Option<usize>, default = None
290        }
291    }
292
293    impl ConfigExtension for TestExtension {
294        const PREFIX: &'static str = "test";
295    }
296
297    #[test]
298    fn task_context_extensions() -> Result<()> {
299        let runtime = Arc::new(RuntimeEnv::default());
300        let mut extensions = Extensions::new();
301        extensions.insert(TestExtension::default());
302
303        let mut config = ConfigOptions::new().with_extensions(extensions);
304        config.set("test.value", "24")?;
305        config.set("test.option_value", "42")?;
306        let session_config = SessionConfig::from(config);
307
308        let task_context = TaskContext::new(
309            Some("task_id".to_string()),
310            "session_id".to_string(),
311            session_config,
312            HashMap::default(),
313            HashMap::default(),
314            HashMap::default(),
315            HashMap::default(),
316            runtime,
317        );
318
319        let test = task_context
320            .session_config()
321            .options()
322            .extensions
323            .get::<TestExtension>();
324        assert!(test.is_some());
325
326        assert_eq!(test.unwrap().value, 24);
327        assert_eq!(test.unwrap().option_value, Some(42));
328
329        Ok(())
330    }
331
332    #[test]
333    fn task_context_extensions_default() -> Result<()> {
334        let runtime = Arc::new(RuntimeEnv::default());
335        let mut extensions = Extensions::new();
336        extensions.insert(TestExtension::default());
337
338        let config = ConfigOptions::new().with_extensions(extensions);
339        let session_config = SessionConfig::from(config);
340
341        let task_context = TaskContext::new(
342            Some("task_id".to_string()),
343            "session_id".to_string(),
344            session_config,
345            HashMap::default(),
346            HashMap::default(),
347            HashMap::default(),
348            HashMap::default(),
349            runtime,
350        );
351
352        let test = task_context
353            .session_config()
354            .options()
355            .extensions
356            .get::<TestExtension>();
357        assert!(test.is_some());
358
359        assert_eq!(test.unwrap().value, 42);
360        assert_eq!(test.unwrap().option_value, None);
361
362        Ok(())
363    }
364}
datafusion_execution/task.rs

datafusion_execution/
task.rs
