jayzhan211 commented on code in PR #22906:
URL: https://github.com/apache/datafusion/pull/22906#discussion_r4155374712
##########
datafusion/expr-common/src/casts.rs:
##########
@@ -76,6 +94,506 @@ 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>> {
+ if let Some(preimage) = maybe_range_preimage(source_type, target_type,
lit_value)? {
+ return Ok(Some(preimage));
+ }
+
+ if is_date_narrowing_cast(source_type, target_type) {
+ return Ok(None);
+ }
+
+ 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(
Review Comment:
Casting a timezone-naive timestamp to a timezone-aware one isn't a pure unit
change: arrow-cast calls `adjust_timestamp_to_timezone` for `(None, Some(tz))`,
which reads the stored value as local time and shifts it to UTC.
`timestamp_narrowing_range_preimage` builds the bucket from the raw literal and
ignores this, so the range is off by the zone offset. On `main` this query
wasn't rewritten. The exact path (`is_exact_cast_safe`, same-unit and widening)
has the same hole, and #22142 lists it as "timezone silently dropped".
```sql
create table t(ts timestamp, l timestamp) as values
('2024-01-01T00:00:00.5'::timestamp, '2024-01-01T00:00:00.5'::timestamp);
-- literal moved into a column, no rewrite: 1
select count(*) from t where arrow_cast(ts, 'Timestamp(Millisecond,
Some("+07:00"))')
= arrow_cast(l, 'Timestamp(Millisecond, Some("+07:00"))');
-- literal, rewritten to `ts >= 1704042000500000000 AND ts <
1704042000501000000`: 0
select count(*) from t where arrow_cast(ts, 'Timestamp(Millisecond,
Some("+07:00"))')
= arrow_cast('2024-01-01T00:00:00.5'::timestamp, 'Timestamp(Millisecond,
Some("+07:00"))');
```
Fix: reject this timezone pair in both places (the second check also covers
the `IN`-list path):
```rs
/// Arrow shifts values when casting a timezone-naive timestamp to a
/// timezone-aware one (local wall clock -> UTC), so the cast has no
/// literal-independent preimage.
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(_)))
)
}
```
```diff
pub fn cast_predicate_preimage(
@@
) -> Result<Option<CastPredicatePreimage>> {
+ 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)? {
@@ fn is_exact_cast_safe
if is_timestamp_cast(src, tgt) {
- return !is_timestamp_precision_narrowing_cast(src, tgt);
+ return !is_timestamp_precision_narrowing_cast(src, tgt)
+ && !is_naive_to_tz_timestamp_cast(src, tgt);
}
```
Please add a result-level slt for this case, not just `EXPLAIN`. With the
fix, `optimizer_ine) -> ms("UTC")`) goes back to keeping the cast, unless you
explicitly allow zero-offset`"UTC"`/`"+00:00"`.
##########
datafusion/expr-common/src/casts.rs:
##########
@@ -156,16 +654,6 @@ pub fn is_timestamp_precision_narrowing_cast(
pub fn is_date_narrowing_cast(from_type: &DataType, to_type: &DataType) ->
bool {
Review Comment:
The `is_date_narrowing_cast` early returns here, at `casts.rs:159` and at
`physical-expr/src/simplifier/unwrap_cast.rs:134` are dead code.
`is_exact_cast_safe` already rejects `Date64 -> Date32` (including the
dictionary-wrapped form, which this bare-type check misses), and the range,
int→string and widening paths can't match date types. Dropping all three leaves
the allowlist as the single gate:
```diff
- if is_date_narrowing_cast(source_type, target_type) {
- return Ok(None);
- }
```
##########
datafusion/physical-expr/src/simplifier/unwrap_cast.rs:
##########
@@ -122,28 +123,84 @@ fn extract_cast_info(
/// 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,
cast_type: &DataType,
+ 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)?;
- if is_timestamp_precision_narrowing_cast(&inner_type, cast_type)
- || is_date_narrowing_cast(&inner_type, cast_type)
- {
+ if is_date_narrowing_cast(&inner_type, cast_type) {
return Ok(None);
}
- // 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)));
+ match cast_predicate_preimage(&inner_type, cast_type, op, literal_value)? {
+ Some(CastPredicatePreimage::Exact(casted_literal)) => {
+ let literal_expr = lit(casted_literal);
+ let binary_expr = BinaryExpr::new(inner_expr, op, literal_expr);
+ Ok(Some(Arc::new(binary_expr)))
+ }
+ Some(CastPredicatePreimage::Range(interval)) => {
+ rewrite_with_preimage(interval, op, inner_expr).map(Some)
+ }
+ None => Ok(None),
}
+}
- Ok(None)
+fn rewrite_with_preimage(
+ interval: datafusion_expr_common::interval_arithmetic::Interval,
+ op: Operator,
+ expr: Arc<dyn PhysicalExpr>,
+) -> Result<Arc<dyn PhysicalExpr>> {
+ let (lower, upper) = interval.into_bounds();
+ let (lower, upper) = (lit(lower), lit(upper));
+
+ let rewritten_expr = match op {
+ Operator::Lt => binary(Arc::clone(&expr), Operator::Lt, lower),
+ Operator::GtEq => binary(Arc::clone(&expr), Operator::GtEq, lower),
+ Operator::Gt => binary(Arc::clone(&expr), Operator::GtEq, upper),
+ Operator::LtEq => binary(Arc::clone(&expr), Operator::Lt, upper),
+ Operator::Eq => binary(
+ binary(Arc::clone(&expr), Operator::GtEq, lower),
+ Operator::And,
+ binary(expr, Operator::Lt, upper),
+ ),
+ Operator::NotEq => binary(
+ binary(Arc::clone(&expr), Operator::Lt, lower),
+ Operator::Or,
+ binary(expr, Operator::GtEq, upper),
+ ),
+ Operator::IsNotDistinctFrom => binary(
+ binary(
+ is_not_null(Arc::clone(&expr))?,
+ Operator::And,
+ binary(Arc::clone(&expr), Operator::GtEq, lower),
+ ),
+ Operator::And,
+ binary(expr, Operator::Lt, upper),
+ ),
+ Operator::IsDistinctFrom => binary(
+ binary(
+ binary(Arc::clone(&expr), Operator::Lt, lower),
+ Operator::Or,
+ binary(Arc::clone(&expr), Operator::GtEq, upper),
+ ),
+ Operator::Or,
+ is_null(expr)?,
+ ),
+ _ => unreachable!("preimage only supports comparison operators"),
Review Comment:
```suggestion
_ => return internal_err!("Expect comparison operators, got {op}"),
```
--
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]
