Skip to content

Preserve timestamp casts when unit or timezone changes #128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
osipovartem merged 1 commit into from
Oct 9, 2026
Merged

Preserve timestamp casts when unit or timezone changes #128

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
167 changes: 153 additions & 14 deletions datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,7 @@ 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_supported_type, try_cast_literal_to_type,
};

pub(super) fn unwrap_cast_in_comparison_for_binary(
Expand Down Expand Up @@ -135,7 +134,7 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary(
return false;
};

if is_timestamp_precision_narrowing_cast(&expr_type, field.data_type())
if changes_timestamp_semantics(&expr_type, field.data_type())
|| is_date_narrowing_cast(&expr_type, field.data_type())
{
return false;
Expand Down Expand Up @@ -180,7 +179,7 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist(
return false;
}

if is_timestamp_precision_narrowing_cast(&expr_type, field.data_type())
if changes_timestamp_semantics(&expr_type, field.data_type())
|| is_date_narrowing_cast(&expr_type, field.data_type())
{
return false;
Expand All @@ -205,6 +204,15 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist(
true
}

fn changes_timestamp_semantics(from_type: &DataType, to_type: &DataType) -> bool {
// Changing units can truncate or overflow; attaching a timezone can shift values.
matches!(
(from_type, to_type),
(DataType::Timestamp(from_unit, from_tz), DataType::Timestamp(to_unit, to_tz))
if from_unit != to_unit || from_tz != to_tz
)
}

///// Tries to move a cast from an expression (such as column) to the literal other side of a comparison operator./
///
/// Specifically, rewrites
Expand Down Expand Up @@ -598,15 +606,57 @@ mod tests {
}

#[test]
/// Basic integration test for unwrapping casts with different timezones
fn test_unwrap_cast_with_timestamp_nanos() {
fn test_not_unwrap_cast_with_timestamp_timezone_change() {
use arrow::array::{Array, TimestampNanosecondArray};
use arrow::compute::{CastOptions, cast_with_options};

let offset: Arc<str> = Arc::from("+01:00");
let offset_type =
DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::clone(&offset)));
let source = TimestampNanosecondArray::from(vec![Some(0)]);
let shifted = cast_with_options(
&source,
&offset_type,
&CastOptions {
safe: false,
..Default::default()
},
)
.unwrap();
let shifted = shifted
.as_any()
.downcast_ref::<TimestampNanosecondArray>()
.unwrap();
assert_ne!(shifted.value(0), 0);

let schema = expr_test_schema();
// cast(ts_nano as Timestamp(Nanosecond, UTC)) < 1666612093000000000::Timestamp(Nanosecond, Utc))
let expr_lt = try_cast(col("ts_nano_none"), timestamp_nano_utc_type())
.lt(lit_timestamp_nano_utc(1666612093000000000));
let widened = try_cast(col("ts_nano_none"), offset_type);
let offset_literal = |value| {
lit(ScalarValue::TimestampNanosecond(
Some(value),
Some(Arc::clone(&offset)),
))
};
let equality = widened.clone().eq(offset_literal(0));
assert_eq!(optimize_test(equality.clone(), &schema), equality);

let list = in_list(
widened,
vec![
offset_literal(0),
offset_literal(1),
offset_literal(2),
offset_literal(3),
],
false,
);
assert_eq!(optimize_test(list.clone(), &schema), list);

let unchanged = try_cast(col("ts_nano_utf"), timestamp_nano_utc_type())
.lt(lit_timestamp_nano_utc(1_666_612_093_000_000_000));
let expected =
col("ts_nano_none").lt(lit_timestamp_nano_none(1666612093000000000));
assert_eq!(optimize_test(expr_lt, &schema), expected);
col("ts_nano_utf").lt(lit_timestamp_nano_utc(1_666_612_093_000_000_000));
assert_eq!(optimize_test(unchanged, &schema), expected);
}

#[test]
Expand All @@ -619,13 +669,102 @@ mod tests {
}

#[test]
fn test_unwrap_cast_timestamp_precision_widening() {
fn test_not_unwrap_cast_timestamp_precision_widening() {
let schema = expr_test_schema();
let expr_input = cast(col("ts_millis_none"), timestamp_nano_none_type())
.eq(lit_timestamp_nano_none(1_000_000));
let expected = col("ts_millis_none").eq(lit_timestamp_millis_none(1));
assert_eq!(optimize_test(expr_input.clone(), &schema), expr_input);

assert_eq!(optimize_test(expr_input, &schema), expected);
let exact_list = in_list(
cast(col("ts_millis_none"), timestamp_nano_none_type()),
vec![
lit_timestamp_nano_none(1_000_000),
lit_timestamp_nano_none(2_000_000),
lit_timestamp_nano_none(3_000_000),
lit_timestamp_nano_none(4_000_000),
],
false,
);
assert_eq!(optimize_test(exact_list.clone(), &schema), exact_list);
}

#[test]
fn test_not_unwrap_try_cast_timestamp_widening_that_can_overflow() {
use arrow::array::{Array, TimestampMillisecondArray};
use arrow::compute::{CastOptions, cast_with_options};

let overflow_ms = i64::MAX / 1_000_000 + 1;
let values = TimestampMillisecondArray::from(vec![Some(overflow_ms)]);
assert!(
cast_with_options(
&values,
&timestamp_nano_none_type(),
&CastOptions {
safe: false,
..Default::default()
},
)
.is_err()
);
let safe_cast = cast_with_options(
&values,
&timestamp_nano_none_type(),
&CastOptions {
safe: true,
..Default::default()
},
)
.unwrap();
assert!(safe_cast.is_null(0));

let schema = expr_test_schema();
let widened = try_cast(col("ts_millis_none"), timestamp_nano_none_type());
let equality = widened.clone().eq(lit_timestamp_nano_none(1_000_000));
assert_eq!(optimize_test(equality.clone(), &schema), equality);

let range = widened.clone().gt(lit_timestamp_nano_none(0));
assert_eq!(optimize_test(range.clone(), &schema), range);

let not_in = in_list(
widened,
vec![
lit_timestamp_nano_none(1_000_000),
lit_timestamp_nano_none(2_000_000),
lit_timestamp_nano_none(3_000_000),
lit_timestamp_nano_none(4_000_000),
],
true,
);
assert_eq!(optimize_test(not_in.clone(), &schema), not_in);
}

#[test]
fn test_not_unwrap_timestamp_precision_widening_with_inexact_literal() {
let schema = expr_test_schema();
let column = cast(col("ts_millis_none"), timestamp_nano_none_type());

for literal in [1_000_001, -1] {
let equality = column.clone().eq(lit_timestamp_nano_none(literal));
assert_eq!(optimize_test(equality.clone(), &schema), equality);

let reversed = lit_timestamp_nano_none(literal).eq(column.clone());
assert_eq!(optimize_test(reversed.clone(), &schema), reversed);

let range = column.clone().lt(lit_timestamp_nano_none(literal));
assert_eq!(optimize_test(range.clone(), &schema), range);
}

let inexact_list = in_list(
column,
vec![
lit_timestamp_nano_none(1_000_000),
lit_timestamp_nano_none(2_000_001),
lit_timestamp_nano_none(3_000_000),
lit_timestamp_nano_none(4_000_000),
],
false,
);
assert_eq!(optimize_test(inexact_list.clone(), &schema), inexact_list);
}

fn optimize_test(expr: Expr, schema: &DFSchemaRef) -> Expr {
Expand Down
Loading