datafusion-physical-expr 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.

//! [`LiteralGuarantee`] predicate analysis to determine if a column is a
//! constant.

use crate::utils::split_disjunction;
use crate::{PhysicalExpr, split_conjunction};
use datafusion_common::{Column, HashMap, ScalarValue};
use datafusion_expr::Operator;
use std::collections::HashSet;
use std::fmt::{self, Display, Formatter};
use std::sync::Arc;

/// Represents a guarantee that must be true for a boolean expression to
/// evaluate to `true`.
///
/// The guarantee takes the form of a column and a set of literal (constant)
/// [`ScalarValue`]s. For the expression to evaluate to `true`, the column *must
/// satisfy* the guarantee(s).
///
/// To satisfy the guarantee, depending on [`Guarantee`], the values in the
/// column must either:
///
/// 1. be ONLY one of that set
/// 2. NOT be ANY of that set
///
/// # Uses `LiteralGuarantee`s
///
/// `LiteralGuarantee`s can be used to simplify filter expressions and skip data
/// files (e.g. row groups in parquet files) by proving expressions can not
/// possibly evaluate to `true`. For example, if we have a guarantee that `a`
/// must be in (`1`) for a filter to evaluate to `true`, then we can skip any
/// partition where we know that `a` never has the value of `1`.
///
/// **Important**: If a `LiteralGuarantee` is not satisfied, the relevant
/// expression is *guaranteed* to evaluate to `false` or `null`. **However**,
/// the opposite does not hold. Even if all `LiteralGuarantee`s are satisfied,
/// that does **not** guarantee that the predicate will actually evaluate to
/// `true`: it may still evaluate to `true`, `false` or `null`.
///
/// # Creating `LiteralGuarantee`s
///
/// Use [`LiteralGuarantee::analyze`] to extract literal guarantees from a
/// filter predicate.
///
/// # Details
/// A guarantee can be one of two forms:
///
/// 1. The column must be one the values for the predicate to be `true`. If the
///    column takes on any other value, the predicate can not evaluate to `true`.
///    For example,
///    `(a = 1)`, `(a = 1 OR a = 2)` or `a IN (1, 2, 3)`
///
/// 2. The column must NOT be one of the values for the predicate to be `true`.
///    If the column can ONLY take one of these values, the predicate can not
///    evaluate to `true`. For example,
///    `(a != 1)`, `(a != 1 AND a != 2)` or `a NOT IN (1, 2, 3)`
#[derive(Debug, Clone, PartialEq)]
pub struct LiteralGuarantee {
    pub column: Column,
    pub guarantee: Guarantee,
    pub literals: HashSet<ScalarValue>,
}

/// What is guaranteed about the values for a [`LiteralGuarantee`]?
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Guarantee {
    /// Guarantee that the expression is `true` if `column` is one of the values. If
    /// `column` is not one of the values, the expression can not be `true`.
    In,
    /// Guarantee that the expression is `true` if `column` is not ANY of the
    /// values. If `column` only takes one of these values, the expression can
    /// not be `true`.
    NotIn,
}

