diff --git a/datafusion/expr-common/src/casts.rs b/datafusion/expr-common/src/casts.rs index 3518c02772672..620e1b60e087c 100644 --- a/datafusion/expr-common/src/casts.rs +++ b/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; diff --git a/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs b/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs index d0165b568a753..3fe7428e748ef 100644 --- a/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs +++ b/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 diff --git a/datafusion/physical-expr/src/simplifier/unwrap_cast.rs b/datafusion/physical-expr/src/simplifier/unwrap_cast.rs index 3e67fc8291a4e..20e92c4a37025 100644 --- a/datafusion/physical-expr/src/simplifier/unwrap_cast.rs +++ b/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,7 +63,7 @@ fn try_unwrap_cast_binary( schema: &Schema, ) -> Result>> { // 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::(), ) && binary.op().supports_propagation() @@ -71,6 +71,7 @@ fn try_unwrap_cast_binary( 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::(), 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, -) -> Option<(&Arc, &DataType)> { +) -> Option<(&Arc, &DataType, bool)> { if let Some(cast) = expr.downcast_ref::() { - Some((cast.expr(), cast.cast_type())) + Some((cast.expr(), cast.cast_type(), false)) } else if let Some(try_cast) = expr.downcast_ref::() { - 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, literal_value: &ScalarValue, cast_type: &DataType, + is_try_cast: bool, op: Operator, schema: &Schema, ) -> Result>> { @@ -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::().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 diff --git a/datafusion/sqllogictest/test_files/cast.slt b/datafusion/sqllogictest/test_files/cast.slt index 414ee2f636db5..6405e8f51c262 100644 --- a/datafusion/sqllogictest/test_files/cast.slt +++ b/datafusion/sqllogictest/test_files/cast.slt @@ -21,6 +21,77 @@ SELECT cast(c as varchar) FROM (SELECT 1 as c) ---- 1 +# Narrowing TRY_CAST comparisons must preserve NULL-on-overflow semantics. +statement ok +CREATE TABLE try_cast_narrowing AS VALUES + (-2147483649::BIGINT), + (-2147483648::BIGINT), + (5::BIGINT), + (2147483647::BIGINT), + (2147483648::BIGINT), + (NULL::BIGINT); + +query IIBB +SELECT column1, + try_cast(column1 AS INT) AS narrowed, + try_cast(column1 AS INT) > 1, + 1 < try_cast(column1 AS INT) +FROM try_cast_narrowing +ORDER BY column1 NULLS LAST; +---- +-2147483649 NULL NULL NULL +-2147483648 -2147483648 false false +5 5 true true +2147483647 2147483647 true true +2147483648 NULL NULL NULL +NULL NULL NULL NULL + +query I +SELECT column1 +FROM try_cast_narrowing +WHERE try_cast(column1 AS INT) > 1 +ORDER BY column1; +---- +5 +2147483647 + +query IBB +SELECT column1, + try_cast(column1 AS INT) IN (1, 2, 3, 4, 5), + try_cast(column1 AS INT) NOT IN (1, 2, 3, 4, 5) +FROM try_cast_narrowing +ORDER BY column1 NULLS LAST; +---- +-2147483649 NULL NULL +-2147483648 false true +5 true false +2147483647 false true +2147483648 NULL NULL +NULL NULL NULL + +statement ok +DROP TABLE try_cast_narrowing; + +# A signed-to-unsigned cast can narrow the domain without changing the width. +statement ok +CREATE TABLE try_cast_signed_to_unsigned AS VALUES + (-1::INT), (0::INT), (1::INT), (NULL::INT); + +query IIB +SELECT column1, + try_cast(column1 AS INTEGER UNSIGNED) AS unsigned_value, + try_cast(column1 AS INTEGER UNSIGNED) > 0::INTEGER UNSIGNED +FROM try_cast_signed_to_unsigned +ORDER BY column1 NULLS LAST; +---- +-1 NULL NULL +0 0 false +1 1 true +NULL NULL NULL + +statement ok +DROP TABLE try_cast_signed_to_unsigned; + # cast tinyint query I SELECT cast(10 as tinyint) diff --git a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt index 1dc66f398d40d..810f894c6cc7b 100644 --- a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt +++ b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt @@ -1099,3 +1099,66 @@ set datafusion.execution.parquet.pushdown_filters = false; statement ok DROP TABLE case_without_else; + +# Narrowing TRY_CAST must keep NULL-on-overflow semantics in the physical +# simplifier, independently of the logical cast-unwrapping rule. +statement ok +COPY (VALUES (5::BIGINT), (7::BIGINT), (9999999999::BIGINT), + (10000000000::BIGINT), (-2147483649::BIGINT), (NULL::BIGINT)) +TO 'test_files/scratch/parquet_filter_pushdown/try_cast_narrowing.parquet' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' '2'); + +statement ok +CREATE EXTERNAL TABLE try_cast_narrowing +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_filter_pushdown/try_cast_narrowing.parquet'; + +statement ok +SET datafusion.optimizer.max_passes = 0; + +statement ok +SET datafusion.execution.parquet.pruning = false; + +statement ok +SET datafusion.execution.parquet.pushdown_filters = false; + +query I +SELECT count(*) FROM try_cast_narrowing WHERE try_cast(column1 AS INT) > 1; +---- +2 + +statement ok +SET datafusion.execution.parquet.pushdown_filters = true; + +query I +SELECT count(*) FROM try_cast_narrowing WHERE try_cast(column1 AS INT) > 1; +---- +2 + +query I +SELECT count(*) FROM try_cast_narrowing WHERE 1 < try_cast(column1 AS INT); +---- +2 + +query I +SELECT count(*) FROM try_cast_narrowing WHERE try_cast(column1 AS INT) != 5; +---- +1 + +query I +SELECT count(*) FROM try_cast_narrowing WHERE try_cast(column1 AS INT) < 1; +---- +0 + +statement ok +RESET datafusion.optimizer.max_passes; + +statement ok +RESET datafusion.execution.parquet.pruning; + +statement ok +RESET datafusion.execution.parquet.pushdown_filters; + +statement ok +DROP TABLE try_cast_narrowing;