blob: c7912bbf70b05f7a72d52edd5d8569871943b7f0 [file]
// 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.
//! Expression simplification API
use arrow::{
array::{new_null_array, AsArray},
datatypes::{DataType, Field, Schema},
record_batch::RecordBatch,
};
use std::borrow::Cow;
use std::collections::HashSet;
use std::ops::Not;
use std::sync::Arc;
use datafusion_common::{
cast::{as_large_list_array, as_list_array},
metadata::FieldMetadata,
tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter},
};
use datafusion_common::{
exec_datafusion_err, internal_err, DFSchema, DataFusionError, Result, ScalarValue,
};
use datafusion_expr::{
and, binary::BinaryTypeCoercer, lit, or, BinaryExpr, Case, ColumnarValue, Expr, Like,
Operator, Volatility,
};
use datafusion_expr::{expr::ScalarFunction, interval_arithmetic::NullableInterval};
use datafusion_expr::{
expr::{InList, InSubquery},
utils::{iter_conjunction, iter_conjunction_owned},
};
use datafusion_expr::{simplify::ExprSimplifyResult, Cast, TryCast};
use datafusion_physical_expr::{create_physical_expr, execution_props::ExecutionProps};
use super::inlist_simplifier::ShortenInListSimplifier;
use super::utils::*;
use crate::analyzer::type_coercion::TypeCoercionRewriter;
use crate::simplify_expressions::guarantees::GuaranteeRewriter;
use crate::simplify_expressions::regex::simplify_regex_expr;
use crate::simplify_expressions::unwrap_cast::{
is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary,
is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist,
unwrap_cast_in_comparison_for_binary,
};
use crate::simplify_expressions::SimplifyInfo;
use datafusion_expr_common::casts::try_cast_literal_to_type;
use indexmap::IndexSet;
use regex::Regex;
/// This structure handles API for expression simplification
///
/// Provides simplification information based on DFSchema and
/// [`ExecutionProps`]. This is the default implementation used by DataFusion
///
/// For example:
/// ```
/// use arrow::datatypes::{DataType, Field, Schema};
/// use datafusion_common::{DataFusionError, ToDFSchema};
/// use datafusion_expr::execution_props::ExecutionProps;
/// use datafusion_expr::simplify::SimplifyContext;
/// use datafusion_expr::{col, lit};
/// use datafusion_optimizer::simplify_expressions::ExprSimplifier;
///
/// // Create the schema
/// let schema = Schema::new(vec![Field::new("i", DataType::Int64, false)])
/// .to_dfschema_ref()
/// .unwrap();
///
/// // Create the simplifier
/// let props = ExecutionProps::new();
/// let context = SimplifyContext::new(&props).with_schema(schema);
/// let simplifier = ExprSimplifier::new(context);
///
/// // Use the simplifier
///
/// // b < 2 or (1 > 3)
/// let expr = col("b").lt(lit(2)).or(lit(1).gt(lit(3)));
///
/// // b < 2
/// let simplified = simplifier.simplify(expr).unwrap();
/// assert_eq!(simplified, col("b").lt(lit(2)));
/// ```
pub struct ExprSimplifier<S> {
info: S,
/// Guarantees about the values of columns. This is provided by the user
/// in [ExprSimplifier::with_guarantees()].
guarantees: Vec<(Expr, NullableInterval)>,
/// Should expressions be canonicalized before simplification? Defaults to
/// true
canonicalize: bool,
/// Maximum number of simplifier cycles
max_simplifier_cycles: u32,
}
pub const THRESHOLD_INLINE_INLIST: usize = 3;
pub const DEFAULT_MAX_SIMPLIFIER_CYCLES: u32 = 3;
impl<S: SimplifyInfo> ExprSimplifier<S> {
/// Create a new `ExprSimplifier` with the given `info` such as an
/// instance of [`SimplifyContext`]. See
/// [`simplify`](Self::simplify) for an example.
///
/// [`SimplifyContext`]: datafusion_expr::simplify::SimplifyContext
pub fn new(info: S) -> Self {
Self {
info,
guarantees: vec![],
canonicalize: true,
max_simplifier_cycles: DEFAULT_MAX_SIMPLIFIER_CYCLES,
}
}
/// Simplifies this [`Expr`] as much as possible, evaluating
/// constants and applying algebraic simplifications.
///
/// The types of the expression must match what operators expect,
/// or else an error may occur trying to evaluate. See
/// [`coerce`](Self::coerce) for a function to help.
///
/// # Example:
///
/// `b > 2 AND b > 2`
///
/// can be written to
///
/// `b > 2`
///
/// ```
/// use arrow::datatypes::DataType;
/// use datafusion_common::DFSchema;
/// use datafusion_common::Result;
/// use datafusion_expr::execution_props::ExecutionProps;
/// use datafusion_expr::simplify::SimplifyContext;
/// use datafusion_expr::simplify::SimplifyInfo;
/// use datafusion_expr::{col, lit, Expr};
/// use datafusion_optimizer::simplify_expressions::ExprSimplifier;
/// use std::sync::Arc;
///
/// /// Simple implementation that provides `Simplifier` the information it needs
/// /// See SimplifyContext for a structure that does this.
/// #[derive(Default)]
/// struct Info {
/// execution_props: ExecutionProps,
/// };
///
/// impl SimplifyInfo for Info {
/// fn is_boolean_type(&self, expr: &Expr) -> Result<bool> {
/// Ok(false)
/// }
/// fn nullable(&self, expr: &Expr) -> Result<bool> {
/// Ok(true)
/// }
/// fn execution_props(&self) -> &ExecutionProps {
/// &self.execution_props
/// }
/// fn get_data_type(&self, expr: &Expr) -> Result<DataType> {
/// Ok(DataType::Int32)
/// }
/// }
///
/// // Create the simplifier
/// let simplifier = ExprSimplifier::new(Info::default());
///
/// // b < 2
/// let b_lt_2 = col("b").gt(lit(2));
///
/// // (b < 2) OR (b < 2)
/// let expr = b_lt_2.clone().or(b_lt_2.clone());
///
/// // (b < 2) OR (b < 2) --> (b < 2)
/// let expr = simplifier.simplify(expr).unwrap();
/// assert_eq!(expr, b_lt_2);
/// ```
pub fn simplify(&self, expr: Expr) -> Result<Expr> {
Ok(self.simplify_with_cycle_count_transformed(expr)?.0.data)
}
/// Like [Self::simplify], simplifies this [`Expr`] as much as possible, evaluating
/// constants and applying algebraic simplifications. Additionally returns a `u32`
/// representing the number of simplification cycles performed, which can be useful for testing
/// optimizations.
///
/// See [Self::simplify] for details and usage examples.
#[deprecated(
since = "48.0.0",
note = "Use `simplify_with_cycle_count_transformed` instead"
)]
#[allow(unused_mut)]
pub fn simplify_with_cycle_count(&self, mut expr: Expr) -> Result<(Expr, u32)> {
let (transformed, cycle_count) =
self.simplify_with_cycle_count_transformed(expr)?;
Ok((transformed.data, cycle_count))
}
/// Like [Self::simplify], simplifies this [`Expr`] as much as possible, evaluating
/// constants and applying algebraic simplifications. Additionally returns a `u32`
/// representing the number of simplification cycles performed, which can be useful for testing
/// optimizations.
///
/// # Returns
///
/// A tuple containing:
/// - The simplified expression wrapped in a `Transformed<Expr>` indicating if changes were made
/// - The number of simplification cycles that were performed
///
/// See [Self::simplify] for details and usage examples.
pub fn simplify_with_cycle_count_transformed(
&self,
mut expr: Expr,
) -> Result<(Transformed<Expr>, u32)> {
let mut simplifier = Simplifier::new(&self.info);
let mut const_evaluator = ConstEvaluator::try_new(self.info.execution_props())?;
let mut shorten_in_list_simplifier = ShortenInListSimplifier::new();
let mut guarantee_rewriter = GuaranteeRewriter::new(&self.guarantees);
if self.canonicalize {
expr = expr.rewrite(&mut Canonicalizer::new()).data()?
}
// Evaluating constants can enable new simplifications and
// simplifications can enable new constant evaluation
// see `Self::with_max_cycles`
let mut num_cycles = 0;
let mut has_transformed = false;
loop {
let Transformed {
data, transformed, ..
} = expr
.rewrite(&mut const_evaluator)?
.transform_data(|expr| expr.rewrite(&mut simplifier))?
.transform_data(|expr| expr.rewrite(&mut guarantee_rewriter))?;
expr = data;
num_cycles += 1;
// Track if any transformation occurred
has_transformed = has_transformed || transformed;
if !transformed || num_cycles >= self.max_simplifier_cycles {
break;
}
}
// shorten inlist should be started after other inlist rules are applied
expr = expr.rewrite(&mut shorten_in_list_simplifier).data()?;
Ok((
Transformed::new_transformed(expr, has_transformed),
num_cycles,
))
}
/// Apply type coercion to an [`Expr`] so that it can be
/// evaluated as a [`PhysicalExpr`](datafusion_physical_expr::PhysicalExpr).
///
/// See the [type coercion module](datafusion_expr::type_coercion)
/// documentation for more details on type coercion
pub fn coerce(&self, expr: Expr, schema: &DFSchema) -> Result<Expr> {
let mut expr_rewrite = TypeCoercionRewriter { schema };
expr.rewrite(&mut expr_rewrite).data()
}
/// Input guarantees about the values of columns.
///
/// The guarantees can simplify expressions. For example, if a column `x` is
/// guaranteed to be `3`, then the expression `x > 1` can be replaced by the
/// literal `true`.
///
/// The guarantees are provided as a `Vec<(Expr, NullableInterval)>`,
/// where the [Expr] is a column reference and the [NullableInterval]
/// is an interval representing the known possible values of that column.
///
/// ```rust
/// use arrow::datatypes::{DataType, Field, Schema};
/// use datafusion_common::{Result, ScalarValue, ToDFSchema};
/// use datafusion_expr::execution_props::ExecutionProps;
/// use datafusion_expr::interval_arithmetic::{Interval, NullableInterval};
/// use datafusion_expr::simplify::SimplifyContext;
/// use datafusion_expr::{col, lit, Expr};
/// use datafusion_optimizer::simplify_expressions::ExprSimplifier;
///
/// let schema = Schema::new(vec![
/// Field::new("x", DataType::Int64, false),
/// Field::new("y", DataType::UInt32, false),
/// Field::new("z", DataType::Int64, false),
/// ])
/// .to_dfschema_ref()
/// .unwrap();
///
/// // Create the simplifier
/// let props = ExecutionProps::new();
/// let context = SimplifyContext::new(&props).with_schema(schema);
///
/// // Expression: (x >= 3) AND (y + 2 < 10) AND (z > 5)
/// let expr_x = col("x").gt_eq(lit(3_i64));
/// let expr_y = (col("y") + lit(2_u32)).lt(lit(10_u32));
/// let expr_z = col("z").gt(lit(5_i64));
/// let expr = expr_x.and(expr_y).and(expr_z.clone());
///
/// let guarantees = vec![
/// // x ∈ [3, 5]
/// (
/// col("x"),
/// NullableInterval::NotNull {
/// values: Interval::make(Some(3_i64), Some(5_i64)).unwrap(),
/// },
/// ),
/// // y = 3
/// (
/// col("y"),
/// NullableInterval::from(ScalarValue::UInt32(Some(3))),
/// ),
/// ];
/// let simplifier = ExprSimplifier::new(context).with_guarantees(guarantees);
/// let output = simplifier.simplify(expr).unwrap();
/// // Expression becomes: true AND true AND (z > 5), which simplifies to
/// // z > 5.
/// assert_eq!(output, expr_z);
/// ```
pub fn with_guarantees(mut self, guarantees: Vec<(Expr, NullableInterval)>) -> Self {
self.guarantees = guarantees;
self
}
/// Should `Canonicalizer` be applied before simplification?
///
/// If true (the default), the expression will be rewritten to canonical
/// form before simplification. This is useful to ensure that the simplifier
/// can apply all possible simplifications.
///
/// Some expressions, such as those in some Joins, can not be canonicalized
/// without changing their meaning. In these cases, canonicalization should
/// be disabled.
///
/// ```rust
/// use arrow::datatypes::{DataType, Field, Schema};
/// use datafusion_common::{Result, ScalarValue, ToDFSchema};
/// use datafusion_expr::execution_props::ExecutionProps;
/// use datafusion_expr::interval_arithmetic::{Interval, NullableInterval};
/// use datafusion_expr::simplify::SimplifyContext;
/// use datafusion_expr::{col, lit, Expr};
/// use datafusion_optimizer::simplify_expressions::ExprSimplifier;
///
/// let schema = Schema::new(vec![
/// Field::new("a", DataType::Int64, false),
/// Field::new("b", DataType::Int64, false),
/// Field::new("c", DataType::Int64, false),
/// ])
/// .to_dfschema_ref()
/// .unwrap();
///
/// // Create the simplifier
/// let props = ExecutionProps::new();
/// let context = SimplifyContext::new(&props).with_schema(schema);
/// let simplifier = ExprSimplifier::new(context);
///
/// // Expression: a = c AND 1 = b
/// let expr = col("a").eq(col("c")).and(lit(1).eq(col("b")));
///
/// // With canonicalization, the expression is rewritten to canonical form
/// // (though it is no simpler in this case):
/// let canonical = simplifier.simplify(expr.clone()).unwrap();
/// // Expression has been rewritten to: (c = a AND b = 1)
/// assert_eq!(canonical, col("c").eq(col("a")).and(col("b").eq(lit(1))));
///
/// // If canonicalization is disabled, the expression is not changed
/// let non_canonicalized = simplifier
/// .with_canonicalize(false)
/// .simplify(expr.clone())
/// .unwrap();
///
/// assert_eq!(non_canonicalized, expr);
/// ```
pub fn with_canonicalize(mut self, canonicalize: bool) -> Self {
self.canonicalize = canonicalize;
self
}
/// Specifies the maximum number of simplification cycles to run.
///
/// The simplifier can perform multiple passes of simplification. This is
/// because the output of one simplification step can allow more optimizations
/// in another simplification step. For example, constant evaluation can allow more
/// expression simplifications, and expression simplifications can allow more constant
/// evaluations.
///
/// This method specifies the maximum number of allowed iteration cycles before the simplifier
/// returns an [Expr] output. However, it does not always perform the maximum number of cycles.
/// The simplifier will attempt to detect when an [Expr] is unchanged by all the simplification
/// passes, and return early. This avoids wasting time on unnecessary [Expr] tree traversals.
///
/// If no maximum is specified, the value of [DEFAULT_MAX_SIMPLIFIER_CYCLES] is used
/// instead.
///
/// ```rust
/// use arrow::datatypes::{DataType, Field, Schema};
/// use datafusion_expr::{col, lit, Expr};
/// use datafusion_common::{Result, ScalarValue, ToDFSchema};
/// use datafusion_expr::execution_props::ExecutionProps;
/// use datafusion_expr::simplify::SimplifyContext;
/// use datafusion_optimizer::simplify_expressions::ExprSimplifier;
///
/// let schema = Schema::new(vec![
/// Field::new("a", DataType::Int64, false),
/// ])
/// .to_dfschema_ref().unwrap();
///
/// // Create the simplifier
/// let props = ExecutionProps::new();
/// let context = SimplifyContext::new(&props)
/// .with_schema(schema);
/// let simplifier = ExprSimplifier::new(context);
///
/// // Expression: a IS NOT NULL
/// let expr = col("a").is_not_null();
///
/// // When using default maximum cycles, 2 cycles will be performed.
/// let (simplified_expr, count) = simplifier.simplify_with_cycle_count_transformed(expr.clone()).unwrap();
/// assert_eq!(simplified_expr.data, lit(true));
/// // 2 cycles were executed, but only 1 was needed
/// assert_eq!(count, 2);
///
/// // Only 1 simplification pass is necessary here, so we can set the maximum cycles to 1.
/// let (simplified_expr, count) = simplifier.with_max_cycles(1).simplify_with_cycle_count_transformed(expr.clone()).unwrap();
/// // Expression has been rewritten to: (c = a AND b = 1)
/// assert_eq!(simplified_expr.data, lit(true));
/// // Only 1 cycle was executed
/// assert_eq!(count, 1);
/// ```
pub fn with_max_cycles(mut self, max_simplifier_cycles: u32) -> Self {
self.max_simplifier_cycles = max_simplifier_cycles;
self
}
}
/// Canonicalize any BinaryExprs that are not in canonical form
///
/// `<literal> <op> <col>` is rewritten to `<col> <op> <literal>`
///
/// `<col1> <op> <col2>` is rewritten so that the name of `col1` sorts higher
/// than `col2` (`a > b` would be canonicalized to `b < a`)
struct Canonicalizer {}
impl Canonicalizer {
fn new() -> Self {
Self {}
}
}
impl TreeNodeRewriter for Canonicalizer {
type Node = Expr;
fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
let Expr::BinaryExpr(BinaryExpr { left, op, right }) = expr else {
return Ok(Transformed::no(expr));
};
match (left.as_ref(), right.as_ref(), op.swap()) {
// <col1> <op> <col2>
(Expr::Column(left_col), Expr::Column(right_col), Some(swapped_op))
if right_col > left_col =>
{
Ok(Transformed::yes(Expr::BinaryExpr(BinaryExpr {
left: right,
op: swapped_op,
right: left,
})))
}
// <literal> <op> <col>
(Expr::Literal(_a, _), Expr::Column(_b), Some(swapped_op)) => {
Ok(Transformed::yes(Expr::BinaryExpr(BinaryExpr {
left: right,
op: swapped_op,
right: left,
})))
}
_ => Ok(Transformed::no(Expr::BinaryExpr(BinaryExpr {
left,
op,
right,
}))),
}
}
}
#[allow(rustdoc::private_intra_doc_links)]
/// Partially evaluate `Expr`s so constant subtrees are evaluated at plan time.
///
/// Note it does not handle algebraic rewrites such as `(a or false)`
/// --> `a`, which is handled by [`Simplifier`]
struct ConstEvaluator<'a> {
/// `can_evaluate` is used during the depth-first-search of the
/// `Expr` tree to track if any siblings (or their descendants) were
/// non evaluatable (e.g. had a column reference or volatile
/// function)
///
/// Specifically, `can_evaluate[N]` represents the state of
/// traversal when we are N levels deep in the tree, one entry for
/// this Expr and each of its parents.
///
/// After visiting all siblings if `can_evaluate.top()` is true, that
/// means there were no non evaluatable siblings (or their
/// descendants) so this `Expr` can be evaluated
can_evaluate: Vec<bool>,
execution_props: &'a ExecutionProps,
input_schema: DFSchema,
input_batch: RecordBatch,
}
#[allow(dead_code)]
/// The simplify result of ConstEvaluator
enum ConstSimplifyResult {
// Expr was simplified and contains the new expression
Simplified(ScalarValue, Option<FieldMetadata>),
// Expr was not simplified and original value is returned
NotSimplified(ScalarValue, Option<FieldMetadata>),
// Evaluation encountered an error, contains the original expression
SimplifyRuntimeError(DataFusionError, Expr),
}
impl TreeNodeRewriter for ConstEvaluator<'_> {
type Node = Expr;
fn f_down(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
// Default to being able to evaluate this node
self.can_evaluate.push(true);
// if this expr is not ok to evaluate, mark entire parent
// stack as not ok (as all parents have at least one child or
// descendant that can not be evaluated
if !Self::can_evaluate(&expr) {
// walk back up stack, marking first parent that is not mutable
let parent_iter = self.can_evaluate.iter_mut().rev();
for p in parent_iter {
if !*p {
// optimization: if we find an element on the
// stack already marked, know all elements above are also marked
break;
}
*p = false;
}
}
// NB: do not short circuit recursion even if we find a non
// evaluatable node (so we can fold other children, args to
// functions, etc.)
Ok(Transformed::no(expr))
}
fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
match self.can_evaluate.pop() {
// Certain expressions such as `CASE` and `COALESCE` are short-circuiting
// and may not evaluate all their sub expressions. Thus, if
// any error is countered during simplification, return the original
// so that normal evaluation can occur
Some(true) => match self.evaluate_to_scalar(expr) {
ConstSimplifyResult::Simplified(s, m) => {
Ok(Transformed::yes(Expr::Literal(s, m)))
}
ConstSimplifyResult::NotSimplified(s, m) => {
Ok(Transformed::no(Expr::Literal(s, m)))
}
ConstSimplifyResult::SimplifyRuntimeError(err, expr) => {
// For CAST expressions with literal inputs, propagate the error at plan time rather than deferring to execution time.
// This provides clearer error messages and fails fast.
if let Expr::Cast(Cast { ref expr, .. })
| Expr::TryCast(TryCast { ref expr, .. }) = expr
{
if matches!(expr.as_ref(), Expr::Literal(_, _)) {
return Err(err);
}
}
// For other expressions (like CASE, COALESCE), preserve the original
// to allow short-circuit evaluation at execution time
Ok(Transformed::yes(expr))
}
},
Some(false) => Ok(Transformed::no(expr)),
_ => internal_err!("Failed to pop can_evaluate"),
}
}
}
impl<'a> ConstEvaluator<'a> {
/// Create a new `ConstantEvaluator`. Session constants (such as
/// the time for `now()` are taken from the passed
/// `execution_props`.
pub fn try_new(execution_props: &'a ExecutionProps) -> Result<Self> {
// The dummy column name is unused and doesn't matter as only
// expressions without column references can be evaluated
static DUMMY_COL_NAME: &str = ".";
let schema = Arc::new(Schema::new(vec![Field::new(
DUMMY_COL_NAME,
DataType::Null,
true,
)]));
let input_schema = DFSchema::try_from(Arc::clone(&schema))?;
// Need a single "input" row to produce a single output row
let col = new_null_array(&DataType::Null, 1);
let input_batch = RecordBatch::try_new(schema, vec![col])?;
Ok(Self {
can_evaluate: vec![],
execution_props,
input_schema,
input_batch,
})
}
/// Can a function of the specified volatility be evaluated?
fn volatility_ok(volatility: Volatility) -> bool {
match volatility {
Volatility::Immutable => true,
// Values for functions such as now() are taken from ExecutionProps
Volatility::Stable => true,
Volatility::Volatile => false,
}
}
/// Can the expression be evaluated at plan time, (assuming all of
/// its children can also be evaluated)?
fn can_evaluate(expr: &Expr) -> bool {
// check for reasons we can't evaluate this node
//
// NOTE all expr types are listed here so when new ones are
// added they can be checked for their ability to be evaluated
// at plan time
match expr {
// TODO: remove the next line after `Expr::Wildcard` is removed
#[expect(deprecated)]
Expr::AggregateFunction { .. }
| Expr::ScalarVariable(_, _)
| Expr::Column(_)
| Expr::OuterReferenceColumn(_, _)
| Expr::Exists { .. }
| Expr::InSubquery(_)
| Expr::ScalarSubquery(_)
| Expr::WindowFunction { .. }
| Expr::GroupingSet(_)
| Expr::Wildcard { .. }
| Expr::Placeholder(_) => false,
Expr::ScalarFunction(ScalarFunction { func, .. }) => {
Self::volatility_ok(func.signature().volatility)
}
Expr::Literal(_, _)
| Expr::Alias(..)
| Expr::Unnest(_)
| Expr::BinaryExpr { .. }
| Expr::Not(_)
| Expr::IsNotNull(_)
| Expr::IsNull(_)
| Expr::IsTrue(_)
| Expr::IsFalse(_)
| Expr::IsUnknown(_)
| Expr::IsNotTrue(_)
| Expr::IsNotFalse(_)
| Expr::IsNotUnknown(_)
| Expr::Negative(_)
| Expr::Between { .. }
| Expr::Like { .. }
| Expr::SimilarTo { .. }
| Expr::Case(_)
| Expr::Cast { .. }
| Expr::TryCast { .. }
| Expr::InList { .. } => true,
}
}
/// Internal helper to evaluates an Expr
pub(crate) fn evaluate_to_scalar(&mut self, expr: Expr) -> ConstSimplifyResult {
if let Expr::Literal(s, m) = expr {
return ConstSimplifyResult::NotSimplified(s, m);
}
let phys_expr =
match create_physical_expr(&expr, &self.input_schema, self.execution_props) {
Ok(e) => e,
Err(err) => return ConstSimplifyResult::SimplifyRuntimeError(err, expr),
};
let metadata = phys_expr
.return_field(self.input_batch.schema_ref())
.ok()
.and_then(|f| {
let m = f.metadata();
match m.is_empty() {
true => None,
false => Some(FieldMetadata::from(m)),
}
});
let col_val = match phys_expr.evaluate(&self.input_batch) {
Ok(v) => v,
Err(err) => return ConstSimplifyResult::SimplifyRuntimeError(err, expr),
};
match col_val {
ColumnarValue::Array(a) => {
if a.len() != 1 {
ConstSimplifyResult::SimplifyRuntimeError(
exec_datafusion_err!("Could not evaluate the expression, found a result of length {}", a.len()),
expr,
)
} else if as_list_array(&a).is_ok() {
ConstSimplifyResult::Simplified(
ScalarValue::List(a.as_list::<i32>().to_owned().into()),
metadata,
)
} else if as_large_list_array(&a).is_ok() {
ConstSimplifyResult::Simplified(
ScalarValue::LargeList(a.as_list::<i64>().to_owned().into()),
metadata,
)
} else {
// Non-ListArray
match ScalarValue::try_from_array(&a, 0) {
Ok(s) => ConstSimplifyResult::Simplified(s, metadata),
Err(err) => ConstSimplifyResult::SimplifyRuntimeError(err, expr),
}
}
}
ColumnarValue::Scalar(s) => ConstSimplifyResult::Simplified(s, metadata),
}
}
}
/// Simplifies [`Expr`]s by applying algebraic transformation rules
///
/// Example transformations that are applied:
/// * `expr = true` and `expr != false` to `expr` when `expr` is of boolean type
/// * `expr = false` and `expr != true` to `!expr` when `expr` is of boolean type
/// * `true = true` and `false = false` to `true`
/// * `false = true` and `true = false` to `false`
/// * `!!expr` to `expr`
/// * `expr = null` and `expr != null` to `null`
struct Simplifier<'a, S> {
info: &'a S,
}
impl<'a, S> Simplifier<'a, S> {
pub fn new(info: &'a S) -> Self {
Self { info }
}
}
impl<S: SimplifyInfo> TreeNodeRewriter for Simplifier<'_, S> {
type Node = Expr;
/// rewrite the expression simplifying any constant expressions
fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
use datafusion_expr::Operator::{
And, BitwiseAnd, BitwiseOr, BitwiseShiftLeft, BitwiseShiftRight, BitwiseXor,
Divide, Eq, Modulo, Multiply, NotEq, Or, RegexIMatch, RegexMatch,
RegexNotIMatch, RegexNotMatch,
};
let info = self.info;
Ok(match expr {
// `value op NULL` -> `NULL`
// `NULL op value` -> `NULL`
// except for few operators that can return non-null value even when one of the operands is NULL
ref expr @ Expr::BinaryExpr(BinaryExpr {
ref left,
ref op,
ref right,
}) if op.returns_null_on_null()
&& (is_null(left.as_ref()) || is_null(right.as_ref())) =>
{
Transformed::yes(Expr::Literal(
ScalarValue::try_new_null(&info.get_data_type(expr)?)?,
None,
))
}
// `NULL {AND, OR} NULL` -> `NULL`
Expr::BinaryExpr(BinaryExpr {
left,
op: And | Or,
right,
}) if is_null(&left) && is_null(&right) => Transformed::yes(lit_bool_null()),
//
// Rules for Eq
//
// true = A --> A
// false = A --> !A
// null = A --> null
Expr::BinaryExpr(BinaryExpr {
left,
op: Eq,
right,
}) if is_bool_lit(&left) && info.is_boolean_type(&right)? => {
Transformed::yes(match as_bool_lit(&left)? {
Some(true) => *right,
Some(false) => Expr::Not(right),
None => lit_bool_null(),
})
}
// A = true --> A
// A = false --> !A
// A = null --> null
Expr::BinaryExpr(BinaryExpr {
left,
op: Eq,
right,
}) if is_bool_lit(&right) && info.is_boolean_type(&left)? => {
Transformed::yes(match as_bool_lit(&right)? {
Some(true) => *left,
Some(false) => Expr::Not(left),
None => lit_bool_null(),
})
}
// According to SQL's null semantics, NULL = NULL evaluates to NULL
// Both sides are the same expression (A = A) and A is non-volatile expression
// A = A --> A IS NOT NULL OR NULL
// A = A --> true (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: Eq,
right,
}) if (left == right) & !left.is_volatile() => {
Transformed::yes(match !info.nullable(&left)? {
true => lit(true),
false => Expr::BinaryExpr(BinaryExpr {
left: Box::new(Expr::IsNotNull(left)),
op: Or,
right: Box::new(lit_bool_null()),
}),
})
}
// Rules for NotEq
//
// true != A --> !A
// false != A --> A
// null != A --> null
Expr::BinaryExpr(BinaryExpr {
left,
op: NotEq,
right,
}) if is_bool_lit(&left) && info.is_boolean_type(&right)? => {
Transformed::yes(match as_bool_lit(&left)? {
Some(true) => Expr::Not(right),
Some(false) => *right,
None => lit_bool_null(),
})
}
// A != true --> !A
// A != false --> A
// A != null --> null,
Expr::BinaryExpr(BinaryExpr {
left,
op: NotEq,
right,
}) if is_bool_lit(&right) && info.is_boolean_type(&left)? => {
Transformed::yes(match as_bool_lit(&right)? {
Some(true) => Expr::Not(left),
Some(false) => *left,
None => lit_bool_null(),
})
}
//
// Rules for OR
//
// true OR A --> true (even if A is null)
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right: _,
}) if is_true(&left) => Transformed::yes(*left),
// false OR A --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if is_false(&left) => Transformed::yes(*right),
// A OR true --> true (even if A is null)
Expr::BinaryExpr(BinaryExpr {
left: _,
op: Or,
right,
}) if is_true(&right) => Transformed::yes(*right),
// A OR false --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if is_false(&right) => Transformed::yes(*left),
// A OR !A ---> true (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if is_not_of(&right, &left) && !info.nullable(&left)? => {
Transformed::yes(lit(true))
}
// !A OR A ---> true (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if is_not_of(&left, &right) && !info.nullable(&right)? => {
Transformed::yes(lit(true))
}
// (..A..) OR A --> (..A..)
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if expr_contains(&left, &right, Or) => Transformed::yes(*left),
// A OR (..A..) --> (..A..)
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if expr_contains(&right, &left, Or) => Transformed::yes(*right),
// A OR (A AND B) --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if is_op_with(And, &right, &left) => Transformed::yes(*left),
// (A AND B) OR A --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if is_op_with(And, &left, &right) => Transformed::yes(*right),
// Eliminate common factors in conjunctions e.g
// (A AND B) OR (A AND C) -> A AND (B OR C)
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if has_common_conjunction(&left, &right) => {
let lhs: IndexSet<Expr> = iter_conjunction_owned(*left).collect();
let (common, rhs): (Vec<_>, Vec<_>) = iter_conjunction_owned(*right)
.partition(|e| lhs.contains(e) && !e.is_volatile());
let new_rhs = rhs.into_iter().reduce(and);
let new_lhs = lhs.into_iter().filter(|e| !common.contains(e)).reduce(and);
let common_conjunction = common.into_iter().reduce(and).unwrap();
let new_expr = match (new_lhs, new_rhs) {
(Some(lhs), Some(rhs)) => and(common_conjunction, or(lhs, rhs)),
(_, _) => common_conjunction,
};
Transformed::yes(new_expr)
}
//
// Rules for AND
//
// true AND A --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if is_true(&left) => Transformed::yes(*right),
// false AND A --> false (even if A is null)
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right: _,
}) if is_false(&left) => Transformed::yes(*left),
// A AND true --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if is_true(&right) => Transformed::yes(*left),
// A AND false --> false (even if A is null)
Expr::BinaryExpr(BinaryExpr {
left: _,
op: And,
right,
}) if is_false(&right) => Transformed::yes(*right),
// A AND !A ---> false (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if is_not_of(&right, &left) && !info.nullable(&left)? => {
Transformed::yes(lit(false))
}
// !A AND A ---> false (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if is_not_of(&left, &right) && !info.nullable(&right)? => {
Transformed::yes(lit(false))
}
// (..A..) AND A --> (..A..)
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if expr_contains(&left, &right, And) => Transformed::yes(*left),
// A AND (..A..) --> (..A..)
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if expr_contains(&right, &left, And) => Transformed::yes(*right),
// A AND (A OR B) --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if is_op_with(Or, &right, &left) => Transformed::yes(*left),
// (A OR B) AND A --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if is_op_with(Or, &left, &right) => Transformed::yes(*right),
// A >= constant AND constant <= A --> A = constant
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if can_reduce_to_equal_statement(&left, &right) => {
if let Expr::BinaryExpr(BinaryExpr {
left: left_left,
right: left_right,
..
}) = *left
{
Transformed::yes(Expr::BinaryExpr(BinaryExpr {
left: left_left,
op: Eq,
right: left_right,
}))
} else {
return internal_err!("can_reduce_to_equal_statement should only be called with a BinaryExpr");
}
}
//
// Rules for Multiply
//
// A * 1 --> A (with type coercion if needed)
Expr::BinaryExpr(BinaryExpr {
left,
op: Multiply,
right,
}) if is_one(&right) => {
simplify_right_is_one_case(info, left, &Multiply, &right)?
}
// 1 * A --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: Multiply,
right,
}) if is_one(&left) => {
// 1 * A is equivalent to A * 1
simplify_right_is_one_case(info, right, &Multiply, &left)?
}
// A * 0 --> 0 (if A is not null and not floating, since NAN * 0 -> NAN)
Expr::BinaryExpr(BinaryExpr {
left,
op: Multiply,
right,
}) if !info.nullable(&left)?
&& !info.get_data_type(&left)?.is_floating()
&& is_zero(&right) =>
{
Transformed::yes(*right)
}
// 0 * A --> 0 (if A is not null and not floating, since 0 * NAN -> NAN)
Expr::BinaryExpr(BinaryExpr {
left,
op: Multiply,
right,
}) if !info.nullable(&right)?
&& !info.get_data_type(&right)?.is_floating()
&& is_zero(&left) =>
{
Transformed::yes(*left)
}
//
// Rules for Divide
//
// A / 1 --> A
Expr::BinaryExpr(BinaryExpr {
left,
op: Divide,
right,
}) if is_one(&right) => {
simplify_right_is_one_case(info, left, &Divide, &right)?
}
//
// Rules for Modulo
//
// A % 1 --> 0 (if A is not nullable and not floating, since NAN % 1 --> NAN)
Expr::BinaryExpr(BinaryExpr {
left,
op: Modulo,
right,
}) if !info.nullable(&left)?
&& !info.get_data_type(&left)?.is_floating()
&& is_one(&right) =>
{
Transformed::yes(Expr::Literal(
ScalarValue::new_zero(&info.get_data_type(&left)?)?,
None,
))
}
//
// Rules for BitwiseAnd
//
// A & 0 -> 0 (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseAnd,
right,
}) if !info.nullable(&left)? && is_zero(&right) => Transformed::yes(*right),
// 0 & A -> 0 (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseAnd,
right,
}) if !info.nullable(&right)? && is_zero(&left) => Transformed::yes(*left),
// !A & A -> 0 (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseAnd,
right,
}) if is_negative_of(&left, &right) && !info.nullable(&right)? => {
Transformed::yes(Expr::Literal(
ScalarValue::new_zero(&info.get_data_type(&left)?)?,
None,
))
}
// A & !A -> 0 (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseAnd,
right,
}) if is_negative_of(&right, &left) && !info.nullable(&left)? => {
Transformed::yes(Expr::Literal(
ScalarValue::new_zero(&info.get_data_type(&left)?)?,
None,
))
}
// (..A..) & A --> (..A..)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseAnd,
right,
}) if expr_contains(&left, &right, BitwiseAnd) => Transformed::yes(*left),
// A & (..A..) --> (..A..)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseAnd,
right,
}) if expr_contains(&right, &left, BitwiseAnd) => Transformed::yes(*right),
// A & (A | B) --> A (if B not null)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseAnd,
right,
}) if !info.nullable(&right)? && is_op_with(BitwiseOr, &right, &left) => {
Transformed::yes(*left)
}
// (A | B) & A --> A (if B not null)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseAnd,
right,
}) if !info.nullable(&left)? && is_op_with(BitwiseOr, &left, &right) => {
Transformed::yes(*right)
}
//
// Rules for BitwiseOr
//
// A | 0 -> A (even if A is null)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseOr,
right,
}) if is_zero(&right) => Transformed::yes(*left),
// 0 | A -> A (even if A is null)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseOr,
right,
}) if is_zero(&left) => Transformed::yes(*right),
// !A | A -> -1 (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseOr,
right,
}) if is_negative_of(&left, &right) && !info.nullable(&right)? => {
Transformed::yes(Expr::Literal(
ScalarValue::new_negative_one(&info.get_data_type(&left)?)?,
None,
))
}
// A | !A -> -1 (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseOr,
right,
}) if is_negative_of(&right, &left) && !info.nullable(&left)? => {
Transformed::yes(Expr::Literal(
ScalarValue::new_negative_one(&info.get_data_type(&left)?)?,
None,
))
}
// (..A..) | A --> (..A..)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseOr,
right,
}) if expr_contains(&left, &right, BitwiseOr) => Transformed::yes(*left),
// A | (..A..) --> (..A..)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseOr,
right,
}) if expr_contains(&right, &left, BitwiseOr) => Transformed::yes(*right),
// A | (A & B) --> A (if B not null)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseOr,
right,
}) if !info.nullable(&right)? && is_op_with(BitwiseAnd, &right, &left) => {
Transformed::yes(*left)
}
// (A & B) | A --> A (if B not null)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseOr,
right,
}) if !info.nullable(&left)? && is_op_with(BitwiseAnd, &left, &right) => {
Transformed::yes(*right)
}
//
// Rules for BitwiseXor
//
// A ^ 0 -> A (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseXor,
right,
}) if !info.nullable(&left)? && is_zero(&right) => Transformed::yes(*left),
// 0 ^ A -> A (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseXor,
right,
}) if !info.nullable(&right)? && is_zero(&left) => Transformed::yes(*right),
// !A ^ A -> -1 (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseXor,
right,
}) if is_negative_of(&left, &right) && !info.nullable(&right)? => {
Transformed::yes(Expr::Literal(
ScalarValue::new_negative_one(&info.get_data_type(&left)?)?,
None,
))
}
// A ^ !A -> -1 (if A not nullable)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseXor,
right,
}) if is_negative_of(&right, &left) && !info.nullable(&left)? => {
Transformed::yes(Expr::Literal(
ScalarValue::new_negative_one(&info.get_data_type(&left)?)?,
None,
))
}
// (..A..) ^ A --> (the expression without A, if number of A is odd, otherwise one A)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseXor,
right,
}) if expr_contains(&left, &right, BitwiseXor) => {
let expr = delete_xor_in_complex_expr(&left, &right, false);
Transformed::yes(if expr == *right {
Expr::Literal(
ScalarValue::new_zero(&info.get_data_type(&right)?)?,
None,
)
} else {
expr
})
}
// A ^ (..A..) --> (the expression without A, if number of A is odd, otherwise one A)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseXor,
right,
}) if expr_contains(&right, &left, BitwiseXor) => {
let expr = delete_xor_in_complex_expr(&right, &left, true);
Transformed::yes(if expr == *left {
Expr::Literal(
ScalarValue::new_zero(&info.get_data_type(&left)?)?,
None,
)
} else {
expr
})
}
//
// Rules for BitwiseShiftRight
//
// A >> 0 -> A (even if A is null)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseShiftRight,
right,
}) if is_zero(&right) => Transformed::yes(*left),
//
// Rules for BitwiseShiftRight
//
// A << 0 -> A (even if A is null)
Expr::BinaryExpr(BinaryExpr {
left,
op: BitwiseShiftLeft,
right,
}) if is_zero(&right) => Transformed::yes(*left),
//
// Rules for Not
//
Expr::Not(inner) => Transformed::yes(negate_clause(*inner)),
//
// Rules for Negative
//
Expr::Negative(inner) => Transformed::yes(distribute_negation(*inner)),
//
// Rules for Case
//
// Inline a comparison to a literal with the case statement into the `THEN` clauses.
// which can enable further simplifications
// CASE WHEN X THEN "a" WHEN Y THEN "b" ... END = "a" --> CASE WHEN X THEN "a" = "a" WHEN Y THEN "b" = "a" END
Expr::BinaryExpr(BinaryExpr {
left,
op: op @ (Eq | NotEq),
right,
}) if is_case_with_literal_outputs(&left) && is_lit(&right) => {
let case = into_case(*left)?;
Transformed::yes(Expr::Case(Case {
expr: None,
when_then_expr: case
.when_then_expr
.into_iter()
.map(|(when, then)| {
(
when,
Box::new(Expr::BinaryExpr(BinaryExpr {
left: then,
op,
right: right.clone(),
})),
)
})
.collect(),
else_expr: case.else_expr.map(|els| {
Box::new(Expr::BinaryExpr(BinaryExpr {
left: els,
op,
right,
}))
}),
}))
}
// CASE WHEN true THEN A ... END --> A
// CASE WHEN X THEN A WHEN TRUE THEN B ... END --> CASE WHEN X THEN A ELSE B END
// CASE WHEN false THEN A END --> NULL
// CASE WHEN false THEN A ELSE B END --> B
// CASE WHEN X THEN A WHEN false THEN B END --> CASE WHEN X THEN A ELSE B END
Expr::Case(Case {
expr: None,
when_then_expr,
mut else_expr,
}) if when_then_expr
.iter()
.any(|(when, _)| is_true(when.as_ref()) || is_false(when.as_ref())) =>
{
let out_type = info.get_data_type(&when_then_expr[0].1)?;
let mut new_when_then_expr = Vec::with_capacity(when_then_expr.len());
for (when, then) in when_then_expr.into_iter() {
if is_true(when.as_ref()) {
// Skip adding the rest of the when-then expressions after WHEN true
// CASE WHEN X THEN A WHEN TRUE THEN B ... END --> CASE WHEN X THEN A ELSE B END
else_expr = Some(then);
break;
} else if !is_false(when.as_ref()) {
new_when_then_expr.push((when, then));
}
// else: skip WHEN false cases
}
// Exclude CASE statement altogether if there are no when-then expressions left
if new_when_then_expr.is_empty() {
// CASE WHEN false THEN A ELSE B END --> B
if let Some(else_expr) = else_expr {
return Ok(Transformed::yes(*else_expr));
// CASE WHEN false THEN A END --> NULL
} else {
let null =
Expr::Literal(ScalarValue::try_new_null(&out_type)?, None);
return Ok(Transformed::yes(null));
}
}
Transformed::yes(Expr::Case(Case {
expr: None,
when_then_expr: new_when_then_expr,
else_expr,
}))
}
// CASE
// WHEN X THEN A
// WHEN Y THEN B
// ...
// ELSE Q
// END
//
// ---> (X AND A) OR (Y AND B AND NOT X) OR ... (NOT (X OR Y) AND Q)
//
// Note: the rationale for this rewrite is that the expr can then be further
// simplified using the existing rules for AND/OR
Expr::Case(Case {
expr: None,
when_then_expr,
else_expr,
}) if !when_then_expr.is_empty()
// The rewrite is O(n²) in general so limit to small number of when-thens that can be true
&& (when_then_expr.len() < 3 // small number of input whens
// or all thens are literal bools and a small number of them are true
|| (when_then_expr.iter().all(|(_, then)| is_bool_lit(then))
&& when_then_expr.iter().filter(|(_, then)| is_true(then)).count() < 3))
&& info.is_boolean_type(&when_then_expr[0].1)? =>
{
// String disjunction of all the when predicates encountered so far. Not nullable.
let mut filter_expr = lit(false);
// The disjunction of all the cases
let mut out_expr = lit(false);
for (when, then) in when_then_expr {
let when = is_exactly_true(*when, info)?;
let case_expr =
when.clone().and(filter_expr.clone().not()).and(*then);
out_expr = out_expr.or(case_expr);
filter_expr = filter_expr.or(when);
}
let else_expr = else_expr.map(|b| *b).unwrap_or_else(lit_bool_null);
let case_expr = filter_expr.not().and(else_expr);
out_expr = out_expr.or(case_expr);
// Do a first pass at simplification
out_expr.rewrite(self)?
}
// CASE
// WHEN X THEN true
// WHEN Y THEN true
// WHEN Z THEN false
// ...
// ELSE true
// END
//
// --->
//
// NOT(CASE
// WHEN X THEN false
// WHEN Y THEN false
// WHEN Z THEN true
// ...
// ELSE false
// END)
//
// Note: the rationale for this rewrite is that the case can then be further
// simplified into a small number of ANDs and ORs
Expr::Case(Case {
expr: None,
when_then_expr,
else_expr,
}) if !when_then_expr.is_empty()
&& when_then_expr
.iter()
.all(|(_, then)| is_bool_lit(then)) // all thens are literal bools
// This simplification is only helpful if we end up with a small number of true thens
&& when_then_expr
.iter()
.filter(|(_, then)| is_false(then))
.count()
< 3
&& else_expr.as_deref().is_none_or(is_bool_lit) =>
{
Transformed::yes(
Expr::Case(Case {
expr: None,
when_then_expr: when_then_expr
.into_iter()
.map(|(when, then)| (when, Box::new(Expr::Not(then))))
.collect(),
else_expr: else_expr
.map(|else_expr| Box::new(Expr::Not(else_expr))),
})
.not(),
)
}
Expr::ScalarFunction(ScalarFunction { func: udf, args }) => {
match udf.simplify(args, info)? {
ExprSimplifyResult::Original(args) => {
Transformed::no(Expr::ScalarFunction(ScalarFunction {
func: udf,
args,
}))
}
ExprSimplifyResult::Simplified(expr) => Transformed::yes(expr),
}
}
Expr::AggregateFunction(datafusion_expr::expr::AggregateFunction {
ref func,
..
}) => match (func.simplify(), expr) {
(Some(simplify_function), Expr::AggregateFunction(af)) => {
Transformed::yes(simplify_function(af, info)?)
}
(_, expr) => Transformed::no(expr),
},
Expr::WindowFunction(ref window_fun) => match (window_fun.simplify(), expr) {
(Some(simplify_function), Expr::WindowFunction(wf)) => {
Transformed::yes(simplify_function(*wf, info)?)
}
(_, expr) => Transformed::no(expr),
},
//
// Rules for Between
//
// a between 3 and 5 --> a >= 3 AND a <=5
// a not between 3 and 5 --> a < 3 OR a > 5
Expr::Between(between) => Transformed::yes(if between.negated {
let l = *between.expr.clone();
let r = *between.expr;
or(l.lt(*between.low), r.gt(*between.high))
} else {
and(
between.expr.clone().gt_eq(*between.low),
between.expr.lt_eq(*between.high),
)
}),
//
// Rules for regexes
//
Expr::BinaryExpr(BinaryExpr {
left,
op: op @ (RegexMatch | RegexNotMatch | RegexIMatch | RegexNotIMatch),
right,
}) => Transformed::yes(simplify_regex_expr(left, op, right)?),
// Rules for Like
Expr::Like(like) => {
// `\` is implicit escape, see https://github.com/apache/datafusion/issues/13291
let escape_char = like.escape_char.unwrap_or('\\');
match as_string_scalar(&like.pattern) {
Some((data_type, pattern_str)) => {
match pattern_str {
None => return Ok(Transformed::yes(lit_bool_null())),
Some(pattern_str) if pattern_str == "%" => {
// exp LIKE '%' is
// - when exp is not NULL, it's true
// - when exp is NULL, it's NULL
// exp NOT LIKE '%' is
// - when exp is not NULL, it's false
// - when exp is NULL, it's NULL
let result_for_non_null = lit(!like.negated);
Transformed::yes(if !info.nullable(&like.expr)? {
result_for_non_null
} else {
Expr::Case(Case {
expr: Some(Box::new(Expr::IsNotNull(like.expr))),
when_then_expr: vec![(
Box::new(lit(true)),
Box::new(result_for_non_null),
)],
else_expr: None,
})
})
}
Some(pattern_str)
if pattern_str.contains("%%")
&& !pattern_str.contains(escape_char) =>
{
// Repeated occurrences of wildcard are redundant so remove them
// exp LIKE '%%' --> exp LIKE '%'
let simplified_pattern = Regex::new("%%+")
.unwrap()
.replace_all(pattern_str, "%")
.to_string();
Transformed::yes(Expr::Like(Like {
pattern: Box::new(to_string_scalar(
&data_type,
Some(simplified_pattern),
)),
..like
}))
}
Some(pattern_str)
if !like.case_insensitive
&& !pattern_str
.contains(['%', '_', escape_char].as_ref()) =>
{
// If the pattern does not contain any wildcards, we can simplify the like expression to an equality expression
// TODO: handle escape characters
Transformed::yes(Expr::BinaryExpr(BinaryExpr {
left: like.expr.clone(),
op: if like.negated { NotEq } else { Eq },
right: like.pattern.clone(),
}))
}
Some(_pattern_str) => Transformed::no(Expr::Like(like)),
}
}
None => Transformed::no(Expr::Like(like)),
}
}
// a is not null/unknown --> true (if a is not nullable)
Expr::IsNotNull(expr) | Expr::IsNotUnknown(expr)
if !info.nullable(&expr)? =>
{
Transformed::yes(lit(true))
}
// a is null/unknown --> false (if a is not nullable)
Expr::IsNull(expr) | Expr::IsUnknown(expr) if !info.nullable(&expr)? => {
Transformed::yes(lit(false))
}
// expr IN () --> false
// expr NOT IN () --> true
Expr::InList(InList {
expr: _,
list,
negated,
}) if list.is_empty() => Transformed::yes(lit(negated)),
// null in (x, y, z) --> null
// null not in (x, y, z) --> null
Expr::InList(InList {
expr,
list,
negated: _,
}) if is_null(expr.as_ref()) && !list.is_empty() => {
Transformed::yes(lit_bool_null())
}
// expr IN ((subquery)) -> expr IN (subquery), see ##5529
Expr::InList(InList {
expr,
mut list,
negated,
}) if list.len() == 1
&& matches!(list.first(), Some(Expr::ScalarSubquery { .. })) =>
{
let Expr::ScalarSubquery(subquery) = list.remove(0) else {
unreachable!()
};
Transformed::yes(Expr::InSubquery(InSubquery::new(
expr, subquery, negated,
)))
}
// Combine multiple OR expressions into a single IN list expression if possible
//
// i.e. `a = 1 OR a = 2 OR a = 3` -> `a IN (1, 2, 3)`
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if are_inlist_and_eq(left.as_ref(), right.as_ref()) => {
let lhs = to_inlist(*left).unwrap();
let rhs = to_inlist(*right).unwrap();
let mut seen: HashSet<Expr> = HashSet::new();
let list = lhs
.list
.into_iter()
.chain(rhs.list)
.filter(|e| seen.insert(e.to_owned()))
.collect::<Vec<_>>();
let merged_inlist = InList {
expr: lhs.expr,
list,
negated: false,
};
Transformed::yes(Expr::InList(merged_inlist))
}
// Simplify expressions that is guaranteed to be true or false to a literal boolean expression
//
// Rules:
// If both expressions are `IN` or `NOT IN`, then we can apply intersection or union on both lists
// Intersection:
// 1. `a in (1,2,3) AND a in (4,5) -> a in (), which is false`
// 2. `a in (1,2,3) AND a in (2,3,4) -> a in (2,3)`
// 3. `a not in (1,2,3) OR a not in (3,4,5,6) -> a not in (3)`
// Union:
// 4. `a not int (1,2,3) AND a not in (4,5,6) -> a not in (1,2,3,4,5,6)`
// # This rule is handled by `or_in_list_simplifier.rs`
// 5. `a in (1,2,3) OR a in (4,5,6) -> a in (1,2,3,4,5,6)`
// If one of the expressions is `IN` and another one is `NOT IN`, then we apply exception on `In` expression
// 6. `a in (1,2,3,4) AND a not in (1,2,3,4,5) -> a in (), which is false`
// 7. `a not in (1,2,3,4) AND a in (1,2,3,4,5) -> a = 5`
// 8. `a in (1,2,3,4) AND a not in (5,6,7,8) -> a in (1,2,3,4)`
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if are_inlist_and_eq_and_match_neg(
left.as_ref(),
right.as_ref(),
false,
false,
) =>
{
match (*left, *right) {
(Expr::InList(l1), Expr::InList(l2)) => {
return inlist_intersection(l1, &l2, false).map(Transformed::yes);
}
// Matched previously once
_ => unreachable!(),
}
}
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if are_inlist_and_eq_and_match_neg(
left.as_ref(),
right.as_ref(),
true,
true,
) =>
{
match (*left, *right) {
(Expr::InList(l1), Expr::InList(l2)) => {
return inlist_union(l1, l2, true).map(Transformed::yes);
}
// Matched previously once
_ => unreachable!(),
}
}
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if are_inlist_and_eq_and_match_neg(
left.as_ref(),
right.as_ref(),
false,
true,
) =>
{
match (*left, *right) {
(Expr::InList(l1), Expr::InList(l2)) => {
return inlist_except(l1, &l2).map(Transformed::yes);
}
// Matched previously once
_ => unreachable!(),
}
}
Expr::BinaryExpr(BinaryExpr {
left,
op: And,
right,
}) if are_inlist_and_eq_and_match_neg(
left.as_ref(),
right.as_ref(),
true,
false,
) =>
{
match (*left, *right) {
(Expr::InList(l1), Expr::InList(l2)) => {
return inlist_except(l2, &l1).map(Transformed::yes);
}
// Matched previously once
_ => unreachable!(),
}
}
Expr::BinaryExpr(BinaryExpr {
left,
op: Or,
right,
}) if are_inlist_and_eq_and_match_neg(
left.as_ref(),
right.as_ref(),
true,
true,
) =>
{
match (*left, *right) {
(Expr::InList(l1), Expr::InList(l2)) => {
return inlist_intersection(l1, &l2, true).map(Transformed::yes);
}
// Matched previously once
_ => unreachable!(),
}
}
// =======================================
// unwrap_cast_in_comparison
// =======================================
//
// For case:
// try_cast/cast(expr as data_type) op literal
Expr::BinaryExpr(BinaryExpr { left, op, right })
if is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary(
info, &left, op, &right,
) && op.supports_propagation() =>
{
unwrap_cast_in_comparison_for_binary(info, *left, *right, op)?
}
// literal op try_cast/cast(expr as data_type)
// -->
// try_cast/cast(expr as data_type) op_swap literal
Expr::BinaryExpr(BinaryExpr { left, op, right })
if is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary(
info, &right, op, &left,
) && op.supports_propagation()
&& op.swap().is_some() =>
{
unwrap_cast_in_comparison_for_binary(
info,
*right,
*left,
op.swap().unwrap(),
)?
}
// For case:
// try_cast/cast(expr as left_type) in (expr1,expr2,expr3)
Expr::InList(InList {
expr: mut left,
list,
negated,
}) if is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist(
info, &left, &list,
) =>
{
let (Expr::TryCast(TryCast {
expr: left_expr, ..
})
| Expr::Cast(Cast {
expr: left_expr, ..
})) = left.as_mut()
else {
return internal_err!("Expect cast expr, but got {:?}", left)?;
};
let expr_type = info.get_data_type(left_expr)?;
let right_exprs = list
.into_iter()
.map(|right| {
match right {
Expr::Literal(right_lit_value, _) => {
// if the right_lit_value can be casted to the type of internal_left_expr
// we need to unwrap the cast for cast/try_cast expr, and add cast to the literal
let Some(value) = try_cast_literal_to_type(&right_lit_value, &expr_type) else {
internal_err!(
"Can't cast the list expr {:?} to type {}",
right_lit_value, &expr_type
)?
};
Ok(lit(value))
}
other_expr => internal_err!(
"Only support literal expr to optimize, but the expr is {:?}",
&other_expr
),
}
})
.collect::<Result<Vec<_>>>()?;
Transformed::yes(Expr::InList(InList {
expr: std::mem::take(left_expr),
list: right_exprs,
negated,
}))
}
// no additional rewrites possible
expr => Transformed::no(expr),
})
}
}
fn as_string_scalar(expr: &Expr) -> Option<(DataType, &Option<String>)> {
match expr {
Expr::Literal(ScalarValue::Utf8(s), _) => Some((DataType::Utf8, s)),
Expr::Literal(ScalarValue::LargeUtf8(s), _) => Some((DataType::LargeUtf8, s)),
Expr::Literal(ScalarValue::Utf8View(s), _) => Some((DataType::Utf8View, s)),
_ => None,
}
}
fn to_string_scalar(data_type: &DataType, value: Option<String>) -> Expr {
match data_type {
DataType::Utf8 => Expr::Literal(ScalarValue::Utf8(value), None),
DataType::LargeUtf8 => Expr::Literal(ScalarValue::LargeUtf8(value), None),
DataType::Utf8View => Expr::Literal(ScalarValue::Utf8View(value), None),
_ => unreachable!(),
}
}
fn has_common_conjunction(lhs: &Expr, rhs: &Expr) -> bool {
let lhs_set: HashSet<&Expr> = iter_conjunction(lhs).collect();
iter_conjunction(rhs).any(|e| lhs_set.contains(&e) && !e.is_volatile())
}
// TODO: We might not need this after defer pattern for Box is stabilized. https://github.com/rust-lang/rust/issues/87121
fn are_inlist_and_eq_and_match_neg(
left: &Expr,
right: &Expr,
is_left_neg: bool,
is_right_neg: bool,
) -> bool {
match (left, right) {
(Expr::InList(l), Expr::InList(r)) => {
l.expr == r.expr && l.negated == is_left_neg && r.negated == is_right_neg
}
_ => false,
}
}
// TODO: We might not need this after defer pattern for Box is stabilized. https://github.com/rust-lang/rust/issues/87121
fn are_inlist_and_eq(left: &Expr, right: &Expr) -> bool {
let left = as_inlist(left);
let right = as_inlist(right);
if let (Some(lhs), Some(rhs)) = (left, right) {
matches!(lhs.expr.as_ref(), Expr::Column(_))
&& matches!(rhs.expr.as_ref(), Expr::Column(_))
&& lhs.expr == rhs.expr
&& !lhs.negated
&& !rhs.negated
} else {
false
}
}
/// Try to convert an expression to an in-list expression
fn as_inlist(expr: &'_ Expr) -> Option<Cow<'_, InList>> {
match expr {
Expr::InList(inlist) => Some(Cow::Borrowed(inlist)),
Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == Operator::Eq => {
match (left.as_ref(), right.as_ref()) {
(Expr::Column(_), Expr::Literal(_, _)) => Some(Cow::Owned(InList {
expr: left.clone(),
list: vec![*right.clone()],
negated: false,
})),
(Expr::Literal(_, _), Expr::Column(_)) => Some(Cow::Owned(InList {
expr: right.clone(),
list: vec![*left.clone()],
negated: false,
})),
_ => None,
}
}
_ => None,
}
}
fn to_inlist(expr: Expr) -> Option<InList> {
match expr {
Expr::InList(inlist) => Some(inlist),
Expr::BinaryExpr(BinaryExpr {
left,
op: Operator::Eq,
right,
}) => match (left.as_ref(), right.as_ref()) {
(Expr::Column(_), Expr::Literal(_, _)) => Some(InList {
expr: left,
list: vec![*right],
negated: false,
}),
(Expr::Literal(_, _), Expr::Column(_)) => Some(InList {
expr: right,
list: vec![*left],
negated: false,
}),
_ => None,
},
_ => None,
}
}
/// Return the union of two inlist expressions
/// maintaining the order of the elements in the two lists
fn inlist_union(mut l1: InList, l2: InList, negated: bool) -> Result<Expr> {
// extend the list in l1 with the elements in l2 that are not already in l1
let l1_items: HashSet<_> = l1.list.iter().collect();
// keep all l2 items that do not also appear in l1
let keep_l2: Vec<_> = l2
.list
.into_iter()
.filter_map(|e| if l1_items.contains(&e) { None } else { Some(e) })
.collect();
l1.list.extend(keep_l2);
l1.negated = negated;
Ok(Expr::InList(l1))
}
/// Return the intersection of two inlist expressions
/// maintaining the order of the elements in the two lists
fn inlist_intersection(mut l1: InList, l2: &InList, negated: bool) -> Result<Expr> {
let l2_items = l2.list.iter().collect::<HashSet<_>>();
// remove all items from l1 that are not in l2
l1.list.retain(|e| l2_items.contains(e));
// e in () is always false
// e not in () is always true
if l1.list.is_empty() {
return Ok(lit(negated));
}
Ok(Expr::InList(l1))
}
/// Return the all items in l1 that are not in l2
/// maintaining the order of the elements in the two lists
fn inlist_except(mut l1: InList, l2: &InList) -> Result<Expr> {
let l2_items = l2.list.iter().collect::<HashSet<_>>();
// keep only items from l1 that are not in l2
l1.list.retain(|e| !l2_items.contains(e));
if l1.list.is_empty() {
return Ok(lit(false));
}
Ok(Expr::InList(l1))
}
/// Returns expression testing a boolean `expr` for being exactly `true` (not `false` or NULL).
fn is_exactly_true(expr: Expr, info: &impl SimplifyInfo) -> Result<Expr> {
if !info.nullable(&expr)? {
Ok(expr)
} else {
Ok(Expr::BinaryExpr(BinaryExpr {
left: Box::new(expr),
op: Operator::IsNotDistinctFrom,
right: Box::new(lit(true)),
}))
}
}
// A * 1 -> A
// A / 1 -> A
//
// Move this function body out of the large match branch avoid stack overflow
fn simplify_right_is_one_case<S: SimplifyInfo>(
info: &S,
left: Box<Expr>,
op: &Operator,
right: &Expr,
) -> Result<Transformed<Expr>> {
// Check if resulting type would be different due to coercion
let left_type = info.get_data_type(&left)?;
let right_type = info.get_data_type(right)?;
match BinaryTypeCoercer::new(&left_type, op, &right_type).get_result_type() {
Ok(result_type) => {
// Only cast if the types differ
if left_type != result_type {
Ok(Transformed::yes(Expr::Cast(Cast::new(left, result_type))))
} else {
Ok(Transformed::yes(*left))
}
}
Err(_) => Ok(Transformed::yes(*left)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::simplify_expressions::SimplifyContext;
use crate::test::test_table_scan_with_name;
use arrow::datatypes::FieldRef;
use datafusion_common::{assert_contains, DFSchemaRef, ToDFSchema};
use datafusion_expr::{
expr::WindowFunction,
function::{
AccumulatorArgs, AggregateFunctionSimplification,
WindowFunctionSimplification,
},
interval_arithmetic::Interval,
*,
};
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
use datafusion_physical_expr::PhysicalExpr;
use std::hash::Hash;
use std::sync::LazyLock;
use std::{
collections::HashMap,
ops::{BitAnd, BitOr, BitXor},
sync::Arc,
};
// ------------------------------
// --- ExprSimplifier tests -----
// ------------------------------
#[test]
fn api_basic() {
let props = ExecutionProps::new();
let simplifier =
ExprSimplifier::new(SimplifyContext::new(&props).with_schema(test_schema()));
let expr = lit(1) + lit(2);
let expected = lit(3);
assert_eq!(expected, simplifier.simplify(expr).unwrap());
}
#[test]
fn basic_coercion() {
let schema = test_schema();
let props = ExecutionProps::new();
let simplifier = ExprSimplifier::new(
SimplifyContext::new(&props).with_schema(Arc::clone(&schema)),
);
// Note expr type is int32 (not int64)
// (1i64 + 2i32) < i
let expr = (lit(1i64) + lit(2i32)).lt(col("i"));
// should fully simplify to 3 < i (though i has been coerced to i64)
let expected = lit(3i64).lt(col("i"));
let expr = simplifier.coerce(expr, &schema).unwrap();
assert_eq!(expected, simplifier.simplify(expr).unwrap());
}
fn test_schema() -> DFSchemaRef {
static TEST_SCHEMA: LazyLock<DFSchemaRef> = LazyLock::new(|| {
Schema::new(vec![
Field::new("i", DataType::Int64, false),
Field::new("b", DataType::Boolean, true),
])
.to_dfschema_ref()
.unwrap()
});
Arc::clone(&TEST_SCHEMA)
}
#[test]
fn simplify_and_constant_prop() {
let props = ExecutionProps::new();
let simplifier =
ExprSimplifier::new(SimplifyContext::new(&props).with_schema(test_schema()));
// should be able to simplify to false
// (i * (1 - 2)) > 0
let expr = (col("i") * (lit(1) - lit(1))).gt(lit(0));
let expected = lit(false);
assert_eq!(expected, simplifier.simplify(expr).unwrap());
}
#[test]
fn simplify_and_constant_prop_with_case() {
let props = ExecutionProps::new();
let simplifier =
ExprSimplifier::new(SimplifyContext::new(&props).with_schema(test_schema()));
// CASE
// WHEN i>5 AND false THEN i > 5
// WHEN i<5 AND true THEN i < 5
// ELSE false
// END
//
// Can be simplified to `i < 5`
let expr = when(col("i").gt(lit(5)).and(lit(false)), col("i").gt(lit(5)))
.when(col("i").lt(lit(5)).and(lit(true)), col("i").lt(lit(5)))
.otherwise(lit(false))
.unwrap();
let expected = col("i").lt(lit(5));
assert_eq!(expected, simplifier.simplify(expr).unwrap());
}
// ------------------------------
// --- Simplifier tests -----
// ------------------------------
#[test]
fn test_simplify_canonicalize() {
{
let expr = lit(1).lt(col("c2")).and(col("c2").gt(lit(1)));
let expected = col("c2").gt(lit(1));
assert_eq!(simplify(expr), expected);
}
{
let expr = col("c1").lt(col("c2")).and(col("c2").gt(col("c1")));
let expected = col("c2").gt(col("c1"));
assert_eq!(simplify(expr), expected);
}
{
let expr = col("c1")
.eq(lit(1))
.and(lit(1).eq(col("c1")))
.and(col("c1").eq(lit(3)));
let expected = col("c1").eq(lit(1)).and(col("c1").eq(lit(3)));
assert_eq!(simplify(expr), expected);
}
{
let expr = col("c1")
.eq(col("c2"))
.and(col("c1").gt(lit(5)))
.and(col("c2").eq(col("c1")));
let expected = col("c2").eq(col("c1")).and(col("c1").gt(lit(5)));
assert_eq!(simplify(expr), expected);
}
{
let expr = col("c1")
.eq(lit(1))
.and(col("c2").gt(lit(3)).or(lit(3).lt(col("c2"))));
let expected = col("c1").eq(lit(1)).and(col("c2").gt(lit(3)));
assert_eq!(simplify(expr), expected);
}
{
let expr = col("c1").lt(lit(5)).and(col("c1").gt_eq(lit(5)));
let expected = col("c1").lt(lit(5)).and(col("c1").gt_eq(lit(5)));
assert_eq!(simplify(expr), expected);
}
{
let expr = col("c1").lt(lit(5)).and(col("c1").gt_eq(lit(5)));
let expected = col("c1").lt(lit(5)).and(col("c1").gt_eq(lit(5)));
assert_eq!(simplify(expr), expected);
}
{
let expr = col("c1").gt(col("c2")).and(col("c1").gt(col("c2")));
let expected = col("c2").lt(col("c1"));
assert_eq!(simplify(expr), expected);
}
}
#[test]
fn test_simplify_eq_not_self() {
// `expr_a`: column `c2` is nullable, so `c2 = c2` simplifies to `c2 IS NOT NULL OR NULL`
// This ensures the expression is only true when `c2` is not NULL, accounting for SQL's NULL semantics.
let expr_a = col("c2").eq(col("c2"));
let expected_a = col("c2").is_not_null().or(lit_bool_null());
// `expr_b`: column `c2_non_null` is explicitly non-nullable, so `c2_non_null = c2_non_null` is always true
let expr_b = col("c2_non_null").eq(col("c2_non_null"));
let expected_b = lit(true);
assert_eq!(simplify(expr_a), expected_a);
assert_eq!(simplify(expr_b), expected_b);
}
#[test]
fn test_simplify_or_true() {
let expr_a = col("c2").or(lit(true));
let expr_b = lit(true).or(col("c2"));
let expected = lit(true);
assert_eq!(simplify(expr_a), expected);
assert_eq!(simplify(expr_b), expected);
}
#[test]
fn test_simplify_or_false() {
let expr_a = lit(false).or(col("c2"));
let expr_b = col("c2").or(lit(false));
let expected = col("c2");
assert_eq!(simplify(expr_a), expected);
assert_eq!(simplify(expr_b), expected);
}
#[test]
fn test_simplify_or_same() {
let expr = col("c2").or(col("c2"));
let expected = col("c2");
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_or_not_self() {
// A OR !A if A is not nullable --> true
// !A OR A if A is not nullable --> true
let expr_a = col("c2_non_null").or(col("c2_non_null").not());
let expr_b = col("c2_non_null").not().or(col("c2_non_null"));
let expected = lit(true);
assert_eq!(simplify(expr_a), expected);
assert_eq!(simplify(expr_b), expected);
}
#[test]
fn test_simplify_and_false() {
let expr_a = lit(false).and(col("c2"));
let expr_b = col("c2").and(lit(false));
let expected = lit(false);
assert_eq!(simplify(expr_a), expected);
assert_eq!(simplify(expr_b), expected);
}
#[test]
fn test_simplify_and_same() {
let expr = col("c2").and(col("c2"));
let expected = col("c2");
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_and_true() {
let expr_a = lit(true).and(col("c2"));
let expr_b = col("c2").and(lit(true));
let expected = col("c2");
assert_eq!(simplify(expr_a), expected);
assert_eq!(simplify(expr_b), expected);
}
#[test]
fn test_simplify_and_not_self() {
// A AND !A if A is not nullable --> false
// !A AND A if A is not nullable --> false
let expr_a = col("c2_non_null").and(col("c2_non_null").not());
let expr_b = col("c2_non_null").not().and(col("c2_non_null"));
let expected = lit(false);
assert_eq!(simplify(expr_a), expected);
assert_eq!(simplify(expr_b), expected);
}
#[test]
fn test_simplify_multiply_by_one() {
let expr_a = col("c2") * lit(1);
let expr_b = lit(1) * col("c2");
let expected = col("c2");
assert_eq!(simplify(expr_a), expected);
assert_eq!(simplify(expr_b), expected);
let expr = col("c2") * lit(ScalarValue::Decimal128(Some(10000000000), 38, 10));
assert_eq!(simplify(expr), expected);
let expr = lit(ScalarValue::Decimal128(Some(10000000000), 31, 10)) * col("c2");
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_multiply_by_null() {
let null = lit(ScalarValue::Int64(None));
// A * null --> null
{
let expr = col("c3") * null.clone();
assert_eq!(simplify(expr), null);
}
// null * A --> null
{
let expr = null.clone() * col("c3");
assert_eq!(simplify(expr), null);
}
}
#[test]
fn test_simplify_multiply_by_zero() {
// cannot optimize A * null (null * A) if A is nullable
{
let expr_a = col("c2") * lit(0);
let expr_b = lit(0) * col("c2");
assert_eq!(simplify(expr_a.clone()), expr_a);
assert_eq!(simplify(expr_b.clone()), expr_b);
}
// 0 * A --> 0 if A is not nullable
{
let expr = lit(0) * col("c2_non_null");
assert_eq!(simplify(expr), lit(0));
}
// A * 0 --> 0 if A is not nullable
{
let expr = col("c2_non_null") * lit(0);
assert_eq!(simplify(expr), lit(0));
}
// A * Decimal128(0) --> 0 if A is not nullable
{
let expr = col("c2_non_null") * lit(ScalarValue::Decimal128(Some(0), 31, 10));
assert_eq!(
simplify(expr),
lit(ScalarValue::Decimal128(Some(0), 31, 10))
);
let expr = binary_expr(
lit(ScalarValue::Decimal128(Some(0), 31, 10)),
Operator::Multiply,
col("c2_non_null"),
);
assert_eq!(
simplify(expr),
lit(ScalarValue::Decimal128(Some(0), 31, 10))
);
}
}
#[test]
fn test_simplify_divide_by_one() {
let expr = binary_expr(col("c2"), Operator::Divide, lit(1));
let expected = col("c2");
assert_eq!(simplify(expr), expected);
let expr = col("c2") / lit(ScalarValue::Decimal128(Some(10000000000), 31, 10));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_divide_null() {
// A / null --> null
let null = lit(ScalarValue::Int64(None));
{
let expr = col("c3") / null.clone();
assert_eq!(simplify(expr), null);
}
// null / A --> null
{
let expr = null.clone() / col("c3");
assert_eq!(simplify(expr), null);
}
}
#[test]
fn test_simplify_divide_by_same() {
let expr = col("c2") / col("c2");
// if c2 is null, c2 / c2 = null, so can't simplify
let expected = expr.clone();
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_modulo_by_null() {
let null = lit(ScalarValue::Int64(None));
// A % null --> null
{
let expr = col("c3") % null.clone();
assert_eq!(simplify(expr), null);
}
// null % A --> null
{
let expr = null.clone() % col("c3");
assert_eq!(simplify(expr), null);
}
}
#[test]
fn test_simplify_modulo_by_one() {
let expr = col("c2") % lit(1);
// if c2 is null, c2 % 1 = null, so can't simplify
let expected = expr.clone();
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_divide_zero_by_zero() {
// because divide by 0 maybe occur in short-circuit expression
// so we should not simplify this, and throw error in runtime
let expr = lit(0) / lit(0);
let expected = expr.clone();
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_divide_by_zero() {
// because divide by 0 maybe occur in short-circuit expression
// so we should not simplify this, and throw error in runtime
let expr = col("c2_non_null") / lit(0);
let expected = expr.clone();
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_modulo_by_one_non_null() {
let expr = col("c3_non_null") % lit(1);
let expected = lit(0_i64);
assert_eq!(simplify(expr), expected);
let expr =
col("c3_non_null") % lit(ScalarValue::Decimal128(Some(10000000000), 31, 10));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_bitwise_xor_by_null() {
let null = lit(ScalarValue::Int64(None));
// A ^ null --> null
{
let expr = col("c3") ^ null.clone();
assert_eq!(simplify(expr), null);
}
// null ^ A --> null
{
let expr = null.clone() ^ col("c3");
assert_eq!(simplify(expr), null);
}
}
#[test]
fn test_simplify_bitwise_shift_right_by_null() {
let null = lit(ScalarValue::Int64(None));
// A >> null --> null
{
let expr = col("c3") >> null.clone();
assert_eq!(simplify(expr), null);
}
// null >> A --> null
{
let expr = null.clone() >> col("c3");
assert_eq!(simplify(expr), null);
}
}
#[test]
fn test_simplify_bitwise_shift_left_by_null() {
let null = lit(ScalarValue::Int64(None));
// A << null --> null
{
let expr = col("c3") << null.clone();
assert_eq!(simplify(expr), null);
}
// null << A --> null
{
let expr = null.clone() << col("c3");
assert_eq!(simplify(expr), null);
}
}
#[test]
fn test_simplify_bitwise_and_by_zero() {
// A & 0 --> 0
{
let expr = col("c2_non_null") & lit(0);
assert_eq!(simplify(expr), lit(0));
}
// 0 & A --> 0
{
let expr = lit(0) & col("c2_non_null");
assert_eq!(simplify(expr), lit(0));
}
}
#[test]
fn test_simplify_bitwise_or_by_zero() {
// A | 0 --> A
{
let expr = col("c2_non_null") | lit(0);
assert_eq!(simplify(expr), col("c2_non_null"));
}
// 0 | A --> A
{
let expr = lit(0) | col("c2_non_null");
assert_eq!(simplify(expr), col("c2_non_null"));
}
}
#[test]
fn test_simplify_bitwise_xor_by_zero() {
// A ^ 0 --> A
{
let expr = col("c2_non_null") ^ lit(0);
assert_eq!(simplify(expr), col("c2_non_null"));
}
// 0 ^ A --> A
{
let expr = lit(0) ^ col("c2_non_null");
assert_eq!(simplify(expr), col("c2_non_null"));
}
}
#[test]
fn test_simplify_bitwise_bitwise_shift_right_by_zero() {
// A >> 0 --> A
{
let expr = col("c2_non_null") >> lit(0);
assert_eq!(simplify(expr), col("c2_non_null"));
}
}
#[test]
fn test_simplify_bitwise_bitwise_shift_left_by_zero() {
// A << 0 --> A
{
let expr = col("c2_non_null") << lit(0);
assert_eq!(simplify(expr), col("c2_non_null"));
}
}
#[test]
fn test_simplify_bitwise_and_by_null() {
let null = Expr::Literal(ScalarValue::Int64(None), None);
// A & null --> null
{
let expr = col("c3") & null.clone();
assert_eq!(simplify(expr), null);
}
// null & A --> null
{
let expr = null.clone() & col("c3");
assert_eq!(simplify(expr), null);
}
}
#[test]
fn test_simplify_composed_bitwise_and() {
// ((c2 > 5) & (c1 < 6)) & (c2 > 5) --> (c2 > 5) & (c1 < 6)
let expr = bitwise_and(
bitwise_and(col("c2").gt(lit(5)), col("c1").lt(lit(6))),
col("c2").gt(lit(5)),
);
let expected = bitwise_and(col("c2").gt(lit(5)), col("c1").lt(lit(6)));
assert_eq!(simplify(expr), expected);
// (c2 > 5) & ((c2 > 5) & (c1 < 6)) --> (c2 > 5) & (c1 < 6)
let expr = bitwise_and(
col("c2").gt(lit(5)),
bitwise_and(col("c2").gt(lit(5)), col("c1").lt(lit(6))),
);
let expected = bitwise_and(col("c2").gt(lit(5)), col("c1").lt(lit(6)));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_composed_bitwise_or() {
// ((c2 > 5) | (c1 < 6)) | (c2 > 5) --> (c2 > 5) | (c1 < 6)
let expr = bitwise_or(
bitwise_or(col("c2").gt(lit(5)), col("c1").lt(lit(6))),
col("c2").gt(lit(5)),
);
let expected = bitwise_or(col("c2").gt(lit(5)), col("c1").lt(lit(6)));
assert_eq!(simplify(expr), expected);
// (c2 > 5) | ((c2 > 5) | (c1 < 6)) --> (c2 > 5) | (c1 < 6)
let expr = bitwise_or(
col("c2").gt(lit(5)),
bitwise_or(col("c2").gt(lit(5)), col("c1").lt(lit(6))),
);
let expected = bitwise_or(col("c2").gt(lit(5)), col("c1").lt(lit(6)));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_composed_bitwise_xor() {
// with an even number of the column "c2"
// c2 ^ ((c2 ^ (c2 | c1)) ^ (c1 & c2)) --> (c2 | c1) ^ (c1 & c2)
let expr = bitwise_xor(
col("c2"),
bitwise_xor(
bitwise_xor(col("c2"), bitwise_or(col("c2"), col("c1"))),
bitwise_and(col("c1"), col("c2")),
),
);
let expected = bitwise_xor(
bitwise_or(col("c2"), col("c1")),
bitwise_and(col("c1"), col("c2")),
);
assert_eq!(simplify(expr), expected);
// with an odd number of the column "c2"
// c2 ^ (c2 ^ (c2 | c1)) ^ ((c1 & c2) ^ c2) --> c2 ^ ((c2 | c1) ^ (c1 & c2))
let expr = bitwise_xor(
col("c2"),
bitwise_xor(
bitwise_xor(col("c2"), bitwise_or(col("c2"), col("c1"))),
bitwise_xor(bitwise_and(col("c1"), col("c2")), col("c2")),
),
);
let expected = bitwise_xor(
col("c2"),
bitwise_xor(
bitwise_or(col("c2"), col("c1")),
bitwise_and(col("c1"), col("c2")),
),
);
assert_eq!(simplify(expr), expected);
// with an even number of the column "c2"
// ((c2 ^ (c2 | c1)) ^ (c1 & c2)) ^ c2 --> (c2 | c1) ^ (c1 & c2)
let expr = bitwise_xor(
bitwise_xor(
bitwise_xor(col("c2"), bitwise_or(col("c2"), col("c1"))),
bitwise_and(col("c1"), col("c2")),
),
col("c2"),
);
let expected = bitwise_xor(
bitwise_or(col("c2"), col("c1")),
bitwise_and(col("c1"), col("c2")),
);
assert_eq!(simplify(expr), expected);
// with an odd number of the column "c2"
// (c2 ^ (c2 | c1)) ^ ((c1 & c2) ^ c2) ^ c2 --> ((c2 | c1) ^ (c1 & c2)) ^ c2
let expr = bitwise_xor(
bitwise_xor(
bitwise_xor(col("c2"), bitwise_or(col("c2"), col("c1"))),
bitwise_xor(bitwise_and(col("c1"), col("c2")), col("c2")),
),
col("c2"),
);
let expected = bitwise_xor(
bitwise_xor(
bitwise_or(col("c2"), col("c1")),
bitwise_and(col("c1"), col("c2")),
),
col("c2"),
);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_negated_bitwise_and() {
// !c4 & c4 --> 0
let expr = (-col("c4_non_null")) & col("c4_non_null");
let expected = lit(0u32);
assert_eq!(simplify(expr), expected);
// c4 & !c4 --> 0
let expr = col("c4_non_null") & (-col("c4_non_null"));
let expected = lit(0u32);
assert_eq!(simplify(expr), expected);
// !c3 & c3 --> 0
let expr = (-col("c3_non_null")) & col("c3_non_null");
let expected = lit(0i64);
assert_eq!(simplify(expr), expected);
// c3 & !c3 --> 0
let expr = col("c3_non_null") & (-col("c3_non_null"));
let expected = lit(0i64);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_negated_bitwise_or() {
// !c4 | c4 --> -1
let expr = (-col("c4_non_null")) | col("c4_non_null");
let expected = lit(-1i32);
assert_eq!(simplify(expr), expected);
// c4 | !c4 --> -1
let expr = col("c4_non_null") | (-col("c4_non_null"));
let expected = lit(-1i32);
assert_eq!(simplify(expr), expected);
// !c3 | c3 --> -1
let expr = (-col("c3_non_null")) | col("c3_non_null");
let expected = lit(-1i64);
assert_eq!(simplify(expr), expected);
// c3 | !c3 --> -1
let expr = col("c3_non_null") | (-col("c3_non_null"));
let expected = lit(-1i64);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_negated_bitwise_xor() {
// !c4 ^ c4 --> -1
let expr = (-col("c4_non_null")) ^ col("c4_non_null");
let expected = lit(-1i32);
assert_eq!(simplify(expr), expected);
// c4 ^ !c4 --> -1
let expr = col("c4_non_null") ^ (-col("c4_non_null"));
let expected = lit(-1i32);
assert_eq!(simplify(expr), expected);
// !c3 ^ c3 --> -1
let expr = (-col("c3_non_null")) ^ col("c3_non_null");
let expected = lit(-1i64);
assert_eq!(simplify(expr), expected);
// c3 ^ !c3 --> -1
let expr = col("c3_non_null") ^ (-col("c3_non_null"));
let expected = lit(-1i64);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_bitwise_and_or() {
// (c2 < 3) & ((c2 < 3) | c1) -> (c2 < 3)
let expr = bitwise_and(
col("c2_non_null").lt(lit(3)),
bitwise_or(col("c2_non_null").lt(lit(3)), col("c1_non_null")),
);
let expected = col("c2_non_null").lt(lit(3));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_bitwise_or_and() {
// (c2 < 3) | ((c2 < 3) & c1) -> (c2 < 3)
let expr = bitwise_or(
col("c2_non_null").lt(lit(3)),
bitwise_and(col("c2_non_null").lt(lit(3)), col("c1_non_null")),
);
let expected = col("c2_non_null").lt(lit(3));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_simple_bitwise_and() {
// (c2 > 5) & (c2 > 5) -> (c2 > 5)
let expr = (col("c2").gt(lit(5))).bitand(col("c2").gt(lit(5)));
let expected = col("c2").gt(lit(5));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_simple_bitwise_or() {
// (c2 > 5) | (c2 > 5) -> (c2 > 5)
let expr = (col("c2").gt(lit(5))).bitor(col("c2").gt(lit(5)));
let expected = col("c2").gt(lit(5));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_simple_bitwise_xor() {
// c4 ^ c4 -> 0
let expr = (col("c4")).bitxor(col("c4"));
let expected = lit(0u32);
assert_eq!(simplify(expr), expected);
// c3 ^ c3 -> 0
let expr = col("c3").bitxor(col("c3"));
let expected = lit(0i64);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_modulo_by_zero_non_null() {
// because modulo by 0 maybe occur in short-circuit expression
// so we should not simplify this, and throw error in runtime.
let expr = col("c2_non_null") % lit(0);
let expected = expr.clone();
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_simple_and() {
// (c2 > 5) AND (c2 > 5) -> (c2 > 5)
let expr = (col("c2").gt(lit(5))).and(col("c2").gt(lit(5)));
let expected = col("c2").gt(lit(5));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_composed_and() {
// ((c2 > 5) AND (c1 < 6)) AND (c2 > 5)
let expr = and(
and(col("c2").gt(lit(5)), col("c1").lt(lit(6))),
col("c2").gt(lit(5)),
);
let expected = and(col("c2").gt(lit(5)), col("c1").lt(lit(6)));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_negated_and() {
// (c2 > 5) AND !(c2 > 5) --> (c2 > 5) AND (c2 <= 5)
let expr = and(col("c2").gt(lit(5)), Expr::not(col("c2").gt(lit(5))));
let expected = col("c2").gt(lit(5)).and(col("c2").lt_eq(lit(5)));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_or_and() {
let l = col("c2").gt(lit(5));
let r = and(col("c1").lt(lit(6)), col("c2").gt(lit(5)));
// (c2 > 5) OR ((c1 < 6) AND (c2 > 5))
let expr = or(l.clone(), r.clone());
let expected = l.clone();
assert_eq!(simplify(expr), expected);
// ((c1 < 6) AND (c2 > 5)) OR (c2 > 5)
let expr = or(r, l);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_or_and_non_null() {
let l = col("c2_non_null").gt(lit(5));
let r = and(col("c1_non_null").lt(lit(6)), col("c2_non_null").gt(lit(5)));
// (c2 > 5) OR ((c1 < 6) AND (c2 > 5)) --> c2 > 5
let expr = or(l.clone(), r.clone());
// This is only true if `c1 < 6` is not nullable / can not be null.
let expected = col("c2_non_null").gt(lit(5));
assert_eq!(simplify(expr), expected);
// ((c1 < 6) AND (c2 > 5)) OR (c2 > 5) --> c2 > 5
let expr = or(l, r);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_and_or() {
let l = col("c2").gt(lit(5));
let r = or(col("c1").lt(lit(6)), col("c2").gt(lit(5)));
// (c2 > 5) AND ((c1 < 6) OR (c2 > 5)) --> c2 > 5
let expr = and(l.clone(), r.clone());
let expected = l.clone();
assert_eq!(simplify(expr), expected);
// ((c1 < 6) OR (c2 > 5)) AND (c2 > 5) --> c2 > 5
let expr = and(r, l);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_and_or_non_null() {
let l = col("c2_non_null").gt(lit(5));
let r = or(col("c1_non_null").lt(lit(6)), col("c2_non_null").gt(lit(5)));
// (c2 > 5) AND ((c1 < 6) OR (c2 > 5)) --> c2 > 5
let expr = and(l.clone(), r.clone());
// This is only true if `c1 < 6` is not nullable / can not be null.
let expected = col("c2_non_null").gt(lit(5));
assert_eq!(simplify(expr), expected);
// ((c1 < 6) OR (c2 > 5)) AND (c2 > 5) --> c2 > 5
let expr = and(l, r);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_by_de_morgan_laws() {
// Laws with logical operations
// !(c3 AND c4) --> !c3 OR !c4
let expr = and(col("c3"), col("c4")).not();
let expected = or(col("c3").not(), col("c4").not());
assert_eq!(simplify(expr), expected);
// !(c3 OR c4) --> !c3 AND !c4
let expr = or(col("c3"), col("c4")).not();
let expected = and(col("c3").not(), col("c4").not());
assert_eq!(simplify(expr), expected);
// !(!c3) --> c3
let expr = col("c3").not().not();
let expected = col("c3");
assert_eq!(simplify(expr), expected);
// Laws with bitwise operations
// !(c3 & c4) --> !c3 | !c4
let expr = -bitwise_and(col("c3"), col("c4"));
let expected = bitwise_or(-col("c3"), -col("c4"));
assert_eq!(simplify(expr), expected);
// !(c3 | c4) --> !c3 & !c4
let expr = -bitwise_or(col("c3"), col("c4"));
let expected = bitwise_and(-col("c3"), -col("c4"));
assert_eq!(simplify(expr), expected);
// !(!c3) --> c3
let expr = -(-col("c3"));
let expected = col("c3");
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_null_and_false() {
let expr = and(lit_bool_null(), lit(false));
let expr_eq = lit(false);
assert_eq!(simplify(expr), expr_eq);
}
#[test]
fn test_simplify_divide_null_by_null() {
let null = lit(ScalarValue::Int32(None));
let expr_plus = null.clone() / null.clone();
let expr_eq = null;
assert_eq!(simplify(expr_plus), expr_eq);
}
#[test]
fn test_simplify_simplify_arithmetic_expr() {
let expr_plus = lit(1) + lit(1);
assert_eq!(simplify(expr_plus), lit(2));
}
#[test]
fn test_simplify_simplify_eq_expr() {
let expr_eq = binary_expr(lit(1), Operator::Eq, lit(1));
assert_eq!(simplify(expr_eq), lit(true));
}
#[test]
fn test_simplify_regex() {
// malformed regex
assert_contains!(
try_simplify(regex_match(col("c1"), lit("foo{")))
.unwrap_err()
.to_string(),
"regex parse error"
);
// unsupported cases
assert_no_change(regex_match(col("c1"), lit("foo.*")));
assert_no_change(regex_match(col("c1"), lit("(foo)")));
assert_no_change(regex_match(col("c1"), lit("%")));
assert_no_change(regex_match(col("c1"), lit("_")));
assert_no_change(regex_match(col("c1"), lit("f%o")));
assert_no_change(regex_match(col("c1"), lit("^f%o")));
assert_no_change(regex_match(col("c1"), lit("f_o")));
// empty cases
assert_change(
regex_match(col("c1"), lit("")),
if_not_null(col("c1"), true),
);
assert_change(
regex_not_match(col("c1"), lit("")),
if_not_null(col("c1"), false),
);
assert_change(
regex_imatch(col("c1"), lit("")),
if_not_null(col("c1"), true),
);
assert_change(
regex_not_imatch(col("c1"), lit("")),
if_not_null(col("c1"), false),
);
// single character
assert_change(regex_match(col("c1"), lit("x")), col("c1").like(lit("%x%")));
// single word
assert_change(
regex_match(col("c1"), lit("foo")),
col("c1").like(lit("%foo%")),
);
// regular expressions that match an exact literal
assert_change(regex_match(col("c1"), lit("^$")), col("c1").eq(lit("")));
assert_change(
regex_not_match(col("c1"), lit("^$")),
col("c1").not_eq(lit("")),
);
assert_change(
regex_match(col("c1"), lit("^foo$")),
col("c1").eq(lit("foo")),
);
assert_change(
regex_not_match(col("c1"), lit("^foo$")),
col("c1").not_eq(lit("foo")),
);
// regular expressions that match exact captured literals
assert_change(
regex_match(col("c1"), lit("^(foo|bar)$")),
col("c1").eq(lit("foo")).or(col("c1").eq(lit("bar"))),
);
assert_change(
regex_not_match(col("c1"), lit("^(foo|bar)$")),
col("c1")
.not_eq(lit("foo"))
.and(col("c1").not_eq(lit("bar"))),
);
assert_change(
regex_match(col("c1"), lit("^(foo)$")),
col("c1").eq(lit("foo")),
);
assert_change(
regex_match(col("c1"), lit("^(foo|bar|baz)$")),
((col("c1").eq(lit("foo"))).or(col("c1").eq(lit("bar"))))
.or(col("c1").eq(lit("baz"))),
);
assert_change(
regex_match(col("c1"), lit("^(foo|bar|baz|qux)$")),
col("c1")
.in_list(vec![lit("foo"), lit("bar"), lit("baz"), lit("qux")], false),
);
assert_change(
regex_match(col("c1"), lit("^(fo_o)$")),
col("c1").eq(lit("fo_o")),
);
assert_change(
regex_match(col("c1"), lit("^(fo_o)$")),
col("c1").eq(lit("fo_o")),
);
assert_change(
regex_match(col("c1"), lit("^(fo_o|ba_r)$")),
col("c1").eq(lit("fo_o")).or(col("c1").eq(lit("ba_r"))),
);
assert_change(
regex_not_match(col("c1"), lit("^(fo_o|ba_r)$")),
col("c1")
.not_eq(lit("fo_o"))
.and(col("c1").not_eq(lit("ba_r"))),
);
assert_change(
regex_match(col("c1"), lit("^(fo_o|ba_r|ba_z)$")),
((col("c1").eq(lit("fo_o"))).or(col("c1").eq(lit("ba_r"))))
.or(col("c1").eq(lit("ba_z"))),
);
assert_change(
regex_match(col("c1"), lit("^(fo_o|ba_r|baz|qu_x)$")),
col("c1").in_list(
vec![lit("fo_o"), lit("ba_r"), lit("baz"), lit("qu_x")],
false,
),
);
// regular expressions that mismatch captured literals
assert_no_change(regex_match(col("c1"), lit("(foo|bar)")));
assert_no_change(regex_match(col("c1"), lit("(foo|bar)*")));
assert_no_change(regex_match(col("c1"), lit("(fo_o|b_ar)")));
assert_no_change(regex_match(col("c1"), lit("(foo|ba_r)*")));
assert_no_change(regex_match(col("c1"), lit("(fo_o|ba_r)*")));
assert_no_change(regex_match(col("c1"), lit("^(foo|bar)*")));
assert_no_change(regex_match(col("c1"), lit("^(foo)(bar)$")));
assert_no_change(regex_match(col("c1"), lit("^")));
assert_no_change(regex_match(col("c1"), lit("$")));
assert_no_change(regex_match(col("c1"), lit("$^")));
assert_no_change(regex_match(col("c1"), lit("$foo^")));
// regular expressions that match a partial literal
assert_change(
regex_match(col("c1"), lit("^foo")),
col("c1").like(lit("foo%")),
);
assert_change(
regex_match(col("c1"), lit("foo$")),
col("c1").like(lit("%foo")),
);
assert_change(
regex_match(col("c1"), lit("^foo|bar$")),
col("c1").like(lit("foo%")).or(col("c1").like(lit("%bar"))),
);
// OR-chain
assert_change(
regex_match(col("c1"), lit("foo|bar|baz")),
col("c1")
.like(lit("%foo%"))
.or(col("c1").like(lit("%bar%")))
.or(col("c1").like(lit("%baz%"))),
);
assert_change(
regex_match(col("c1"), lit("foo|x|baz")),
col("c1")
.like(lit("%foo%"))
.or(col("c1").like(lit("%x%")))
.or(col("c1").like(lit("%baz%"))),
);
assert_change(
regex_not_match(col("c1"), lit("foo|bar|baz")),
col("c1")
.not_like(lit("%foo%"))
.and(col("c1").not_like(lit("%bar%")))
.and(col("c1").not_like(lit("%baz%"))),
);
// both anchored expressions (translated to equality) and unanchored
assert_change(
regex_match(col("c1"), lit("foo|^x$|baz")),
col("c1")
.like(lit("%foo%"))
.or(col("c1").eq(lit("x")))
.or(col("c1").like(lit("%baz%"))),
);
assert_change(
regex_not_match(col("c1"), lit("foo|^bar$|baz")),
col("c1")
.not_like(lit("%foo%"))
.and(col("c1").not_eq(lit("bar")))
.and(col("c1").not_like(lit("%baz%"))),
);
// Too many patterns (MAX_REGEX_ALTERNATIONS_EXPANSION)
assert_no_change(regex_match(col("c1"), lit("foo|bar|baz|blarg|bozo|etc")));
}
#[track_caller]
fn assert_no_change(expr: Expr) {
let optimized = simplify(expr.clone());
assert_eq!(expr, optimized);
}
#[track_caller]
fn assert_change(expr: Expr, expected: Expr) {
let optimized = simplify(expr);
assert_eq!(optimized, expected);
}
fn regex_match(left: Expr, right: Expr) -> Expr {
Expr::BinaryExpr(BinaryExpr {
left: Box::new(left),
op: Operator::RegexMatch,
right: Box::new(right),
})
}
fn regex_not_match(left: Expr, right: Expr) -> Expr {
Expr::BinaryExpr(BinaryExpr {
left: Box::new(left),
op: Operator::RegexNotMatch,
right: Box::new(right),
})
}
fn regex_imatch(left: Expr, right: Expr) -> Expr {
Expr::BinaryExpr(BinaryExpr {
left: Box::new(left),
op: Operator::RegexIMatch,
right: Box::new(right),
})
}
fn regex_not_imatch(left: Expr, right: Expr) -> Expr {
Expr::BinaryExpr(BinaryExpr {
left: Box::new(left),
op: Operator::RegexNotIMatch,
right: Box::new(right),
})
}
// ------------------------------
// ----- Simplifier tests -------
// ------------------------------
fn try_simplify(expr: Expr) -> Result<Expr> {
let schema = expr_test_schema();
let execution_props = ExecutionProps::new();
let simplifier = ExprSimplifier::new(
SimplifyContext::new(&execution_props).with_schema(schema),
);
simplifier.simplify(expr)
}
fn coerce(expr: Expr) -> Expr {
let schema = expr_test_schema();
let execution_props = ExecutionProps::new();
let simplifier = ExprSimplifier::new(
SimplifyContext::new(&execution_props).with_schema(Arc::clone(&schema)),
);
simplifier.coerce(expr, schema.as_ref()).unwrap()
}
fn simplify(expr: Expr) -> Expr {
try_simplify(expr).unwrap()
}
fn try_simplify_with_cycle_count(expr: Expr) -> Result<(Expr, u32)> {
let schema = expr_test_schema();
let execution_props = ExecutionProps::new();
let simplifier = ExprSimplifier::new(
SimplifyContext::new(&execution_props).with_schema(schema),
);
let (expr, count) = simplifier.simplify_with_cycle_count_transformed(expr)?;
Ok((expr.data, count))
}
fn simplify_with_cycle_count(expr: Expr) -> (Expr, u32) {
try_simplify_with_cycle_count(expr).unwrap()
}
fn simplify_with_guarantee(
expr: Expr,
guarantees: Vec<(Expr, NullableInterval)>,
) -> Expr {
let schema = expr_test_schema();
let execution_props = ExecutionProps::new();
let simplifier = ExprSimplifier::new(
SimplifyContext::new(&execution_props).with_schema(schema),
)
.with_guarantees(guarantees);
simplifier.simplify(expr).unwrap()
}
fn expr_test_schema() -> DFSchemaRef {
static EXPR_TEST_SCHEMA: LazyLock<DFSchemaRef> = LazyLock::new(|| {
Arc::new(
DFSchema::from_unqualified_fields(
vec![
Field::new("c1", DataType::Utf8, true),
Field::new("c2", DataType::Boolean, true),
Field::new("c3", DataType::Int64, true),
Field::new("c4", DataType::UInt32, true),
Field::new("c1_non_null", DataType::Utf8, false),
Field::new("c2_non_null", DataType::Boolean, false),
Field::new("c3_non_null", DataType::Int64, false),
Field::new("c4_non_null", DataType::UInt32, false),
Field::new("c5", DataType::FixedSizeBinary(3), true),
]
.into(),
HashMap::new(),
)
.unwrap(),
)
});
Arc::clone(&EXPR_TEST_SCHEMA)
}
#[test]
fn simplify_expr_null_comparison() {
// x = null is always null
assert_eq!(
simplify(lit(true).eq(lit(ScalarValue::Boolean(None)))),
lit(ScalarValue::Boolean(None)),
);
// null != null is always null
assert_eq!(
simplify(
lit(ScalarValue::Boolean(None)).not_eq(lit(ScalarValue::Boolean(None)))
),
lit(ScalarValue::Boolean(None)),
);
// x != null is always null
assert_eq!(
simplify(col("c2").not_eq(lit(ScalarValue::Boolean(None)))),
lit(ScalarValue::Boolean(None)),
);
// null = x is always null
assert_eq!(
simplify(lit(ScalarValue::Boolean(None)).eq(col("c2"))),
lit(ScalarValue::Boolean(None)),
);
}
#[test]
fn simplify_expr_is_not_null() {
assert_eq!(
simplify(Expr::IsNotNull(Box::new(col("c1")))),
Expr::IsNotNull(Box::new(col("c1")))
);
// 'c1_non_null IS NOT NULL' is always true
assert_eq!(
simplify(Expr::IsNotNull(Box::new(col("c1_non_null")))),
lit(true)
);
}
#[test]
fn simplify_expr_is_null() {
assert_eq!(
simplify(Expr::IsNull(Box::new(col("c1")))),
Expr::IsNull(Box::new(col("c1")))
);
// 'c1_non_null IS NULL' is always false
assert_eq!(
simplify(Expr::IsNull(Box::new(col("c1_non_null")))),
lit(false)
);
}
#[test]
fn simplify_expr_is_unknown() {
assert_eq!(simplify(col("c2").is_unknown()), col("c2").is_unknown(),);
// 'c2_non_null is unknown' is always false
assert_eq!(simplify(col("c2_non_null").is_unknown()), lit(false));
}
#[test]
fn simplify_expr_is_not_known() {
assert_eq!(
simplify(col("c2").is_not_unknown()),
col("c2").is_not_unknown()
);
// 'c2_non_null is not unknown' is always true
assert_eq!(simplify(col("c2_non_null").is_not_unknown()), lit(true));
}
#[test]
fn simplify_expr_eq() {
let schema = expr_test_schema();
assert_eq!(col("c2").get_type(&schema).unwrap(), DataType::Boolean);
// true = true -> true
assert_eq!(simplify(lit(true).eq(lit(true))), lit(true));
// true = false -> false
assert_eq!(simplify(lit(true).eq(lit(false))), lit(false),);
// c2 = true -> c2
assert_eq!(simplify(col("c2").eq(lit(true))), col("c2"));
// c2 = false => !c2
assert_eq!(simplify(col("c2").eq(lit(false))), col("c2").not(),);
}
#[test]
fn simplify_expr_eq_skip_nonboolean_type() {
let schema = expr_test_schema();
// When one of the operand is not of boolean type, folding the
// other boolean constant will change return type of
// expression to non-boolean.
//
// Make sure c1 column to be used in tests is not boolean type
assert_eq!(col("c1").get_type(&schema).unwrap(), DataType::Utf8);
// don't fold c1 = foo
assert_eq!(simplify(col("c1").eq(lit("foo"))), col("c1").eq(lit("foo")),);
}
#[test]
fn simplify_expr_not_eq() {
let schema = expr_test_schema();
assert_eq!(col("c2").get_type(&schema).unwrap(), DataType::Boolean);
// c2 != true -> !c2
assert_eq!(simplify(col("c2").not_eq(lit(true))), col("c2").not(),);
// c2 != false -> c2
assert_eq!(simplify(col("c2").not_eq(lit(false))), col("c2"),);
// test constant
assert_eq!(simplify(lit(true).not_eq(lit(true))), lit(false),);
assert_eq!(simplify(lit(true).not_eq(lit(false))), lit(true),);
}
#[test]
fn simplify_expr_not_eq_skip_nonboolean_type() {
let schema = expr_test_schema();
// when one of the operand is not of boolean type, folding the
// other boolean constant will change return type of
// expression to non-boolean.
assert_eq!(col("c1").get_type(&schema).unwrap(), DataType::Utf8);
assert_eq!(
simplify(col("c1").not_eq(lit("foo"))),
col("c1").not_eq(lit("foo")),
);
}
#[test]
fn simplify_literal_case_equality() {
// CASE WHEN c2 != false THEN "ok" ELSE "not_ok"
let simple_case = Expr::Case(Case::new(
None,
vec![(
Box::new(col("c2_non_null").not_eq(lit(false))),
Box::new(lit("ok")),
)],
Some(Box::new(lit("not_ok"))),
));
// CASE WHEN c2 != false THEN "ok" ELSE "not_ok" == "ok"
// -->
// CASE WHEN c2 != false THEN "ok" == "ok" ELSE "not_ok" == "ok"
// -->
// CASE WHEN c2 != false THEN true ELSE false
// -->
// c2
assert_eq!(
simplify(binary_expr(simple_case.clone(), Operator::Eq, lit("ok"),)),
col("c2_non_null"),
);
// CASE WHEN c2 != false THEN "ok" ELSE "not_ok" != "ok"
// -->
// NOT(CASE WHEN c2 != false THEN "ok" == "ok" ELSE "not_ok" == "ok")
// -->
// NOT(CASE WHEN c2 != false THEN true ELSE false)
// -->
// NOT(c2)
assert_eq!(
simplify(binary_expr(simple_case, Operator::NotEq, lit("ok"),)),
not(col("c2_non_null")),
);
let complex_case = Expr::Case(Case::new(
None,
vec![
(
Box::new(col("c1").eq(lit("inboxed"))),
Box::new(lit("pending")),
),
(
Box::new(col("c1").eq(lit("scheduled"))),
Box::new(lit("pending")),
),
(
Box::new(col("c1").eq(lit("completed"))),
Box::new(lit("completed")),
),
(
Box::new(col("c1").eq(lit("paused"))),
Box::new(lit("paused")),
),
(Box::new(col("c2")), Box::new(lit("running"))),
(
Box::new(col("c1").eq(lit("invoked")).and(col("c3").gt(lit(0)))),
Box::new(lit("backing-off")),
),
],
Some(Box::new(lit("ready"))),
));
assert_eq!(
simplify(binary_expr(
complex_case.clone(),
Operator::Eq,
lit("completed"),
)),
not_distinct_from(col("c1").eq(lit("completed")), lit(true)).and(
distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.and(distinct_from(col("c1").eq(lit("scheduled")), lit(true)))
)
);
assert_eq!(
simplify(binary_expr(
complex_case.clone(),
Operator::NotEq,
lit("completed"),
)),
distinct_from(col("c1").eq(lit("completed")), lit(true))
.or(not_distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.or(not_distinct_from(col("c1").eq(lit("scheduled")), lit(true))))
);
assert_eq!(
simplify(binary_expr(
complex_case.clone(),
Operator::Eq,
lit("running"),
)),
not_distinct_from(col("c2"), lit(true)).and(
distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.and(distinct_from(col("c1").eq(lit("scheduled")), lit(true)))
.and(distinct_from(col("c1").eq(lit("completed")), lit(true)))
.and(distinct_from(col("c1").eq(lit("paused")), lit(true)))
)
);
assert_eq!(
simplify(binary_expr(
complex_case.clone(),
Operator::Eq,
lit("ready"),
)),
distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.and(distinct_from(col("c1").eq(lit("scheduled")), lit(true)))
.and(distinct_from(col("c1").eq(lit("completed")), lit(true)))
.and(distinct_from(col("c1").eq(lit("paused")), lit(true)))
.and(distinct_from(col("c2"), lit(true)))
.and(distinct_from(
col("c1").eq(lit("invoked")).and(col("c3").gt(lit(0))),
lit(true)
))
);
assert_eq!(
simplify(binary_expr(
complex_case.clone(),
Operator::NotEq,
lit("ready"),
)),
not_distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.or(not_distinct_from(col("c1").eq(lit("scheduled")), lit(true)))
.or(not_distinct_from(col("c1").eq(lit("completed")), lit(true)))
.or(not_distinct_from(col("c1").eq(lit("paused")), lit(true)))
.or(not_distinct_from(col("c2"), lit(true)))
.or(not_distinct_from(
col("c1").eq(lit("invoked")).and(col("c3").gt(lit(0))),
lit(true)
))
);
}
#[test]
fn simplify_expr_case_when_then_else() {
// CASE WHEN c2 != false THEN "ok" == "not_ok" ELSE c2 == true
// -->
// CASE WHEN c2 THEN false ELSE c2
// -->
// false
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![(
Box::new(col("c2_non_null").not_eq(lit(false))),
Box::new(lit("ok").eq(lit("not_ok"))),
)],
Some(Box::new(col("c2_non_null").eq(lit(true)))),
))),
lit(false) // #1716
);
// CASE WHEN c2 != false THEN "ok" == "ok" ELSE c2
// -->
// CASE WHEN c2 THEN true ELSE c2
// -->
// c2
//
// Need to call simplify 2x due to
// https://github.com/apache/datafusion/issues/1160
assert_eq!(
simplify(simplify(Expr::Case(Case::new(
None,
vec![(
Box::new(col("c2_non_null").not_eq(lit(false))),
Box::new(lit("ok").eq(lit("ok"))),
)],
Some(Box::new(col("c2_non_null").eq(lit(true)))),
)))),
col("c2_non_null")
);
// CASE WHEN ISNULL(c2) THEN true ELSE c2
// -->
// ISNULL(c2) OR c2
//
// Need to call simplify 2x due to
// https://github.com/apache/datafusion/issues/1160
assert_eq!(
simplify(simplify(Expr::Case(Case::new(
None,
vec![(Box::new(col("c2").is_null()), Box::new(lit(true)),)],
Some(Box::new(col("c2"))),
)))),
col("c2")
.is_null()
.or(col("c2").is_not_null().and(col("c2")))
);
// CASE WHEN c1 then true WHEN c2 then false ELSE true
// --> c1 OR (NOT(c1) AND c2 AND FALSE) OR (NOT(c1 OR c2) AND TRUE)
// --> c1 OR (NOT(c1) AND NOT(c2))
// --> c1 OR NOT(c2)
//
// Need to call simplify 2x due to
// https://github.com/apache/datafusion/issues/1160
assert_eq!(
simplify(simplify(Expr::Case(Case::new(
None,
vec![
(Box::new(col("c1_non_null")), Box::new(lit(true)),),
(Box::new(col("c2_non_null")), Box::new(lit(false)),),
],
Some(Box::new(lit(true))),
)))),
col("c1_non_null").or(col("c1_non_null").not().and(col("c2_non_null").not()))
);
// CASE WHEN c1 then true WHEN c2 then true ELSE false
// --> c1 OR (NOT(c1) AND c2 AND TRUE) OR (NOT(c1 OR c2) AND FALSE)
// --> c1 OR (NOT(c1) AND c2)
// --> c1 OR c2
//
// Need to call simplify 2x due to
// https://github.com/apache/datafusion/issues/1160
assert_eq!(
simplify(simplify(Expr::Case(Case::new(
None,
vec![
(Box::new(col("c1_non_null")), Box::new(lit(true)),),
(Box::new(col("c2_non_null")), Box::new(lit(false)),),
],
Some(Box::new(lit(true))),
)))),
col("c1_non_null").or(col("c1_non_null").not().and(col("c2_non_null").not()))
);
// CASE WHEN c > 0 THEN true END AS c1
assert_eq!(
simplify(simplify(Expr::Case(Case::new(
None,
vec![(Box::new(col("c3").gt(lit(0_i64))), Box::new(lit(true)))],
None,
)))),
not_distinct_from(col("c3").gt(lit(0_i64)), lit(true)).or(distinct_from(
col("c3").gt(lit(0_i64)),
lit(true)
)
.and(lit_bool_null()))
);
// CASE WHEN c > 0 THEN true ELSE false END AS c1
assert_eq!(
simplify(simplify(Expr::Case(Case::new(
None,
vec![(Box::new(col("c3").gt(lit(0_i64))), Box::new(lit(true)))],
Some(Box::new(lit(false))),
)))),
not_distinct_from(col("c3").gt(lit(0_i64)), lit(true))
);
}
#[test]
fn simplify_expr_case_when_first_true() {
// CASE WHEN true THEN 1 ELSE c1 END --> 1
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![(Box::new(lit(true)), Box::new(lit(1)),)],
Some(Box::new(col("c1"))),
))),
lit(1)
);
// CASE WHEN true THEN col('a') ELSE col('b') END --> col('a')
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![(Box::new(lit(true)), Box::new(lit("a")),)],
Some(Box::new(lit("b"))),
))),
lit("a")
);
// CASE WHEN true THEN col('a') WHEN col('x') > 5 THEN col('b') ELSE col('c') END --> col('a')
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![
(Box::new(lit(true)), Box::new(lit("a"))),
(Box::new(lit("x").gt(lit(5))), Box::new(lit("b"))),
],
Some(Box::new(lit("c"))),
))),
lit("a")
);
// CASE WHEN true THEN col('a') END --> col('a') (no else clause)
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![(Box::new(lit(true)), Box::new(lit("a")),)],
None,
))),
lit("a")
);
// Negative test: CASE WHEN c2 THEN 1 ELSE 2 END should not be simplified
let expr = Expr::Case(Case::new(
None,
vec![(Box::new(col("c2")), Box::new(lit(1)))],
Some(Box::new(lit(2))),
));
assert_eq!(simplify(expr.clone()), expr);
// Negative test: CASE WHEN false THEN 1 ELSE 2 END should not use this rule
let expr = Expr::Case(Case::new(
None,
vec![(Box::new(lit(false)), Box::new(lit(1)))],
Some(Box::new(lit(2))),
));
assert_ne!(simplify(expr), lit(1));
// Negative test: CASE WHEN col('c1') > 5 THEN 1 ELSE 2 END should not be simplified
let expr = Expr::Case(Case::new(
None,
vec![(Box::new(col("c1").gt(lit(5))), Box::new(lit(1)))],
Some(Box::new(lit(2))),
));
assert_eq!(simplify(expr.clone()), expr);
}
#[test]
fn simplify_expr_case_when_any_true() {
// CASE WHEN c3 > 0 THEN 'a' WHEN true THEN 'b' ELSE 'c' END --> CASE WHEN c3 > 0 THEN 'a' ELSE 'b' END
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![
(Box::new(col("c3").gt(lit(0))), Box::new(lit("a"))),
(Box::new(lit(true)), Box::new(lit("b"))),
],
Some(Box::new(lit("c"))),
))),
Expr::Case(Case::new(
None,
vec![(Box::new(col("c3").gt(lit(0))), Box::new(lit("a")))],
Some(Box::new(lit("b"))),
))
);
// CASE WHEN c3 > 0 THEN 'a' WHEN c4 < 0 THEN 'b' WHEN true THEN 'c' WHEN c3 = 0 THEN 'd' ELSE 'e' END
// --> CASE WHEN c3 > 0 THEN 'a' WHEN c4 < 0 THEN 'b' ELSE 'c' END
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![
(Box::new(col("c3").gt(lit(0))), Box::new(lit("a"))),
(Box::new(col("c4").lt(lit(0))), Box::new(lit("b"))),
(Box::new(lit(true)), Box::new(lit("c"))),
(Box::new(col("c3").eq(lit(0))), Box::new(lit("d"))),
],
Some(Box::new(lit("e"))),
))),
Expr::Case(Case::new(
None,
vec![
(Box::new(col("c3").gt(lit(0))), Box::new(lit("a"))),
(Box::new(col("c4").lt(lit(0))), Box::new(lit("b"))),
],
Some(Box::new(lit("c"))),
))
);
// CASE WHEN c3 > 0 THEN 1 WHEN c4 < 0 THEN 2 WHEN true THEN 3 END (no else)
// --> CASE WHEN c3 > 0 THEN 1 WHEN c4 < 0 THEN 2 ELSE 3 END
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![
(Box::new(col("c3").gt(lit(0))), Box::new(lit(1))),
(Box::new(col("c4").lt(lit(0))), Box::new(lit(2))),
(Box::new(lit(true)), Box::new(lit(3))),
],
None,
))),
Expr::Case(Case::new(
None,
vec![
(Box::new(col("c3").gt(lit(0))), Box::new(lit(1))),
(Box::new(col("c4").lt(lit(0))), Box::new(lit(2))),
],
Some(Box::new(lit(3))),
))
);
// Negative test: CASE WHEN c3 > 0 THEN c3 WHEN c4 < 0 THEN 2 ELSE 3 END should not be simplified
let expr = Expr::Case(Case::new(
None,
vec![
(Box::new(col("c3").gt(lit(0))), Box::new(col("c3"))),
(Box::new(col("c4").lt(lit(0))), Box::new(lit(2))),
],
Some(Box::new(lit(3))),
));
assert_eq!(simplify(expr.clone()), expr);
}
#[test]
fn simplify_expr_case_when_any_false() {
// CASE WHEN false THEN 'a' END --> NULL
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![(Box::new(lit(false)), Box::new(lit("a")))],
None,
))),
Expr::Literal(ScalarValue::Utf8(None), None)
);
// CASE WHEN false THEN 2 ELSE 1 END --> 1
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![(Box::new(lit(false)), Box::new(lit(2)))],
Some(Box::new(lit(1))),
))),
lit(1),
);
// CASE WHEN c3 < 10 THEN 'b' WHEN false then c3 ELSE c4 END --> CASE WHEN c3 < 10 THEN b ELSE c4 END
assert_eq!(
simplify(Expr::Case(Case::new(
None,
vec![
(Box::new(col("c3").lt(lit(10))), Box::new(lit("b"))),
(Box::new(lit(false)), Box::new(col("c3"))),
],
Some(Box::new(col("c4"))),
))),
Expr::Case(Case::new(
None,
vec![(Box::new(col("c3").lt(lit(10))), Box::new(lit("b")))],
Some(Box::new(col("c4"))),
))
);
// Negative test: CASE WHEN c3 = 4 THEN 1 ELSE 2 END should not be simplified
let expr = Expr::Case(Case::new(
None,
vec![(Box::new(col("c3").eq(lit(4))), Box::new(lit(1)))],
Some(Box::new(lit(2))),
));
assert_eq!(simplify(expr.clone()), expr);
}
fn distinct_from(left: impl Into<Expr>, right: impl Into<Expr>) -> Expr {
Expr::BinaryExpr(BinaryExpr {
left: Box::new(left.into()),
op: Operator::IsDistinctFrom,
right: Box::new(right.into()),
})
}
fn not_distinct_from(left: impl Into<Expr>, right: impl Into<Expr>) -> Expr {
Expr::BinaryExpr(BinaryExpr {
left: Box::new(left.into()),
op: Operator::IsNotDistinctFrom,
right: Box::new(right.into()),
})
}
#[test]
fn simplify_expr_bool_or() {
// col || true is always true
assert_eq!(simplify(col("c2").or(lit(true))), lit(true),);
// col || false is always col
assert_eq!(simplify(col("c2").or(lit(false))), col("c2"),);
// true || null is always true
assert_eq!(simplify(lit(true).or(lit_bool_null())), lit(true),);
// null || true is always true
assert_eq!(simplify(lit_bool_null().or(lit(true))), lit(true),);
// false || null is always null
assert_eq!(simplify(lit(false).or(lit_bool_null())), lit_bool_null(),);
// null || false is always null
assert_eq!(simplify(lit_bool_null().or(lit(false))), lit_bool_null(),);
// ( c1 BETWEEN Int32(0) AND Int32(10) ) OR Boolean(NULL)
// it can be either NULL or TRUE depending on the value of `c1 BETWEEN Int32(0) AND Int32(10)`
// and should not be rewritten
let expr = col("c1").between(lit(0), lit(10));
let expr = expr.or(lit_bool_null());
let result = simplify(expr);
let expected_expr = or(
and(col("c1").gt_eq(lit(0)), col("c1").lt_eq(lit(10))),
lit_bool_null(),
);
assert_eq!(expected_expr, result);
}
#[test]
fn simplify_inlist() {
assert_eq!(simplify(in_list(col("c1"), vec![], false)), lit(false));
assert_eq!(simplify(in_list(col("c1"), vec![], true)), lit(true));
// null in (...) --> null
assert_eq!(
simplify(in_list(lit_bool_null(), vec![col("c1"), lit(1)], false)),
lit_bool_null()
);
// null not in (...) --> null
assert_eq!(
simplify(in_list(lit_bool_null(), vec![col("c1"), lit(1)], true)),
lit_bool_null()
);
assert_eq!(
simplify(in_list(col("c1"), vec![lit(1)], false)),
col("c1").eq(lit(1))
);
assert_eq!(
simplify(in_list(col("c1"), vec![lit(1)], true)),
col("c1").not_eq(lit(1))
);
// more complex expressions can be simplified if list contains
// one element only
assert_eq!(
simplify(in_list(col("c1") * lit(10), vec![lit(2)], false)),
(col("c1") * lit(10)).eq(lit(2))
);
assert_eq!(
simplify(in_list(col("c1"), vec![lit(1), lit(2)], false)),
col("c1").eq(lit(1)).or(col("c1").eq(lit(2)))
);
assert_eq!(
simplify(in_list(col("c1"), vec![lit(1), lit(2)], true)),
col("c1").not_eq(lit(1)).and(col("c1").not_eq(lit(2)))
);
let subquery = Arc::new(test_table_scan_with_name("test").unwrap());
assert_eq!(
simplify(in_list(
col("c1"),
vec![scalar_subquery(Arc::clone(&subquery))],
false
)),
in_subquery(col("c1"), Arc::clone(&subquery))
);
assert_eq!(
simplify(in_list(
col("c1"),
vec![scalar_subquery(Arc::clone(&subquery))],
true
)),
not_in_subquery(col("c1"), subquery)
);
let subquery1 =
scalar_subquery(Arc::new(test_table_scan_with_name("test1").unwrap()));
let subquery2 =
scalar_subquery(Arc::new(test_table_scan_with_name("test2").unwrap()));
// c1 NOT IN (<subquery1>, <subquery2>) -> c1 != <subquery1> AND c1 != <subquery2>
assert_eq!(
simplify(in_list(
col("c1"),
vec![subquery1.clone(), subquery2.clone()],
true
)),
col("c1")
.not_eq(subquery1.clone())
.and(col("c1").not_eq(subquery2.clone()))
);
// c1 IN (<subquery1>, <subquery2>) -> c1 == <subquery1> OR c1 == <subquery2>
assert_eq!(
simplify(in_list(
col("c1"),
vec![subquery1.clone(), subquery2.clone()],
false
)),
col("c1").eq(subquery1).or(col("c1").eq(subquery2))
);
// 1. c1 IN (1,2,3,4) AND c1 IN (5,6,7,8) -> false
let expr = in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], false).and(
in_list(col("c1"), vec![lit(5), lit(6), lit(7), lit(8)], false),
);
assert_eq!(simplify(expr), lit(false));
// 2. c1 IN (1,2,3,4) AND c1 IN (4,5,6,7) -> c1 = 4
let expr = in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], false).and(
in_list(col("c1"), vec![lit(4), lit(5), lit(6), lit(7)], false),
);
assert_eq!(simplify(expr), col("c1").eq(lit(4)));
// 3. c1 NOT IN (1, 2, 3, 4) OR c1 NOT IN (5, 6, 7, 8) -> true
let expr = in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], true).or(
in_list(col("c1"), vec![lit(5), lit(6), lit(7), lit(8)], true),
);
assert_eq!(simplify(expr), lit(true));
// 3.5 c1 NOT IN (1, 2, 3, 4) OR c1 NOT IN (4, 5, 6, 7) -> c1 != 4 (4 overlaps)
let expr = in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], true).or(
in_list(col("c1"), vec![lit(4), lit(5), lit(6), lit(7)], true),
);
assert_eq!(simplify(expr), col("c1").not_eq(lit(4)));
// 4. c1 NOT IN (1,2,3,4) AND c1 NOT IN (4,5,6,7) -> c1 NOT IN (1,2,3,4,5,6,7)
let expr = in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], true).and(
in_list(col("c1"), vec![lit(4), lit(5), lit(6), lit(7)], true),
);
assert_eq!(
simplify(expr),
in_list(
col("c1"),
vec![lit(1), lit(2), lit(3), lit(4), lit(5), lit(6), lit(7)],
true
)
);
// 5. c1 IN (1,2,3,4) OR c1 IN (2,3,4,5) -> c1 IN (1,2,3,4,5)
let expr = in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], false).or(
in_list(col("c1"), vec![lit(2), lit(3), lit(4), lit(5)], false),
);
assert_eq!(
simplify(expr),
in_list(
col("c1"),
vec![lit(1), lit(2), lit(3), lit(4), lit(5)],
false
)
);
// 6. c1 IN (1,2,3) AND c1 NOT INT (1,2,3,4,5) -> false
let expr = in_list(col("c1"), vec![lit(1), lit(2), lit(3)], false).and(in_list(
col("c1"),
vec![lit(1), lit(2), lit(3), lit(4), lit(5)],
true,
));
assert_eq!(simplify(expr), lit(false));
// 7. c1 NOT IN (1,2,3,4) AND c1 IN (1,2,3,4,5) -> c1 = 5
let expr =
in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], true).and(in_list(
col("c1"),
vec![lit(1), lit(2), lit(3), lit(4), lit(5)],
false,
));
assert_eq!(simplify(expr), col("c1").eq(lit(5)));
// 8. c1 IN (1,2,3,4) AND c1 NOT IN (5,6,7,8) -> c1 IN (1,2,3,4)
let expr = in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], false).and(
in_list(col("c1"), vec![lit(5), lit(6), lit(7), lit(8)], true),
);
assert_eq!(
simplify(expr),
in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], false)
);
// inlist with more than two expressions
// c1 IN (1,2,3,4,5,6) AND c1 IN (1,3,5,6) AND c1 IN (3,6) -> c1 = 3 OR c1 = 6
let expr = in_list(
col("c1"),
vec![lit(1), lit(2), lit(3), lit(4), lit(5), lit(6)],
false,
)
.and(in_list(
col("c1"),
vec![lit(1), lit(3), lit(5), lit(6)],
false,
))
.and(in_list(col("c1"), vec![lit(3), lit(6)], false));
assert_eq!(
simplify(expr),
col("c1").eq(lit(3)).or(col("c1").eq(lit(6)))
);
// c1 NOT IN (1,2,3,4) AND c1 IN (5,6,7,8) AND c1 NOT IN (3,4,5,6) AND c1 IN (8,9,10) -> c1 = 8
let expr = in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], true).and(
in_list(col("c1"), vec![lit(5), lit(6), lit(7), lit(8)], false)
.and(in_list(
col("c1"),
vec![lit(3), lit(4), lit(5), lit(6)],
true,
))
.and(in_list(col("c1"), vec![lit(8), lit(9), lit(10)], false)),
);
assert_eq!(simplify(expr), col("c1").eq(lit(8)));
// Contains non-InList expression
// c1 NOT IN (1,2,3,4) OR c1 != 5 OR c1 NOT IN (6,7,8,9) -> c1 NOT IN (1,2,3,4) OR c1 != 5 OR c1 NOT IN (6,7,8,9)
let expr =
in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], true).or(col("c1")
.not_eq(lit(5))
.or(in_list(
col("c1"),
vec![lit(6), lit(7), lit(8), lit(9)],
true,
)));
// TODO: Further simplify this expression
// https://github.com/apache/datafusion/issues/8970
// assert_eq!(simplify(expr.clone()), lit(true));
assert_eq!(simplify(expr.clone()), expr);
}
#[test]
fn simplify_null_in_empty_inlist() {
// `NULL::boolean IN ()` == `NULL::boolean IN (SELECT foo FROM empty)` == false
let expr = in_list(lit_bool_null(), vec![], false);
assert_eq!(simplify(expr), lit(false));
// `NULL::boolean NOT IN ()` == `NULL::boolean NOT IN (SELECT foo FROM empty)` == true
let expr = in_list(lit_bool_null(), vec![], true);
assert_eq!(simplify(expr), lit(true));
// `NULL IN ()` == `NULL IN (SELECT foo FROM empty)` == false
let null_null = || Expr::Literal(ScalarValue::Null, None);
let expr = in_list(null_null(), vec![], false);
assert_eq!(simplify(expr), lit(false));
// `NULL NOT IN ()` == `NULL NOT IN (SELECT foo FROM empty)` == true
let expr = in_list(null_null(), vec![], true);
assert_eq!(simplify(expr), lit(true));
}
#[test]
fn just_simplifier_simplify_null_in_empty_inlist() {
let simplify = |expr: Expr| -> Expr {
let schema = expr_test_schema();
let execution_props = ExecutionProps::new();
let info = SimplifyContext::new(&execution_props).with_schema(schema);
let simplifier = &mut Simplifier::new(&info);
expr.rewrite(simplifier)
.expect("Failed to simplify expression")
.data
};
// `NULL::boolean IN ()` == `NULL::boolean IN (SELECT foo FROM empty)` == false
let expr = in_list(lit_bool_null(), vec![], false);
assert_eq!(simplify(expr), lit(false));
// `NULL::boolean NOT IN ()` == `NULL::boolean NOT IN (SELECT foo FROM empty)` == true
let expr = in_list(lit_bool_null(), vec![], true);
assert_eq!(simplify(expr), lit(true));
// `NULL IN ()` == `NULL IN (SELECT foo FROM empty)` == false
let null_null = || Expr::Literal(ScalarValue::Null, None);
let expr = in_list(null_null(), vec![], false);
assert_eq!(simplify(expr), lit(false));
// `NULL NOT IN ()` == `NULL NOT IN (SELECT foo FROM empty)` == true
let expr = in_list(null_null(), vec![], true);
assert_eq!(simplify(expr), lit(true));
}
#[test]
fn simplify_large_or() {
let expr = (0..5)
.map(|i| col("c1").eq(lit(i)))
.fold(lit(false), |acc, e| acc.or(e));
assert_eq!(
simplify(expr),
in_list(col("c1"), (0..5).map(lit).collect(), false),
);
}
#[test]
fn simplify_expr_bool_and() {
// col & true is always col
assert_eq!(simplify(col("c2").and(lit(true))), col("c2"),);
// col & false is always false
assert_eq!(simplify(col("c2").and(lit(false))), lit(false),);
// true && null is always null
assert_eq!(simplify(lit(true).and(lit_bool_null())), lit_bool_null(),);
// null && true is always null
assert_eq!(simplify(lit_bool_null().and(lit(true))), lit_bool_null(),);
// false && null is always false
assert_eq!(simplify(lit(false).and(lit_bool_null())), lit(false),);
// null && false is always false
assert_eq!(simplify(lit_bool_null().and(lit(false))), lit(false),);
// c1 BETWEEN Int32(0) AND Int32(10) AND Boolean(NULL)
// it can be either NULL or FALSE depending on the value of `c1 BETWEEN Int32(0) AND Int32(10)`
// and the Boolean(NULL) should remain
let expr = col("c1").between(lit(0), lit(10));
let expr = expr.and(lit_bool_null());
let result = simplify(expr);
let expected_expr = and(
and(col("c1").gt_eq(lit(0)), col("c1").lt_eq(lit(10))),
lit_bool_null(),
);
assert_eq!(expected_expr, result);
}
#[test]
fn simplify_expr_between() {
// c2 between 3 and 4 is c2 >= 3 and c2 <= 4
let expr = col("c2").between(lit(3), lit(4));
assert_eq!(
simplify(expr),
and(col("c2").gt_eq(lit(3)), col("c2").lt_eq(lit(4)))
);
// c2 not between 3 and 4 is c2 < 3 or c2 > 4
let expr = col("c2").not_between(lit(3), lit(4));
assert_eq!(
simplify(expr),
or(col("c2").lt(lit(3)), col("c2").gt(lit(4)))
);
}
#[test]
fn test_like_and_ilike() {
let null = lit(ScalarValue::Utf8(None));
// expr [NOT] [I]LIKE NULL
let expr = col("c1").like(null.clone());
assert_eq!(simplify(expr), lit_bool_null());
let expr = col("c1").not_like(null.clone());
assert_eq!(simplify(expr), lit_bool_null());
let expr = col("c1").ilike(null.clone());
assert_eq!(simplify(expr), lit_bool_null());
let expr = col("c1").not_ilike(null.clone());
assert_eq!(simplify(expr), lit_bool_null());
// expr [NOT] [I]LIKE '%'
let expr = col("c1").like(lit("%"));
assert_eq!(simplify(expr), if_not_null(col("c1"), true));
let expr = col("c1").not_like(lit("%"));
assert_eq!(simplify(expr), if_not_null(col("c1"), false));
let expr = col("c1").ilike(lit("%"));
assert_eq!(simplify(expr), if_not_null(col("c1"), true));
let expr = col("c1").not_ilike(lit("%"));
assert_eq!(simplify(expr), if_not_null(col("c1"), false));
// expr [NOT] [I]LIKE '%%'
let expr = col("c1").like(lit("%%"));
assert_eq!(simplify(expr), if_not_null(col("c1"), true));
let expr = col("c1").not_like(lit("%%"));
assert_eq!(simplify(expr), if_not_null(col("c1"), false));
let expr = col("c1").ilike(lit("%%"));
assert_eq!(simplify(expr), if_not_null(col("c1"), true));
let expr = col("c1").not_ilike(lit("%%"));
assert_eq!(simplify(expr), if_not_null(col("c1"), false));
// not_null_expr [NOT] [I]LIKE '%'
let expr = col("c1_non_null").like(lit("%"));
assert_eq!(simplify(expr), lit(true));
let expr = col("c1_non_null").not_like(lit("%"));
assert_eq!(simplify(expr), lit(false));
let expr = col("c1_non_null").ilike(lit("%"));
assert_eq!(simplify(expr), lit(true));
let expr = col("c1_non_null").not_ilike(lit("%"));
assert_eq!(simplify(expr), lit(false));
// not_null_expr [NOT] [I]LIKE '%%'
let expr = col("c1_non_null").like(lit("%%"));
assert_eq!(simplify(expr), lit(true));
let expr = col("c1_non_null").not_like(lit("%%"));
assert_eq!(simplify(expr), lit(false));
let expr = col("c1_non_null").ilike(lit("%%"));
assert_eq!(simplify(expr), lit(true));
let expr = col("c1_non_null").not_ilike(lit("%%"));
assert_eq!(simplify(expr), lit(false));
// null_constant [NOT] [I]LIKE '%'
let expr = null.clone().like(lit("%"));
assert_eq!(simplify(expr), lit_bool_null());
let expr = null.clone().not_like(lit("%"));
assert_eq!(simplify(expr), lit_bool_null());
let expr = null.clone().ilike(lit("%"));
assert_eq!(simplify(expr), lit_bool_null());
let expr = null.clone().not_ilike(lit("%"));
assert_eq!(simplify(expr), lit_bool_null());
// null_constant [NOT] [I]LIKE '%%'
let expr = null.clone().like(lit("%%"));
assert_eq!(simplify(expr), lit_bool_null());
let expr = null.clone().not_like(lit("%%"));
assert_eq!(simplify(expr), lit_bool_null());
let expr = null.clone().ilike(lit("%%"));
assert_eq!(simplify(expr), lit_bool_null());
let expr = null.clone().not_ilike(lit("%%"));
assert_eq!(simplify(expr), lit_bool_null());
// null_constant [NOT] [I]LIKE 'a%'
let expr = null.clone().like(lit("a%"));
assert_eq!(simplify(expr), lit_bool_null());
let expr = null.clone().not_like(lit("a%"));
assert_eq!(simplify(expr), lit_bool_null());
let expr = null.clone().ilike(lit("a%"));
assert_eq!(simplify(expr), lit_bool_null());
let expr = null.clone().not_ilike(lit("a%"));
assert_eq!(simplify(expr), lit_bool_null());
// expr [NOT] [I]LIKE with pattern without wildcards
let expr = col("c1").like(lit("a"));
assert_eq!(simplify(expr), col("c1").eq(lit("a")));
let expr = col("c1").not_like(lit("a"));
assert_eq!(simplify(expr), col("c1").not_eq(lit("a")));
let expr = col("c1").like(lit("a_"));
assert_eq!(simplify(expr), col("c1").like(lit("a_")));
let expr = col("c1").not_like(lit("a_"));
assert_eq!(simplify(expr), col("c1").not_like(lit("a_")));
let expr = col("c1").ilike(lit("a"));
assert_eq!(simplify(expr), col("c1").ilike(lit("a")));
let expr = col("c1").not_ilike(lit("a"));
assert_eq!(simplify(expr), col("c1").not_ilike(lit("a")));
}
#[test]
fn test_simplify_with_guarantee() {
// (c3 >= 3) AND (c4 + 2 < 10 OR (c1 NOT IN ("a", "b")))
let expr_x = col("c3").gt(lit(3_i64));
let expr_y = (col("c4") + lit(2_u32)).lt(lit(10_u32));
let expr_z = col("c1").in_list(vec![lit("a"), lit("b")], true);
let expr = expr_x.clone().and(expr_y.or(expr_z));
// All guaranteed null
let guarantees = vec![
(col("c3"), NullableInterval::from(ScalarValue::Int64(None))),
(col("c4"), NullableInterval::from(ScalarValue::UInt32(None))),
(col("c1"), NullableInterval::from(ScalarValue::Utf8(None))),
];
let output = simplify_with_guarantee(expr.clone(), guarantees);
assert_eq!(output, lit_bool_null());
// All guaranteed false
let guarantees = vec![
(
col("c3"),
NullableInterval::NotNull {
values: Interval::make(Some(0_i64), Some(2_i64)).unwrap(),
},
),
(
col("c4"),
NullableInterval::from(ScalarValue::UInt32(Some(9))),
),
(col("c1"), NullableInterval::from(ScalarValue::from("a"))),
];
let output = simplify_with_guarantee(expr.clone(), guarantees);
assert_eq!(output, lit(false));
// Guaranteed false or null -> no change.
let guarantees = vec![
(
col("c3"),
NullableInterval::MaybeNull {
values: Interval::make(Some(0_i64), Some(2_i64)).unwrap(),
},
),
(
col("c4"),
NullableInterval::MaybeNull {
values: Interval::make(Some(9_u32), Some(9_u32)).unwrap(),
},
),
(
col("c1"),
NullableInterval::NotNull {
values: Interval::try_new(
ScalarValue::from("d"),
ScalarValue::from("f"),
)
.unwrap(),
},
),
];
let output = simplify_with_guarantee(expr.clone(), guarantees);
assert_eq!(&output, &expr_x);
// Sufficient true guarantees
let guarantees = vec![
(
col("c3"),
NullableInterval::from(ScalarValue::Int64(Some(9))),
),
(
col("c4"),
NullableInterval::from(ScalarValue::UInt32(Some(3))),
),
];
let output = simplify_with_guarantee(expr.clone(), guarantees);
assert_eq!(output, lit(true));
// Only partially simplify
let guarantees = vec![(
col("c4"),
NullableInterval::from(ScalarValue::UInt32(Some(3))),
)];
let output = simplify_with_guarantee(expr, guarantees);
assert_eq!(&output, &expr_x);
}
#[test]
fn test_expression_partial_simplify_1() {
// (1 + 2) + (4 / 0) -> 3 + (4 / 0)
let expr = (lit(1) + lit(2)) + (lit(4) / lit(0));
let expected = (lit(3)) + (lit(4) / lit(0));
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_expression_partial_simplify_2() {
// (1 > 2) and (4 / 0) -> false
let expr = (lit(1).gt(lit(2))).and(lit(4) / lit(0));
let expected = lit(false);
assert_eq!(simplify(expr), expected);
}
#[test]
fn test_simplify_cycles() {
// TRUE
let expr = lit(true);
let expected = lit(true);
let (expr, num_iter) = simplify_with_cycle_count(expr);
assert_eq!(expr, expected);
assert_eq!(num_iter, 1);
// (true != NULL) OR (5 > 10)
let expr = lit(true).not_eq(lit_bool_null()).or(lit(5).gt(lit(10)));
let expected = lit_bool_null();
let (expr, num_iter) = simplify_with_cycle_count(expr);
assert_eq!(expr, expected);
assert_eq!(num_iter, 2);
// NOTE: this currently does not simplify
// (((c4 - 10) + 10) *100) / 100
let expr = (((col("c4") - lit(10)) + lit(10)) * lit(100)) / lit(100);
let expected = expr.clone();
let (expr, num_iter) = simplify_with_cycle_count(expr);
assert_eq!(expr, expected);
assert_eq!(num_iter, 1);
// ((c4<1 or c3<2) and c3_non_null<3) and false
let expr = col("c4")
.lt(lit(1))
.or(col("c3").lt(lit(2)))
.and(col("c3_non_null").lt(lit(3)))
.and(lit(false));
let expected = lit(false);
let (expr, num_iter) = simplify_with_cycle_count(expr);
assert_eq!(expr, expected);
assert_eq!(num_iter, 2);
}
fn boolean_test_schema() -> DFSchemaRef {
static BOOLEAN_TEST_SCHEMA: LazyLock<DFSchemaRef> = LazyLock::new(|| {
Schema::new(vec![
Field::new("A", DataType::Boolean, false),
Field::new("B", DataType::Boolean, false),
Field::new("C", DataType::Boolean, false),
Field::new("D", DataType::Boolean, false),
])
.to_dfschema_ref()
.unwrap()
});
Arc::clone(&BOOLEAN_TEST_SCHEMA)
}
#[test]
fn simplify_common_factor_conjunction_in_disjunction() {
let props = ExecutionProps::new();
let schema = boolean_test_schema();
let simplifier =
ExprSimplifier::new(SimplifyContext::new(&props).with_schema(schema));
let a = || col("A");
let b = || col("B");
let c = || col("C");
let d = || col("D");
// (A AND B) OR (A AND C) -> A AND (B OR C)
let expr = a().and(b()).or(a().and(c()));
let expected = a().and(b().or(c()));
assert_eq!(expected, simplifier.simplify(expr).unwrap());
// (A AND B) OR (A AND C) OR (A AND D) -> A AND (B OR C OR D)
let expr = a().and(b()).or(a().and(c())).or(a().and(d()));
let expected = a().and(b().or(c()).or(d()));
assert_eq!(expected, simplifier.simplify(expr).unwrap());
// A OR (B AND C AND A) -> A
let expr = a().or(b().and(c().and(a())));
let expected = a();
assert_eq!(expected, simplifier.simplify(expr).unwrap());
}
#[test]
fn test_simplify_udaf() {
let udaf = AggregateUDF::new_from_impl(SimplifyMockUdaf::new_with_simplify());
let aggregate_function_expr =
Expr::AggregateFunction(expr::AggregateFunction::new_udf(
udaf.into(),
vec![],
false,
None,
vec![],
None,
));
let expected = col("result_column");
assert_eq!(simplify(aggregate_function_expr), expected);
let udaf = AggregateUDF::new_from_impl(SimplifyMockUdaf::new_without_simplify());
let aggregate_function_expr =
Expr::AggregateFunction(expr::AggregateFunction::new_udf(
udaf.into(),
vec![],
false,
None,
vec![],
None,
));
let expected = aggregate_function_expr.clone();
assert_eq!(simplify(aggregate_function_expr), expected);
}
/// A Mock UDAF which defines `simplify` to be used in tests
/// related to UDAF simplification
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct SimplifyMockUdaf {
simplify: bool,
}
impl SimplifyMockUdaf {
/// make simplify method return new expression
fn new_with_simplify() -> Self {
Self { simplify: true }
}
/// make simplify method return no change
fn new_without_simplify() -> Self {
Self { simplify: false }
}
}
impl AggregateUDFImpl for SimplifyMockUdaf {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
"mock_simplify"
}
fn signature(&self) -> &Signature {
unimplemented!()
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
unimplemented!("not needed for tests")
}
fn accumulator(
&self,
_acc_args: AccumulatorArgs,
) -> Result<Box<dyn Accumulator>> {
unimplemented!("not needed for tests")
}
fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
unimplemented!("not needed for testing")
}
fn create_groups_accumulator(
&self,
_args: AccumulatorArgs,
) -> Result<Box<dyn GroupsAccumulator>> {
unimplemented!("not needed for testing")
}
fn simplify(&self) -> Option<AggregateFunctionSimplification> {
if self.simplify {
Some(Box::new(|_, _| Ok(col("result_column"))))
} else {
None
}
}
}
#[test]
fn test_simplify_udwf() {
let udwf = WindowFunctionDefinition::WindowUDF(
WindowUDF::new_from_impl(SimplifyMockUdwf::new_with_simplify()).into(),
);
let window_function_expr = Expr::from(WindowFunction::new(udwf, vec![]));
let expected = col("result_column");
assert_eq!(simplify(window_function_expr), expected);
let udwf = WindowFunctionDefinition::WindowUDF(
WindowUDF::new_from_impl(SimplifyMockUdwf::new_without_simplify()).into(),
);
let window_function_expr = Expr::from(WindowFunction::new(udwf, vec![]));
let expected = window_function_expr.clone();
assert_eq!(simplify(window_function_expr), expected);
}
/// A Mock UDWF which defines `simplify` to be used in tests
/// related to UDWF simplification
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct SimplifyMockUdwf {
simplify: bool,
}
impl SimplifyMockUdwf {
/// make simplify method return new expression
fn new_with_simplify() -> Self {
Self { simplify: true }
}
/// make simplify method return no change
fn new_without_simplify() -> Self {
Self { simplify: false }
}
}
impl WindowUDFImpl for SimplifyMockUdwf {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
"mock_simplify"
}
fn signature(&self) -> &Signature {
unimplemented!()
}
fn simplify(&self) -> Option<WindowFunctionSimplification> {
if self.simplify {
Some(Box::new(|_, _| Ok(col("result_column"))))
} else {
None
}
}
fn partition_evaluator(
&self,
_partition_evaluator_args: PartitionEvaluatorArgs,
) -> Result<Box<dyn PartitionEvaluator>> {
unimplemented!("not needed for tests")
}
fn field(&self, _field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
unimplemented!("not needed for tests")
}
fn limit_effect(&self, _args: &[Arc<dyn PhysicalExpr>]) -> LimitEffect {
LimitEffect::Unknown
}
}
#[derive(Debug, PartialEq, Eq, Hash)]
struct VolatileUdf {
signature: Signature,
}
impl VolatileUdf {
pub fn new() -> Self {
Self {
signature: Signature::exact(vec![], Volatility::Volatile),
}
}
}
impl ScalarUDFImpl for VolatileUdf {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
"VolatileUdf"
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(DataType::Int16)
}
fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
panic!("dummy - not implemented")
}
}
#[test]
fn test_optimize_volatile_conditions() {
let fun = Arc::new(ScalarUDF::new_from_impl(VolatileUdf::new()));
let rand = Expr::ScalarFunction(ScalarFunction::new_udf(fun, vec![]));
{
let expr = rand
.clone()
.eq(lit(0))
.or(col("column1").eq(lit(2)).and(rand.clone().eq(lit(0))));
assert_eq!(simplify(expr.clone()), expr);
}
{
let expr = col("column1")
.eq(lit(2))
.or(col("column1").eq(lit(2)).and(rand.clone().eq(lit(0))));
assert_eq!(simplify(expr), col("column1").eq(lit(2)));
}
{
let expr = (col("column1").eq(lit(2)).and(rand.clone().eq(lit(0)))).or(col(
"column1",
)
.eq(lit(2))
.and(rand.clone().eq(lit(0))));
assert_eq!(
simplify(expr),
col("column1")
.eq(lit(2))
.and((rand.clone().eq(lit(0))).or(rand.clone().eq(lit(0))))
);
}
}
#[test]
fn simplify_fixed_size_binary_eq_lit() {
let bytes = [1u8, 2, 3].as_slice();
// The expression starts simple.
let expr = col("c5").eq(lit(bytes));
// The type coercer introduces a cast.
let coerced = coerce(expr.clone());
let schema = expr_test_schema();
assert_eq!(
coerced,
col("c5")
.cast_to(&DataType::Binary, schema.as_ref())
.unwrap()
.eq(lit(bytes))
);
// The simplifier removes the cast.
assert_eq!(
simplify(coerced),
col("c5").eq(Expr::Literal(
ScalarValue::FixedSizeBinary(3, Some(bytes.to_vec()),),
None
))
);
}
#[test]
fn simplify_cast_literal() {
// Test that CAST(literal) expressions are evaluated at plan time
// CAST(123 AS Int64) should become 123i64
let expr = Expr::Cast(Cast::new(Box::new(lit(123i32)), DataType::Int64));
let expected = lit(123i64);
assert_eq!(simplify(expr), expected);
// CAST(1761630189642 AS Timestamp(Nanosecond, Some("+00:00")))
// Integer to timestamp cast
let expr = Expr::Cast(Cast::new(
Box::new(lit(1761630189642i64)),
DataType::Timestamp(
arrow::datatypes::TimeUnit::Nanosecond,
Some("+00:00".into()),
),
));
// Should evaluate to a timestamp literal
let result = simplify(expr);
match result {
Expr::Literal(ScalarValue::TimestampNanosecond(Some(val), tz), _) => {
assert_eq!(val, 1761630189642i64);
assert_eq!(tz.as_deref(), Some("+00:00"));
}
other => panic!("Expected TimestampNanosecond literal, got: {other:?}"),
}
// Test CAST of invalid string to timestamp - should return an error at plan time
// This represents the case from the issue: CAST(Utf8("1761630189642") AS Timestamp)
// "1761630189642" is NOT a valid timestamp string format
let expr = Expr::Cast(Cast::new(
Box::new(lit("1761630189642")),
DataType::Timestamp(
arrow::datatypes::TimeUnit::Nanosecond,
Some("+00:00".into()),
),
));
// The simplification should now fail with an error at plan time
let schema = test_schema();
let props = ExecutionProps::new();
let simplifier =
ExprSimplifier::new(SimplifyContext::new(&props).with_schema(schema));
let result = simplifier.simplify(expr);
assert!(result.is_err(), "Expected error for invalid cast");
let err_msg = result.unwrap_err().to_string();
assert_contains!(err_msg, "Error parsing timestamp");
}
fn if_not_null(expr: Expr, then: bool) -> Expr {
Expr::Case(Case {
expr: Some(expr.is_not_null().into()),
when_then_expr: vec![(lit(true).into(), lit(then).into())],
else_expr: None,
})
}
}