impl LiteralGuarantee {
    /// Create a new instance of the guarantee if the provided operator is
    /// supported. Returns None otherwise. See [`LiteralGuarantee::analyze`] to
    /// create these structures from an predicate (boolean expression).
    fn new<'a>(
        column_name: impl Into<String>,
        guarantee: Guarantee,
        literals: impl IntoIterator<Item = &'a ScalarValue>,
    ) -> Self {
        let literals: HashSet<_> = literals.into_iter().cloned().collect();

        Self {
            column: Column::from_name(column_name),
            guarantee,
            literals,
        }
    }

    /// Return a list of [`LiteralGuarantee`]s that must be satisfied for `expr`
    /// to evaluate to `true`.
    ///
    /// If more than one `LiteralGuarantee` is returned, they must **all** hold
    /// for the expression to possibly be `true`. If any is not satisfied, the
    /// expression is guaranteed to be `null` or `false`.
    ///
    /// # Notes:
    /// 1. `expr` must be a boolean expression or inlist expression.
    /// 2. `expr` is not simplified prior to analysis.
    pub fn analyze(expr: &Arc<dyn PhysicalExpr>) -> Vec<LiteralGuarantee> {
        // split conjunction: <expr> AND <expr> AND ...
        split_conjunction(expr)
            .into_iter()
            // for an `AND` conjunction to be true, all terms individually must be true
            .fold(GuaranteeBuilder::new(), |builder, expr| {
                if let Some(cel) = ColOpLit::try_new(expr) {
                    builder.aggregate_conjunct(&cel)
                } else if let Some(inlist) = expr
                    .as_any()
                    .downcast_ref::<crate::expressions::InListExpr>()
                {
                    if let Some(inlist) = ColInList::try_new(inlist) {
                        builder.aggregate_multi_conjunct(
                            inlist.col,
                            inlist.guarantee,
                            inlist.list.iter().map(|lit| lit.value()),
                        )
                    } else {
                        builder
                    }
                } else {
                    // split disjunction: <expr> OR <expr> OR ...
                    let disjunctions = split_disjunction(expr);

                    // We are trying to add a guarantee that a column must be
                    // in/not in a particular set of values for the expression
                    // to evaluate to true.
                    //
                    // A disjunction is true, if at least one of the terms is be
                    // true.
                    //
                    // Thus, we can infer a guarantee if all terms are of the
                    // form `(col <op> literal) OR (col <op> literal) OR ...`.
                    //
                    // For example, we can infer that `a = 1 OR a = 2 OR a = 3`
                    // is guaranteed to be true ONLY if a is in (`1`, `2` or `3`).
                    //
                    // However, for something like  `a = 1 OR a = 2 OR a < 0` we
                    // **can't** guarantee that the predicate is only true if a
                    // is in (`1`, `2`), as it could also be true if `a` were less
                    // than zero.
                    let terms = disjunctions
                        .iter()
                        .filter_map(|expr| ColOpLit::try_new(expr))
                        .collect::<Vec<_>>();

                    // if all terms are 'col <op> literal' with the same column
                    // and operation we can infer any guarantees
                    //
                    // For those like (a != foo AND (a != bar OR a != baz)).
                    // We can't combine the (a != bar OR a != baz) part, but
                    // it also doesn't invalidate our knowledge that a !=
                    // foo is required for the expression to be true.
                    // So we can only create a multi value guarantee for `=`
                    // (or a single value). (e.g. ignore `a != foo OR a != bar`)
                    let first_term = terms.first();
                    if !terms.is_empty()
                        && terms.len() == disjunctions.len()
                        && terms.iter().all(|term| {
                            term.col.name() == first_term.unwrap().col.name()
                                && term.guarantee == Guarantee::In
                        })
                    {
                        builder.aggregate_multi_conjunct(
                            first_term.unwrap().col,
                            Guarantee::In,
                            terms.iter().map(|term| term.lit.value()),
                        )
                    } else {
                        // Handle disjunctions with conjunctions like (a = 1 AND b = 2) OR (a = 2 AND b = 3)
                        // Extract termsets from each disjunction
                        // if in each termset, they have same column, and the guarantee is In,
                        // we can infer a guarantee for the column
                        // e.g. (a = 1 AND b = 2) OR (a = 2 AND b = 3) is `a IN (1, 2) AND b IN (2, 3)`
                        // otherwise, we can't infer a guarantee
                        let termsets: Vec<Vec<ColOpLitOrInList>> = disjunctions
                            .iter()
                            .map(|expr| {
                                split_conjunction(expr)
                                    .into_iter()
                                    .filter_map(ColOpLitOrInList::try_new)
                                    .filter(|term| term.guarantee() == Guarantee::In)
                                    .collect()
                            })
                            .collect();

                        // Early return if any termset is empty (can't infer guarantees)
                        if termsets.iter().any(|terms| terms.is_empty()) {
                            return builder;
                        }

                        // Find columns that appear in all termsets
                        let common_cols = find_common_columns(&termsets);
                        if common_cols.is_empty() {
                            return builder;
                        }

                        // Build guarantees for common columns
                        let mut builder = builder;
                        for col in common_cols {
                            let literals: Vec<_> = termsets
                                .iter()
                                .filter_map(|terms| {
                                    terms.iter().find(|term| term.col() == col).map(
                                        |term| {
                                            term.lits().into_iter().map(|lit| lit.value())
                                        },
                                    )
                                })
                                .flatten()
                                .collect();

                            builder = builder.aggregate_multi_conjunct(
                                col,
                                Guarantee::In,
                                literals.into_iter(),
                            );
                        }

                        builder
                    }
                }
            })
            .build()
    }
}

impl Display for LiteralGuarantee {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let mut sorted_literals: Vec<_> =
            self.literals.iter().map(|lit| lit.to_string()).collect();
        sorted_literals.sort();
        match self.guarantee {
            Guarantee::In => write!(
                f,
                "{} in ({})",
                self.column.name,
                sorted_literals.join(", ")
            ),
            Guarantee::NotIn => write!(
                f,
                "{} not in ({})",
                self.column.name,
                sorted_literals.join(", ")
            ),
        }
    }
}

/// Combines conjuncts (aka terms `AND`ed together) into [`LiteralGuarantee`]s,
/// preserving insert order
#[derive(Debug, Default)]
struct GuaranteeBuilder<'a> {
    /// List of guarantees that have been created so far
    /// if we have determined a subsequent conjunct invalidates a guarantee
    /// e.g. `a = foo AND a = bar` then the relevant guarantee will be None
    guarantees: Vec<Option<LiteralGuarantee>>,

    /// Key is the (column name, guarantee type)
    /// Value is the index into `guarantees`
    map: HashMap<(&'a crate::expressions::Column, Guarantee), usize>,
}

impl<'a> GuaranteeBuilder<'a> {
    fn new() -> Self {
        Default::default()
    }

