datafusion_expr/
simplify.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
18//! Structs to provide the information needed for expression simplification.
19
20use std::sync::Arc;
21
22use arrow::datatypes::DataType;
23use chrono::{DateTime, Utc};
24use datafusion_common::config::ConfigOptions;
25use datafusion_common::{DFSchema, DFSchemaRef, Result};
26
27use crate::{Expr, ExprSchemable};
28
29/// Provides simplification information based on schema, query execution time,
30/// and configuration options.
31///
32/// # Example
33/// See the `simplify_demo` in the [`expr_api` example]
34///
35/// [`expr_api` example]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/query_planning/expr_api.rs
36#[derive(Debug, Clone)]
37pub struct SimplifyContext {
38    schema: DFSchemaRef,
39    query_execution_start_time: Option<DateTime<Utc>>,
40    config_options: Arc<ConfigOptions>,
41}
42
43impl Default for SimplifyContext {
44    fn default() -> Self {
45        Self {
46            schema: Arc::new(DFSchema::empty()),
47            query_execution_start_time: None,
48            config_options: Arc::new(ConfigOptions::default()),
49        }
50    }
51}
52
53impl SimplifyContext {
54    /// Set the [`ConfigOptions`] for this context
55    pub fn with_config_options(mut self, config_options: Arc<ConfigOptions>) -> Self {
56        self.config_options = config_options;
57        self
58    }
59
60    /// Set the schema for this context
61    pub fn with_schema(mut self, schema: DFSchemaRef) -> Self {
62        self.schema = schema;
63        self
64    }
65
66    /// Set the query execution start time
67    pub fn with_query_execution_start_time(
68        mut self,
69        query_execution_start_time: Option<DateTime<Utc>>,
70    ) -> Self {
71        self.query_execution_start_time = query_execution_start_time;
72        self
73    }
74
75    /// Set the query execution start to the current time
76    pub fn with_current_time(mut self) -> Self {
77        self.query_execution_start_time = Some(Utc::now());
78        self
79    }
80
81    /// Returns the schema
82    pub fn schema(&self) -> &DFSchemaRef {
83        &self.schema
84    }
85
86    /// Returns true if this Expr has boolean type
87    pub fn is_boolean_type(&self, expr: &Expr) -> Result<bool> {
88        Ok(expr.get_type(&self.schema)? == DataType::Boolean)
89    }
90
91    /// Returns true if expr is nullable
92    pub fn nullable(&self, expr: &Expr) -> Result<bool> {
93        expr.nullable(self.schema.as_ref())
94    }
95
96    /// Returns data type of this expr needed for determining optimized int type of a value
97    pub fn get_data_type(&self, expr: &Expr) -> Result<DataType> {
98        expr.get_type(&self.schema)
99    }
100
101    /// Returns the time at which the query execution started.
102    /// If `None`, time-dependent functions like `now()` will not be simplified.
103    pub fn query_execution_start_time(&self) -> Option<DateTime<Utc>> {
104        self.query_execution_start_time
105    }
106
107    /// Returns the configuration options for the session.
108    pub fn config_options(&self) -> &Arc<ConfigOptions> {
109        &self.config_options
110    }
111}
112
113/// Was the expression simplified?
114#[derive(Debug)]
115pub enum ExprSimplifyResult {
116    /// The function call was simplified to an entirely new Expr
117    Simplified(Expr),
118    /// The function call could not be simplified, and the arguments
119    /// are return unmodified.
120    Original(Vec<Expr>),
121}
datafusion_expr/simplify.rs

datafusion_expr/
simplify.rs
