datafusion-optimizer 53.1.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Utility functions for expression simplification

use arrow::datatypes::i256;
use datafusion_common::{Result, ScalarValue, internal_err};
use datafusion_expr::{
    Case, Expr, Like, Operator,
    expr::{Between, BinaryExpr, InList},
    expr_fn::{and, bitwise_and, bitwise_or, or},
};

pub static POWS_OF_TEN: [i128; 38] = [
    1,
    10,
    100,
    1000,
    10000,
    100000,
    1000000,
    10000000,
    100000000,
    1000000000,
    10000000000,
    100000000000,
    1000000000000,
    10000000000000,
    100000000000000,
    1000000000000000,
    10000000000000000,
    100000000000000000,
    1000000000000000000,
    10000000000000000000,
    100000000000000000000,
    1000000000000000000000,
    10000000000000000000000,
    100000000000000000000000,
    1000000000000000000000000,
    10000000000000000000000000,
    100000000000000000000000000,
    1000000000000000000000000000,
    10000000000000000000000000000,
    100000000000000000000000000000,
    1000000000000000000000000000000,
    10000000000000000000000000000000,
    100000000000000000000000000000000,
    1000000000000000000000000000000000,
    10000000000000000000000000000000000,
    100000000000000000000000000000000000,
    1000000000000000000000000000000000000,
    10000000000000000000000000000000000000,
];

/// returns true if `needle` is found in a chain of search_op
/// expressions. Such as: (A AND B) AND C
fn expr_contains_inner(expr: &Expr, needle: &Expr, search_op: Operator) -> bool {
    match expr {
        Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == search_op => {
            expr_contains_inner(left, needle, search_op)
                || expr_contains_inner(right, needle, search_op)
        }
        _ => expr == needle,
    }
}

/// check volatile calls and return if expr contains needle
pub fn expr_contains(expr: &Expr, needle: &Expr, search_op: Operator) -> bool {
    expr_contains_inner(expr, needle, search_op) && !needle.is_volatile()
}

/// Deletes all 'needles' or remains one 'needle' that are found in a chain of xor
/// expressions. Such as: A ^ (A ^ (B ^ A))
pub fn delete_xor_in_complex_expr(expr: &Expr, needle: &Expr, is_left: bool) -> Expr {
    /// Deletes recursively 'needles' in a chain of xor expressions
    fn recursive_delete_xor_in_expr(
        expr: &Expr,
        needle: &Expr,
        xor_counter: &mut i32,
    ) -> Expr {
        match expr {
            Expr::BinaryExpr(BinaryExpr { left, op, right })
                if *op == Operator::BitwiseXor =>
            {
                let left_expr = recursive_delete_xor_in_expr(left, needle, xor_counter);
                let right_expr = recursive_delete_xor_in_expr(right, needle, xor_counter);
                if left_expr == *needle {
                    *xor_counter += 1;
                    return right_expr;
                } else if right_expr == *needle {
                    *xor_counter += 1;
                    return left_expr;
                }

                Expr::BinaryExpr(BinaryExpr::new(
                    Box::new(left_expr),
                    *op,
                    Box::new(right_expr),
                ))
            }
            _ => expr.clone(),
        }
    }

    let mut xor_counter: i32 = 0;
    let result_expr = recursive_delete_xor_in_expr(expr, needle, &mut xor_counter);
    if result_expr == *needle {
        return needle.clone();
    } else if xor_counter % 2 == 0 {
        if is_left {
            return Expr::BinaryExpr(BinaryExpr::new(
                Box::new(needle.clone()),
                Operator::BitwiseXor,
                Box::new(result_expr),
            ));
        } else {
            return Expr::BinaryExpr(BinaryExpr::new(
                Box::new(result_expr),
                Operator::BitwiseXor,
                Box::new(needle.clone()),
            ));
        }
    }
    result_expr
}