    /// Aggregate a new single `AND col <op> literal` term to this builder
    /// combining with existing guarantees if possible.
    ///
    /// # Examples
    /// * `AND (a = 1)`: `a` is guaranteed to be 1
    /// * `AND (a != 1)`: a is guaranteed to not be 1
    fn aggregate_conjunct(self, col_op_lit: &ColOpLit<'a>) -> Self {
        self.aggregate_multi_conjunct(
            col_op_lit.col,
            col_op_lit.guarantee,
            [col_op_lit.lit.value()],
        )
    }

    /// Aggregates a new single column, multi literal term to this builder
    /// combining with previously known guarantees if possible.
    ///
    /// # Examples
    /// For the following examples, we can guarantee the expression is `true` if:
    /// * `AND (a = 1 OR a = 2 OR a = 3)`: a is in (1, 2, or 3)
    /// * `AND (a IN (1,2,3))`: a is in (1, 2, or 3)
    /// * `AND (a != 1 OR a != 2 OR a != 3)`: a is not in (1, 2, or 3)
    /// * `AND (a NOT IN (1,2,3))`: a is not in (1, 2, or 3)
    fn aggregate_multi_conjunct(
        mut self,
        col: &'a crate::expressions::Column,
        guarantee: Guarantee,
        new_values: impl IntoIterator<Item = &'a ScalarValue>,
    ) -> Self {
        let key = (col, guarantee);
        if let Some(index) = self.map.get(&key) {
            // already have a guarantee for this column
            let entry = &mut self.guarantees[*index];

            let Some(existing) = entry else {
                // determined the previous guarantee for this column has been
                // invalidated, nothing to do
                return self;
            };

            // Combine conjuncts if we have `a != foo AND a != bar`. `a = foo
            // AND a = bar` doesn't make logical sense so we don't optimize this
            // case
            match existing.guarantee {
                // knew that the column could not be a set of values
                //
                // For example, if we previously had `a != 5` and now we see
                // another `AND a != 6` we know that a must not be either 5 or 6
                // for the expression to be true
                Guarantee::NotIn => {
                    let new_values: HashSet<_> = new_values.into_iter().collect();
                    existing.literals.extend(new_values.into_iter().cloned());
                }
                Guarantee::In => {
                    let intersection = new_values
                        .into_iter()
                        .filter(|new_value| existing.literals.contains(*new_value))
                        .collect::<Vec<_>>();
                    // for an In guarantee, if the intersection is not empty,  we can extend the guarantee
                    // e.g. `a IN (1,2,3) AND a IN (2,3,4)` is `a IN (2,3)`
                    // otherwise, we invalidate the guarantee
                    // e.g. `a IN (1,2,3) AND a IN (4,5,6)` is `a IN ()`, which is invalid
                    if !intersection.is_empty() {
                        existing.literals = intersection.into_iter().cloned().collect();
                    } else {
                        // at least one was not, so invalidate the guarantee
                        *entry = None;
                    }
                }
            }
        } else {
            // This is a new guarantee
            let new_values: HashSet<_> = new_values.into_iter().collect();

            let guarantee = LiteralGuarantee::new(col.name(), guarantee, new_values);
            // add it to the list of guarantees
            self.guarantees.push(Some(guarantee));
            self.map.insert(key, self.guarantees.len() - 1);
        }

        self
    }

    /// Return all guarantees that have been created so far
    fn build(self) -> Vec<LiteralGuarantee> {
        // filter out any guarantees that have been invalidated
        self.guarantees.into_iter().flatten().collect()
    }
}

/// Represents a single `col [not]in literal` expression
struct ColOpLit<'a> {
    col: &'a crate::expressions::Column,
    guarantee: Guarantee,
    lit: &'a crate::expressions::Literal,
}

impl<'a> ColOpLit<'a> {
    /// Returns Some(ColOpLit) if the expression is either:
    /// 1. `col <op> literal`
    /// 2. `literal <op> col`
    /// 3. operator is `=` or `!=`
    ///
    /// Returns None otherwise
    fn try_new(expr: &'a Arc<dyn PhysicalExpr>) -> Option<Self> {
        let binary_expr = expr
            .as_any()
            .downcast_ref::<crate::expressions::BinaryExpr>()?;

        let (left, op, right) = (
            binary_expr.left().as_any(),
            binary_expr.op(),
            binary_expr.right().as_any(),
        );
        let guarantee = match op {
            Operator::Eq => Guarantee::In,
            Operator::NotEq => Guarantee::NotIn,
            _ => return None,
        };
        // col <op> literal
        if let (Some(col), Some(lit)) = (
            left.downcast_ref::<crate::expressions::Column>(),
            right.downcast_ref::<crate::expressions::Literal>(),
        ) {
            Some(Self {
                col,
                guarantee,
                lit,
            })
        }
        // literal <op> col
        else if let (Some(lit), Some(col)) = (
            left.downcast_ref::<crate::expressions::Literal>(),
            right.downcast_ref::<crate::expressions::Column>(),
        ) {
            Some(Self {
                col,
                guarantee,
                lit,
            })
        } else {
            None
        }
    }
}

/// Represents a single `col [not]in literal` expression
struct ColInList<'a> {
    col: &'a crate::expressions::Column,
    guarantee: Guarantee,
    list: Vec<&'a crate::expressions::Literal>,
}

