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


##########
datafusion/expr-common/src/casts.rs:
##########
@@ -76,6 +109,530 @@ 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, 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 safety gate: the source→target cast must normally 
be
+    // value-preserving over 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 (injective) and order-preserving at the family level.
+/// Timestamp widening is classified family-safe but retains its pre-existing
+/// overflow limitation at extreme source values; it is not a full-domain
+/// guarantee. Other accepted casts can be round-tripped without information
+/// loss and preserve comparison ordering.
+///
+/// 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 the widening overflow limitation 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 only: target integer digits and scale must
+        // be at least as large as source.
+        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!(),
+    }
+}
+
+/// Target must have at least as many integer digits and at least as large a
+/// scale as source.
+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);
+    tgt_int >= src_int && s_tgt >= s_src

Review Comment:
   `decimal(38,2)_col = 4.00` (Float64 literal) is coerced to `CAST(col AS 
Decimal128(38, 15))`. Main unwraps this to `col = Decimal128(4.00,38,2)`; here 
`tgt_int (23) < src_int (36)` keeps the cast, so Bloom filters stop being 
evaluated. That's why `prune_decimal_eq` had to go from 
`matched_by_bloom_filter(Some(2))` to `Some(0)`. The only hazard is a 
strict-cast overflow error for values ≥ 1e23, which is the same trade-off this 
PR already accepts for timestamp widening. Suggest allowing the coercion-capped 
shape:
   
   ```diff
        let src_int = (p_src as i16) - (s_src as i16);
        let tgt_int = (p_tgt as i16) - (s_tgt as i16);
   -    tgt_int >= src_int && s_tgt >= s_src
   +    // Coercion caps precision at the type maximum and can drop integer 
digits.
   +    // As with timestamp widening, the only difference is the strict cast's
   +    // overflow error.
   +    let capped = match tgt {
   +        DataType::Decimal128(p, _) => *p == DECIMAL128_MAX_PRECISION,
   +        DataType::Decimal64(p, _) => *p == DECIMAL64_MAX_PRECISION,
   +        DataType::Decimal32(p, _) => *p == DECIMAL32_MAX_PRECISION,
   +        _ => false,
   +    };
   +    s_tgt >= s_src && (tgt_int >= src_int || capped)
   ```



-- 
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