/Users/andrewlamb/Software/datafusion/datafusion/expr/src/execution_props.rs
Line | Count | Source (jump to first uncovered line) |
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 | | |
18 | | use crate::var_provider::{VarProvider, VarType}; |
19 | | use chrono::{DateTime, TimeZone, Utc}; |
20 | | use datafusion_common::alias::AliasGenerator; |
21 | | use std::collections::HashMap; |
22 | | use std::sync::Arc; |
23 | | |
24 | | /// Holds per-query execution properties and data (such as statement |
25 | | /// starting timestamps). |
26 | | /// |
27 | | /// An [`ExecutionProps`] is created each time a `LogicalPlan` is |
28 | | /// prepared for execution (optimized). If the same plan is optimized |
29 | | /// multiple times, a new `ExecutionProps` is created each time. |
30 | | /// |
31 | | /// It is important that this structure be cheap to create as it is |
32 | | /// done so during predicate pruning and expression simplification |
33 | | #[derive(Clone, Debug)] |
34 | | pub struct ExecutionProps { |
35 | | pub query_execution_start_time: DateTime<Utc>, |
36 | | /// Alias generator used by subquery optimizer rules |
37 | | pub alias_generator: Arc<AliasGenerator>, |
38 | | /// Providers for scalar variables |
39 | | pub var_providers: Option<HashMap<VarType, Arc<dyn VarProvider + Send + Sync>>>, |
40 | | } |
41 | | |
42 | | impl Default for ExecutionProps { |
43 | 0 | fn default() -> Self { |
44 | 0 | Self::new() |
45 | 0 | } |
46 | | } |
47 | | |
48 | | impl ExecutionProps { |
49 | | /// Creates a new execution props |
50 | 0 | pub fn new() -> Self { |
51 | 0 | ExecutionProps { |
52 | 0 | // Set this to a fixed sentinel to make it obvious if this is |
53 | 0 | // not being updated / propagated correctly |
54 | 0 | query_execution_start_time: Utc.timestamp_nanos(0), |
55 | 0 | alias_generator: Arc::new(AliasGenerator::new()), |
56 | 0 | var_providers: None, |
57 | 0 | } |
58 | 0 | } |
59 | | |
60 | | /// Set the query execution start time to use |
61 | 0 | pub fn with_query_execution_start_time( |
62 | 0 | mut self, |
63 | 0 | query_execution_start_time: DateTime<Utc>, |
64 | 0 | ) -> Self { |
65 | 0 | self.query_execution_start_time = query_execution_start_time; |
66 | 0 | self |
67 | 0 | } |
68 | | |
69 | | /// Marks the execution of query started timestamp. |
70 | | /// This also instantiates a new alias generator. |
71 | 0 | pub fn start_execution(&mut self) -> &Self { |
72 | 0 | self.query_execution_start_time = Utc::now(); |
73 | 0 | self.alias_generator = Arc::new(AliasGenerator::new()); |
74 | 0 | &*self |
75 | 0 | } |
76 | | |
77 | | /// Registers a variable provider, returning the existing |
78 | | /// provider, if any |
79 | 0 | pub fn add_var_provider( |
80 | 0 | &mut self, |
81 | 0 | var_type: VarType, |
82 | 0 | provider: Arc<dyn VarProvider + Send + Sync>, |
83 | 0 | ) -> Option<Arc<dyn VarProvider + Send + Sync>> { |
84 | 0 | let mut var_providers = self.var_providers.take().unwrap_or_default(); |
85 | 0 |
|
86 | 0 | let old_provider = var_providers.insert(var_type, provider); |
87 | 0 |
|
88 | 0 | self.var_providers = Some(var_providers); |
89 | 0 |
|
90 | 0 | old_provider |
91 | 0 | } |
92 | | |
93 | | /// Returns the provider for the `var_type`, if any |
94 | 0 | pub fn get_var_provider( |
95 | 0 | &self, |
96 | 0 | var_type: VarType, |
97 | 0 | ) -> Option<Arc<dyn VarProvider + Send + Sync>> { |
98 | 0 | self.var_providers |
99 | 0 | .as_ref() |
100 | 0 | .and_then(|var_providers| var_providers.get(&var_type).cloned()) |
101 | 0 | } |
102 | | } |
103 | | |
104 | | #[cfg(test)] |
105 | | mod test { |
106 | | use super::*; |
107 | | #[test] |
108 | | fn debug() { |
109 | | let props = ExecutionProps::new(); |
110 | | assert_eq!("ExecutionProps { query_execution_start_time: 1970-01-01T00:00:00Z, alias_generator: AliasGenerator { next_id: 1 }, var_providers: None }", format!("{props:?}")); |
111 | | } |
112 | | } |