impl<'a> ColInList<'a> {
    /// Returns Some(ColInList) if the expression is either:
    /// 1. `col <op> (literal1, literal2, ...)`
    /// 3. operator is `in` or `not in`
    ///
    /// Returns None otherwise
    fn try_new(inlist: &'a crate::expressions::InListExpr) -> Option<Self> {
        // Only support single-column inlist currently, multi-column inlist is not supported
        let col = inlist
            .expr()
            .as_any()
            .downcast_ref::<crate::expressions::Column>()?;

        let literals = inlist
            .list()
            .iter()
            .map(|e| e.as_any().downcast_ref::<crate::expressions::Literal>())
            .collect::<Option<Vec<_>>>()?;

        let guarantee = if inlist.negated() {
            Guarantee::NotIn
        } else {
            Guarantee::In
        };

        Some(Self {
            col,
            guarantee,
            list: literals,
        })
    }
}

/// Represents a single `col [not]in literal` expression or a single `col <op> literal` expression
enum ColOpLitOrInList<'a> {
    ColOpLit(ColOpLit<'a>),
    ColInList(ColInList<'a>),
}

impl<'a> ColOpLitOrInList<'a> {
    fn try_new(expr: &'a Arc<dyn PhysicalExpr>) -> Option<Self> {
        match expr
            .as_any()
            .downcast_ref::<crate::expressions::InListExpr>()
        {
            Some(inlist) => Some(Self::ColInList(ColInList::try_new(inlist)?)),
            None => ColOpLit::try_new(expr).map(Self::ColOpLit),
        }
    }

    fn guarantee(&self) -> Guarantee {
        match self {
            Self::ColOpLit(col_op_lit) => col_op_lit.guarantee,
            Self::ColInList(col_in_list) => col_in_list.guarantee,
        }
    }

    fn col(&self) -> &'a crate::expressions::Column {
        match self {
            Self::ColOpLit(col_op_lit) => col_op_lit.col,
            Self::ColInList(col_in_list) => col_in_list.col,
        }
    }

    fn lits(&self) -> Vec<&'a crate::expressions::Literal> {
        match self {
            Self::ColOpLit(col_op_lit) => vec![col_op_lit.lit],
            Self::ColInList(col_in_list) => col_in_list.list.clone(),
        }
    }
}

/// Find columns that appear in all termsets
fn find_common_columns<'a>(
    termsets: &[Vec<ColOpLitOrInList<'a>>],
) -> Vec<&'a crate::expressions::Column> {
    if termsets.is_empty() {
        return Vec::new();
    }

    // Start with columns from the first termset
    let mut common_cols: HashSet<_> = termsets[0].iter().map(|term| term.col()).collect();

    // check if any common_col in one termset occur many times
    // e.g. (a = 1 AND a = 2) OR (a = 2 AND b = 3), should not infer a guarantee
    // TODO: for above case, we can infer a IN (2) AND b IN (3)
    if common_cols.len() != termsets[0].len() {
        return Vec::new();
    }

    // Intersect with columns from remaining termsets
    for termset in termsets.iter().skip(1) {
        let termset_cols: HashSet<_> = termset.iter().map(|term| term.col()).collect();
        if termset_cols.len() != termset.len() {
            return Vec::new();
        }
        common_cols = common_cols.intersection(&termset_cols).cloned().collect();
    }

    common_cols.into_iter().collect()
}

#[cfg(test)]
mod test {
    use std::sync::LazyLock;

    use super::*;
    use crate::planner::logical2physical;

    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
    use datafusion_expr::expr_fn::*;
    use datafusion_expr::{Expr, lit};

    use itertools::Itertools;

    #[test]
    fn test_literal() {
        // a single literal offers no guarantee
        test_analyze(lit(true), vec![])
    }

    #[test]
    fn test_single() {
        // a = "foo"
        test_analyze(col("a").eq(lit("foo")), vec![in_guarantee("a", ["foo"])]);
        // "foo" = a
        test_analyze(lit("foo").eq(col("a")), vec![in_guarantee("a", ["foo"])]);
        // a != "foo"
        test_analyze(
            col("a").not_eq(lit("foo")),
            vec![not_in_guarantee("a", ["foo"])],
        );
        // "foo" != a
        test_analyze(
            lit("foo").not_eq(col("a")),
            vec![not_in_guarantee("a", ["foo"])],
        );
    }