pub fn is_zero(s: &Expr) -> bool {
    match s {
        Expr::Literal(ScalarValue::Int8(Some(0)), _)
        | Expr::Literal(ScalarValue::Int16(Some(0)), _)
        | Expr::Literal(ScalarValue::Int32(Some(0)), _)
        | Expr::Literal(ScalarValue::Int64(Some(0)), _)
        | Expr::Literal(ScalarValue::UInt8(Some(0)), _)
        | Expr::Literal(ScalarValue::UInt16(Some(0)), _)
        | Expr::Literal(ScalarValue::UInt32(Some(0)), _)
        | Expr::Literal(ScalarValue::UInt64(Some(0)), _) => true,
        Expr::Literal(ScalarValue::Float32(Some(v)), _) if *v == 0. => true,
        Expr::Literal(ScalarValue::Float64(Some(v)), _) if *v == 0. => true,
        Expr::Literal(ScalarValue::Decimal128(Some(v), _p, _s), _) if *v == 0 => true,
        Expr::Literal(ScalarValue::Decimal256(Some(v), _p, _s), _)
            if *v == i256::ZERO =>
        {
            true
        }
        _ => false,
    }
}

pub fn is_one(s: &Expr) -> bool {
    match s {
        Expr::Literal(ScalarValue::Int8(Some(1)), _)
        | Expr::Literal(ScalarValue::Int16(Some(1)), _)
        | Expr::Literal(ScalarValue::Int32(Some(1)), _)
        | Expr::Literal(ScalarValue::Int64(Some(1)), _)
        | Expr::Literal(ScalarValue::UInt8(Some(1)), _)
        | Expr::Literal(ScalarValue::UInt16(Some(1)), _)
        | Expr::Literal(ScalarValue::UInt32(Some(1)), _)
        | Expr::Literal(ScalarValue::UInt64(Some(1)), _) => true,
        Expr::Literal(ScalarValue::Float32(Some(v)), _) if *v == 1. => true,
        Expr::Literal(ScalarValue::Float64(Some(v)), _) if *v == 1. => true,
        Expr::Literal(ScalarValue::Decimal128(Some(v), _p, s), _) => {
            *s >= 0
                && POWS_OF_TEN
                    .get(*s as usize)
                    .map(|x| x == v)
                    .unwrap_or_default()
        }
        Expr::Literal(ScalarValue::Decimal256(Some(v), _p, s), _) => {
            *s >= 0
                && match i256::from(10).checked_pow(*s as u32) {
                    Some(res) => res == *v,
                    None => false,
                }
        }
        _ => false,
    }
}

pub fn is_true(expr: &Expr) -> bool {
    match expr {
        Expr::Literal(ScalarValue::Boolean(Some(v)), _) => *v,
        _ => false,
    }
}

/// returns true if expr is a
/// `Expr::Literal(ScalarValue::Boolean(v))` , false otherwise
pub fn is_bool_lit(expr: &Expr) -> bool {
    matches!(expr, Expr::Literal(ScalarValue::Boolean(_), _))
}

/// Return a literal NULL value of Boolean data type
pub fn lit_bool_null() -> Expr {
    Expr::Literal(ScalarValue::Boolean(None), None)
}

pub fn is_null(expr: &Expr) -> bool {
    match expr {
        Expr::Literal(v, _) => v.is_null(),
        _ => false,
    }
}

pub fn is_false(expr: &Expr) -> bool {
    match expr {
        Expr::Literal(ScalarValue::Boolean(Some(v)), _) => !(*v),
        _ => false,
    }
}

/// returns true if `haystack` looks like (needle OP X) or (X OP needle)
pub fn is_op_with(target_op: Operator, haystack: &Expr, needle: &Expr) -> bool {
    matches!(haystack, Expr::BinaryExpr(BinaryExpr { left, op, right }) if op == &target_op && (needle == left.as_ref() || needle == right.as_ref()) && !needle.is_volatile())
}

pub fn can_reduce_to_equal_statement(haystack: &Expr, needle: &Expr) -> bool {
    match (haystack, needle) {
        // a >= constant and constant <= a => a = constant
        (
            Expr::BinaryExpr(BinaryExpr {
                left,
                op: Operator::GtEq,
                right,
            }),
            Expr::BinaryExpr(BinaryExpr {
                left: n_left,
                op: Operator::LtEq,
                right: n_right,
            }),
        ) if left == n_left && right == n_right => true,
        _ => false,
    }
}

/// returns true if `not_expr` is !`expr` (not)
pub fn is_not_of(not_expr: &Expr, expr: &Expr) -> bool {
    matches!(not_expr, Expr::Not(inner) if expr == inner.as_ref())
}

