alamb commented on code in PR #16530:
URL: https://github.com/apache/datafusion/pull/16530#discussion_r2176033909
##########
datafusion/datasource-parquet/src/opener.rs:
##########
@@ -233,7 +234,13 @@ impl FileOpener for ParquetOpener {
)
.rewrite(p)
.map_err(ArrowError::from)
+ .map(|p| {
+ PhysicalExprSimplifier::new(&physical_file_schema)
Review Comment:
I think it is worth worth mentioning here that we are optimizing the
physical expression here as it may have been altered due to differences in
table and file schema
Perhaps something like this
```suggestion
// After rewriting to the file schema, further
simplifications may
// be possible.
PhysicalExprSimplifier::new(&physical_file_schema)
```
##########
datafusion/expr-common/src/casts.rs:
##########
@@ -0,0 +1,1227 @@
+// 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.
+
+//! Utilities for casting scalar literals to different data types
Review Comment:
👍
##########
datafusion/physical-expr/src/simplifier/unwrap_cast.rs:
##########
@@ -0,0 +1,621 @@
+// 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.
+
+//! Unwrap casts in binary comparisons for physical expressions
+//!
+//! This module provides optimization for physical expressions similar to the
logical
+//! optimizer's unwrap_cast module. It attempts to remove casts from
comparisons to
+//! literals by applying the casts to the literals if possible.
+//!
+//! The optimization improves performance by:
+//! 1. Reducing runtime cast operations on column data
+//! 2. Enabling better predicate pushdown opportunities
+//! 3. Optimizing filter expressions in physical plans
+//!
+//! # Example
+//!
+//! Physical expression: `cast(column as INT64) > INT64(10)`
+//! Optimized to: `column > INT32(10)` (assuming column is INT32)
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Schema};
+use datafusion_common::{
+ tree_node::{Transformed, TreeNode},
+ Result, ScalarValue,
+};
+use datafusion_expr::Operator;
+use datafusion_expr_common::casts::try_cast_literal_to_type;
+
+use crate::expressions::{lit, BinaryExpr, CastExpr, Literal, TryCastExpr};
+use crate::PhysicalExpr;
+
+/// Attempts to unwrap casts in comparison expressions.
+pub(crate) fn unwrap_cast_in_comparison(
+ expr: Arc<dyn PhysicalExpr>,
+ schema: &Schema,
+) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ expr.transform_down(|e| {
+ if let Some(binary) = e.as_any().downcast_ref::<BinaryExpr>() {
+ if let Some(unwrapped) = try_unwrap_cast_binary(binary, schema)? {
+ return Ok(Transformed::yes(unwrapped));
+ }
+ }
+ Ok(Transformed::no(e))
+ })
+}
+
+/// Try to unwrap casts in binary expressions
+fn try_unwrap_cast_binary(
+ binary: &BinaryExpr,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Case 1: cast(left_expr) op literal
+ if let (Some((inner_expr, _cast_type)), Some(literal)) = (
+ extract_cast_info(binary.left()),
+ binary.right().as_any().downcast_ref::<Literal>(),
+ ) {
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ *binary.op(),
+ schema,
+ )? {
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ // Case 2: literal op cast(right_expr)
+ if let (Some(literal), Some((inner_expr, _cast_type))) = (
+ binary.left().as_any().downcast_ref::<Literal>(),
+ extract_cast_info(binary.right()),
+ ) {
+ // For literal op cast(expr), we need to swap the operator
+ let swapped_op = swap_operator(*binary.op());
Review Comment:
If we can't swap the operator I think the rewrite should stop (not just
unwrap the comparison)
##########
datafusion/physical-expr/src/simplifier/unwrap_cast.rs:
##########
@@ -0,0 +1,621 @@
+// 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.
+
+//! Unwrap casts in binary comparisons for physical expressions
+//!
+//! This module provides optimization for physical expressions similar to the
logical
+//! optimizer's unwrap_cast module. It attempts to remove casts from
comparisons to
+//! literals by applying the casts to the literals if possible.
+//!
+//! The optimization improves performance by:
+//! 1. Reducing runtime cast operations on column data
+//! 2. Enabling better predicate pushdown opportunities
+//! 3. Optimizing filter expressions in physical plans
+//!
+//! # Example
+//!
+//! Physical expression: `cast(column as INT64) > INT64(10)`
+//! Optimized to: `column > INT32(10)` (assuming column is INT32)
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Schema};
+use datafusion_common::{
+ tree_node::{Transformed, TreeNode},
+ Result, ScalarValue,
+};
+use datafusion_expr::Operator;
+use datafusion_expr_common::casts::try_cast_literal_to_type;
+
+use crate::expressions::{lit, BinaryExpr, CastExpr, Literal, TryCastExpr};
+use crate::PhysicalExpr;
+
+/// Attempts to unwrap casts in comparison expressions.
+pub(crate) fn unwrap_cast_in_comparison(
+ expr: Arc<dyn PhysicalExpr>,
+ schema: &Schema,
+) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ expr.transform_down(|e| {
+ if let Some(binary) = e.as_any().downcast_ref::<BinaryExpr>() {
+ if let Some(unwrapped) = try_unwrap_cast_binary(binary, schema)? {
+ return Ok(Transformed::yes(unwrapped));
+ }
+ }
+ Ok(Transformed::no(e))
+ })
+}
+
+/// Try to unwrap casts in binary expressions
+fn try_unwrap_cast_binary(
+ binary: &BinaryExpr,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Case 1: cast(left_expr) op literal
+ if let (Some((inner_expr, _cast_type)), Some(literal)) = (
+ extract_cast_info(binary.left()),
+ binary.right().as_any().downcast_ref::<Literal>(),
+ ) {
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ *binary.op(),
+ schema,
+ )? {
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ // Case 2: literal op cast(right_expr)
+ if let (Some(literal), Some((inner_expr, _cast_type))) = (
+ binary.left().as_any().downcast_ref::<Literal>(),
+ extract_cast_info(binary.right()),
+ ) {
+ // For literal op cast(expr), we need to swap the operator
+ let swapped_op = swap_operator(*binary.op());
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ swapped_op,
+ schema,
+ )? {
+ // unwrapped is already the correct expression: inner_expr
swapped_op casted_literal
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ Ok(None)
+}
+
+/// Extract cast information from a physical expression
+fn extract_cast_info(
+ expr: &Arc<dyn PhysicalExpr>,
+) -> Option<(&Arc<dyn PhysicalExpr>, &DataType)> {
+ if let Some(cast) = expr.as_any().downcast_ref::<CastExpr>() {
+ Some((cast.expr(), cast.cast_type()))
+ } else if let Some(try_cast) = expr.as_any().downcast_ref::<TryCastExpr>()
{
+ Some((try_cast.expr(), try_cast.cast_type()))
+ } else {
+ None
+ }
+}
+
+/// Try to unwrap a cast in comparison by moving the cast to the literal
+fn try_unwrap_cast_comparison(
+ inner_expr: Arc<dyn PhysicalExpr>,
+ literal_value: &ScalarValue,
+ op: Operator,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Get the data type of the inner expression
+ let inner_type = inner_expr.data_type(schema)?;
+
+ // Try to cast the literal to the inner expression's type
+ if let Some(casted_literal) = try_cast_literal_to_type(literal_value,
&inner_type) {
+ let literal_expr = lit(casted_literal);
+ let binary_expr = BinaryExpr::new(inner_expr, op, literal_expr);
+ return Ok(Some(Arc::new(binary_expr)));
+ }
+
+ Ok(None)
+}
+
+/// Swap comparison operators for right-side cast unwrapping
+fn swap_operator(op: Operator) -> Operator {
+ match op {
+ Operator::Lt => Operator::Gt,
+ Operator::LtEq => Operator::GtEq,
+ Operator::Gt => Operator::Lt,
+ Operator::GtEq => Operator::LtEq,
+ // Symmetric operators remain the same
+ Operator::Eq | Operator::NotEq => op,
+ // Other operators are not swappable
+ _ => op,
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::expressions::{col, lit};
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::Operator;
+
+ /// Check if an expression is a cast expression
+ fn is_cast_expr(expr: &Arc<dyn PhysicalExpr>) -> bool {
+ expr.as_any().downcast_ref::<CastExpr>().is_some()
+ || expr.as_any().downcast_ref::<TryCastExpr>().is_some()
+ }
+
+ /// Check if a binary expression is suitable for cast unwrapping
+ fn is_binary_expr_with_cast_and_literal(binary: &BinaryExpr) -> bool {
+ // Check if left is cast and right is literal
+ let left_cast_right_literal = is_cast_expr(binary.left())
+ && binary.right().as_any().downcast_ref::<Literal>().is_some();
+
+ // Check if left is literal and right is cast
+ let left_literal_right_cast =
+ binary.left().as_any().downcast_ref::<Literal>().is_some()
+ && is_cast_expr(binary.right());
+
+ left_cast_right_literal || left_literal_right_cast
+ }
+
+ fn test_schema() -> Schema {
+ Schema::new(vec![
+ Field::new("c1", DataType::Int32, false),
+ Field::new("c2", DataType::Int64, false),
+ Field::new("c3", DataType::Utf8, false),
+ ])
+ }
+
+ #[test]
+ fn test_unwrap_cast_in_binary_comparison() {
+ let schema = test_schema();
+
+ // Create: cast(c1 as INT64) > INT64(10)
+ let column_expr = col("c1", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(column_expr, DataType::Int64,
None));
+ let literal_expr = lit(ScalarValue::Int64(Some(10)));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(cast_expr, Operator::Gt, literal_expr));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap();
Review Comment:
As a a minor comment it would be nice if we could reduce the boiler plate in
these tests -- specifically, you could perhaps factor out the unwrapping and
verifying that it was transformed
##########
datafusion/physical-expr/src/simplifier/unwrap_cast.rs:
##########
@@ -0,0 +1,621 @@
+// 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.
+
+//! Unwrap casts in binary comparisons for physical expressions
+//!
+//! This module provides optimization for physical expressions similar to the
logical
+//! optimizer's unwrap_cast module. It attempts to remove casts from
comparisons to
+//! literals by applying the casts to the literals if possible.
+//!
+//! The optimization improves performance by:
+//! 1. Reducing runtime cast operations on column data
+//! 2. Enabling better predicate pushdown opportunities
+//! 3. Optimizing filter expressions in physical plans
+//!
+//! # Example
+//!
+//! Physical expression: `cast(column as INT64) > INT64(10)`
+//! Optimized to: `column > INT32(10)` (assuming column is INT32)
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Schema};
+use datafusion_common::{
+ tree_node::{Transformed, TreeNode},
+ Result, ScalarValue,
+};
+use datafusion_expr::Operator;
+use datafusion_expr_common::casts::try_cast_literal_to_type;
+
+use crate::expressions::{lit, BinaryExpr, CastExpr, Literal, TryCastExpr};
+use crate::PhysicalExpr;
+
+/// Attempts to unwrap casts in comparison expressions.
+pub(crate) fn unwrap_cast_in_comparison(
+ expr: Arc<dyn PhysicalExpr>,
+ schema: &Schema,
+) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ expr.transform_down(|e| {
+ if let Some(binary) = e.as_any().downcast_ref::<BinaryExpr>() {
+ if let Some(unwrapped) = try_unwrap_cast_binary(binary, schema)? {
+ return Ok(Transformed::yes(unwrapped));
+ }
+ }
+ Ok(Transformed::no(e))
+ })
+}
+
+/// Try to unwrap casts in binary expressions
+fn try_unwrap_cast_binary(
+ binary: &BinaryExpr,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Case 1: cast(left_expr) op literal
+ if let (Some((inner_expr, _cast_type)), Some(literal)) = (
+ extract_cast_info(binary.left()),
+ binary.right().as_any().downcast_ref::<Literal>(),
+ ) {
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ *binary.op(),
+ schema,
+ )? {
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ // Case 2: literal op cast(right_expr)
+ if let (Some(literal), Some((inner_expr, _cast_type))) = (
+ binary.left().as_any().downcast_ref::<Literal>(),
+ extract_cast_info(binary.right()),
+ ) {
+ // For literal op cast(expr), we need to swap the operator
+ let swapped_op = swap_operator(*binary.op());
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ swapped_op,
+ schema,
+ )? {
+ // unwrapped is already the correct expression: inner_expr
swapped_op casted_literal
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ Ok(None)
+}
+
+/// Extract cast information from a physical expression
Review Comment:
Can we please update this to explain what info is extracted?
it seems like it looks for `CAST(expr, datatype)` and returns the `expr` and
`datatype` if if finds a match
##########
datafusion/physical-expr/src/simplifier/unwrap_cast.rs:
##########
@@ -0,0 +1,621 @@
+// 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.
+
+//! Unwrap casts in binary comparisons for physical expressions
+//!
+//! This module provides optimization for physical expressions similar to the
logical
+//! optimizer's unwrap_cast module. It attempts to remove casts from
comparisons to
+//! literals by applying the casts to the literals if possible.
+//!
+//! The optimization improves performance by:
+//! 1. Reducing runtime cast operations on column data
+//! 2. Enabling better predicate pushdown opportunities
+//! 3. Optimizing filter expressions in physical plans
+//!
+//! # Example
+//!
+//! Physical expression: `cast(column as INT64) > INT64(10)`
+//! Optimized to: `column > INT32(10)` (assuming column is INT32)
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Schema};
+use datafusion_common::{
+ tree_node::{Transformed, TreeNode},
+ Result, ScalarValue,
+};
+use datafusion_expr::Operator;
+use datafusion_expr_common::casts::try_cast_literal_to_type;
+
+use crate::expressions::{lit, BinaryExpr, CastExpr, Literal, TryCastExpr};
+use crate::PhysicalExpr;
+
+/// Attempts to unwrap casts in comparison expressions.
+pub(crate) fn unwrap_cast_in_comparison(
+ expr: Arc<dyn PhysicalExpr>,
+ schema: &Schema,
+) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ expr.transform_down(|e| {
+ if let Some(binary) = e.as_any().downcast_ref::<BinaryExpr>() {
+ if let Some(unwrapped) = try_unwrap_cast_binary(binary, schema)? {
+ return Ok(Transformed::yes(unwrapped));
+ }
+ }
+ Ok(Transformed::no(e))
+ })
+}
+
+/// Try to unwrap casts in binary expressions
+fn try_unwrap_cast_binary(
+ binary: &BinaryExpr,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Case 1: cast(left_expr) op literal
+ if let (Some((inner_expr, _cast_type)), Some(literal)) = (
+ extract_cast_info(binary.left()),
+ binary.right().as_any().downcast_ref::<Literal>(),
+ ) {
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ *binary.op(),
+ schema,
+ )? {
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ // Case 2: literal op cast(right_expr)
+ if let (Some(literal), Some((inner_expr, _cast_type))) = (
+ binary.left().as_any().downcast_ref::<Literal>(),
+ extract_cast_info(binary.right()),
+ ) {
+ // For literal op cast(expr), we need to swap the operator
+ let swapped_op = swap_operator(*binary.op());
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ swapped_op,
+ schema,
+ )? {
+ // unwrapped is already the correct expression: inner_expr
swapped_op casted_literal
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ Ok(None)
+}
+
+/// Extract cast information from a physical expression
+fn extract_cast_info(
+ expr: &Arc<dyn PhysicalExpr>,
+) -> Option<(&Arc<dyn PhysicalExpr>, &DataType)> {
+ if let Some(cast) = expr.as_any().downcast_ref::<CastExpr>() {
+ Some((cast.expr(), cast.cast_type()))
+ } else if let Some(try_cast) = expr.as_any().downcast_ref::<TryCastExpr>()
{
+ Some((try_cast.expr(), try_cast.cast_type()))
+ } else {
+ None
+ }
+}
+
+/// Try to unwrap a cast in comparison by moving the cast to the literal
+fn try_unwrap_cast_comparison(
+ inner_expr: Arc<dyn PhysicalExpr>,
+ literal_value: &ScalarValue,
+ op: Operator,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Get the data type of the inner expression
+ let inner_type = inner_expr.data_type(schema)?;
+
+ // Try to cast the literal to the inner expression's type
+ if let Some(casted_literal) = try_cast_literal_to_type(literal_value,
&inner_type) {
+ let literal_expr = lit(casted_literal);
+ let binary_expr = BinaryExpr::new(inner_expr, op, literal_expr);
+ return Ok(Some(Arc::new(binary_expr)));
+ }
+
+ Ok(None)
+}
+
+/// Swap comparison operators for right-side cast unwrapping
+fn swap_operator(op: Operator) -> Operator {
+ match op {
+ Operator::Lt => Operator::Gt,
+ Operator::LtEq => Operator::GtEq,
+ Operator::Gt => Operator::Lt,
+ Operator::GtEq => Operator::LtEq,
+ // Symmetric operators remain the same
+ Operator::Eq | Operator::NotEq => op,
+ // Other operators are not swappable
+ _ => op,
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::expressions::{col, lit};
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::Operator;
+
+ /// Check if an expression is a cast expression
+ fn is_cast_expr(expr: &Arc<dyn PhysicalExpr>) -> bool {
+ expr.as_any().downcast_ref::<CastExpr>().is_some()
+ || expr.as_any().downcast_ref::<TryCastExpr>().is_some()
+ }
+
+ /// Check if a binary expression is suitable for cast unwrapping
+ fn is_binary_expr_with_cast_and_literal(binary: &BinaryExpr) -> bool {
+ // Check if left is cast and right is literal
+ let left_cast_right_literal = is_cast_expr(binary.left())
+ && binary.right().as_any().downcast_ref::<Literal>().is_some();
+
+ // Check if left is literal and right is cast
+ let left_literal_right_cast =
+ binary.left().as_any().downcast_ref::<Literal>().is_some()
+ && is_cast_expr(binary.right());
+
+ left_cast_right_literal || left_literal_right_cast
+ }
+
+ fn test_schema() -> Schema {
+ Schema::new(vec![
+ Field::new("c1", DataType::Int32, false),
+ Field::new("c2", DataType::Int64, false),
+ Field::new("c3", DataType::Utf8, false),
+ ])
+ }
+
+ #[test]
+ fn test_unwrap_cast_in_binary_comparison() {
+ let schema = test_schema();
+
+ // Create: cast(c1 as INT64) > INT64(10)
+ let column_expr = col("c1", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(column_expr, DataType::Int64,
None));
+ let literal_expr = lit(ScalarValue::Int64(Some(10)));
Review Comment:
I think you can write this (and similar expressions) below much more
sucinctly like this:
```suggestion
let literal_expr = lit(10i64);
```
##########
datafusion/expr-common/src/casts.rs:
##########
@@ -0,0 +1,1227 @@
+// 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.
+
+//! Utilities for casting scalar literals to different data types
+//!
+//! This module contains functions for casting ScalarValue literals
+//! to different data types, originally extracted from the optimizer's
+//! unwrap_cast module to be shared between logical and physical layers.
+
+use std::cmp::Ordering;
+
+use arrow::datatypes::{
+ DataType, TimeUnit, MAX_DECIMAL128_FOR_EACH_PRECISION,
+ MIN_DECIMAL128_FOR_EACH_PRECISION,
+};
+use arrow::temporal_conversions::{MICROSECONDS, MILLISECONDS, NANOSECONDS};
+use datafusion_common::ScalarValue;
+
+/// Convert a literal value from one data type to another
+pub fn try_cast_literal_to_type(
Review Comment:
This almost seems like it is / should be a method on `ScalarValue` like
`ScalarValue::try_cast` or something -- might make it more discoverable
We could do this as a follow on PR -- I see you just moved this to a new
place
##########
datafusion/physical-expr/src/simplifier/unwrap_cast.rs:
##########
@@ -0,0 +1,621 @@
+// 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.
+
+//! Unwrap casts in binary comparisons for physical expressions
+//!
+//! This module provides optimization for physical expressions similar to the
logical
+//! optimizer's unwrap_cast module. It attempts to remove casts from
comparisons to
+//! literals by applying the casts to the literals if possible.
+//!
+//! The optimization improves performance by:
+//! 1. Reducing runtime cast operations on column data
+//! 2. Enabling better predicate pushdown opportunities
+//! 3. Optimizing filter expressions in physical plans
+//!
+//! # Example
+//!
+//! Physical expression: `cast(column as INT64) > INT64(10)`
+//! Optimized to: `column > INT32(10)` (assuming column is INT32)
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Schema};
+use datafusion_common::{
+ tree_node::{Transformed, TreeNode},
+ Result, ScalarValue,
+};
+use datafusion_expr::Operator;
+use datafusion_expr_common::casts::try_cast_literal_to_type;
+
+use crate::expressions::{lit, BinaryExpr, CastExpr, Literal, TryCastExpr};
+use crate::PhysicalExpr;
+
+/// Attempts to unwrap casts in comparison expressions.
+pub(crate) fn unwrap_cast_in_comparison(
+ expr: Arc<dyn PhysicalExpr>,
+ schema: &Schema,
+) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ expr.transform_down(|e| {
+ if let Some(binary) = e.as_any().downcast_ref::<BinaryExpr>() {
+ if let Some(unwrapped) = try_unwrap_cast_binary(binary, schema)? {
+ return Ok(Transformed::yes(unwrapped));
+ }
+ }
+ Ok(Transformed::no(e))
+ })
+}
+
+/// Try to unwrap casts in binary expressions
+fn try_unwrap_cast_binary(
+ binary: &BinaryExpr,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Case 1: cast(left_expr) op literal
+ if let (Some((inner_expr, _cast_type)), Some(literal)) = (
+ extract_cast_info(binary.left()),
+ binary.right().as_any().downcast_ref::<Literal>(),
+ ) {
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ *binary.op(),
+ schema,
+ )? {
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ // Case 2: literal op cast(right_expr)
+ if let (Some(literal), Some((inner_expr, _cast_type))) = (
+ binary.left().as_any().downcast_ref::<Literal>(),
+ extract_cast_info(binary.right()),
+ ) {
+ // For literal op cast(expr), we need to swap the operator
+ let swapped_op = swap_operator(*binary.op());
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ swapped_op,
+ schema,
+ )? {
+ // unwrapped is already the correct expression: inner_expr
swapped_op casted_literal
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ Ok(None)
+}
+
+/// Extract cast information from a physical expression
+fn extract_cast_info(
+ expr: &Arc<dyn PhysicalExpr>,
+) -> Option<(&Arc<dyn PhysicalExpr>, &DataType)> {
+ if let Some(cast) = expr.as_any().downcast_ref::<CastExpr>() {
+ Some((cast.expr(), cast.cast_type()))
+ } else if let Some(try_cast) = expr.as_any().downcast_ref::<TryCastExpr>()
{
+ Some((try_cast.expr(), try_cast.cast_type()))
+ } else {
+ None
+ }
+}
+
+/// Try to unwrap a cast in comparison by moving the cast to the literal
+fn try_unwrap_cast_comparison(
+ inner_expr: Arc<dyn PhysicalExpr>,
+ literal_value: &ScalarValue,
+ op: Operator,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Get the data type of the inner expression
+ let inner_type = inner_expr.data_type(schema)?;
+
+ // Try to cast the literal to the inner expression's type
+ if let Some(casted_literal) = try_cast_literal_to_type(literal_value,
&inner_type) {
+ let literal_expr = lit(casted_literal);
+ let binary_expr = BinaryExpr::new(inner_expr, op, literal_expr);
+ return Ok(Some(Arc::new(binary_expr)));
+ }
+
+ Ok(None)
+}
+
+/// Swap comparison operators for right-side cast unwrapping
+fn swap_operator(op: Operator) -> Operator {
+ match op {
+ Operator::Lt => Operator::Gt,
+ Operator::LtEq => Operator::GtEq,
+ Operator::Gt => Operator::Lt,
+ Operator::GtEq => Operator::LtEq,
+ // Symmetric operators remain the same
+ Operator::Eq | Operator::NotEq => op,
+ // Other operators are not swappable
+ _ => op,
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::expressions::{col, lit};
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::Operator;
+
+ /// Check if an expression is a cast expression
+ fn is_cast_expr(expr: &Arc<dyn PhysicalExpr>) -> bool {
+ expr.as_any().downcast_ref::<CastExpr>().is_some()
+ || expr.as_any().downcast_ref::<TryCastExpr>().is_some()
+ }
+
+ /// Check if a binary expression is suitable for cast unwrapping
+ fn is_binary_expr_with_cast_and_literal(binary: &BinaryExpr) -> bool {
+ // Check if left is cast and right is literal
+ let left_cast_right_literal = is_cast_expr(binary.left())
+ && binary.right().as_any().downcast_ref::<Literal>().is_some();
+
+ // Check if left is literal and right is cast
+ let left_literal_right_cast =
+ binary.left().as_any().downcast_ref::<Literal>().is_some()
+ && is_cast_expr(binary.right());
+
+ left_cast_right_literal || left_literal_right_cast
+ }
+
+ fn test_schema() -> Schema {
+ Schema::new(vec![
+ Field::new("c1", DataType::Int32, false),
+ Field::new("c2", DataType::Int64, false),
+ Field::new("c3", DataType::Utf8, false),
+ ])
+ }
+
+ #[test]
+ fn test_unwrap_cast_in_binary_comparison() {
+ let schema = test_schema();
+
+ // Create: cast(c1 as INT64) > INT64(10)
+ let column_expr = col("c1", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(column_expr, DataType::Int64,
None));
+ let literal_expr = lit(ScalarValue::Int64(Some(10)));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(cast_expr, Operator::Gt, literal_expr));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap();
+
+ // Should be transformed
+ assert!(result.transformed);
+
+ // The result should be: c1 > INT32(10)
+ let optimized = result.data;
+ let optimized_binary =
optimized.as_any().downcast_ref::<BinaryExpr>().unwrap();
+
+ // Check that left side is no longer a cast
+ assert!(!is_cast_expr(optimized_binary.left()));
+
+ // Check that right side is a literal with the correct type and value
+ let right_literal = optimized_binary
+ .right()
+ .as_any()
+ .downcast_ref::<Literal>()
+ .unwrap();
+ assert_eq!(right_literal.value(), &ScalarValue::Int32(Some(10)));
+ }
+
+ #[test]
+ fn test_unwrap_cast_with_literal_on_left() {
+ let schema = test_schema();
+
+ // Create: INT64(10) < cast(c1 as INT64)
+ let column_expr = col("c1", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(column_expr, DataType::Int64,
None));
+ let literal_expr = lit(ScalarValue::Int64(Some(10)));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(literal_expr, Operator::Lt, cast_expr));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap();
+
+ // Should be transformed
+ assert!(result.transformed);
+
+ // The result should be equivalent to: c1 > INT32(10)
+ let optimized = result.data;
+ let optimized_binary =
optimized.as_any().downcast_ref::<BinaryExpr>().unwrap();
+
+ // Check the operator was swapped
+ assert_eq!(*optimized_binary.op(), Operator::Gt);
+ }
+
+ #[test]
+ fn test_no_unwrap_when_types_unsupported() {
+ let schema = Schema::new(vec![Field::new("f1", DataType::Float32,
false)]);
+
+ // Create: cast(f1 as FLOAT64) > FLOAT64(10.5)
+ let column_expr = col("f1", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(column_expr, DataType::Float64,
None));
+ let literal_expr = lit(ScalarValue::Float64(Some(10.5)));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(cast_expr, Operator::Gt, literal_expr));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap();
+
+ // Should NOT be transformed (floating point types not supported)
+ assert!(!result.transformed);
+ }
+
+ #[test]
+ fn test_is_binary_expr_with_cast_and_literal() {
+ let schema = test_schema();
+
+ let column_expr = col("c1", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(column_expr, DataType::Int64,
None));
+ let literal_expr = lit(ScalarValue::Int64(Some(10)));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(cast_expr, Operator::Gt, literal_expr));
+ let binary_ref =
binary_expr.as_any().downcast_ref::<BinaryExpr>().unwrap();
+
+ assert!(is_binary_expr_with_cast_and_literal(binary_ref));
+ }
+
+ #[test]
+ fn test_swap_operator() {
+ assert_eq!(swap_operator(Operator::Lt), Operator::Gt);
+ assert_eq!(swap_operator(Operator::LtEq), Operator::GtEq);
+ assert_eq!(swap_operator(Operator::Gt), Operator::Lt);
+ assert_eq!(swap_operator(Operator::GtEq), Operator::LtEq);
+ assert_eq!(swap_operator(Operator::Eq), Operator::Eq);
+ assert_eq!(swap_operator(Operator::NotEq), Operator::NotEq);
+ }
+
+ #[test]
+ fn test_unwrap_cast_literal_on_left_side() {
+ // Test case for: literal <= cast(column)
+ // This was the specific case that caused the bug
+ let schema = Schema::new(vec![Field::new(
+ "decimal_col",
+ DataType::Decimal128(9, 2),
+ true,
+ )]);
+
+ // Create: Decimal128(400) <= cast(decimal_col as Decimal128(22, 2))
+ let column_expr = col("decimal_col", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(
+ column_expr,
+ DataType::Decimal128(22, 2),
+ None,
+ ));
+ let literal_expr = lit(ScalarValue::Decimal128(Some(400), 22, 2));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(literal_expr, Operator::LtEq, cast_expr));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap();
+
+ // Should be transformed
+ assert!(result.transformed);
+
+ // The result should be: decimal_col >= Decimal128(400, 9, 2)
+ let optimized = result.data;
+ let optimized_binary =
optimized.as_any().downcast_ref::<BinaryExpr>().unwrap();
+
+ // Check operator was swapped correctly
+ assert_eq!(*optimized_binary.op(), Operator::GtEq);
+
+ // Check that left side is the column without cast
+ assert!(!is_cast_expr(optimized_binary.left()));
+
+ // Check that right side is a literal with the correct type
+ let right_literal = optimized_binary
+ .right()
+ .as_any()
+ .downcast_ref::<Literal>()
+ .unwrap();
+ assert_eq!(
+ right_literal.value().data_type(),
+ DataType::Decimal128(9, 2)
+ );
+ }
+
+ #[test]
+ fn test_unwrap_cast_with_different_comparison_operators() {
+ let schema = Schema::new(vec![Field::new("int_col", DataType::Int32,
false)]);
+
+ // Test all comparison operators with literal on the left
+ let operators = vec![
+ (Operator::Lt, Operator::Gt),
+ (Operator::LtEq, Operator::GtEq),
+ (Operator::Gt, Operator::Lt),
+ (Operator::GtEq, Operator::LtEq),
+ (Operator::Eq, Operator::Eq),
+ (Operator::NotEq, Operator::NotEq),
+ ];
+
+ for (original_op, expected_op) in operators {
+ // Create: INT64(100) op cast(int_col as INT64)
+ let column_expr = col("int_col", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(column_expr,
DataType::Int64, None));
+ let literal_expr = lit(ScalarValue::Int64(Some(100)));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(literal_expr, original_op,
cast_expr));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(binary_expr,
&schema).unwrap();
+
+ // Should be transformed
+ assert!(result.transformed);
+
+ let optimized = result.data;
+ let optimized_binary =
+ optimized.as_any().downcast_ref::<BinaryExpr>().unwrap();
+
+ // Check the operator was swapped correctly
+ assert_eq!(
+ *optimized_binary.op(),
+ expected_op,
+ "Failed for operator {original_op:?} -> {expected_op:?}"
+ );
+
+ // Check that left side has no cast
+ assert!(!is_cast_expr(optimized_binary.left()));
+
+ // Check that the literal was cast to the column type
+ let right_literal = optimized_binary
+ .right()
+ .as_any()
+ .downcast_ref::<Literal>()
+ .unwrap();
+ assert_eq!(right_literal.value(), &ScalarValue::Int32(Some(100)));
+ }
+ }
+
+ #[test]
+ fn test_unwrap_cast_with_decimal_types() {
+ // Test various decimal precision/scale combinations
+ let test_cases = vec![
+ // (column_precision, column_scale, cast_precision, cast_scale,
value)
+ (9, 2, 22, 2, 400),
+ (10, 3, 20, 3, 1000),
+ (5, 1, 10, 1, 99),
+ ];
+
+ for (col_p, col_s, cast_p, cast_s, value) in test_cases {
+ let schema = Schema::new(vec![Field::new(
+ "decimal_col",
+ DataType::Decimal128(col_p, col_s),
+ true,
+ )]);
+
+ // Test both: cast(column) op literal AND literal op cast(column)
+
+ // Case 1: cast(column) > literal
+ let column_expr = col("decimal_col", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(
+ Arc::clone(&column_expr),
+ DataType::Decimal128(cast_p, cast_s),
+ None,
+ ));
+ let literal_expr = lit(ScalarValue::Decimal128(Some(value),
cast_p, cast_s));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(cast_expr, Operator::Gt,
literal_expr));
+
+ let result = unwrap_cast_in_comparison(binary_expr,
&schema).unwrap();
+ assert!(result.transformed);
+
+ // Case 2: literal < cast(column)
+ let cast_expr = Arc::new(CastExpr::new(
+ column_expr,
+ DataType::Decimal128(cast_p, cast_s),
+ None,
+ ));
+ let literal_expr = lit(ScalarValue::Decimal128(Some(value),
cast_p, cast_s));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(literal_expr, Operator::Lt,
cast_expr));
+
+ let result = unwrap_cast_in_comparison(binary_expr,
&schema).unwrap();
+ assert!(result.transformed);
+ }
+ }
+
+ #[test]
+ fn test_unwrap_cast_with_null_literals() {
+ // Test with NULL literals to ensure they're handled correctly
+ let schema = Schema::new(vec![Field::new("int_col", DataType::Int32,
true)]);
+
+ // Create: cast(int_col as INT64) = NULL
+ let column_expr = col("int_col", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(column_expr, DataType::Int64,
None));
+ let null_literal = lit(ScalarValue::Int64(None));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(cast_expr, Operator::Eq, null_literal));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap();
+
+ // Should be transformed
+ assert!(result.transformed);
+
+ // Verify the NULL was cast to the column type
+ let optimized = result.data;
+ let optimized_binary =
optimized.as_any().downcast_ref::<BinaryExpr>().unwrap();
+ let right_literal = optimized_binary
+ .right()
+ .as_any()
+ .downcast_ref::<Literal>()
+ .unwrap();
+ assert_eq!(right_literal.value(), &ScalarValue::Int32(None));
+ }
+
+ #[test]
+ fn test_unwrap_cast_with_try_cast() {
+ // Test that TryCast expressions are also unwrapped correctly
+ let schema = Schema::new(vec![Field::new("str_col", DataType::Utf8,
true)]);
+
+ // Create: try_cast(str_col as INT64) > INT64(100)
+ let column_expr = col("str_col", &schema).unwrap();
+ let try_cast_expr = Arc::new(TryCastExpr::new(column_expr,
DataType::Int64));
+ let literal_expr = lit(ScalarValue::Int64(Some(100)));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(try_cast_expr, Operator::Gt,
literal_expr));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap();
+
+ // Should NOT be transformed (string to int cast not supported)
+ assert!(!result.transformed);
+ }
+
+ #[test]
+ fn test_unwrap_cast_preserves_non_comparison_operators() {
+ // Test that non-comparison operators in AND/OR expressions are
preserved
+ let schema = Schema::new(vec![Field::new("int_col", DataType::Int32,
false)]);
+
+ // Create: cast(int_col as INT64) > INT64(10) AND cast(int_col as
INT64) < INT64(20)
+ let column_expr = col("int_col", &schema).unwrap();
+
+ let cast1 = Arc::new(CastExpr::new(
+ Arc::clone(&column_expr),
+ DataType::Int64,
+ None,
+ ));
+ let lit1 = lit(ScalarValue::Int64(Some(10)));
+ let compare1 = Arc::new(BinaryExpr::new(cast1, Operator::Gt, lit1));
+
+ let cast2 = Arc::new(CastExpr::new(column_expr, DataType::Int64,
None));
+ let lit2 = lit(ScalarValue::Int64(Some(20)));
+ let compare2 = Arc::new(BinaryExpr::new(cast2, Operator::Lt, lit2));
+
+ let and_expr = Arc::new(BinaryExpr::new(compare1, Operator::And,
compare2));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(and_expr, &schema).unwrap();
+
+ // Should be transformed
+ assert!(result.transformed);
+
+ // Verify the AND operator is preserved
+ let optimized = result.data;
+ let and_binary =
optimized.as_any().downcast_ref::<BinaryExpr>().unwrap();
+ assert_eq!(*and_binary.op(), Operator::And);
+
+ // Both sides should have their casts unwrapped
Review Comment:
It might make the tests easier to write / read if you just made the expected
output directly
```rust
let expected = lit(10i32)...;
assert_eq!(expected, optimized)
```
##########
datafusion/physical-expr/src/simplifier/mod.rs:
##########
@@ -0,0 +1,184 @@
+// 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.
+
+//! Simplifier for Physical Expressions
+
+use arrow::datatypes::Schema;
+use datafusion_common::{
+ tree_node::{Transformed, TreeNode, TreeNodeRewriter},
+ Result,
+};
+use std::sync::Arc;
+
+use crate::PhysicalExpr;
+
+pub mod unwrap_cast;
+
+/// Simplifies physical expressions by applying various optimizations
Review Comment:
```suggestion
/// Simplifies physical expressions by applying various optimizations
///
/// This can be useful after adapting expressions from a table schema
/// to a file schema. For example, casts added to match the types may
/// potentially be unwrapped.
```
##########
datafusion/physical-expr/src/simplifier/unwrap_cast.rs:
##########
@@ -0,0 +1,621 @@
+// 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.
+
+//! Unwrap casts in binary comparisons for physical expressions
+//!
+//! This module provides optimization for physical expressions similar to the
logical
+//! optimizer's unwrap_cast module. It attempts to remove casts from
comparisons to
+//! literals by applying the casts to the literals if possible.
+//!
+//! The optimization improves performance by:
+//! 1. Reducing runtime cast operations on column data
+//! 2. Enabling better predicate pushdown opportunities
+//! 3. Optimizing filter expressions in physical plans
+//!
+//! # Example
+//!
+//! Physical expression: `cast(column as INT64) > INT64(10)`
+//! Optimized to: `column > INT32(10)` (assuming column is INT32)
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Schema};
+use datafusion_common::{
+ tree_node::{Transformed, TreeNode},
+ Result, ScalarValue,
+};
+use datafusion_expr::Operator;
+use datafusion_expr_common::casts::try_cast_literal_to_type;
+
+use crate::expressions::{lit, BinaryExpr, CastExpr, Literal, TryCastExpr};
+use crate::PhysicalExpr;
+
+/// Attempts to unwrap casts in comparison expressions.
+pub(crate) fn unwrap_cast_in_comparison(
+ expr: Arc<dyn PhysicalExpr>,
+ schema: &Schema,
+) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ expr.transform_down(|e| {
+ if let Some(binary) = e.as_any().downcast_ref::<BinaryExpr>() {
+ if let Some(unwrapped) = try_unwrap_cast_binary(binary, schema)? {
+ return Ok(Transformed::yes(unwrapped));
+ }
+ }
+ Ok(Transformed::no(e))
+ })
+}
+
+/// Try to unwrap casts in binary expressions
+fn try_unwrap_cast_binary(
+ binary: &BinaryExpr,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Case 1: cast(left_expr) op literal
+ if let (Some((inner_expr, _cast_type)), Some(literal)) = (
+ extract_cast_info(binary.left()),
+ binary.right().as_any().downcast_ref::<Literal>(),
+ ) {
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ *binary.op(),
+ schema,
+ )? {
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ // Case 2: literal op cast(right_expr)
+ if let (Some(literal), Some((inner_expr, _cast_type))) = (
+ binary.left().as_any().downcast_ref::<Literal>(),
+ extract_cast_info(binary.right()),
+ ) {
+ // For literal op cast(expr), we need to swap the operator
+ let swapped_op = swap_operator(*binary.op());
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ swapped_op,
+ schema,
+ )? {
+ // unwrapped is already the correct expression: inner_expr
swapped_op casted_literal
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ Ok(None)
+}
+
+/// Extract cast information from a physical expression
+fn extract_cast_info(
+ expr: &Arc<dyn PhysicalExpr>,
+) -> Option<(&Arc<dyn PhysicalExpr>, &DataType)> {
+ if let Some(cast) = expr.as_any().downcast_ref::<CastExpr>() {
+ Some((cast.expr(), cast.cast_type()))
+ } else if let Some(try_cast) = expr.as_any().downcast_ref::<TryCastExpr>()
{
+ Some((try_cast.expr(), try_cast.cast_type()))
+ } else {
+ None
+ }
+}
+
+/// Try to unwrap a cast in comparison by moving the cast to the literal
+fn try_unwrap_cast_comparison(
+ inner_expr: Arc<dyn PhysicalExpr>,
+ literal_value: &ScalarValue,
+ op: Operator,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Get the data type of the inner expression
+ let inner_type = inner_expr.data_type(schema)?;
+
+ // Try to cast the literal to the inner expression's type
+ if let Some(casted_literal) = try_cast_literal_to_type(literal_value,
&inner_type) {
+ let literal_expr = lit(casted_literal);
+ let binary_expr = BinaryExpr::new(inner_expr, op, literal_expr);
+ return Ok(Some(Arc::new(binary_expr)));
+ }
+
+ Ok(None)
+}
+
+/// Swap comparison operators for right-side cast unwrapping
+fn swap_operator(op: Operator) -> Operator {
Review Comment:
This could use `Operator::swap`, right? I don't think there is any reason to
have it specially here
##########
datafusion/physical-expr/src/simplifier/unwrap_cast.rs:
##########
@@ -0,0 +1,621 @@
+// 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.
+
+//! Unwrap casts in binary comparisons for physical expressions
+//!
+//! This module provides optimization for physical expressions similar to the
logical
+//! optimizer's unwrap_cast module. It attempts to remove casts from
comparisons to
+//! literals by applying the casts to the literals if possible.
+//!
+//! The optimization improves performance by:
+//! 1. Reducing runtime cast operations on column data
+//! 2. Enabling better predicate pushdown opportunities
+//! 3. Optimizing filter expressions in physical plans
+//!
+//! # Example
+//!
+//! Physical expression: `cast(column as INT64) > INT64(10)`
+//! Optimized to: `column > INT32(10)` (assuming column is INT32)
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Schema};
+use datafusion_common::{
+ tree_node::{Transformed, TreeNode},
+ Result, ScalarValue,
+};
+use datafusion_expr::Operator;
+use datafusion_expr_common::casts::try_cast_literal_to_type;
+
+use crate::expressions::{lit, BinaryExpr, CastExpr, Literal, TryCastExpr};
+use crate::PhysicalExpr;
+
+/// Attempts to unwrap casts in comparison expressions.
+pub(crate) fn unwrap_cast_in_comparison(
+ expr: Arc<dyn PhysicalExpr>,
+ schema: &Schema,
+) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ expr.transform_down(|e| {
+ if let Some(binary) = e.as_any().downcast_ref::<BinaryExpr>() {
+ if let Some(unwrapped) = try_unwrap_cast_binary(binary, schema)? {
+ return Ok(Transformed::yes(unwrapped));
+ }
+ }
+ Ok(Transformed::no(e))
+ })
+}
+
+/// Try to unwrap casts in binary expressions
+fn try_unwrap_cast_binary(
+ binary: &BinaryExpr,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Case 1: cast(left_expr) op literal
+ if let (Some((inner_expr, _cast_type)), Some(literal)) = (
+ extract_cast_info(binary.left()),
+ binary.right().as_any().downcast_ref::<Literal>(),
+ ) {
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ *binary.op(),
+ schema,
+ )? {
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ // Case 2: literal op cast(right_expr)
+ if let (Some(literal), Some((inner_expr, _cast_type))) = (
+ binary.left().as_any().downcast_ref::<Literal>(),
+ extract_cast_info(binary.right()),
+ ) {
+ // For literal op cast(expr), we need to swap the operator
+ let swapped_op = swap_operator(*binary.op());
+ if let Some(unwrapped) = try_unwrap_cast_comparison(
+ Arc::clone(inner_expr),
+ literal.value(),
+ swapped_op,
+ schema,
+ )? {
+ // unwrapped is already the correct expression: inner_expr
swapped_op casted_literal
+ return Ok(Some(unwrapped));
+ }
+ }
+
+ Ok(None)
+}
+
+/// Extract cast information from a physical expression
+fn extract_cast_info(
+ expr: &Arc<dyn PhysicalExpr>,
+) -> Option<(&Arc<dyn PhysicalExpr>, &DataType)> {
+ if let Some(cast) = expr.as_any().downcast_ref::<CastExpr>() {
+ Some((cast.expr(), cast.cast_type()))
+ } else if let Some(try_cast) = expr.as_any().downcast_ref::<TryCastExpr>()
{
+ Some((try_cast.expr(), try_cast.cast_type()))
+ } else {
+ None
+ }
+}
+
+/// Try to unwrap a cast in comparison by moving the cast to the literal
+fn try_unwrap_cast_comparison(
+ inner_expr: Arc<dyn PhysicalExpr>,
+ literal_value: &ScalarValue,
+ op: Operator,
+ schema: &Schema,
+) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+ // Get the data type of the inner expression
+ let inner_type = inner_expr.data_type(schema)?;
+
+ // Try to cast the literal to the inner expression's type
+ if let Some(casted_literal) = try_cast_literal_to_type(literal_value,
&inner_type) {
+ let literal_expr = lit(casted_literal);
+ let binary_expr = BinaryExpr::new(inner_expr, op, literal_expr);
+ return Ok(Some(Arc::new(binary_expr)));
+ }
+
+ Ok(None)
+}
+
+/// Swap comparison operators for right-side cast unwrapping
+fn swap_operator(op: Operator) -> Operator {
+ match op {
+ Operator::Lt => Operator::Gt,
+ Operator::LtEq => Operator::GtEq,
+ Operator::Gt => Operator::Lt,
+ Operator::GtEq => Operator::LtEq,
+ // Symmetric operators remain the same
+ Operator::Eq | Operator::NotEq => op,
+ // Other operators are not swappable
+ _ => op,
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::expressions::{col, lit};
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::Operator;
+
+ /// Check if an expression is a cast expression
+ fn is_cast_expr(expr: &Arc<dyn PhysicalExpr>) -> bool {
+ expr.as_any().downcast_ref::<CastExpr>().is_some()
+ || expr.as_any().downcast_ref::<TryCastExpr>().is_some()
+ }
+
+ /// Check if a binary expression is suitable for cast unwrapping
+ fn is_binary_expr_with_cast_and_literal(binary: &BinaryExpr) -> bool {
+ // Check if left is cast and right is literal
+ let left_cast_right_literal = is_cast_expr(binary.left())
+ && binary.right().as_any().downcast_ref::<Literal>().is_some();
+
+ // Check if left is literal and right is cast
+ let left_literal_right_cast =
+ binary.left().as_any().downcast_ref::<Literal>().is_some()
+ && is_cast_expr(binary.right());
+
+ left_cast_right_literal || left_literal_right_cast
+ }
+
+ fn test_schema() -> Schema {
+ Schema::new(vec![
+ Field::new("c1", DataType::Int32, false),
+ Field::new("c2", DataType::Int64, false),
+ Field::new("c3", DataType::Utf8, false),
+ ])
+ }
+
+ #[test]
+ fn test_unwrap_cast_in_binary_comparison() {
+ let schema = test_schema();
+
+ // Create: cast(c1 as INT64) > INT64(10)
+ let column_expr = col("c1", &schema).unwrap();
+ let cast_expr = Arc::new(CastExpr::new(column_expr, DataType::Int64,
None));
+ let literal_expr = lit(ScalarValue::Int64(Some(10)));
+ let binary_expr =
+ Arc::new(BinaryExpr::new(cast_expr, Operator::Gt, literal_expr));
+
+ // Apply unwrap cast optimization
+ let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap();
+
+ // Should be transformed
+ assert!(result.transformed);
+
+ // The result should be: c1 > INT32(10)
+ let optimized = result.data;
+ let optimized_binary =
optimized.as_any().downcast_ref::<BinaryExpr>().unwrap();
+
+ // Check that left side is no longer a cast
+ assert!(!is_cast_expr(optimized_binary.left()));
+
+ // Check that right side is a literal with the correct type and value
+ let right_literal = optimized_binary
+ .right()
+ .as_any()
+ .downcast_ref::<Literal>()
+ .unwrap();
+ assert_eq!(right_literal.value(), &ScalarValue::Int32(Some(10)));
+ }
+
+ #[test]
+ fn test_unwrap_cast_with_literal_on_left() {
+ let schema = test_schema();
+
+ // Create: INT64(10) < cast(c1 as INT64)
Review Comment:
Can you also please test something where the cast doesn't work? Something
with a column that is `uint8` for example, and the literal can not fit into the
column: `Int64(-5) < CAST(c4 as Int64)
Though I realize in this case the predicate would always be true then 🤔
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
