Coverage Report

Created: 2024-10-13 08:39

/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
}