/// returns true if `not_expr` is !`expr` (bitwise not)
pub fn is_negative_of(not_expr: &Expr, expr: &Expr) -> bool {
    matches!(not_expr, Expr::Negative(inner) if expr == inner.as_ref())
}

/// returns the contained boolean value in `expr` as
/// `Expr::Literal(ScalarValue::Boolean(v))`.
pub fn as_bool_lit(expr: &Expr) -> Result<Option<bool>> {
    match expr {
        Expr::Literal(ScalarValue::Boolean(v), _) => Ok(*v),
        _ => internal_err!("Expected boolean literal, got {expr:?}"),
    }
}

pub fn is_case_with_literal_outputs(expr: &Expr) -> bool {
    match expr {
        Expr::Case(Case {
            expr: None,
            when_then_expr,
            else_expr,
        }) => {
            when_then_expr.iter().all(|(_, then)| is_lit(then))
                && else_expr.as_deref().is_none_or(is_lit)
        }
        _ => false,
    }
}

pub fn into_case(expr: Expr) -> Result<Case> {
    match expr {
        Expr::Case(case) => Ok(case),
        _ => internal_err!("Expected case, got {expr:?}"),
    }
}

pub fn is_lit(expr: &Expr) -> bool {
    matches!(expr, Expr::Literal(_, _))
}

/// Checks if `eq_expr` is `A = L1` and `ne_expr` is `A != L2` where L1 != L2.
/// This pattern can be simplified to just `A = L1` since if A equals L1
/// and L1 is different from L2, then A is automatically not equal to L2.
pub fn is_eq_and_ne_with_different_literal(eq_expr: &Expr, ne_expr: &Expr) -> bool {
    fn extract_var_and_literal(expr: &Expr) -> Option<(&Expr, &Expr)> {
        match expr {
            Expr::BinaryExpr(BinaryExpr {
                left,
                op: Operator::Eq,
                right,
            })
            | Expr::BinaryExpr(BinaryExpr {
                left,
                op: Operator::NotEq,
                right,
            }) => match (left.as_ref(), right.as_ref()) {
                (Expr::Literal(_, _), var) => Some((var, left)),
                (var, Expr::Literal(_, _)) => Some((var, right)),
                _ => None,
            },
            _ => None,
        }
    }
    match (eq_expr, ne_expr) {
        (
            Expr::BinaryExpr(BinaryExpr {
                op: Operator::Eq, ..
            }),
            Expr::BinaryExpr(BinaryExpr {
                op: Operator::NotEq,
                ..
            }),
        ) => {
            // Check if both compare the same expression against different literals
            if let (Some((var1, lit1)), Some((var2, lit2))) = (
                extract_var_and_literal(eq_expr),
                extract_var_and_literal(ne_expr),
            ) && var1 == var2
                && lit1 != lit2
            {
                return true;
            }
            false
        }
        _ => false,
    }
}

/// negate a Not clause
/// input is the clause to be negated.(args of Not clause)
/// For BinaryExpr, use the negation of op instead.
///    not ( A > B) ===> (A <= B)
/// For BoolExpr, not (A and B) ===> (not A) or (not B)
///     not (A or B) ===> (not A) and (not B)
///     not (not A) ===> A
/// For NullExpr, not (A is not null) ===> A is null
///     not (A is null) ===> A is not null
/// For InList, not (A not in (..)) ===> A in (..)
///     not (A in (..)) ===> A not in (..)
/// For Between, not (A between B and C) ===> (A not between B and C)
///     not (A not between B and C) ===> (A between B and C)
/// For others, use Not clause
pub fn negate_clause(expr: Expr) -> Expr {
    match expr {
        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
            if let Some(negated_op) = op.negate() {
                return Expr::BinaryExpr(BinaryExpr::new(left, negated_op, right));
            }
            match op {
                // not (A and B) ===> (not A) or (not B)
                Operator::And => {
                    let left = negate_clause(*left);
                    let right = negate_clause(*right);

                    or(left, right)
                }
                // not (A or B) ===> (not A) and (not B)
                Operator::Or => {
                    let left = negate_clause(*left);
                    let right = negate_clause(*right);

                    and(left, right)
                }
                // use not clause
                _ => Expr::Not(Box::new(Expr::BinaryExpr(BinaryExpr::new(
                    left, op, right,
                )))),
            }
        }
        // not (not A) ===> A
        Expr::Not(expr) => *expr,
        // not (A is not null) ===> A is null
        Expr::IsNotNull(expr) => expr.is_null(),
        // not (A is null) ===> A is not null
        Expr::IsNull(expr) => expr.is_not_null(),
        // not (A not in (..)) ===> A in (..)
        // not (A in (..)) ===> A not in (..)
        Expr::InList(InList {
            expr,
            list,
            negated,
        }) => expr.in_list(list, !negated),
        // not (A between B and C) ===> (A not between B and C)
        // not (A not between B and C) ===> (A between B and C)
        Expr::Between(between) => Expr::Between(Between::new(
            between.expr,
            !between.negated,
            between.low,
            between.high,
        )),
        // not (A like B) ===> A not like B
        Expr::Like(like) => Expr::Like(Like::new(
            !like.negated,
            like.expr,
            like.pattern,
            like.escape_char,
            like.case_insensitive,
        )),
        // use not clause
        _ => Expr::Not(Box::new(expr)),
    }
}

