fix: preserve narrowing TRY_CAST comparison semantics

datafusion/expr-common/src/casts.rs

@@ -157,6 +157,38 @@ pub fn is_date_narrowing_cast(from_type: &DataType, to_type: &DataType) -> bool
     matches!((from_type, to_type), (DataType::Date64, DataType::Date32))
 }
 
+/// Returns true when the source integer domain is not fully representable in
+/// the target integer type.
+///
+/// Comparison cast unwrapping must retain a narrowing `TRY_CAST`: source values
+/// outside the target domain become `NULL`, while the unwrapped comparison
+/// would evaluate those values directly. This check includes signed/unsigned
+/// transitions as well as casts to a smaller integer width.
+pub fn is_integer_narrowing_cast(from_type: &DataType, to_type: &DataType) -> bool {
+    let Some((from_min, from_max)) = integer_bounds(from_type) else {
+        return false;
+    };
+    let Some((to_min, to_max)) = integer_bounds(to_type) else {
+        return false;
+    };
+
+    from_min < to_min || from_max > to_max
+}
+
+fn integer_bounds(data_type: &DataType) -> Option<(i128, i128)> {
+    match data_type {
+        DataType::Int8 => Some((i8::MIN.into(), i8::MAX.into())),
+        DataType::Int16 => Some((i16::MIN.into(), i16::MAX.into())),
+        DataType::Int32 => Some((i32::MIN.into(), i32::MAX.into())),
+        DataType::Int64 => Some((i64::MIN.into(), i64::MAX.into())),
+        DataType::UInt8 => Some((u8::MIN.into(), u8::MAX.into())),
+        DataType::UInt16 => Some((u16::MIN.into(), u16::MAX.into())),
+        DataType::UInt32 => Some((u32::MIN.into(), u32::MAX.into())),
+        DataType::UInt64 => Some((u64::MIN.into(), u64::MAX.into())),
+        _ => None,
+    }
+}
+
 fn timestamp_unit_scale(unit: &TimeUnit) -> i128 {
     match unit {
         TimeUnit::Second => 1,
@@ -1036,6 +1068,34 @@ mod tests {
         assert!(!is_date_narrowing_cast(&DataType::Int64, &DataType::Date32));
     }
 
+    #[test]
+    fn test_is_integer_narrowing_cast() {
+        assert!(is_integer_narrowing_cast(
+            &DataType::Int64,
+            &DataType::Int32
+        ));
+        assert!(is_integer_narrowing_cast(
+            &DataType::Int32,
+            &DataType::UInt32
+        ));
+        assert!(is_integer_narrowing_cast(
+            &DataType::UInt32,
+            &DataType::Int32
+        ));
+        assert!(!is_integer_narrowing_cast(
+            &DataType::Int32,
+            &DataType::Int64
+        ));
+        assert!(!is_integer_narrowing_cast(
+            &DataType::UInt32,
+            &DataType::Int64
+        ));
+        assert!(!is_integer_narrowing_cast(
+            &DataType::Utf8,
+            &DataType::Int32
+        ));
+    }
+
     #[test]
     fn test_scale_date_literal_exactness_and_overflow() {
         const MS_PER_DAY: i128 = MILLISECONDS_IN_DAY as i128;

datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs

@@ -60,8 +60,8 @@ use datafusion_common::{internal_err, tree_node::Transformed};
 use datafusion_expr::{BinaryExpr, lit};
 use datafusion_expr::{Cast, Expr, Operator, TryCast, simplify::SimplifyContext};
 use datafusion_expr_common::casts::{
-    is_date_narrowing_cast, is_supported_type, is_timestamp_precision_narrowing_cast,
-    try_cast_literal_to_type,
+    is_date_narrowing_cast, is_integer_narrowing_cast, is_supported_type,
+    is_timestamp_precision_narrowing_cast, try_cast_literal_to_type,
 };
 
 pub(super) fn unwrap_cast_in_comparison_for_binary(
@@ -113,18 +113,20 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary(
     op: Operator,
     literal: &Expr,
 ) -> bool {
-    match (expr, literal) {
-        (
-            Expr::TryCast(TryCast {
-                expr: left_expr,
-                field,
-            })
-            | Expr::Cast(Cast {
-                expr: left_expr,
-                field,
-            }),
-            Expr::Literal(lit_val, _),
-        ) => {
+    let (left_expr, field, is_try_cast) = match expr {
+        Expr::TryCast(TryCast {
+            expr: left_expr,
+            field,
+        }) => (left_expr, field, true),
+        Expr::Cast(Cast {
+            expr: left_expr,
+            field,
+        }) => (left_expr, field, false),
+        _ => return false,
+    };
+
+    match literal {
+        Expr::Literal(lit_val, _) => {
             let Ok(expr_type) = info.get_data_type(left_expr) else {
                 return false;
             };
@@ -135,6 +137,8 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary(
 
             if is_timestamp_precision_narrowing_cast(&expr_type, field.data_type())
                 || is_date_narrowing_cast(&expr_type, field.data_type())
+                || (is_try_cast
+                    && is_integer_narrowing_cast(&expr_type, field.data_type()))
             {
                 return false;
             }
@@ -156,16 +160,16 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist(
     expr: &Expr,
     list: &[Expr],
 ) -> bool {
-    let (Expr::TryCast(TryCast {
-        expr: left_expr,
-        field,
-    })
-    | Expr::Cast(Cast {
-        expr: left_expr,
-        field,
-    })) = expr
-    else {
-        return false;
+    let (left_expr, field, is_try_cast) = match expr {
+        Expr::TryCast(TryCast {
+            expr: left_expr,
+            field,
+        }) => (left_expr, field, true),
+        Expr::Cast(Cast {
+            expr: left_expr,
+            field,
+        }) => (left_expr, field, false),
+        _ => return false,
     };
 
     let Ok(expr_type) = info.get_data_type(left_expr) else {
@@ -178,6 +182,7 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist(
 
     if is_timestamp_precision_narrowing_cast(&expr_type, field.data_type())
         || is_date_narrowing_cast(&expr_type, field.data_type())
+        || (is_try_cast && is_integer_narrowing_cast(&expr_type, field.data_type()))
     {
         return false;
     }
@@ -336,6 +341,21 @@ mod tests {
         assert_eq!(optimize_test(expr_input, &schema), expected);
     }
 
+    #[test]
+    fn test_not_unwrap_narrowing_integer_try_cast() {
+        let schema = expr_test_schema();
+
+        let comparison = try_cast(col("c2"), DataType::Int32).gt(lit(1i32));
+        assert_eq!(optimize_test(comparison.clone(), &schema), comparison);
+
+        let in_list_expr = in_list(
+            try_cast(col("c2"), DataType::Int32),
+            vec![lit(1i32), lit(2i32), lit(3i32), lit(4i32), lit(5i32)],
+            true,
+        );
+        assert_eq!(optimize_test(in_list_expr.clone(), &schema), in_list_expr);
+    }
+
     #[test]
     fn test_unwrap_cast_comparison_unsigned() {
         // "cast(c6, UINT64) = 0u64 => c6 = 0u32

datafusion/physical-expr/src/simplifier/unwrap_cast.rs

@@ -37,8 +37,8 @@ use arrow::datatypes::{DataType, Schema};
 use datafusion_common::{Result, ScalarValue, tree_node::Transformed};
 use datafusion_expr::Operator;
 use datafusion_expr_common::casts::{
-    is_date_narrowing_cast, is_timestamp_precision_narrowing_cast,
-    try_cast_literal_to_type,
+    is_date_narrowing_cast, is_integer_narrowing_cast,
+    is_timestamp_precision_narrowing_cast, try_cast_literal_to_type,
 };
 
 use crate::PhysicalExpr;
@@ -63,14 +63,15 @@ fn try_unwrap_cast_binary(
     schema: &Schema,
 ) -> Result<Option<Arc<dyn PhysicalExpr>>> {
     // Case 1: cast(left_expr) op literal
-    if let (Some((inner_expr, cast_type)), Some(literal)) = (
+    if let (Some((inner_expr, cast_type, is_try_cast)), Some(literal)) = (
         extract_cast_info(binary.left()),
         binary.right().downcast_ref::<Literal>(),
     ) && binary.op().supports_propagation()
         && let Some(unwrapped) = try_unwrap_cast_comparison(
             Arc::clone(inner_expr),
             literal.value(),
             cast_type,
+            is_try_cast,
             *binary.op(),
             schema,
         )?
@@ -79,7 +80,7 @@ fn try_unwrap_cast_binary(
     }
 
     // Case 2: literal op cast(right_expr)
-    if let (Some(literal), Some((inner_expr, cast_type))) = (
+    if let (Some(literal), Some((inner_expr, cast_type, is_try_cast))) = (
         binary.left().downcast_ref::<Literal>(),
         extract_cast_info(binary.right()),
     ) {
@@ -90,6 +91,7 @@ fn try_unwrap_cast_binary(
                 Arc::clone(inner_expr),
                 literal.value(),
                 cast_type,
+                is_try_cast,
                 swapped_op,
                 schema,
             )?
@@ -106,14 +108,15 @@ fn try_unwrap_cast_binary(
 /// Extract cast information from a physical expression
 ///
 /// If the expression is a CAST(expr, datatype) or TRY_CAST(expr, datatype),
-/// returns Some((inner_expr, target_datatype)). Otherwise returns None.
+/// returns the inner expression, target datatype, and whether the cast is a
+/// TRY_CAST. Otherwise returns None.
 fn extract_cast_info(
     expr: &Arc<dyn PhysicalExpr>,
-) -> Option<(&Arc<dyn PhysicalExpr>, &DataType)> {
+) -> Option<(&Arc<dyn PhysicalExpr>, &DataType, bool)> {
     if let Some(cast) = expr.downcast_ref::<CastExpr>() {
-        Some((cast.expr(), cast.cast_type()))
+        Some((cast.expr(), cast.cast_type(), false))
     } else if let Some(try_cast) = expr.downcast_ref::<TryCastExpr>() {
-        Some((try_cast.expr(), try_cast.cast_type()))
+        Some((try_cast.expr(), try_cast.cast_type(), true))
     } else {
         None
     }
@@ -124,6 +127,7 @@ fn try_unwrap_cast_comparison(
     inner_expr: Arc<dyn PhysicalExpr>,
     literal_value: &ScalarValue,
     cast_type: &DataType,
+    is_try_cast: bool,
     op: Operator,
     schema: &Schema,
 ) -> Result<Option<Arc<dyn PhysicalExpr>>> {
@@ -132,6 +136,7 @@ fn try_unwrap_cast_comparison(
 
     if is_timestamp_precision_narrowing_cast(&inner_type, cast_type)
         || is_date_narrowing_cast(&inner_type, cast_type)
+        || (is_try_cast && is_integer_narrowing_cast(&inner_type, cast_type))
     {
         return Ok(None);
     }
@@ -466,6 +471,84 @@ mod tests {
         assert!(!result.transformed);
     }
 
+    #[test]
+    fn test_no_unwrap_narrowing_integer_try_cast() -> Result<()> {
+        use arrow::array::{BooleanArray, Int64Array};
+        use arrow::record_batch::RecordBatch;
+
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "int_col",
+            DataType::Int64,
+            true,
+        )]));
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int64Array::from(vec![
+                Some(0),
+                Some(1),
+                Some(2),
+                Some(i64::from(i32::MIN) - 1),
+                Some(i64::from(i32::MAX) + 1),
+                None,
+            ]))],
+        )?;
+
+        for (literal, expected) in [
+            (
+                ScalarValue::Int32(Some(1)),
+                BooleanArray::from(vec![
+                    Some(false),
+                    Some(false),
+                    Some(true),
+                    None,
+                    None,
+                    None,
+                ]),
+            ),
+            (
+                ScalarValue::UInt64(Some(1)),
+                BooleanArray::from(vec![
+                    Some(false),
+                    Some(false),
+                    Some(true),
+                    None,
+                    Some(true),
+                    None,
+                ]),
+            ),
+        ] {
+            // Width narrowing and signed/unsigned transitions must both preserve
+            // NULL results for values outside the target integer domain.
+            let cast_type = literal.data_type();
+            for literal_on_left in [false, true] {
+                let column_expr = col("int_col", &schema)?;
+                let try_cast_expr =
+                    Arc::new(TryCastExpr::new(column_expr, cast_type.clone()));
+                let literal_expr = lit(literal.clone());
+                let binary_expr = if literal_on_left {
+                    BinaryExpr::new(literal_expr, Operator::Lt, try_cast_expr)
+                } else {
+                    BinaryExpr::new(try_cast_expr, Operator::Gt, literal_expr)
+                };
+
+                let result = unwrap_cast_in_comparison(Arc::new(binary_expr), &schema)?;
+                assert!(
+                    !result.transformed,
+                    "{cast_type:?}, literal_on_left={literal_on_left}"
+                );
+                let actual =
+                    result.data.evaluate(&batch)?.into_array(batch.num_rows())?;
+                assert_eq!(
+                    actual.as_any().downcast_ref::<BooleanArray>().unwrap(),
+                    &expected,
+                    "{cast_type:?}, literal_on_left={literal_on_left}"
+                );
+            }
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn test_unwrap_cast_preserves_non_comparison_operators() {
         // Test that non-comparison operators in AND/OR expressions are preserved