    #[test]
    fn test_conjunction_single_column() {
        // b = 1 AND b = 2. This is impossible. Ideally this expression could be simplified to false
        test_analyze(col("b").eq(lit(1)).and(col("b").eq(lit(2))), vec![]);
        // b = 1 AND b != 2 . In theory, this could be simplified to `b = 1`.
        test_analyze(
            col("b").eq(lit(1)).and(col("b").not_eq(lit(2))),
            vec![
                // can only be true of b is 1 and b is not 2 (even though it is redundant)
                in_guarantee("b", [1]),
                not_in_guarantee("b", [2]),
            ],
        );
        // b != 1 AND b = 2. In theory, this could be simplified to `b = 2`.
        test_analyze(
            col("b").not_eq(lit(1)).and(col("b").eq(lit(2))),
            vec![
                // can only be true of b is not 1 and b is 2 (even though it is redundant)
                not_in_guarantee("b", [1]),
                in_guarantee("b", [2]),
            ],
        );
        // b != 1 AND b != 2
        test_analyze(
            col("b").not_eq(lit(1)).and(col("b").not_eq(lit(2))),
            vec![not_in_guarantee("b", [1, 2])],
        );
        // b != 1 AND b != 2 and b != 3
        test_analyze(
            col("b")
                .not_eq(lit(1))
                .and(col("b").not_eq(lit(2)))
                .and(col("b").not_eq(lit(3))),
            vec![not_in_guarantee("b", [1, 2, 3])],
        );
        // b != 1 AND b = 2 and b != 3. Can only be true if b is 2 and b is not in (1, 3)
        test_analyze(
            col("b")
                .not_eq(lit(1))
                .and(col("b").eq(lit(2)))
                .and(col("b").not_eq(lit(3))),
            vec![not_in_guarantee("b", [1, 3]), in_guarantee("b", [2])],
        );
        // b != 1 AND b != 2 and b = 3 (in theory could determine b = 3)
        test_analyze(
            col("b")
                .not_eq(lit(1))
                .and(col("b").not_eq(lit(2)))
                .and(col("b").eq(lit(3))),
            vec![not_in_guarantee("b", [1, 2]), in_guarantee("b", [3])],
        );
        // b != 1 AND b != 2 and b > 3 (to be true, b can't be either 1 or 2
        test_analyze(
            col("b")
                .not_eq(lit(1))
                .and(col("b").not_eq(lit(2)))
                .and(col("b").gt(lit(3))),
            vec![not_in_guarantee("b", [1, 2])],
        );
    }

    #[test]
    fn test_conjunction_multi_column() {
        // a = "foo" AND b = 1
        test_analyze(
            col("a").eq(lit("foo")).and(col("b").eq(lit(1))),
            vec![
                // should find both column guarantees
                in_guarantee("a", ["foo"]),
                in_guarantee("b", [1]),
            ],
        );
        // a != "foo" AND b != 1
        test_analyze(
            col("a").not_eq(lit("foo")).and(col("b").not_eq(lit(1))),
            // should find both column guarantees
            vec![not_in_guarantee("a", ["foo"]), not_in_guarantee("b", [1])],
        );
        // a = "foo" AND a = "bar"
        test_analyze(
            col("a").eq(lit("foo")).and(col("a").eq(lit("bar"))),
            // this predicate is impossible ( can't be both foo and bar),
            vec![],
        );
        // a = "foo" AND b != "bar"
        test_analyze(
            col("a").eq(lit("foo")).and(col("a").not_eq(lit("bar"))),
            vec![in_guarantee("a", ["foo"]), not_in_guarantee("a", ["bar"])],
        );
        // a != "foo" AND a != "bar"
        test_analyze(
            col("a").not_eq(lit("foo")).and(col("a").not_eq(lit("bar"))),
            // know it isn't "foo" or "bar"
            vec![not_in_guarantee("a", ["foo", "bar"])],
        );
        // a != "foo" AND a != "bar" and a != "baz"
        test_analyze(
            col("a")
                .not_eq(lit("foo"))
                .and(col("a").not_eq(lit("bar")))
                .and(col("a").not_eq(lit("baz"))),
            // know it isn't "foo" or "bar" or "baz"
            vec![not_in_guarantee("a", ["foo", "bar", "baz"])],
        );
        // a = "foo" AND a = "foo"
        let expr = col("a").eq(lit("foo"));
        test_analyze(expr.clone().and(expr), vec![in_guarantee("a", ["foo"])]);
        // b > 5 AND b = 10 (should get an b = 10 guarantee)
        test_analyze(
            col("b").gt(lit(5)).and(col("b").eq(lit(10))),
            vec![in_guarantee("b", [10])],
        );
        // b > 10 AND b = 10 (this is impossible)
        test_analyze(
            col("b").gt(lit(10)).and(col("b").eq(lit(10))),
            vec![
                //  if b isn't 10, it can not be true (though the expression actually can never be true)
                in_guarantee("b", [10]),
            ],
        );
        // a != "foo" and (a != "bar" OR a != "baz")
        test_analyze(
            col("a")
                .not_eq(lit("foo"))
                .and(col("a").not_eq(lit("bar")).or(col("a").not_eq(lit("baz")))),
            // a is not foo (we can't represent other knowledge about a)
            vec![not_in_guarantee("a", ["foo"])],
        );
    }

    #[test]
    fn test_conjunction_and_disjunction_single_column() {
        // b != 1 AND (b > 2)
        test_analyze(
            col("b").not_eq(lit(1)).and(col("b").gt(lit(2))),
            vec![
                // for the expression to be true, b can not be one
                not_in_guarantee("b", [1]),
            ],
        );

        // b = 1 AND (b = 2 OR b = 3). Could be simplified to false.
        test_analyze(
            col("b")
                .eq(lit(1))
                .and(col("b").eq(lit(2)).or(col("b").eq(lit(3)))),
            vec![
                // in theory, b must be 1 and one of 2,3 for this expression to be true
                // which is a logical contradiction
            ],
        );
    }

