`SessionContext::create_physical_expr` does not unwrap casts (and thus is not always optimal)

[alamb]

Is your feature request related to a problem or challenge?

As described by @ion-elgreco in #14944

Given a dataset with an Int64 column named month, when a predicate such as the following is created

month_id = '202502'

In sql / dataframe queries, this will be simplified to the following. Note the constant is cast and the column is not cast:

month_id = cast('202502', 'Int64')

However, when using SessionContext::create_physical_expr to create a physical expression directly, as is done in delta.rs and other systems like LanceDB, the expression looks like this (cast on the column)

cast(month_id, 'Int64') = '202502'

This is bad for two reasons:

1. PruningPredicate can't handle this type of expression (and thus it can't be used to prune Parquet row groups)
2. Evaluating this filter is substantially slower as it has to apply a transformation to all values of month_id before it can evaluate the filter. And furthermore it does slow string comparison compared to faster int63 comparison

The reason this happens is that the conversion from cast(month_id, 'Int64') = '202502' to month_id = Cast('202502', Int64) happens in the Analyzer, specifically here:

datafusion/datafusion/optimizer/src/unwrap_cast_in_comparison.rs

Lines 39 to 77 in 2fcab2e

	/// [`UnwrapCastInComparison`] attempts to remove casts from
	/// comparisons to literals ([`ScalarValue`]s) by applying the casts
	/// to the literals if possible. It is inspired by the optimizer rule
	/// `UnwrapCastInBinaryComparison` of Spark.
	///
	/// Removing casts often improves performance because:
	/// 1. The cast is done once (to the literal) rather than to every value
	/// 2. Can enable other optimizations such as predicate pushdown that
	/// don't support casting
	///
	/// The rule is applied to expressions of the following forms:
	///
	/// 1. `cast(left_expr as data_type) comparison_op literal_expr`
	/// 2. `literal_expr comparison_op cast(left_expr as data_type)`
	/// 3. `cast(literal_expr) IN (expr1, expr2, ...)`
	/// 4. `literal_expr IN (cast(expr1) , cast(expr2), ...)`
	///
	/// If the expression matches one of the forms above, the rule will
	/// ensure the value of `literal` is in range(min, max) of the
	/// expr's data_type, and if the scalar is within range, the literal
	/// will be casted to the data type of expr on the other side, and the
	/// cast will be removed from the other side.
	///
	/// # Example
	///
	/// If the DataType of c1 is INT32. Given the filter
	///
	/// ```text
	/// Filter: cast(c1 as INT64) > INT64(10)`
	/// ```
	///
	/// This rule will remove the cast and rewrite the expression to:
	///
	/// ```text
	/// Filter: c1 > INT32(10)
	/// ```
	///
	#[derive(Default, Debug)]
	pub struct UnwrapCastInComparison {}

However, this pass is not run as part of SessionContext::create_physical_expr

Describe the solution you'd like

I would like the expressions crated by SessionContext::create_physical_expr to have had their casts unwrapped as well

Describe alternatives you've considered

THe ideal solution in my mind is to remove the entire Analyzer pass and instead do the unwrap in comparisons as part of the expression simplification

https://docs.rs/datafusion/latest/datafusion/optimizer/simplify_expressions/expr_simplifier/struct.ExprSimplifier.html

Additional context

• Partition pruning does not work when values in predicate do not match column datatype delta-io/delta-rs#3278
• Should PruningPredicate coerce? #14944

[jayzhan211]

THe ideal solution in my mind is to remove the entire Analyzer pass and instead do the unwrap in comparisons as part of the expression simplification

Why don't the user apply UnwrapCastInComparison follow on SimplifyExpressions like what DataFusion did?

[alamb]

THe ideal solution in my mind is to remove the entire Analyzer pass and instead do the unwrap in comparisons as part of the expression simplification

Why don't the user apply UnwrapCastInComparison follow on SimplifyExpressions like what DataFusion did?

I don't quite follow what you are suggesting. I am not sure the code to apply unwrap in comparsion for exprs is publically visible, though I haven't checked 🤔

[jayzhan211]