/// bitwise negate a Negative clause
/// input is the clause to be bitwise negated.(args for Negative clause)
/// For BinaryExpr:
///    ~(A & B) ===> ~A | ~B
///    ~(A | B) ===> ~A & ~B
/// For Negative:
///    ~(~A) ===> A
/// For others, use Negative clause
pub fn distribute_negation(expr: Expr) -> Expr {
    match expr {
        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
            match op {
                // ~(A & B) ===> ~A | ~B
                Operator::BitwiseAnd => {
                    let left = distribute_negation(*left);
                    let right = distribute_negation(*right);

                    bitwise_or(left, right)
                }
                // ~(A | B) ===> ~A & ~B
                Operator::BitwiseOr => {
                    let left = distribute_negation(*left);
                    let right = distribute_negation(*right);

                    bitwise_and(left, right)
                }
                // use negative clause
                _ => Expr::Negative(Box::new(Expr::BinaryExpr(BinaryExpr::new(
                    left, op, right,
                )))),
            }
        }
        // ~(~A) ===> A
        Expr::Negative(expr) => *expr,
        // use negative clause
        _ => Expr::Negative(Box::new(expr)),
    }
}

#[cfg(test)]
mod tests {
    use super::{is_one, is_zero};
    use arrow::datatypes::i256;
    use datafusion_common::ScalarValue;
    use datafusion_expr::lit;

    #[test]
    fn test_is_zero() {
        assert!(is_zero(&lit(ScalarValue::Int8(Some(0)))));
        assert!(is_zero(&lit(ScalarValue::Float32(Some(0.0)))));
        assert!(is_zero(&lit(ScalarValue::Decimal128(
            Some(i128::from(0)),
            9,
            0
        ))));
        assert!(is_zero(&lit(ScalarValue::Decimal128(
            Some(i128::from(0)),
            9,
            5
        ))));
        assert!(is_zero(&lit(ScalarValue::Decimal256(
            Some(i256::ZERO),
            9,
            0
        ))));
        assert!(is_zero(&lit(ScalarValue::Decimal256(
            Some(i256::ZERO),
            9,
            5
        ))));
    }

    #[test]
    fn test_is_one() {
        assert!(is_one(&lit(ScalarValue::Int8(Some(1)))));
        assert!(is_one(&lit(ScalarValue::Float32(Some(1.0)))));
        assert!(is_one(&lit(ScalarValue::Decimal128(
            Some(i128::from(1)),
            9,
            0
        ))));
        assert!(is_one(&lit(ScalarValue::Decimal128(
            Some(i128::from(10)),
            9,
            1
        ))));
        assert!(is_one(&lit(ScalarValue::Decimal128(
            Some(i128::from(100)),
            9,
            2
        ))));
        assert!(is_one(&lit(ScalarValue::Decimal256(
            Some(i256::from(1)),
            9,
            0
        ))));
        assert!(is_one(&lit(ScalarValue::Decimal256(
            Some(i256::from(10)),
            9,
            1
        ))));
        assert!(is_one(&lit(ScalarValue::Decimal256(
            Some(i256::from(100)),
            9,
            2
        ))));
        assert!(!is_one(&lit(ScalarValue::Decimal256(
            Some(i256::from(100)),
            9,
            -1
        ))));
    }
}