    #[test]
    fn test_disjunction_single_column() {
        // b = 1 OR b = 2
        test_analyze(
            col("b").eq(lit(1)).or(col("b").eq(lit(2))),
            vec![in_guarantee("b", [1, 2])],
        );
        // b != 1 OR b = 2
        test_analyze(col("b").not_eq(lit(1)).or(col("b").eq(lit(2))), vec![]);
        // b = 1 OR b != 2
        test_analyze(col("b").eq(lit(1)).or(col("b").not_eq(lit(2))), vec![]);
        // b != 1 OR b != 2
        test_analyze(col("b").not_eq(lit(1)).or(col("b").not_eq(lit(2))), vec![]);
        // b != 1 OR b != 2 OR b = 3 -- in theory could guarantee that b = 3
        test_analyze(
            col("b")
                .not_eq(lit(1))
                .or(col("b").not_eq(lit(2)))
                .or(lit("b").eq(lit(3))),
            vec![],
        );
        // b = 1 OR b = 2 OR b = 3
        test_analyze(
            col("b")
                .eq(lit(1))
                .or(col("b").eq(lit(2)))
                .or(col("b").eq(lit(3))),
            vec![in_guarantee("b", [1, 2, 3])],
        );
        // b = 1 OR b = 2 OR b > 3 -- can't guarantee that the expression is only true if a is in (1, 2)
        test_analyze(
            col("b")
                .eq(lit(1))
                .or(col("b").eq(lit(2)))
                .or(lit("b").eq(lit(3))),
            vec![],
        );
    }

    #[test]
    fn test_disjunction_multi_column() {
        // a = "foo" OR b = 1
        test_analyze(
            col("a").eq(lit("foo")).or(col("b").eq(lit(1))),
            // no can't have a single column guarantee (if a = "foo" then b != 1) etc
            vec![],
        );
        // a != "foo" OR b != 1
        test_analyze(
            col("a").not_eq(lit("foo")).or(col("b").not_eq(lit(1))),
            // No single column guarantee
            vec![],
        );
        // a = "foo" OR a = "bar"
        test_analyze(
            col("a").eq(lit("foo")).or(col("a").eq(lit("bar"))),
            vec![in_guarantee("a", ["foo", "bar"])],
        );
        // a = "foo" OR a = "foo"
        test_analyze(
            col("a").eq(lit("foo")).or(col("a").eq(lit("foo"))),
            vec![in_guarantee("a", ["foo"])],
        );
        // a != "foo" OR a != "bar"
        test_analyze(
            col("a").not_eq(lit("foo")).or(col("a").not_eq(lit("bar"))),
            // can't represent knowledge about a in this case
            vec![],
        );
        // a = "foo" OR a = "bar" OR a = "baz"
        test_analyze(
            col("a")
                .eq(lit("foo"))
                .or(col("a").eq(lit("bar")))
                .or(col("a").eq(lit("baz"))),
            vec![in_guarantee("a", ["foo", "bar", "baz"])],
        );
        // (a = "foo" OR a = "bar") AND (a = "baz)"
        test_analyze(
            (col("a").eq(lit("foo")).or(col("a").eq(lit("bar"))))
                .and(col("a").eq(lit("baz"))),
            // this could potentially be represented as 2 constraints with a more
            // sophisticated analysis
            vec![],
        );
        // (a = "foo" OR a = "bar") AND (b = 1)
        test_analyze(
            (col("a").eq(lit("foo")).or(col("a").eq(lit("bar"))))
                .and(col("b").eq(lit(1))),
            vec![in_guarantee("a", ["foo", "bar"]), in_guarantee("b", [1])],
        );
        // (a = "foo" OR a = "bar") OR (b = 1)
        test_analyze(
            col("a")
                .eq(lit("foo"))
                .or(col("a").eq(lit("bar")))
                .or(col("b").eq(lit(1))),
            // can't represent knowledge about a or b in this case
            vec![],
        );
    }

    #[test]
    fn test_single_inlist() {
        // b IN (1, 2, 3)
        test_analyze(
            col("b").in_list(vec![lit(1), lit(2), lit(3)], false),
            vec![in_guarantee("b", [1, 2, 3])],
        );
        // b NOT IN (1, 2, 3)
        test_analyze(
            col("b").in_list(vec![lit(1), lit(2), lit(3)], true),
            vec![not_in_guarantee("b", [1, 2, 3])],
        );
        // b IN (1,2,3,4...24)
        test_analyze(
            col("b").in_list((1..25).map(lit).collect_vec(), false),
            vec![in_guarantee("b", 1..25)],
        );
    }

