jayzhan211 commented on code in PR #22906:
URL: https://github.com/apache/datafusion/pull/22906#discussion_r4294788543


##########
datafusion/expr-common/src/casts.rs:
##########
@@ -76,6 +112,546 @@ pub fn try_cast_literal_to_type(
         .or_else(|| try_cast_binary(lit_value, target_type))
 }
 
+/// Computes a source-domain preimage for `CAST(source AS target_type) OP 
literal`.
+///
+/// This is the shared semantic core for logical and physical cast-predicate
+/// rewrites. It returns an exact same-operator source bound or literal as
+/// [`CastPredicatePreimage::Exact`] where moving the cast preserves comparison
+/// semantics subject to the documented family-level overflow exceptions, and a
+/// [`CastPredicatePreimage::Range`] for many-to-one casts with known preimages
+/// such as timestamp precision narrowing.
+pub fn cast_predicate_preimage(
+    source_type: &DataType,
+    target_type: &DataType,
+    op: Operator,
+    lit_value: &ScalarValue,
+) -> Result<Option<CastPredicatePreimage>> {
+    // Arrow adjusts a timezone-naive value to UTC when the target is
+    // timezone-aware (`adjust_timestamp_to_timezone`). The precision-only
+    // arithmetic here does not model that adjustment, so retain the cast. This
+    // dispatcher check covers the range, exact and widening paths below.
+    if is_naive_to_tz_timestamp_cast(source_type, target_type) {
+        return Ok(None);
+    }
+
+    if let Some(preimage) = maybe_range_preimage(source_type, target_type, 
lit_value)? {
+        return Ok(Some(preimage));
+    }
+
+    if let Some(value) =
+        exact_preimage_int_to_str_eq_like(source_type, target_type, op, 
lit_value)
+    {
+        return Ok(Some(CastPredicatePreimage::Exact(value)));
+    }
+
+    if let Some(value) = exact_preimage_cast(source_type, target_type, 
lit_value) {
+        return Ok(Some(CastPredicatePreimage::Exact(value)));
+    }
+
+    Ok(
+        timestamp_widening_ordered_preimage(source_type, target_type, op, 
lit_value)
+            .map(CastPredicatePreimage::Exact),
+    )
+}
+
+fn maybe_range_preimage(
+    source_type: &DataType,
+    target_type: &DataType,
+    lit_value: &ScalarValue,
+) -> Result<Option<CastPredicatePreimage>> {
+    if !is_timestamp_precision_narrowing_cast(source_type, target_type) {
+        return Ok(None);
+    }
+
+    Ok(
+        timestamp_narrowing_range_preimage(source_type, target_type, 
lit_value)?
+            .map(CastPredicatePreimage::Range),
+    )
+}
+
+/// Computes a singleton source-domain literal for exact cast-predicate 
rewrites.
+///
+/// This intentionally returns `None` for timestamp precision narrowing: those
+/// casts are many-to-one and need range preimages instead.
+pub fn exact_preimage_cast(
+    source_type: &DataType,
+    target_type: &DataType,
+    lit_value: &ScalarValue,
+) -> Option<ScalarValue> {
+    // Apply a family-level gate: the source→target cast must normally preserve
+    // the full source domain. Timestamp precision narrowing is handled as a
+    // Range preimage, not an Exact preimage.
+    // Timestamp widening / equal-unit casts still require the literal
+    // round-trip check below; widening retains its pre-existing overflow
+    // limitation at extreme source values.
+    if !is_exact_cast_safe(source_type, target_type) {
+        return None;
+    }
+
+    let source_value = try_cast_literal_to_type(lit_value, source_type)?;
+    if is_timestamp_cast(source_type, target_type) {
+        let round_tripped = try_cast_literal_to_type(&source_value, 
target_type)?;
+        if &round_tripped != lit_value {
+            return None;
+        }
+    }
+
+    Some(source_value)
+}
+
+/// Returns true when casting a timestamp from `source_type` to `target_type`
+/// loses timestamp precision.
+pub fn is_timestamp_precision_narrowing_cast(
+    source_type: &DataType,
+    target_type: &DataType,
+) -> bool {
+    let (DataType::Timestamp(source_unit, _), DataType::Timestamp(target_unit, 
_)) =
+        (source_type, target_type)
+    else {
+        return false;
+    };
+
+    timestamp_unit_scale(source_unit) > timestamp_unit_scale(target_unit)
+}
+
+fn is_timestamp_cast(source_type: &DataType, target_type: &DataType) -> bool {
+    matches!(
+        (source_type, target_type),
+        (DataType::Timestamp(_, _), DataType::Timestamp(_, _))
+    )
+}
+
+/// Returns true when casting a timezone-naive timestamp to a timezone-aware 
one.
+///
+/// Arrow adjusts such values to UTC (`adjust_timestamp_to_timezone`), which 
the
+/// precision-only preimage arithmetic here does not model, so the cast is 
kept.
+fn is_naive_to_tz_timestamp_cast(source_type: &DataType, target_type: 
&DataType) -> bool {
+    matches!(
+        (source_type, target_type),
+        (
+            DataType::Timestamp(_, None),
+            DataType::Timestamp(_, Some(_))
+        )
+    )
+}
+
+/// Returns `true` when the cast from `source_type` to `target_type` is
+/// value-preserving and order-preserving at the family level. Timestamp
+/// widening retains its pre-existing overflow limitation. Same-variant
+/// maximum-precision decimal widening is an explicit exception: it may lose
+/// integer digits and replace cast overflow errors or `TRY_CAST` NULLs with
+/// booleans for out-of-range sources. This family-level rule is type-based and
+/// does not distinguish implicit coercions from explicit casts.
+///
+/// This is a conservative gate used by [`exact_preimage_cast`]. Timestamp
+/// precision narrowing is not exact-safe (it is handled by
+/// [`CastPredicatePreimage::Range`]); timestamp widening / equal-unit casts 
are
+/// family-safe subject to their documented overflow limitations and still
+/// require the caller's literal round-trip check.
+fn is_exact_cast_safe(source_type: &DataType, target_type: &DataType) -> bool {
+    // Unwrap at most one level of dictionary for family-level checking.
+    let (src, tgt) = match (source_type, target_type) {
+        (DataType::Dictionary(_, v_src), DataType::Dictionary(_, v_tgt)) => {
+            (v_src.as_ref(), v_tgt.as_ref())
+        }
+        (DataType::Dictionary(_, v_src), tgt) => (v_src.as_ref(), tgt),
+        (src, DataType::Dictionary(_, v_tgt)) => (src, v_tgt.as_ref()),
+        _ => (source_type, target_type),
+    };
+
+    if src == tgt {
+        return true;
+    }
+    if is_timestamp_cast(src, tgt) {
+        // Timestamp precision narrowing is handled as a Range preimage, and a
+        // naive → timezone-aware cast reads the raw value as local wall-clock
+        // time and shifts it to UTC; neither is an exact preimage.
+        return !is_timestamp_precision_narrowing_cast(src, tgt)
+            && !is_naive_to_tz_timestamp_cast(src, tgt);
+    }
+
+    // Whitelist of family-level safe casts.  Anything not listed here is
+    // considered potentially information-losing and blocked.
+    if is_integer_type(src) && is_integer_type(tgt) {
+        return is_safe_integer_cast(src, tgt);
+    }
+    if matches!(
+        (src, tgt),
+        (DataType::Date32, DataType::Int32)
+            | (DataType::Int32, DataType::Date32)
+            | (DataType::Date64, DataType::Int64)
+            | (DataType::Int64, DataType::Date64)
+    ) {
+        // Date32 ↔ Int32 and Date64 ↔ Int64 share the same internal
+        // representation.
+        return true;
+    }
+    if matches!((src, tgt), (DataType::Date32, DataType::Date64)) {
+        // Date32 → Date64 is injective: every source day maps to its midnight
+        // millisecond value. The literal conversion below enforces the exact
+        // day-boundary requirement in the reverse direction.
+        return true;
+    }
+    if is_supported_string_type(src) && is_supported_string_type(tgt) {
+        // String width casts among Utf8 / LargeUtf8 / Utf8View.
+        return true;
+    }
+    if is_decimal_type(src) && is_decimal_type(tgt) {
+        // Decimal widening normally retains integer digits and increases
+        // scale; same-variant maximum precision permits the narrow overflow
+        // exception documented by is_safe_decimal_widening.
+        return is_safe_decimal_widening(src, tgt);
+    }
+    if is_integer_type(src) && is_decimal_type(tgt) {
+        // int / uint → decimal: allow only when the target can represent
+        // the full source integer domain.
+        return is_safe_integer_to_decimal(src, tgt);
+    }
+    if matches!((src, tgt), (DataType::FixedSizeBinary(_), DataType::Binary)) {
+        // FixedSizeBinary(n) → Binary is always widening.
+        return true;
+    }
+    if matches!(
+        (src, tgt),
+        (DataType::FixedSizeBinary(n1), DataType::FixedSizeBinary(n2)) if n1 
== n2
+    ) {
+        return true;
+    }
+    false
+}
+
+/// Returns `true` for integer types (signed and unsigned, excluding Date and
+/// Decimal).
+fn is_integer_type(dt: &DataType) -> bool {
+    matches!(
+        dt,
+        DataType::Int8
+            | DataType::Int16
+            | DataType::Int32
+            | DataType::Int64
+            | DataType::UInt8
+            | DataType::UInt16
+            | DataType::UInt32
+            | DataType::UInt64
+    )
+}
+
+/// Integer cast that is value-preserving over the full source domain.
+fn is_safe_integer_cast(src: &DataType, tgt: &DataType) -> bool {
+    use DataType::*;
+    match (src, tgt) {
+        // Signed → wider signed
+        (Int8, Int16 | Int32 | Int64) => true,
+        (Int16, Int32 | Int64) => true,
+        (Int32, Int64) => true,
+        // Unsigned → wider unsigned
+        (UInt8, UInt16 | UInt32 | UInt64) => true,
+        (UInt16, UInt32 | UInt64) => true,
+        (UInt32, UInt64) => true,
+        // Unsigned → wider signed (full source range fits in target)
+        (UInt8, Int16 | Int32 | Int64) => true,
+        (UInt16, Int32 | Int64) => true,
+        (UInt32, Int64) => true,
+        // Anything else is narrowing, partial, or crosses signedness
+        // without guaranteed domain containment.
+        _ => false,
+    }
+}
+
+fn is_decimal_type(dt: &DataType) -> bool {
+    matches!(
+        dt,
+        DataType::Decimal128(_, _)
+            | DataType::Decimal64(_, _)
+            | DataType::Decimal32(_, _)
+    )
+}
+
+fn decimal_precision_scale(dt: &DataType) -> (u8, i8) {
+    match dt {
+        DataType::Decimal128(p, s)
+        | DataType::Decimal64(p, s)
+        | DataType::Decimal32(p, s) => (*p, *s),
+        _ => unreachable!(),
+    }
+}
+
+/// Requires nonnegative scales and normally preserves integer digits while
+/// increasing scale. Same-variant maximum-precision decimals may lose integer
+/// digits: this type-level exception can turn cast overflow errors or
+/// `TRY_CAST` NULLs into booleans for out-of-range source values.
+fn is_safe_decimal_widening(src: &DataType, tgt: &DataType) -> bool {
+    let (p_src, s_src) = decimal_precision_scale(src);
+    let (p_tgt, s_tgt) = decimal_precision_scale(tgt);
+    if s_src < 0 || s_tgt < 0 {
+        return false;
+    }
+    let src_int = (p_src as i16) - (s_src as i16);
+    let tgt_int = (p_tgt as i16) - (s_tgt as i16);
+    let both_max_same_variant = match (src, tgt) {

Review Comment:
   Coercion caps the target at the variant's max precision whatever the source 
precision is, so requiring the source to be max too still misses the case. 
`decimal(30,2)` / `(28,0)` / `(36,4)` vs `4.00` all coerce to `CAST(col AS 
Decimal128(38,15))`, `cast_predicate_preimage` returns `None`, and the cast 
stays (main unwraps all three → Bloom pruning lost). The overflow hazard is the 
same as for a `(38,2)` source, so the source check adds no safety.
   
   Fix (only the `(37,0) -> (38,2)` blocked assertion flips; move it to the 
admitted cases):
   ```diff
   -    let both_max_same_variant = match (src, tgt) {
   -        (DataType::Decimal32(p_src, _), DataType::Decimal32(p_tgt, _)) => {
   -            *p_src == DECIMAL32_MAX_PRECISION && *p_tgt == 
DECIMAL32_MAX_PRECISION
   -        }
   -        (DataType::Decimal64(p_src, _), DataType::Decimal64(p_tgt, _)) => {
   -            *p_src == DECIMAL64_MAX_PRECISION && *p_tgt == 
DECIMAL64_MAX_PRECISION
   -        }
   -        (DataType::Decimal128(p_src, _), DataType::Decimal128(p_tgt, _)) => 
{
   -            *p_src == DECIMAL128_MAX_PRECISION && *p_tgt == 
DECIMAL128_MAX_PRECISION
   -        }
   -        _ => false,
   -    };
   -    s_tgt >= s_src && (tgt_int >= src_int || both_max_same_variant)
   +    // Coercion caps the target at the variant's maximum precision, which 
can
   +    // lose integer digits; same overflow trade-off as timestamp widening.
   +    let capped = match (src, tgt) {
   +        (DataType::Decimal32(_, _), DataType::Decimal32(p, _)) => {
   +            *p == DECIMAL32_MAX_PRECISION
   +        }
   +        (DataType::Decimal64(_, _), DataType::Decimal64(p, _)) => {
   +            *p == DECIMAL64_MAX_PRECISION
   +        }
   +        (DataType::Decimal128(_, _), DataType::Decimal128(p, _)) => {
   +            *p == DECIMAL128_MAX_PRECISION
   +        }
   +        _ => false,
   +    };
   +    s_tgt >= s_src && (tgt_int >= src_int || capped)
   ```
   
   Worth adding a `decimal(30,2)` column to the new `simplify_expr.slt` block 
(`explain ... where price = 4.00` should show `price = Decimal128(4.00,30,2)`).



-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to