datafusion_expr/
execution_props.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::var_provider::{VarProvider, VarType};
19use chrono::{DateTime, TimeZone, Utc};
20use datafusion_common::alias::AliasGenerator;
21use datafusion_common::HashMap;
22use 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)]
34pub 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
42impl Default for ExecutionProps {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl ExecutionProps {
49    /// Creates a new execution props
50    pub fn new() -> Self {
51        ExecutionProps {
52            // Set this to a fixed sentinel to make it obvious if this is
53            // not being updated / propagated correctly
54            query_execution_start_time: Utc.timestamp_nanos(0),
55            alias_generator: Arc::new(AliasGenerator::new()),
56            var_providers: None,
57        }
58    }
59
60    /// Set the query execution start time to use
61    pub fn with_query_execution_start_time(
62        mut self,
63        query_execution_start_time: DateTime<Utc>,
64    ) -> Self {
65        self.query_execution_start_time = query_execution_start_time;
66        self
67    }
68
69    /// Marks the execution of query started timestamp.
70    /// This also instantiates a new alias generator.
71    pub fn start_execution(&mut self) -> &Self {
72        self.query_execution_start_time = Utc::now();
73        self.alias_generator = Arc::new(AliasGenerator::new());
74        &*self
75    }
76
77    /// Registers a variable provider, returning the existing
78    /// provider, if any
79    pub fn add_var_provider(
80        &mut self,
81        var_type: VarType,
82        provider: Arc<dyn VarProvider + Send + Sync>,
83    ) -> Option<Arc<dyn VarProvider + Send + Sync>> {
84        let mut var_providers = self.var_providers.take().unwrap_or_default();
85
86        let old_provider = var_providers.insert(var_type, provider);
87
88        self.var_providers = Some(var_providers);
89
90        old_provider
91    }
92
93    /// Returns the provider for the `var_type`, if any
94    pub fn get_var_provider(
95        &self,
96        var_type: VarType,
97    ) -> Option<Arc<dyn VarProvider + Send + Sync>> {
98        self.var_providers
99            .as_ref()
100            .and_then(|var_providers| var_providers.get(&var_type).cloned())
101    }
102}
103
104#[cfg(test)]
105mod 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}
datafusion_expr/execution_props.rs

datafusion_expr/
execution_props.rs
