/Users/andrewlamb/Software/datafusion/datafusion/expr/src/simplify.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 | | //! Structs and traits to provide the information needed for expression simplification. |
19 | | |
20 | | use arrow::datatypes::DataType; |
21 | | use datafusion_common::{DFSchemaRef, DataFusionError, Result}; |
22 | | |
23 | | use crate::{execution_props::ExecutionProps, Expr, ExprSchemable}; |
24 | | |
25 | | /// Provides the information necessary to apply algebraic simplification to an |
26 | | /// [Expr]. See [SimplifyContext] for one concrete implementation. |
27 | | /// |
28 | | /// This trait exists so that other systems can plug schema |
29 | | /// information in without having to create `DFSchema` objects. If you |
30 | | /// have a [`DFSchemaRef`] you can use [`SimplifyContext`] |
31 | | pub trait SimplifyInfo { |
32 | | /// returns true if this Expr has boolean type |
33 | | fn is_boolean_type(&self, expr: &Expr) -> Result<bool>; |
34 | | |
35 | | /// returns true of this expr is nullable (could possibly be NULL) |
36 | | fn nullable(&self, expr: &Expr) -> Result<bool>; |
37 | | |
38 | | /// Returns details needed for partial expression evaluation |
39 | | fn execution_props(&self) -> &ExecutionProps; |
40 | | |
41 | | /// Returns data type of this expr needed for determining optimized int type of a value |
42 | | fn get_data_type(&self, expr: &Expr) -> Result<DataType>; |
43 | | } |
44 | | |
45 | | /// Provides simplification information based on DFSchema and |
46 | | /// [`ExecutionProps`]. This is the default implementation used by DataFusion |
47 | | /// |
48 | | /// # Example |
49 | | /// See the `simplify_demo` in the [`expr_api` example] |
50 | | /// |
51 | | /// [`expr_api` example]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/expr_api.rs |
52 | | #[derive(Debug, Clone)] |
53 | | pub struct SimplifyContext<'a> { |
54 | | schema: Option<DFSchemaRef>, |
55 | | props: &'a ExecutionProps, |
56 | | } |
57 | | |
58 | | impl<'a> SimplifyContext<'a> { |
59 | | /// Create a new SimplifyContext |
60 | | pub fn new(props: &'a ExecutionProps) -> Self { |
61 | | Self { |
62 | | schema: None, |
63 | | props, |
64 | | } |
65 | | } |
66 | | |
67 | | /// Register a [`DFSchemaRef`] with this context |
68 | | pub fn with_schema(mut self, schema: DFSchemaRef) -> Self { |
69 | | self.schema = Some(schema); |
70 | | self |
71 | | } |
72 | | } |
73 | | |
74 | | impl<'a> SimplifyInfo for SimplifyContext<'a> { |
75 | | /// returns true if this Expr has boolean type |
76 | | fn is_boolean_type(&self, expr: &Expr) -> Result<bool> { |
77 | | if let Some(schema) = &self.schema { |
78 | | if let Ok(DataType::Boolean) = expr.get_type(schema) { |
79 | | return Ok(true); |
80 | | } |
81 | | } |
82 | | |
83 | | Ok(false) |
84 | | } |
85 | | |
86 | | /// Returns true if expr is nullable |
87 | | fn nullable(&self, expr: &Expr) -> Result<bool> { |
88 | 0 | let schema = self.schema.as_ref().ok_or_else(|| { |
89 | 0 | DataFusionError::Internal( |
90 | 0 | "attempt to get nullability without schema".to_string(), |
91 | 0 | ) |
92 | 0 | })?; |
93 | | expr.nullable(schema.as_ref()) |
94 | | } |
95 | | |
96 | | /// Returns data type of this expr needed for determining optimized int type of a value |
97 | | fn get_data_type(&self, expr: &Expr) -> Result<DataType> { |
98 | 0 | let schema = self.schema.as_ref().ok_or_else(|| { |
99 | 0 | DataFusionError::Internal( |
100 | 0 | "attempt to get data type without schema".to_string(), |
101 | 0 | ) |
102 | 0 | })?; |
103 | | expr.get_type(schema) |
104 | | } |
105 | | |
106 | | fn execution_props(&self) -> &ExecutionProps { |
107 | | self.props |
108 | | } |
109 | | } |
110 | | |
111 | | /// Was the expression simplified? |
112 | | #[derive(Debug)] |
113 | | pub enum ExprSimplifyResult { |
114 | | /// The function call was simplified to an entirely new Expr |
115 | | Simplified(Expr), |
116 | | /// the function call could not be simplified, and the arguments |
117 | | /// are return unmodified. |
118 | | Original(Vec<Expr>), |
119 | | } |