impl Optimizer {
    /// Create a new optimizer using the recommended list of rules
    pub fn new() -> Self {
        let rules: Vec<Arc<dyn OptimizerRule + Sync + Send>> = vec![
            Arc::new(EliminateNestedUnion::new()),
            Arc::new(SimplifyExpressions::new()),
            Arc::new(UnwrapCastInComparison::new()),
            Arc::new(ReplaceDistinctWithAggregate::new()),
            Arc::new(EliminateJoin::new()),
            Arc::new(DecorrelatePredicateSubquery::new()),
            Arc::new(ScalarSubqueryToJoin::new()),
            Arc::new(ExtractEquijoinPredicate::new()),
            Arc::new(EliminateDuplicatedExpr::new()),
            Arc::new(EliminateFilter::new()),
            Arc::new(EliminateCrossJoin::new()),
            Arc::new(CommonSubexprEliminate::new()),
            Arc::new(EliminateLimit::new()),
            Arc::new(PropagateEmptyRelation::new()),
            // Must be after PropagateEmptyRelation
            Arc::new(EliminateOneUnion::new()),
            Arc::new(FilterNullJoinKeys::default()),
            Arc::new(EliminateOuterJoin::new()),
            // Filters can't be pushed down past Limits, we should do PushDownFilter after PushDownLimit
            Arc::new(PushDownLimit::new()),
            Arc::new(PushDownFilter::new()),
            Arc::new(SingleDistinctToGroupBy::new()),
            // The previous optimizations added expressions and projections,
            // that might benefit from the following rules
            Arc::new(SimplifyExpressions::new()),
            Arc::new(UnwrapCastInComparison::new()),
            Arc::new(CommonSubexprEliminate::new()),
            Arc::new(EliminateGroupByConstant::new()),
            Arc::new(OptimizeProjections::new()),
        ];

        Self::with_rules(rules)
    }

    /// Create a new optimizer with the given rules
    pub fn with_rules(rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>) -> Self {
        Self { rules }
    }
}

let rules = vec![SimplifyExpressions, UnwrapCastInComparison, 3rd parties rule]
let optimizer = Optimizer::with_rules(rules);

Not sure how they determine the rules they need, but I guess they can add their rule and datafusion builtin rule for their optimizer?

If they only apply SimplifyExpressions then it results like what mentioned above, they need UnwrapCastInComparison as well

[ion-elgreco]

@jayzhan211 the optimizer is used on the plan level?

[jayzhan211]

@jayzhan211 the optimizer is used on the plan level?

Yes, in Logical Plan level, it is called after the LogicalPlan created.

[ion-elgreco]

@jayzhan211 the optimizer is used on the plan level?

Yes, in Logical Plan level, it is called after the LogicalPlan created.

I see but that's not how we use the expression in our case :)

We get an expression predicate, we process it to add more info based on the source statistics when possible, then we simplify and then we prune our TableSnapshot statistics with the PruningPredicate.

Here there is no logical plan optimization involved

[jayzhan211]

What do you mean "simplify", is it SimplifyExpressions or other logic. If it is SimplifyExpressions, then you have logical plan optimization involved, if not and you go directly to create_physical_expr then I think this is the issue similar to Comet where they don't have coerced types for create_physical_expr and they need to done the coercion part themselves.

Alternative idea is that we need to introduce another create_physical_expr_from_non_resolved that allows inputs that are "not resolved" (so we need to do coercion or type related things) and keep current create_physical_expr for inputs that are "resolved" and used for DataFusion or other projects that has DataFusion logical plan optimization. But since create_physical_expr_from_non_resolved will not be used in DataFusion, I'm not sure what is the best way to maintain this

[ion-elgreco]

@jayzhan211 its using the expression simplifier:

        let logical_filter = self.filter.map(|expr| {
            // Simplify the expression first
            let props = ExecutionProps::new();
            let simplify_context =
                SimplifyContext::new(&props).with_schema(df_schema.clone().into());
            let simplifier = ExprSimplifier::new(simplify_context).with_max_cycles(10);
            let simplified = simplifier.simplify(expr).unwrap();

            context
                .create_physical_expr(simplified, &df_schema)
                .unwrap()
        });

[jayzhan211]

I see. We can try moving UnwrapCastInComparison into Simplifier large pattern matching rule