datafusion_optimizer/simplify_expressions/
inlist_simplifier.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//! This module implements a rule that simplifies the values for `InList`s
19
20use super::THRESHOLD_INLINE_INLIST;
21
22use datafusion_common::Result;
23use datafusion_common::tree_node::{Transformed, TreeNodeRewriter};
24use datafusion_expr::Expr;
25use datafusion_expr::expr::InList;
26
27pub(super) struct ShortenInListSimplifier {}
28
29impl ShortenInListSimplifier {
30    pub(super) fn new() -> Self {
31        Self {}
32    }
33}
34
35impl TreeNodeRewriter for ShortenInListSimplifier {
36    type Node = Expr;
37
38    fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
39        // if expr is a single column reference:
40        // expr IN (A, B, ...) --> (expr = A) OR (expr = B) OR (expr = C)
41        if let Expr::InList(InList {
42            ref expr,
43            ref list,
44            negated,
45        }) = expr
46            && !list.is_empty()
47            && (
48                // For lists with only 1 value we allow more complex expressions to be simplified
49                // e.g SUBSTR(c1, 2, 3) IN ('1') -> SUBSTR(c1, 2, 3) = '1'
50                // for more than one we avoid repeating this potentially expensive
51                // expressions
52                list.len() == 1
53                    || list.len() <= THRESHOLD_INLINE_INLIST
54                        && expr.try_as_col().is_some()
55            )
56        {
57            let first_val = list[0].clone();
58            if negated {
59                return Ok(Transformed::yes(list.iter().skip(1).cloned().fold(
60                    (*expr.clone()).not_eq(first_val),
61                    |acc, y| {
62                        // Note that `A and B and C and D` is a left-deep tree structure
63                        // as such we want to maintain this structure as much as possible
64                        // to avoid reordering the expression during each optimization
65                        // pass.
66                        //
67                        // Left-deep tree structure for `A and B and C and D`:
68                        // ```
69                        //        &
70                        //       / \
71                        //      &   D
72                        //     / \
73                        //    &   C
74                        //   / \
75                        //  A   B
76                        // ```
77                        //
78                        // The code below maintain the left-deep tree structure.
79                        acc.and((*expr.clone()).not_eq(y))
80                    },
81                )));
82            } else {
83                return Ok(Transformed::yes(list.iter().skip(1).cloned().fold(
84                    (*expr.clone()).eq(first_val),
85                    |acc, y| {
86                        // Same reasoning as above
87                        acc.or((*expr.clone()).eq(y))
88                    },
89                )));
90            }
91        }
92
93        Ok(Transformed::no(expr))
94    }
95}
datafusion_optimizer/simplify_expressions/inlist_simplifier.rs

datafusion_optimizer/simplify_expressions/
inlist_simplifier.rs