    #[test]
    fn test_inlist_conjunction() {
        // b IN (1, 2, 3) AND b IN (2, 3, 4)
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], false)
                .and(col("b").in_list(vec![lit(2), lit(3), lit(4)], false)),
            vec![in_guarantee("b", [2, 3])],
        );
        // b NOT IN (1, 2, 3) AND b IN (2, 3, 4)
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], true)
                .and(col("b").in_list(vec![lit(2), lit(3), lit(4)], false)),
            vec![
                not_in_guarantee("b", [1, 2, 3]),
                in_guarantee("b", [2, 3, 4]),
            ],
        );
        // b NOT IN (1, 2, 3) AND b NOT IN (2, 3, 4)
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], true)
                .and(col("b").in_list(vec![lit(2), lit(3), lit(4)], true)),
            vec![not_in_guarantee("b", [1, 2, 3, 4])],
        );
        // b IN (1, 2, 3) AND b = 4
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], false)
                .and(col("b").eq(lit(4))),
            vec![],
        );
        // b IN (1, 2, 3) AND b = 2
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], false)
                .and(col("b").eq(lit(2))),
            vec![in_guarantee("b", [2])],
        );
        // b IN (1, 2, 3) AND b != 2
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], false)
                .and(col("b").not_eq(lit(2))),
            vec![in_guarantee("b", [1, 2, 3]), not_in_guarantee("b", [2])],
        );
        // b NOT IN (1, 2, 3) AND b != 4
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], true)
                .and(col("b").not_eq(lit(4))),
            vec![not_in_guarantee("b", [1, 2, 3, 4])],
        );
        // b NOT IN (1, 2, 3) AND b != 2
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], true)
                .and(col("b").not_eq(lit(2))),
            vec![not_in_guarantee("b", [1, 2, 3])],
        );
    }

    #[test]
    fn test_inlist_with_disjunction() {
        // b IN (1, 2, 3) AND (b = 3 OR b = 4)
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], false)
                .and(col("b").eq(lit(3)).or(col("b").eq(lit(4)))),
            vec![in_guarantee("b", [3])],
        );
        // b IN (1, 2, 3) AND (b = 4 OR b = 5)
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], false)
                .and(col("b").eq(lit(4)).or(col("b").eq(lit(5)))),
            vec![],
        );
        // b NOT IN (1, 2, 3) AND (b = 3 OR b = 4)
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], true)
                .and(col("b").eq(lit(3)).or(col("b").eq(lit(4)))),
            vec![not_in_guarantee("b", [1, 2, 3]), in_guarantee("b", [3, 4])],
        );
        // b IN (1, 2, 3) OR b = 2
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], false)
                .or(col("b").eq(lit(2))),
            vec![in_guarantee("b", [1, 2, 3])],
        );
        // b IN (1, 2, 3) OR b != 3
        test_analyze(
            col("b")
                .in_list(vec![lit(1), lit(2), lit(3)], false)
                .or(col("b").not_eq(lit(3))),
            vec![],
        );
    }

    #[test]
    fn test_disjunction_and_conjunction_multi_column() {
        // (a = "foo" AND b = 1) OR (a = "bar" AND b = 2)
        test_analyze(
            (col("a").eq(lit("foo")).and(col("b").eq(lit(1))))
                .or(col("a").eq(lit("bar")).and(col("b").eq(lit(2)))),
            vec![in_guarantee("a", ["foo", "bar"]), in_guarantee("b", [1, 2])],
        );
        // (a = "foo" AND b = 1) OR (a = "bar" AND b = 2) OR (b = 3)
        test_analyze(
            (col("a").eq(lit("foo")).and(col("b").eq(lit(1))))
                .or(col("a").eq(lit("bar")).and(col("b").eq(lit(2))))
                .or(col("b").eq(lit(3))),
            vec![in_guarantee("b", [1, 2, 3])],
        );
        // (a = "foo" AND b = 1) OR (a = "bar" AND b = 2) OR (c = 3)
        test_analyze(
            (col("a").eq(lit("foo")).and(col("b").eq(lit(1))))
                .or(col("a").eq(lit("bar")).and(col("b").eq(lit(2))))
                .or(col("c").eq(lit(3))),
            vec![],
        );
        // (a = "foo" AND b > 1) OR (a = "bar" AND b = 2)
        test_analyze(
            (col("a").eq(lit("foo")).and(col("b").gt(lit(1))))
                .or(col("a").eq(lit("bar")).and(col("b").eq(lit(2)))),
            vec![in_guarantee("a", ["foo", "bar"])],
        );
        // (a = "foo" AND b = 1) OR (b = 1 AND c = 2) OR (c = 3 AND a = "bar")
        test_analyze(
            (col("a").eq(lit("foo")).and(col("b").eq(lit(1))))
                .or(col("b").eq(lit(1)).and(col("c").eq(lit(2))))
                .or(col("c").eq(lit(3)).and(col("a").eq(lit("bar")))),
            vec![],
        );
        // (a = "foo" AND a = "bar") OR (a = "good" AND b = 1)
        // TODO: this should be `a IN ("good") AND b IN (1)`
        test_analyze(
            (col("a").eq(lit("foo")).and(col("a").eq(lit("bar"))))
                .or(col("a").eq(lit("good")).and(col("b").eq(lit(1)))),
            vec![],
        );
        // (a = "foo" AND a = "foo") OR (a = "good" AND b = 1)
        // TODO: this should be `a IN ("foo", "good")`
        test_analyze(
            (col("a").eq(lit("foo")).and(col("a").eq(lit("foo"))))
                .or(col("a").eq(lit("good")).and(col("b").eq(lit(1)))),
            vec![],
        );
        // (a = "foo" AND b = 3) OR (b = 4 AND b = 1) OR (b = 2 AND a = "bar")
        test_analyze(
            (col("a").eq(lit("foo")).and(col("b").eq(lit(3))))
                .or(col("b").eq(lit(4)).and(col("b").eq(lit(1))))
                .or(col("b").eq(lit(2)).and(col("a").eq(lit("bar")))),
            vec![],
        );
        // (b = 1 AND b > 3) OR (a = "foo" AND b = 4)
        test_analyze(
            (col("b").eq(lit(1)).and(col("b").gt(lit(3))))
                .or(col("a").eq(lit("foo")).and(col("b").eq(lit(4)))),
            // if b isn't 1 or 4, it can not be true (though the expression actually can never be true)
            vec![in_guarantee("b", [1, 4])],
        );
        // (a = "foo" AND b = 1) OR (a != "bar" AND b = 2)
        test_analyze(
            (col("a").eq(lit("foo")).and(col("b").eq(lit(1))))
                .or(col("a").not_eq(lit("bar")).and(col("b").eq(lit(2)))),
            vec![in_guarantee("b", [1, 2])],
        );
        // (a = "foo" AND b = 1) OR (a LIKE "%bar" AND b = 2)
        test_analyze(
            (col("a").eq(lit("foo")).and(col("b").eq(lit(1))))
                .or(col("a").like(lit("%bar")).and(col("b").eq(lit(2)))),
            vec![in_guarantee("b", [1, 2])],
        );
        // (a IN ("foo", "bar") AND b = 5) OR (a IN ("foo", "bar") AND b = 6)
        test_analyze(
            (col("a")
                .in_list(vec![lit("foo"), lit("bar")], false)
                .and(col("b").eq(lit(5))))
            .or(col("a")
                .in_list(vec![lit("foo"), lit("bar")], false)
                .and(col("b").eq(lit(6)))),
            vec![in_guarantee("a", ["foo", "bar"]), in_guarantee("b", [5, 6])],
        );
        // (a IN ("foo", "bar") AND b = 5) OR (a IN ("foo") AND b = 6)
        test_analyze(
            (col("a")
                .in_list(vec![lit("foo"), lit("bar")], false)
                .and(col("b").eq(lit(5))))
            .or(col("a")
                .in_list(vec![lit("foo")], false)
                .and(col("b").eq(lit(6)))),
            vec![in_guarantee("a", ["foo", "bar"]), in_guarantee("b", [5, 6])],
        );
        // (a NOT IN ("foo", "bar") AND b = 5) OR (a NOT IN ("foo") AND b = 6)
        test_analyze(
            (col("a")
                .in_list(vec![lit("foo"), lit("bar")], true)
                .and(col("b").eq(lit(5))))
            .or(col("a")
                .in_list(vec![lit("foo")], true)
                .and(col("b").eq(lit(6)))),
            vec![in_guarantee("b", [5, 6])],
        );
    }

    /// Tests that analyzing expr results in the expected guarantees
    fn test_analyze(expr: Expr, expected: Vec<LiteralGuarantee>) {
        println!("Begin analyze of {expr}");
        let schema = schema();
        let physical_expr = logical2physical(&expr, &schema);

        let actual = LiteralGuarantee::analyze(&physical_expr)
            .into_iter()
            .sorted_by_key(|g| g.column.name().to_string())
            .collect::<Vec<_>>();
        assert_eq!(
            expected, actual,
            "expr: {expr}\
               \n\nexpected: {expected:#?}\
               \n\nactual: {actual:#?}\
               \n\nexpr: {expr:#?}\
               \n\nphysical_expr: {physical_expr:#?}"
        );
    }

    /// Guarantee that the expression is true if the column is one of the specified values
    fn in_guarantee<'a, I, S>(column: &str, literals: I) -> LiteralGuarantee
    where
        I: IntoIterator<Item = S>,
        S: Into<ScalarValue> + 'a,
    {
        let literals: Vec<_> = literals.into_iter().map(|s| s.into()).collect();
        LiteralGuarantee::new(column, Guarantee::In, literals.iter())
    }

    /// Guarantee that the expression is true if the column is NOT any of the specified values
    fn not_in_guarantee<'a, I, S>(column: &str, literals: I) -> LiteralGuarantee
    where
        I: IntoIterator<Item = S>,
        S: Into<ScalarValue> + 'a,
    {
        let literals: Vec<_> = literals.into_iter().map(|s| s.into()).collect();
        LiteralGuarantee::new(column, Guarantee::NotIn, literals.iter())
    }

    // Schema for testing
    fn schema() -> SchemaRef {
        static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
            Arc::new(Schema::new(vec![
                Field::new("a", DataType::Utf8, false),
                Field::new("b", DataType::Int32, false),
                Field::new("c", DataType::Int32, false),
            ]))
        });
        Arc::clone(&SCHEMA)
    }
}