From e582791b2820c62a09050a8df69720fe30ccec39 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:32:49 -0500 Subject: [PATCH 1/6] fix: coerce timezone-naive minus timezone-aware timestamps at equal units `Timestamp(u, Some(tz)) - Timestamp(u, None)` at *equal* units is the one mixed timezone-aware/naive pair that arrow can subtract directly, so `BinaryTypeCoercer` inserted no cast at all and arrow subtracted the raw values, reading the naive operand as UTC. Every other unit pairing falls through to `temporal_coercion_strict_timezone` and reads the naive operand in the aware operand's zone, which is what comparisons (`=`, `<`, ...) already do. Coerce the equal-unit case the same way so that all unit pairings agree. Co-Authored-By: Claude Fable 5.1 --- .../expr-common/src/type_coercion/binary.rs | 28 ++++++++++++++++++ .../type_coercion/binary/tests/arithmetic.rs | 29 +++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/datafusion/expr-common/src/type_coercion/binary.rs b/datafusion/expr-common/src/type_coercion/binary.rs index e7c20dde101b0..4ea03943a667c 100644 --- a/datafusion/expr-common/src/type_coercion/binary.rs +++ b/datafusion/expr-common/src/type_coercion/binary.rs @@ -285,6 +285,34 @@ impl<'a> BinaryTypeCoercer<'a> { return Ok(Signature { lhs, rhs, ret }); } Plus | Minus | Multiply | Divide | Modulo => { + // `Timestamp(u, Some(tz)) - Timestamp(u, None)` at *equal* units is + // the one mixed timezone-aware/naive pair arrow can subtract + // directly, so the `self.get_result` probe below would accept it + // and no cast would be inserted: arrow then subtracts the raw + // values, i.e. reads the naive operand as UTC. Every other unit + // pairing falls through to `temporal_coercion_strict_timezone` + // below and reads the naive operand in the aware operand's zone, + // as comparisons already do. Coerce here so that all unit pairings + // agree. + if self.op == &Minus + && matches!( + (lhs, rhs), + (Timestamp(_, Some(_)), Timestamp(_, None)) + | (Timestamp(_, None), Timestamp(_, Some(_))) + ) + && let Some(coerced) = temporal_coercion_strict_timezone(lhs, rhs) + { + let ret = self.get_result(&coerced, &coerced).map_err(|e| { + plan_datafusion_err!( + "Cannot get result type for temporal operation {coerced} {} {coerced}: {e}", self.op + ) + })?; + return Ok(Signature { + lhs: coerced.clone(), + rhs: coerced, + ret, + }); + } if let Ok(ret) = self.get_result(lhs, rhs) { // Temporal arithmetic, e.g. Date32 + Interval diff --git a/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs b/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs index 70a8fc0e35a15..5b5a8c94384d8 100644 --- a/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs +++ b/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs @@ -475,3 +475,32 @@ fn test_decimal_precision_overflow_cross_variant() -> Result<()> { Ok(()) } + +#[test] +fn test_timestamp_minus_mixed_timezone_awareness() -> Result<()> { + let aware_tz: Arc = Arc::from("America/New_York"); + let aware_ms = DataType::Timestamp(Millisecond, Some(Arc::clone(&aware_tz))); + let aware_ns = DataType::Timestamp(Nanosecond, Some(Arc::clone(&aware_tz))); + let naive_ns = DataType::Timestamp(Nanosecond, None); + + // Differing units already coerced to the aware zone; equal units did not. + // Both must now coerce both operands to the aware zone at the finer unit, + // in either operand order. + test_coercion_binary_rule!(aware_ms, naive_ns, Operator::Minus, aware_ns); + test_coercion_binary_rule!(naive_ns, aware_ms, Operator::Minus, aware_ns); + test_coercion_binary_rule!(aware_ns, naive_ns, Operator::Minus, aware_ns); + test_coercion_binary_rule!(naive_ns, aware_ns, Operator::Minus, aware_ns); + + // Pairs that are both aware or both naive are untouched. + test_coercion_binary_rule!(aware_ns, aware_ns, Operator::Minus, aware_ns); + test_coercion_binary_rule!(naive_ns, naive_ns, Operator::Minus, naive_ns); + // Two different aware zones keep their own types (arrow subtracts the + // instants directly); unchanged by this rule. + let other_tz = DataType::Timestamp(Nanosecond, Some(Arc::from("+08:00"))); + test_coercion_binary_rule!(aware_ns, other_tz, Operator::Minus, aware_ns, other_tz); + + // For contrast, comparisons already coerced the mixed pair this way. + test_coercion_binary_rule!(aware_ms, naive_ns, Operator::Eq, aware_ns); + + Ok(()) +} From 8e2ed76c6035cd7b8e588ead0da71790ca95bc76 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:22:59 -0500 Subject: [PATCH 2/6] feat: add session-time-zone aware cast helpers A `Timestamp(unit, Some(tz))` is an instant and `tz` only labels it for display; a `Timestamp(unit, None)` is a wall clock reading. Converting the second to the first has to pick a zone to read the wall clock in, and arrow uses the target type's zone. PostgreSQL and DuckDB instead read it in the session time zone at every *implicit* conversion, so DataFusion currently returns a different instant than either of them whenever coercion inserts such a cast. Add the two helpers the planner and analyzer will use to insert those casts, emitting `CAST(CAST(ts AS ) AS )` when the target zone differs from `datafusion.execution.time_zone`: the inner cast reads the wall clock in the session zone and the outer one only relabels, so the coerced type is unchanged and only the instant differs. Nested naive timestamps (`List`, `LargeList`, `FixedSizeList`, `Struct`, `Map`, `Dictionary`) are re-zoned leaf by leaf. Explicit conversions are deliberately not routed through these helpers: `AT TIME ZONE`, `arrow_cast`, the DataFrame API and Substrait all keep arrow's semantics, as they do in PostgreSQL and DuckDB. With `datafusion.execution.time_zone` unset (the default) both helpers degrade to a plain `cast_to` and plans are unchanged. Co-Authored-By: Claude Fable 5.1 --- datafusion/expr/src/expr_rewriter/mod.rs | 55 ++- datafusion/expr/src/expr_schema.rs | 41 +- datafusion/expr/src/logical_plan/builder.rs | 87 ++-- datafusion/expr/src/type_coercion/mod.rs | 1 + .../src/type_coercion/session_time_zone.rs | 441 ++++++++++++++++++ 5 files changed, 588 insertions(+), 37 deletions(-) create mode 100644 datafusion/expr/src/type_coercion/session_time_zone.rs diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index 4e9839e2f7479..7f756d2b4deb3 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -24,6 +24,7 @@ use std::sync::Arc; use crate::expr::{Alias, Sort, Unnest}; use crate::logical_plan::Projection; +use crate::type_coercion::session_time_zone::cast_to_with_session_time_zone; use crate::{Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder}; use datafusion_common::TableReference; @@ -221,20 +222,46 @@ pub fn strip_outer_reference(expr: Expr) -> Expr { /// Returns plan with expressions coerced to types compatible with /// schema types +/// +/// Timezone-naive timestamps are read in the zone of the type they are coerced +/// to. Use [`coerce_plan_expr_for_schema_with_session_time_zone`] to read them +/// in `datafusion.execution.time_zone` instead. pub fn coerce_plan_expr_for_schema( plan: LogicalPlan, schema: &DFSchema, +) -> Result { + coerce_plan_expr_for_schema_with_session_time_zone(plan, schema, None) +} + +/// Returns plan with expressions coerced to types compatible with schema types, +/// reading a timezone-naive timestamp in `session_time_zone` wherever a cast +/// this function inserts makes one timezone-aware. +/// +/// `session_time_zone` is `datafusion.execution.time_zone`; see +/// [`cast_to_with_session_time_zone`] for exactly what changes. Passing `None` +/// is [`coerce_plan_expr_for_schema`]: every plan is then byte-identical to the +/// one DataFusion produced before the session time zone existed. +/// +/// Only the casts inserted here are affected. An expression that already has +/// the target type is returned untouched, so a cast the *user* wrote — `AT TIME +/// ZONE`, `arrow_cast` — keeps arrow's semantics. +pub fn coerce_plan_expr_for_schema_with_session_time_zone( + plan: LogicalPlan, + schema: &DFSchema, + session_time_zone: Option<&str>, ) -> Result { match plan { // special case Projection to avoid adding multiple projections LogicalPlan::Projection(Projection { expr, input, .. }) => { - let new_exprs = coerce_exprs_for_schema(expr, input.schema(), schema)?; + let new_exprs = + coerce_exprs_for_schema(expr, input.schema(), schema, session_time_zone)?; let projection = Projection::try_new(new_exprs, input)?; Ok(LogicalPlan::Projection(projection)) } _ => { let exprs: Vec = plan.schema().iter().map(Expr::from).collect(); - let new_exprs = coerce_exprs_for_schema(exprs, plan.schema(), schema)?; + let new_exprs = + coerce_exprs_for_schema(exprs, plan.schema(), schema, session_time_zone)?; let add_project = new_exprs.iter().any(|expr| expr.try_as_col().is_none()); if add_project { let projection = Projection::try_new(new_exprs, Arc::new(plan))?; @@ -250,6 +277,7 @@ fn coerce_exprs_for_schema( exprs: Vec, src_schema: &DFSchema, dst_schema: &DFSchema, + session_time_zone: Option<&str>, ) -> Result> { exprs .into_iter() @@ -259,7 +287,13 @@ fn coerce_exprs_for_schema( if new_type != &expr.get_type(src_schema)? { match expr { Expr::Alias(Alias { expr, name, .. }) => { - Ok(expr.cast_to(new_type, src_schema)?.alias(name)) + Ok(cast_to_with_session_time_zone( + *expr, + new_type, + src_schema, + session_time_zone, + )? + .alias(name)) } #[expect(deprecated)] Expr::Wildcard { .. } => Ok(expr), @@ -270,9 +304,20 @@ fn coerce_exprs_for_schema( // (see: https://github.com/apache/datafusion/issues/18818) Expr::Column(ref column) => { let name = column.name().to_owned(); - Ok(expr.cast_to(new_type, src_schema)?.alias(name)) + Ok(cast_to_with_session_time_zone( + expr, + new_type, + src_schema, + session_time_zone, + )? + .alias(name)) } - _ => Ok(expr.cast_to(new_type, src_schema)?), + _ => cast_to_with_session_time_zone( + expr, + new_type, + src_schema, + session_time_zone, + ), } } } diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 75b3a60af1465..e90c49cf5a317 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -26,6 +26,7 @@ use crate::expr::{FieldMetadata, LambdaVariable}; use crate::higher_order_function::HigherOrderReturnFieldArgs; use crate::type_coercion::functions::value_fields_with_higher_order_udf_and_lambdas; use crate::type_coercion::functions::{UDFCoercionExt, fields_with_udf}; +use crate::type_coercion::session_time_zone::cast_to_with_session_time_zone; use crate::udf::ReturnFieldArgs; use crate::{ LogicalPlan, Operator, Projection, Subquery, WindowFunctionDefinition, utils, @@ -822,7 +823,30 @@ fn scalar_subquery_nullable(subquery: &Subquery) -> bool { /// new projection with the casted expression. /// 2. **Non-projection plan**: If the subquery isn't a projection, it adds a projection to the plan /// with the casted first column. +/// +/// A timezone-naive timestamp is read in the zone of `cast_to_type`. Use +/// [`cast_subquery_with_session_time_zone`] to read it in +/// `datafusion.execution.time_zone` instead. pub fn cast_subquery(subquery: Subquery, cast_to_type: &DataType) -> Result { + cast_subquery_with_session_time_zone(subquery, cast_to_type, None) +} + +/// [`cast_subquery`], reading a timezone-naive timestamp in +/// `session_time_zone` if the cast it inserts makes one timezone-aware. +/// +/// `session_time_zone` is `datafusion.execution.time_zone`; see +/// [`cast_to_with_session_time_zone`] for exactly what changes. Passing `None` +/// is [`cast_subquery`]: the subquery is then byte-identical to the one +/// DataFusion produced before the session time zone existed. +/// +/// Only the cast inserted here is affected. A subquery whose projected +/// expression already has `cast_to_type` is returned untouched, so a cast the +/// *user* wrote — `AT TIME ZONE`, `arrow_cast` — keeps arrow's semantics. +pub fn cast_subquery_with_session_time_zone( + subquery: Subquery, + cast_to_type: &DataType, + session_time_zone: Option<&str>, +) -> Result { if subquery.subquery.schema().field(0).data_type() == cast_to_type { return Ok(subquery); } @@ -830,17 +854,24 @@ pub fn cast_subquery(subquery: Subquery, cast_to_type: &DataType) -> Result { - let cast_expr = projection.expr[0] - .clone() - .cast_to(cast_to_type, projection.input.schema())?; + let cast_expr = cast_to_with_session_time_zone( + projection.expr[0].clone(), + cast_to_type, + projection.input.schema().as_ref(), + session_time_zone, + )?; LogicalPlan::Projection(Projection::try_new( vec![cast_expr], Arc::clone(&projection.input), )?) } _ => { - let cast_expr = Expr::Column(Column::from(plan.schema().qualified_field(0))) - .cast_to(cast_to_type, subquery.subquery.schema())?; + let cast_expr = cast_to_with_session_time_zone( + Expr::Column(Column::from(plan.schema().qualified_field(0))), + cast_to_type, + subquery.subquery.schema().as_ref(), + session_time_zone, + )?; LogicalPlan::Projection(Projection::try_new( vec![cast_expr], subquery.subquery, diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 36aa67bbe7e3d..23f2b33b5afb1 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -37,6 +37,7 @@ use crate::logical_plan::{ Union, Unnest, Values, Window, }; use crate::select_expr::SelectExpr; +use crate::type_coercion::session_time_zone::cast_to_with_session_time_zone; use crate::utils::{ can_hash, check_all_columns_from_schema, columnize_expr, compare_sort_expr, expand_qualified_wildcard, expand_wildcard, expr_to_columns, @@ -210,27 +211,12 @@ impl LogicalPlanBuilder { /// so it's usually better to override the default names with a table alias list. /// /// If the values include params/binders such as $1, $2, $3, etc, then the `param_data_types` should be provided. + /// + /// Timezone-naive timestamps are read in the zone of the column type they + /// are cast to. Use [`Self::values_with_session_time_zone`] to read them in + /// `datafusion.execution.time_zone` instead. pub fn values(values: Vec>) -> Result { - if values.is_empty() { - return plan_err!("Values list cannot be empty"); - } - let n_cols = values[0].len(); - if n_cols == 0 { - return plan_err!("Values list cannot be zero length"); - } - for (i, row) in values.iter().enumerate() { - if row.len() != n_cols { - return plan_err!( - "Inconsistent data length across values list: got {} values in row {} but expected {}", - row.len(), - i, - n_cols - ); - } - } - - // Infer from data itself - Self::infer_data(values) + Self::values_with_session_time_zone(values, None, None) } /// Create a values list based relation, and the schema is inferred from data itself or table schema if provided, consuming @@ -242,14 +228,45 @@ impl LogicalPlanBuilder { /// so it's usually better to override the default names with a table alias list. /// /// If the values include params/binders such as $1, $2, $3, etc, then the `param_data_types` should be provided. + /// + /// Timezone-naive timestamps are read in the zone of the column type they + /// are cast to. Use [`Self::values_with_session_time_zone`] to read them in + /// `datafusion.execution.time_zone` instead. pub fn values_with_schema( values: Vec>, schema: &DFSchemaRef, + ) -> Result { + Self::values_with_session_time_zone(values, Some(schema.as_ref()), None) + } + + /// Create a values list based relation, reading a timezone-naive timestamp + /// in `session_time_zone` wherever a cast this function inserts makes one + /// timezone-aware. + /// + /// The column types are taken from `schema` when it is `Some` + /// ([`Self::values_with_schema`]) and inferred from the values themselves + /// when it is `None` ([`Self::values`]). + /// + /// `session_time_zone` is `datafusion.execution.time_zone`; see + /// [`cast_to_with_session_time_zone`] for exactly what changes. Passing + /// `None` gives the plan DataFusion produced before the session time zone + /// existed. + /// + /// Only the casts inserted here are affected. A cell that already has its + /// column's type is left untouched, so a cast the *user* wrote — `AT TIME + /// ZONE`, `arrow_cast` — keeps arrow's semantics. + pub fn values_with_session_time_zone( + values: Vec>, + schema: Option<&DFSchema>, + session_time_zone: Option<&str>, ) -> Result { if values.is_empty() { return plan_err!("Values list cannot be empty"); } - let n_cols = schema.fields().len(); + let n_cols = match schema { + Some(schema) => schema.fields().len(), + None => values[0].len(), + }; if n_cols == 0 { return plan_err!("Values list cannot be zero length"); } @@ -264,13 +281,20 @@ impl LogicalPlanBuilder { } } - // Check the type of value against the schema - Self::infer_values_from_schema(values, schema) + match schema { + // Check the type of value against the schema + Some(schema) => { + Self::infer_values_from_schema(values, schema, session_time_zone) + } + // Infer from data itself + None => Self::infer_data(values, session_time_zone), + } } fn infer_values_from_schema( values: Vec>, schema: &DFSchema, + session_time_zone: Option<&str>, ) -> Result { let n_cols = values[0].len(); let mut fields = ValuesFields::new(); @@ -294,10 +318,13 @@ impl LogicalPlanBuilder { fields.push(field_type.to_owned(), field_nullable); } - Self::infer_inner(values, fields, schema) + Self::infer_inner(values, fields, schema, session_time_zone) } - fn infer_data(values: Vec>) -> Result { + fn infer_data( + values: Vec>, + session_time_zone: Option<&str>, + ) -> Result { let n_cols = values[0].len(); let schema = DFSchema::empty(); let mut fields = ValuesFields::new(); @@ -351,13 +378,14 @@ impl LogicalPlanBuilder { fields.push_with_metadata(data_type, nullable, common_metadata); } - Self::infer_inner(values, fields, &schema) + Self::infer_inner(values, fields, &schema, session_time_zone) } fn infer_inner( mut values: Vec>, fields: ValuesFields, schema: &DFSchema, + session_time_zone: Option<&str>, ) -> Result { let fields = fields.into_fields(); // wrap cast if data type is not same as common type. @@ -369,7 +397,12 @@ impl LogicalPlanBuilder { metadata.clone(), ); } else { - row[j] = std::mem::take(&mut row[j]).cast_to(field_type, schema)?; + row[j] = cast_to_with_session_time_zone( + std::mem::take(&mut row[j]), + field_type, + schema, + session_time_zone, + )?; } } } diff --git a/datafusion/expr/src/type_coercion/mod.rs b/datafusion/expr/src/type_coercion/mod.rs index c92d434e34abe..52f1347911062 100644 --- a/datafusion/expr/src/type_coercion/mod.rs +++ b/datafusion/expr/src/type_coercion/mod.rs @@ -36,6 +36,7 @@ pub mod aggregates { } pub mod functions; pub mod other; +pub mod session_time_zone; pub use datafusion_expr_common::type_coercion::binary; diff --git a/datafusion/expr/src/type_coercion/session_time_zone.rs b/datafusion/expr/src/type_coercion/session_time_zone.rs new file mode 100644 index 0000000000000..2a8af90dfb648 --- /dev/null +++ b/datafusion/expr/src/type_coercion/session_time_zone.rs @@ -0,0 +1,441 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Reading timezone-naive timestamps in the session time zone when they are +//! *implicitly* converted to a timezone-aware type. +//! +//! # The model +//! +//! A `Timestamp(unit, Some(tz))` value is an instant; `tz` is only a *display +//! label* attached to it. A `Timestamp(unit, None)` value is a wall clock +//! reading with no instant attached. Turning the second into the first requires +//! picking a zone to read the wall clock in, and that choice is what decides +//! which instant you get. +//! +//! PostgreSQL and DuckDB always read it in the *session* time zone, never in +//! the zone that happens to label the other operand. DataFusion's type +//! coercion, by contrast, produces a plain +//! `CAST(naive AS Timestamp(unit, Some(other_tz)))`, and arrow's cast reads the +//! naive value in `other_tz`. [`cast_to_with_session_time_zone`] restores the +//! PostgreSQL behaviour by splitting that cast in two: +//! +//! ```text +//! CAST(ts AS Timestamp(ns, "America/New_York")) +//! -- becomes, with datafusion.execution.time_zone = '+08:00' +//! CAST(CAST(ts AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York")) +//! ``` +//! +//! The inner cast reads the wall clock in the session zone, fixing the instant. +//! The outer cast is instant-preserving (an aware -> aware cast only changes the +//! display label), so the coerced *type* is unchanged and only the chosen +//! instant differs. +//! +//! # Where it applies +//! +//! Only where the planner or the analyzer *inserts* a cast on the user's +//! behalf: comparisons, arithmetic, `IN`, `BETWEEN`, `CASE`, `UNION`, function +//! arguments, `VALUES`, `INSERT`, ... . Every one of those call sites builds its +//! cast with [`cast_to_with_session_time_zone`] rather than calling +//! [`ExprSchemable::cast_to`] directly. The three library functions that insert +//! a cast into a whole plan instead of into one expression have a sibling that +//! does the same: +//! [`coerce_plan_expr_for_schema_with_session_time_zone`](crate::expr_rewriter::coerce_plan_expr_for_schema_with_session_time_zone), +//! [`cast_subquery_with_session_time_zone`](crate::expr_schema::cast_subquery_with_session_time_zone) +//! and +//! [`LogicalPlanBuilder::values_with_session_time_zone`](crate::LogicalPlanBuilder::values_with_session_time_zone). +//! +//! Splitting the cast where it is inserted, rather than looking for inserted +//! casts afterwards, is what makes the implicit/explicit distinction exact: an +//! expression that already has the target type is never cast, so a cast the user +//! wrote is never mistaken for one DataFusion inserted. +//! +//! # What it does *not* touch +//! +//! Explicit conversions keep arrow's semantics, which is also what PostgreSQL +//! and DuckDB do for their explicit forms: +//! +//! * `ts AT TIME ZONE 'America/New_York'` reads the wall clock in New York, +//! whatever the session zone is. +//! * `arrow_cast(ts, 'Timestamp(Nanosecond, Some("America/New_York"))')`, the +//! DataFrame API's `cast_to`/`cast` and Substrait casts all keep meaning +//! exactly the arrow cast they name. +//! * SQL `CAST(ts AS TIMESTAMPTZ)` is already planned with the session time zone +//! as its target, so the target zone equals the session zone and there is +//! nothing to split. +//! +//! # A no-op by default +//! +//! `datafusion.execution.time_zone` is unset by default. With no session zone +//! there is no wall clock zone to prefer, so every function here degrades to a +//! plain [`ExprSchemable::cast_to`] and plans are unchanged. + +use std::sync::Arc; + +use arrow::datatypes::{DataType, FieldRef}; +use datafusion_common::{ExprSchema, Result}; + +use crate::{Expr, ExprSchemable}; + +/// Wraps `expr` in a cast to `cast_to_type`, reading a timezone-naive timestamp +/// in `session_time_zone` if the cast makes one timezone-aware. +/// +/// This is the [`ExprSchemable::cast_to`] to use for a cast that DataFusion +/// inserts on the user's behalf. It behaves exactly like `cast_to` except when +/// all of the following hold, in which case it emits the two step cast +/// described in the [module documentation](self): +/// +/// * `session_time_zone` is set (`datafusion.execution.time_zone`), +/// * the cast turns a timezone-naive timestamp — possibly nested inside a +/// `List`, `LargeList`, `FixedSizeList`, `Struct`, `Map` or `Dictionary` — +/// into a timezone-aware one, +/// * and the target zone is not already the session zone. +/// +/// The returned expression always has type `cast_to_type`; only the instant the +/// naive wall clock is read at changes. +/// +/// # Errors +/// +/// As [`ExprSchemable::cast_to`]: when `expr` cannot be typed against `schema`, +/// or cannot be cast to `cast_to_type`. +pub fn cast_to_with_session_time_zone( + expr: Expr, + cast_to_type: &DataType, + schema: &dyn ExprSchema, + session_time_zone: Option<&str>, +) -> Result { + let Some(time_zone) = session_time_zone else { + return expr.cast_to(cast_to_type, schema); + }; + let source_type = expr.get_type(schema)?; + let Some(intermediate) = session_zoned_type(&source_type, cast_to_type, time_zone) + else { + return expr.cast_to(cast_to_type, schema); + }; + expr.cast_to(&intermediate, schema)? + .cast_to(cast_to_type, schema) +} + +/// The intermediate cast target: `target` with every timezone-naive leaf of +/// `source` that `target` makes timezone-aware re-zoned to `time_zone`. +/// +/// `None` when there is no such leaf, i.e. when the cast already reads the naive +/// value in the session time zone or does not read one at all. +fn session_zoned_type( + source: &DataType, + target: &DataType, + time_zone: &str, +) -> Option { + match (source, target) { + (DataType::Timestamp(_, None), DataType::Timestamp(unit, Some(target_tz))) => { + (target_tz.as_ref() != time_zone) + .then(|| DataType::Timestamp(*unit, Some(Arc::from(time_zone)))) + } + (DataType::List(source), DataType::List(target)) => { + Some(DataType::List(rezoned_field(source, target, time_zone)?)) + } + (DataType::LargeList(source), DataType::LargeList(target)) => Some( + DataType::LargeList(rezoned_field(source, target, time_zone)?), + ), + (DataType::ListView(source), DataType::ListView(target)) => Some( + DataType::ListView(rezoned_field(source, target, time_zone)?), + ), + (DataType::LargeListView(source), DataType::LargeListView(target)) => Some( + DataType::LargeListView(rezoned_field(source, target, time_zone)?), + ), + ( + DataType::FixedSizeList(source, source_len), + DataType::FixedSizeList(target, target_len), + ) if source_len == target_len => Some(DataType::FixedSizeList( + rezoned_field(source, target, time_zone)?, + *target_len, + )), + (DataType::Struct(source), DataType::Struct(target)) + if source.len() == target.len() => + { + let mut changed = false; + let fields = source + .iter() + .zip(target.iter()) + .map( + |(source, target)| match rezoned_field(source, target, time_zone) { + Some(field) => { + changed = true; + field + } + None => Arc::clone(target), + }, + ) + .collect::>(); + changed.then(|| DataType::Struct(fields.into())) + } + // A `Map` is a list of a two field (key, value) struct; recursing into + // that entries field covers both the keys and the values. + ( + DataType::Map(source_entries, _), + DataType::Map(target_entries, target_sorted), + ) => Some(DataType::Map( + rezoned_field(source_entries, target_entries, time_zone)?, + *target_sorted, + )), + ( + DataType::Dictionary(_, source_value), + DataType::Dictionary(target_key, target_value), + ) => Some(DataType::Dictionary( + target_key.clone(), + Box::new(session_zoned_type(source_value, target_value, time_zone)?), + )), + // Other nested types (`Union`, `RunEndEncoded`, ...) are not rewritten: + // type coercion does not produce a naive -> aware cast through them, and + // an unhandled type simply keeps DataFusion's pre-existing behaviour. + _ => None, + } +} + +fn rezoned_field( + source: &FieldRef, + target: &FieldRef, + time_zone: &str, +) -> Option { + let data_type = + session_zoned_type(source.data_type(), target.data_type(), time_zone)?; + Some(Arc::new(target.as_ref().clone().with_data_type(data_type))) +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::collections::HashMap; + + use arrow::datatypes::{Field, Fields, TimeUnit}; + use datafusion_common::{DFSchema, ScalarValue}; + use insta::assert_snapshot; + + use crate::expr::Cast; + use crate::{col, lit}; + + const NY: &str = "America/New_York"; + const SESSION: &str = "+08:00"; + + fn naive(unit: TimeUnit) -> DataType { + DataType::Timestamp(unit, None) + } + + fn aware(unit: TimeUnit, tz: &str) -> DataType { + DataType::Timestamp(unit, Some(Arc::from(tz))) + } + + fn schema_with(name: &str, data_type: DataType) -> DFSchema { + DFSchema::from_unqualified_fields( + vec![Field::new(name, data_type, true)].into(), + HashMap::new(), + ) + .unwrap() + } + + /// Casts `col("ts")`, typed as `source`, to `target`. + fn cast(source: DataType, target: &DataType, tz: Option<&str>) -> Expr { + let schema = schema_with("ts", source); + cast_to_with_session_time_zone(col("ts"), target, &schema, tz).unwrap() + } + + /// The data type of the cast nested inside `expr`, which must be a cast of + /// a cast. + fn inner_cast_type(expr: &Expr) -> &DataType { + let Expr::Cast(Cast { expr: inner, .. }) = expr else { + panic!("expected a cast, got {expr}"); + }; + let Expr::Cast(Cast { field, .. }) = inner.as_ref() else { + panic!("expected a nested cast, got {inner}"); + }; + field.data_type() + } + + #[test] + fn naive_to_aware_cast_is_split() { + assert_snapshot!( + cast(naive(TimeUnit::Nanosecond), &aware(TimeUnit::Nanosecond, NY), Some(SESSION)), + @r#"CAST(CAST(ts AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York"))"# + ); + } + + #[test] + fn unset_session_time_zone_is_a_no_op() { + assert_snapshot!( + cast(naive(TimeUnit::Nanosecond), &aware(TimeUnit::Nanosecond, NY), None), + @r#"CAST(ts AS Timestamp(ns, "America/New_York"))"# + ); + } + + #[test] + fn cast_to_session_zone_is_untouched() { + assert_snapshot!( + cast(naive(TimeUnit::Nanosecond), &aware(TimeUnit::Nanosecond, SESSION), Some(SESSION)), + @r#"CAST(ts AS Timestamp(ns, "+08:00"))"# + ); + } + + #[test] + fn aware_to_aware_cast_is_untouched() { + assert_snapshot!( + cast(aware(TimeUnit::Nanosecond, "UTC"), &aware(TimeUnit::Nanosecond, NY), Some(SESSION)), + @r#"CAST(ts AS Timestamp(ns, "America/New_York"))"# + ); + } + + #[test] + fn aware_to_naive_cast_is_untouched() { + assert_snapshot!( + cast(aware(TimeUnit::Nanosecond, NY), &naive(TimeUnit::Nanosecond), Some(SESSION)), + @"CAST(ts AS Timestamp(ns))" + ); + } + + #[test] + fn non_timestamp_cast_is_untouched() { + assert_snapshot!( + cast(DataType::Int32, &DataType::Int64, Some(SESSION)), + @"CAST(ts AS Int64)" + ); + } + + /// A cast whose source already has the target type is not a cast at all. + #[test] + fn identity_cast_is_untouched() { + assert_snapshot!( + cast(aware(TimeUnit::Nanosecond, NY), &aware(TimeUnit::Nanosecond, NY), Some(SESSION)), + @"ts" + ); + } + + /// The naive value is truncated to the target unit by the inner cast, as a + /// single cast would have done. + #[test] + fn unit_change_happens_in_the_inner_cast() { + assert_snapshot!( + cast(naive(TimeUnit::Nanosecond), &aware(TimeUnit::Second, NY), Some(SESSION)), + @r#"CAST(CAST(ts AS Timestamp(s, "+08:00")) AS Timestamp(s, "America/New_York"))"# + ); + } + + #[test] + fn literal_cast_is_split() { + let schema = DFSchema::empty(); + let expr = lit(ScalarValue::TimestampNanosecond(Some(0), None)); + assert_snapshot!( + cast_to_with_session_time_zone( + expr, + &aware(TimeUnit::Nanosecond, NY), + &schema, + Some(SESSION), + ) + .unwrap(), + @r#"CAST(CAST(TimestampNanosecond(0, None) AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York"))"# + ); + } + + #[test] + fn nested_list_is_split() { + let list = + |inner: DataType| DataType::List(Arc::new(Field::new("item", inner, true))); + assert_snapshot!( + cast( + list(naive(TimeUnit::Nanosecond)), + &list(aware(TimeUnit::Nanosecond, NY)), + Some(SESSION), + ), + @r#"CAST(CAST(ts AS List(Timestamp(ns, "+08:00"))) AS List(Timestamp(ns, "America/New_York")))"# + ); + } + + #[test] + fn nested_fixed_size_list_is_split() { + let list = |inner: DataType| { + DataType::FixedSizeList(Arc::new(Field::new("item", inner, true)), 2) + }; + let expr = cast( + list(naive(TimeUnit::Nanosecond)), + &list(aware(TimeUnit::Nanosecond, NY)), + Some(SESSION), + ); + assert_eq!( + inner_cast_type(&expr), + &list(aware(TimeUnit::Nanosecond, SESSION)) + ); + } + + #[test] + fn nested_struct_is_split() { + let strukt = |inner: DataType| { + DataType::Struct(Fields::from(vec![ + Field::new("i", DataType::Int32, true), + Field::new("ts", inner, true), + ])) + }; + let expr = cast( + strukt(naive(TimeUnit::Nanosecond)), + &strukt(aware(TimeUnit::Nanosecond, NY)), + Some(SESSION), + ); + assert_eq!( + inner_cast_type(&expr), + &strukt(aware(TimeUnit::Nanosecond, SESSION)) + ); + } + + #[test] + fn nested_dictionary_is_split() { + let dict = |inner: DataType| { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(inner)) + }; + let expr = cast( + dict(naive(TimeUnit::Nanosecond)), + &dict(aware(TimeUnit::Nanosecond, NY)), + Some(SESSION), + ); + assert_eq!( + inner_cast_type(&expr), + &dict(aware(TimeUnit::Nanosecond, SESSION)) + ); + } + + #[test] + fn nested_map_is_split() { + let map = |inner: DataType| { + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", inner, true), + ])), + false, + )), + false, + ) + }; + let expr = cast( + map(naive(TimeUnit::Nanosecond)), + &map(aware(TimeUnit::Nanosecond, NY)), + Some(SESSION), + ); + assert_eq!( + inner_cast_type(&expr), + &map(aware(TimeUnit::Nanosecond, SESSION)) + ); + } +} From acbdaade845700eee3d558654e9e07e8384ee386 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:33:04 -0500 Subject: [PATCH 3/6] feat: coerce naive timestamps in the session time zone MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type coercion inserts `CAST(naive AS Timestamp(u, Some(tz)))` wherever a timezone-naive timestamp meets a timezone-aware one, and arrow reads the naive wall clock in `tz` — the zone that happens to label the *other* operand. PostgreSQL and DuckDB read it in the session time zone instead, so DataFusion answers a query like `tstz - ts` with a different instant than either of them. Route every cast the analyzer inserts through `cast_to_with_session_time_zone`, so that with `datafusion.execution.time_zone` set the naive side is read in the session zone and then relabelled to the coerced type. The coerced types are unchanged; only the instant differs. A scalar function argument that is already a literal was folded to the coerced type outright rather than wrapped in a cast, which read its wall clock in the coerced type's own zone; it now keeps the split cast when there is one, and is folded as before when there is not. `coerce_plan_expr_for_schema` (UNION branch alignment, view output coercion) and `cast_subquery` build a whole projection rather than one cast, so their casts are split afterwards — only at the top level of each projected expression, since anything deeper was written by the user. The projection keeps the schema it already had so that splitting a cast never renames a column. `TypeCoercionRewriter::new` alone still coerces the way it always did; `with_session_time_zone` opts in, and `TypeCoercion` passes the session zone for SQL. `ExprSimplifier::coerce` passes it too so that `SessionContext::create_physical_expr` matches SQL planning. The two optimizer rules that build a rewriter over an already coerced plan deliberately do not. Co-Authored-By: Claude Fable 5.1 --- .../optimizer/src/analyzer/type_coercion.rs | 665 +++++++++++++++--- .../optimizer/src/scalar_subquery_to_join.rs | 7 +- .../simplify_expressions/expr_simplifier.rs | 4 +- datafusion/optimizer/src/utils.rs | 4 +- 4 files changed, 593 insertions(+), 87 deletions(-) diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 07d3173d63342..f06e2333a0cea 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -39,8 +39,8 @@ use datafusion_expr::expr::{ HigherOrderFunction, InList, InSubquery, Like, ScalarFunction, SetComparison, Sort, WindowFunction, }; -use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema; -use datafusion_expr::expr_schema::cast_subquery; +use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema_with_session_time_zone; +use datafusion_expr::expr_schema::cast_subquery_with_session_time_zone; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::type_coercion::binary::{ comparison_coercion, like_coercion, regex_coercion, type_union_coercion, @@ -52,6 +52,7 @@ use datafusion_expr::type_coercion::other::{ get_coerce_type_for_case_expression, get_coerce_type_for_case_when, get_coerce_type_for_list, }; +use datafusion_expr::type_coercion::session_time_zone::cast_to_with_session_time_zone; use datafusion_expr::type_coercion::{ is_datetime, is_interval, is_signed_numeric, is_timestamp, }; @@ -86,7 +87,11 @@ fn coerce_output(plan: LogicalPlan, config: &ConfigOptions) -> Result Result { static EMPTY_SCHEMA: LazyLock = LazyLock::new(DFSchema::empty); + let session_time_zone = config.execution.time_zone.as_deref(); + // recurse let transformed_plan = plan - .transform_up_with_subqueries(|plan| analyze_internal(&EMPTY_SCHEMA, plan))? + .transform_up_with_subqueries(|plan| { + analyze_internal(&EMPTY_SCHEMA, plan, session_time_zone) + })? .data; // finish @@ -113,9 +122,13 @@ impl AnalyzerRule for TypeCoercion { /// use the external schema to handle the correlated subqueries case /// /// Assumes that children have already been optimized +/// +/// `session_time_zone` is `datafusion.execution.time_zone`; see +/// [`cast_to_with_session_time_zone`]. fn analyze_internal( external_schema: &DFSchema, plan: LogicalPlan, + session_time_zone: Option<&str>, ) -> Result> { // get schema representing all available input fields. This is used for data type // resolution only, so order does not matter here @@ -151,13 +164,19 @@ fn analyze_internal( // Coerce filter predicates to boolean (handles `WHERE NULL`) let plan = if let LogicalPlan::Filter(mut filter) = plan { - filter.predicate = filter.predicate.cast_to(&DataType::Boolean, &schema)?; + filter.predicate = cast_to_with_session_time_zone( + filter.predicate, + &DataType::Boolean, + &schema, + session_time_zone, + )?; LogicalPlan::Filter(filter) } else { plan }; - let mut expr_rewrite = TypeCoercionRewriter::new(&schema); + let mut expr_rewrite = + TypeCoercionRewriter::new(&schema).with_session_time_zone(session_time_zone); let name_preserver = NamePreserver::new(&plan); // apply coercion rewrite all expressions in the plan individually @@ -175,13 +194,71 @@ fn analyze_internal( /// Rewrite expressions to apply type coercion. pub struct TypeCoercionRewriter<'a> { pub(crate) schema: &'a DFSchema, + /// `datafusion.execution.time_zone`, when set. + /// + /// Timezone-naive timestamps are read in this zone wherever coercion turns + /// one into a timezone-aware timestamp. `None`, the default, keeps arrow's + /// behaviour of reading them in the target type's own zone. + session_time_zone: Option<&'a str>, } impl<'a> TypeCoercionRewriter<'a> { /// Create a new [`TypeCoercionRewriter`] with a provided schema /// representing both the inputs and output of the [`LogicalPlan`] node. + /// + /// The session time zone is unset; see + /// [`with_session_time_zone`](Self::with_session_time_zone). pub fn new(schema: &'a DFSchema) -> Self { - Self { schema } + Self { + schema, + session_time_zone: None, + } + } + + /// Read timezone-naive timestamps in `session_time_zone` wherever coercion + /// converts one to a timezone-aware type. + /// + /// Pass `datafusion.execution.time_zone` to coerce the way SQL planning + /// does, and so the way PostgreSQL and DuckDB do. See + /// [`cast_to_with_session_time_zone`] for what changes. + pub fn with_session_time_zone(mut self, session_time_zone: Option<&'a str>) -> Self { + self.session_time_zone = session_time_zone; + self + } + + /// Wrap `expr` in a cast to `to`, reading timezone-naive timestamps in the + /// session time zone. + /// + /// Every cast this rewriter inserts goes through here. + fn cast(&self, expr: Expr, to: &DataType, schema: &DFSchema) -> Result { + cast_to_with_session_time_zone(expr, to, schema, self.session_time_zone) + } + + /// Support the `IsTrue` `IsNotTrue` `IsFalse` `IsNotFalse` type coercion. + /// The above op will be rewrite to the binary op when creating the physical + /// op. + fn casted_expr_for_bool_op(&self, expr: Expr) -> Result { + let left_type = expr.get_type(self.schema)?; + BinaryTypeCoercer::new(&left_type, &Operator::IsDistinctFrom, &DataType::Boolean) + .get_input_types()?; + self.cast(expr, &DataType::Boolean, self.schema) + } + + /// [`cast_subquery_with_session_time_zone`], reading timezone-naive + /// timestamps in the session time zone. + /// + /// The cast lands in the subquery's projection rather than on an expression + /// of this plan, so it is built there rather than through [`Self::cast`]. + fn cast_subquery( + &self, + subquery: Subquery, + cast_to_type: &DataType, + ) -> Result { + cast_subquery_with_session_time_zone( + subquery, + cast_to_type, + self.session_time_zone, + ) } /// Coerce the [`LogicalPlan`]. @@ -192,8 +269,10 @@ impl<'a> TypeCoercionRewriter<'a> { match plan { LogicalPlan::Join(join) => self.coerce_join(join), LogicalPlan::AsOfJoin(join) => self.coerce_asof_join(join), - LogicalPlan::Union(union) => Self::coerce_union(union), - LogicalPlan::Limit(limit) => Self::coerce_limit(limit), + LogicalPlan::Union(union) => { + Self::coerce_union_in_time_zone(union, self.session_time_zone) + } + LogicalPlan::Limit(limit) => self.coerce_limit(limit), LogicalPlan::Dml(dml) => self.coerce_dml(dml), _ => Ok(plan), } @@ -221,7 +300,8 @@ impl<'a> TypeCoercionRewriter<'a> { datafusion_expr::dml::MergeIntoAction::Update(assignments) => { for (column, value) in assignments { let field = target_schema.field_with_unqualified_name(column)?; - *value = value.clone().cast_to(field.data_type(), self.schema)?; + *value = + self.cast(value.clone(), field.data_type(), self.schema)?; } } datafusion_expr::dml::MergeIntoAction::Insert { columns, values } => { @@ -230,14 +310,14 @@ impl<'a> TypeCoercionRewriter<'a> { values.iter_mut().zip(target_schema.fields()) { *value = - value.clone().cast_to(field.data_type(), self.schema)?; + self.cast(value.clone(), field.data_type(), self.schema)?; } } else { for (column, value) in columns.iter().zip(values) { let field = target_schema.field_with_unqualified_name(column)?; *value = - value.clone().cast_to(field.data_type(), self.schema)?; + self.cast(value.clone(), field.data_type(), self.schema)?; } } } @@ -317,7 +397,19 @@ impl<'a> TypeCoercionRewriter<'a> { /// Coerce the union’s inputs to a common schema compatible with all inputs. /// This occurs after wildcard expansion and the coercion of the input expressions. + /// + /// Timezone-naive timestamps are read in the zone of the branch they are + /// aligned to. Use + /// [`with_session_time_zone`](Self::with_session_time_zone) and coerce the + /// whole plan to read them in the session time zone instead. pub fn coerce_union(union_plan: Union) -> Result { + Self::coerce_union_in_time_zone(union_plan, None) + } + + fn coerce_union_in_time_zone( + union_plan: Union, + session_time_zone: Option<&str>, + ) -> Result { let union_schema = Arc::new(coerce_union_schema_with_schema( &union_plan.inputs, &union_plan.schema, @@ -326,8 +418,11 @@ impl<'a> TypeCoercionRewriter<'a> { .inputs .into_iter() .map(|p| { - let plan = - coerce_plan_expr_for_schema(Arc::unwrap_or_clone(p), &union_schema)?; + let plan = coerce_plan_expr_for_schema_with_session_time_zone( + Arc::unwrap_or_clone(p), + &union_schema, + session_time_zone, + )?; match plan { LogicalPlan::Projection(Projection { expr, input, .. }) => { Ok(Arc::new(project_with_column_index( @@ -347,15 +442,16 @@ impl<'a> TypeCoercionRewriter<'a> { } /// Coerce the fetch and skip expression to Int64 type. - fn coerce_limit(limit: Limit) -> Result { + fn coerce_limit(&self, limit: Limit) -> Result { fn coerce_limit_expr( + rewriter: &TypeCoercionRewriter<'_>, expr: Expr, schema: &DFSchema, expr_name: &str, ) -> Result { let dt = expr.get_type(schema)?; if dt.is_integer() || dt.is_null() { - expr.cast_to(&DataType::Int64, schema) + rewriter.cast(expr, &DataType::Int64, schema) } else { plan_err!("Expected {expr_name} to be an integer or null, but got {dt}") } @@ -364,11 +460,11 @@ impl<'a> TypeCoercionRewriter<'a> { let empty_schema = DFSchema::empty(); let new_fetch = limit .fetch - .map(|expr| coerce_limit_expr(*expr, &empty_schema, "LIMIT")) + .map(|expr| coerce_limit_expr(self, *expr, &empty_schema, "LIMIT")) .transpose()?; let new_skip = limit .skip - .map(|expr| coerce_limit_expr(*expr, &empty_schema, "OFFSET")) + .map(|expr| coerce_limit_expr(self, *expr, &empty_schema, "OFFSET")) .transpose()?; Ok(LogicalPlan::Limit(Limit { input: limit.input, @@ -381,7 +477,7 @@ impl<'a> TypeCoercionRewriter<'a> { let expr_type = expr.get_type(self.schema)?; match expr_type { DataType::Boolean => Ok(expr), - DataType::Null => expr.cast_to(&DataType::Boolean, self.schema), + DataType::Null => self.cast(expr, &DataType::Boolean, self.schema), other => { plan_err!("{description} must be boolean type, but got {other:?}") } @@ -416,7 +512,7 @@ impl<'a> TypeCoercionRewriter<'a> { &right_type, )? } else { - left.cast_to(&left_type, left_schema)? + self.cast(left, &left_type, left_schema)? }; let right_expr = if !right_cast_ok { @@ -428,7 +524,7 @@ impl<'a> TypeCoercionRewriter<'a> { &left_type, )? } else { - right.cast_to(&right_type, right_schema)? + self.cast(right, &right_type, right_schema)? }; Ok((left_expr, right_expr)) @@ -568,9 +664,9 @@ impl<'a> TypeCoercionRewriter<'a> { { Box::new(expr) } - _ => Box::new(expr.cast_to(&coerced_type, self.schema)?), + _ => Box::new(self.cast(expr, &coerced_type, self.schema)?), }; - let pattern = Box::new(pattern.cast_to(&coerced_type, self.schema)?); + let pattern = Box::new(self.cast(pattern, &coerced_type, self.schema)?); Ok((expr, pattern)) } } @@ -588,8 +684,12 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { outer_ref_columns, spans, }) => { - let new_plan = - analyze_internal(self.schema, Arc::unwrap_or_clone(subquery))?.data; + let new_plan = analyze_internal( + self.schema, + Arc::unwrap_or_clone(subquery), + self.session_time_zone, + )? + .data; Ok(Transformed::yes(Expr::ScalarSubquery(Subquery { subquery: Arc::new(new_plan), outer_ref_columns, @@ -600,6 +700,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { let new_plan = analyze_internal( self.schema, Arc::unwrap_or_clone(subquery.subquery), + self.session_time_zone, )? .data; Ok(Transformed::yes(Expr::Exists(Exists { @@ -619,6 +720,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { let new_plan = analyze_internal( self.schema, Arc::unwrap_or_clone(subquery.subquery), + self.session_time_zone, )? .data; let expr_type = expr.get_type(self.schema)?; @@ -634,8 +736,8 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { spans: subquery.spans, }; Ok(Transformed::yes(Expr::InSubquery(InSubquery::new( - Box::new(expr.cast_to(&common_type, self.schema)?), - cast_subquery(new_subquery, &common_type)?, + Box::new(self.cast(*expr, &common_type, self.schema)?), + self.cast_subquery(new_subquery, &common_type)?, negated, )))) } @@ -648,6 +750,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { let new_plan = analyze_internal( self.schema, Arc::unwrap_or_clone(subquery.subquery), + self.session_time_zone, )? .data; let expr_type = expr.get_type(self.schema)?; @@ -670,33 +773,32 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { spans: subquery.spans, }; Ok(Transformed::yes(Expr::SetComparison(SetComparison::new( - Box::new(expr.cast_to(&common_type, self.schema)?), - cast_subquery(new_subquery, &common_type)?, + Box::new(self.cast(*expr, &common_type, self.schema)?), + self.cast_subquery(new_subquery, &common_type)?, op, quantifier, )))) } - Expr::Not(expr) => Ok(Transformed::yes(not(get_casted_expr_for_bool_op( - *expr, - self.schema, - )?))), + Expr::Not(expr) => { + Ok(Transformed::yes(not(self.casted_expr_for_bool_op(*expr)?))) + } Expr::IsTrue(expr) => Ok(Transformed::yes(is_true( - get_casted_expr_for_bool_op(*expr, self.schema)?, + self.casted_expr_for_bool_op(*expr)?, ))), Expr::IsNotTrue(expr) => Ok(Transformed::yes(is_not_true( - get_casted_expr_for_bool_op(*expr, self.schema)?, + self.casted_expr_for_bool_op(*expr)?, ))), Expr::IsFalse(expr) => Ok(Transformed::yes(is_false( - get_casted_expr_for_bool_op(*expr, self.schema)?, + self.casted_expr_for_bool_op(*expr)?, ))), Expr::IsNotFalse(expr) => Ok(Transformed::yes(is_not_false( - get_casted_expr_for_bool_op(*expr, self.schema)?, + self.casted_expr_for_bool_op(*expr)?, ))), Expr::IsUnknown(expr) => Ok(Transformed::yes(is_unknown( - get_casted_expr_for_bool_op(*expr, self.schema)?, + self.casted_expr_for_bool_op(*expr)?, ))), Expr::IsNotUnknown(expr) => Ok(Transformed::yes(is_not_unknown( - get_casted_expr_for_bool_op(*expr, self.schema)?, + self.casted_expr_for_bool_op(*expr)?, ))), Expr::Negative(expr) => { let data_type = expr.get_type(self.schema)?; @@ -799,10 +901,10 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { ) })?; Ok(Transformed::yes(Expr::Between(Between::new( - Box::new(expr.cast_to(&coercion_type, self.schema)?), + Box::new(self.cast(*expr, &coercion_type, self.schema)?), negated, - Box::new(low.cast_to(&coercion_type, self.schema)?), - Box::new(high.cast_to(&coercion_type, self.schema)?), + Box::new(self.cast(*low, &coercion_type, self.schema)?), + Box::new(self.cast(*high, &coercion_type, self.schema)?), )))) } Expr::InList(InList { @@ -824,11 +926,11 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { ), Some(coerced_type) => { // find the coerced type - let cast_expr = expr.cast_to(&coerced_type, self.schema)?; + let cast_expr = self.cast(*expr, &coerced_type, self.schema)?; let cast_list_expr = list .into_iter() .map(|list_expr| { - list_expr.cast_to(&coerced_type, self.schema) + self.cast(list_expr, &coerced_type, self.schema) }) .collect::>>()?; Ok(Transformed::yes(Expr::InList(InList::new( @@ -840,7 +942,8 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { } } Expr::Case(case) => { - let case = coerce_case_expression(case, self.schema)?; + let case = + coerce_case_expression(case, self.schema, self.session_time_zone)?; Ok(Transformed::yes(Expr::Case(case))) } Expr::ScalarFunction(ScalarFunction { func, args }) => { @@ -848,6 +951,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { args, self.schema, func.as_ref(), + self.session_time_zone, )?; Ok(Transformed::yes(Expr::ScalarFunction( ScalarFunction::new_udf(func, new_expr), @@ -864,11 +968,15 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { null_treatment, }, }) => { - let new_expr = - coerce_arguments_for_signature(args, self.schema, func.as_ref())?; + let new_expr = coerce_arguments_for_signature( + args, + self.schema, + func.as_ref(), + self.session_time_zone, + )?; let filter = filter - .map(|filter| filter.cast_to(&DataType::Boolean, self.schema)) + .map(|filter| self.cast(*filter, &DataType::Boolean, self.schema)) .transpose()? .map(Box::new); @@ -902,15 +1010,25 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { let args = match &fun { expr::WindowFunctionDefinition::AggregateUDF(udf) => { - coerce_arguments_for_signature(args, self.schema, udf.as_ref())? + coerce_arguments_for_signature( + args, + self.schema, + udf.as_ref(), + self.session_time_zone, + )? } expr::WindowFunctionDefinition::WindowUDF(udf) => { - coerce_arguments_for_signature(args, self.schema, udf.as_ref())? + coerce_arguments_for_signature( + args, + self.schema, + udf.as_ref(), + self.session_time_zone, + )? } }; let filter = filter - .map(|filter| filter.cast_to(&DataType::Boolean, self.schema)) + .map(|filter| self.cast(*filter, &DataType::Boolean, self.schema)) .transpose()? .map(Box::new); @@ -948,7 +1066,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { .map(|(arg, new_field)| match (&arg, new_field) { (Expr::Lambda(_lambda), ValueOrLambda::Lambda(_)) => Ok(arg), (Expr::Lambda(_lambda), ValueOrLambda::Value(_)) => internal_err!("value_fields_with_higher_order_udf returned a value for a lambda argument"), - (_, ValueOrLambda::Value(new_field)) => arg.cast_to(new_field.data_type(), self.schema), + (_, ValueOrLambda::Value(new_field)) => self.cast(arg, new_field.data_type(), self.schema), (_, ValueOrLambda::Lambda(_)) => internal_err!("value_fields_with_higher_order_udf returned a lambda for a value argument"), }) .collect::>()?; @@ -1169,15 +1287,6 @@ fn coerce_window_frame( Ok(window_frame) } -// Support the `IsTrue` `IsNotTrue` `IsFalse` `IsNotFalse` type coercion. -// The above op will be rewrite to the binary op when creating the physical op. -fn get_casted_expr_for_bool_op(expr: Expr, schema: &DFSchema) -> Result { - let left_type = expr.get_type(schema)?; - BinaryTypeCoercer::new(&left_type, &Operator::IsDistinctFrom, &DataType::Boolean) - .get_input_types()?; - expr.cast_to(&DataType::Boolean, schema) -} - /// Returns `expressions` coerced to types compatible with /// `signature`, if possible. /// @@ -1186,13 +1295,16 @@ fn coerce_arguments_for_signature( expressions: Vec, schema: &DFSchema, func: &F, + session_time_zone: Option<&str>, ) -> Result> { let coerced_types = coerced_argument_types(&expressions, schema, func)?; expressions .into_iter() .zip(coerced_types) - .map(|(expr, data_type)| expr.cast_to(&data_type, schema)) + .map(|(expr, data_type)| { + cast_to_with_session_time_zone(expr, &data_type, schema, session_time_zone) + }) .collect() } @@ -1204,6 +1316,7 @@ fn coerce_scalar_function_arguments_for_signature( expressions: Vec, schema: &DFSchema, func: &F, + session_time_zone: Option<&str>, ) -> Result> { let coerced_types = coerced_argument_types(&expressions, schema, func)?; @@ -1211,7 +1324,7 @@ fn coerce_scalar_function_arguments_for_signature( .into_iter() .zip(coerced_types) .map(|(expr, data_type)| { - coerce_scalar_function_argument(expr, &data_type, schema) + coerce_scalar_function_argument(expr, &data_type, schema, session_time_zone) }) .collect() } @@ -1238,6 +1351,7 @@ fn coerce_scalar_function_argument( expr: Expr, data_type: &DataType, schema: &DFSchema, + session_time_zone: Option<&str>, ) -> Result { if matches!(&expr, Expr::Cast(_) | Expr::TryCast(_)) && expr.get_type(schema)? == *data_type @@ -1245,23 +1359,40 @@ fn coerce_scalar_function_argument( return Ok(expr); } - let Expr::Literal(value, metadata) = expr else { - return expr.cast_to(data_type, schema); - }; + if !matches!(expr, Expr::Literal(_, _)) { + return cast_to_with_session_time_zone( + expr, + data_type, + schema, + session_time_zone, + ); + } - if value.data_type() != *data_type + let coerced = + cast_to_with_session_time_zone(expr, data_type, schema, session_time_zone)?; + + // Materialize the cast into the literal, unless the session time zone split + // it in two: casting the value would read its wall clock in `data_type`'s + // own time zone, which is the reading the split is there to avoid. + if let Expr::Cast(Cast { expr, .. }) = &coerced + && let Expr::Literal(value, metadata) = expr.as_ref() + && value.data_type() != *data_type && let Ok(value) = value.cast_to(data_type) { - return Ok(Expr::Literal(value, metadata)); + return Ok(Expr::Literal(value, metadata.clone())); } // A failed value cast remains an expression cast so execution produces the // same error as before. Since it is no longer a literal at the coerced type, // it is reported as `None` in `ReturnFieldArgs::scalar_arguments`. - Expr::Literal(value, metadata).cast_to(data_type, schema) + Ok(coerced) } -fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { +fn coerce_case_expression( + case: Case, + schema: &DFSchema, + session_time_zone: Option<&str>, +) -> Result { // Given expressions like: // // CASE a1 @@ -1351,7 +1482,14 @@ fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { let case_expr = case .expr .zip(case_when_coerce_type.as_ref()) - .map(|(case_expr, coercible_type)| case_expr.cast_to(coercible_type, schema)) + .map(|(case_expr, coercible_type)| { + cast_to_with_session_time_zone( + *case_expr, + coercible_type, + schema, + session_time_zone, + ) + }) .transpose()? .map(Box::new); let when_then = case @@ -1359,7 +1497,13 @@ fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { .into_iter() .map(|(when, then)| { let when_type = case_when_coerce_type.as_ref().unwrap_or(&DataType::Boolean); - let when = when.cast_to(when_type, schema).map_err(|e| { + let when = cast_to_with_session_time_zone( + *when, + when_type, + schema, + session_time_zone, + ) + .map_err(|e| { DataFusionError::Context( format!( "WHEN expressions in CASE couldn't be \ @@ -1368,13 +1512,25 @@ fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { Box::new(e), ) })?; - let then = then.cast_to(&then_else_coerce_type, schema)?; + let then = cast_to_with_session_time_zone( + *then, + &then_else_coerce_type, + schema, + session_time_zone, + )?; Ok((Box::new(when), Box::new(then))) }) .collect::>>()?; let else_expr = case .else_expr - .map(|expr| expr.cast_to(&then_else_coerce_type, schema)) + .map(|expr| { + cast_to_with_session_time_zone( + *expr, + &then_else_coerce_type, + schema, + session_time_zone, + ) + }) .transpose()? .map(Box::new); @@ -2721,7 +2877,7 @@ mod test { vec![Field::new("a", DataType::Int64, true)].into(), std::collections::HashMap::new(), )?); - let mut rewriter = TypeCoercionRewriter { schema: &schema }; + let mut rewriter = TypeCoercionRewriter::new(&schema); let expr = is_true(lit(12i32).gt(lit(13i64))); let expected = is_true(cast(lit(12i32), DataType::Int64).gt(lit(13i64))); let result = expr.rewrite(&mut rewriter).data()?; @@ -2732,7 +2888,7 @@ mod test { vec![Field::new("a", DataType::Int64, true)].into(), std::collections::HashMap::new(), )?); - let mut rewriter = TypeCoercionRewriter { schema: &schema }; + let mut rewriter = TypeCoercionRewriter::new(&schema); let expr = is_true(lit(12i32).eq(lit(13i64))); let expected = is_true(cast(lit(12i32), DataType::Int64).eq(lit(13i64))); let result = expr.rewrite(&mut rewriter).data()?; @@ -2743,7 +2899,7 @@ mod test { vec![Field::new("a", DataType::Int64, true)].into(), std::collections::HashMap::new(), )?); - let mut rewriter = TypeCoercionRewriter { schema: &schema }; + let mut rewriter = TypeCoercionRewriter::new(&schema); let expr = is_true(lit(12i32).lt(lit(13i64))); let expected = is_true(cast(lit(12i32), DataType::Int64).lt(lit(13i64))); let result = expr.rewrite(&mut rewriter).data()?; @@ -2860,7 +3016,7 @@ mod test { &then_else_common_type, &schema, ); - let actual = coerce_case_expression(case, &schema)?; + let actual = coerce_case_expression(case, &schema, None)?; assert_eq!(expected, actual); // CASE string WHEN float/integer/string: comparison coercion @@ -2883,7 +3039,7 @@ mod test { &then_else_common_type, &schema, ); - let actual = coerce_case_expression(case, &schema)?; + let actual = coerce_case_expression(case, &schema, None)?; assert_eq!(expected, actual); let case = Case { @@ -2895,7 +3051,7 @@ mod test { ], else_expr: Some(Box::new(col("string"))), }; - let err = coerce_case_expression(case, &schema).unwrap_err(); + let err = coerce_case_expression(case, &schema, None).unwrap_err(); assert_snapshot!( err.strip_backtrace(), @"Error during planning: Failed to coerce case (Interval(MonthDayNano)) and when (Float32, Binary, Utf8) to common types in CASE WHEN expression" @@ -2910,7 +3066,7 @@ mod test { ], else_expr: Some(Box::new(col("timestamp"))), }; - let err = coerce_case_expression(case, &schema).unwrap_err(); + let err = coerce_case_expression(case, &schema, None).unwrap_err(); assert_snapshot!( err.strip_backtrace(), @"Error during planning: Failed to coerce then (Date32, Float32, Binary) and else (Timestamp(ns)) to common types in CASE WHEN expression" @@ -2930,7 +3086,7 @@ mod test { let expected = cast_helper(case.clone(), &$case_when_type, &$then_else_type, &$schema); - let actual = coerce_case_expression(case, &$schema)?; + let actual = coerce_case_expression(case, &$schema, None)?; assert_eq!(expected, actual); }; } @@ -3286,4 +3442,349 @@ mod test { " ) } + + /// Tests for reading timezone-naive timestamps in + /// `datafusion.execution.time_zone` wherever coercion converts one to a + /// timezone-aware type. + /// + /// Each test also pins the plan with no session time zone, which must stay + /// exactly what DataFusion produced before. + mod session_time_zone { + use super::*; + + use datafusion_expr::expr::{SetComparison, SetQuantifier}; + use datafusion_expr::{Case, Union}; + + /// The session time zone: a fixed offset, so that the tests do not + /// depend on any time zone database. + const SESSION: &str = "+08:00"; + const NY: &str = "America/New_York"; + + fn naive() -> DataType { + DataType::Timestamp(TimeUnit::Nanosecond, None) + } + + fn aware() -> DataType { + DataType::Timestamp(TimeUnit::Second, Some(NY.into())) + } + + /// `tstz`, a timezone-aware timestamp, and `ts`, a timezone-naive one. + fn timestamps() -> Arc { + Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new( + DFSchema::from_unqualified_fields( + vec![ + Field::new("tstz", aware(), true), + Field::new("ts", naive(), true), + ] + .into(), + std::collections::HashMap::new(), + ) + .unwrap(), + ), + })) + } + + fn one_column(name: &str, data_type: DataType) -> Arc { + Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new( + DFSchema::from_unqualified_fields( + vec![Field::new(name, data_type, true)].into(), + std::collections::HashMap::new(), + ) + .unwrap(), + ), + })) + } + + macro_rules! assert_analyzed_plan_in_time_zone { + ( + $time_zone: expr, + $plan: expr, + @ $expected: literal $(,)? + ) => {{ + let mut options = ConfigOptions::default(); + options.execution.time_zone = + Option::<&str>::map($time_zone, str::to_string); + let rule = Arc::new(TypeCoercion::new()); + assert_analyzed_plan_with_config_eq_snapshot!( + options, + rule, + $plan, + @ $expected, + ) + }}; + } + + fn subtraction_plan() -> Result { + let subtract = |left: &str, right: &str| { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col(left)), + Operator::Minus, + Box::new(col(right)), + )) + }; + Ok(LogicalPlan::Projection(Projection::try_new( + vec![subtract("tstz", "ts"), subtract("ts", "tstz")], + timestamps(), + )?)) + } + + /// The naive operand of a subtraction is read in the session time zone; + /// the aware operand only has its unit changed. + #[test] + fn subtraction() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + Some(SESSION), + subtraction_plan()?, + @r#" + Projection: CAST(tstz AS Timestamp(ns, "America/New_York")) - CAST(CAST(ts AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York")), CAST(CAST(ts AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York")) - CAST(tstz AS Timestamp(ns, "America/New_York")) + EmptyRelation: rows=0 + "# + ) + } + + #[test] + fn subtraction_without_session_time_zone() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + None, + subtraction_plan()?, + @r#" + Projection: CAST(tstz AS Timestamp(ns, "America/New_York")) - CAST(ts AS Timestamp(ns, "America/New_York")), CAST(ts AS Timestamp(ns, "America/New_York")) - CAST(tstz AS Timestamp(ns, "America/New_York")) + EmptyRelation: rows=0 + "# + ) + } + + /// `= ANY ()` casts inside the subquery's projection. + fn set_comparison_plan() -> Result { + let set_comparison = Expr::SetComparison(SetComparison::new( + Box::new(col("a")), + Subquery { + subquery: one_column("a", naive()), + outer_ref_columns: vec![], + spans: Spans::new(), + }, + Operator::Eq, + SetQuantifier::Any, + )); + Ok(LogicalPlan::Filter(Filter::try_new( + set_comparison, + one_column("a", aware()), + )?)) + } + + #[test] + fn set_comparison() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + Some(SESSION), + set_comparison_plan()?, + @r#" + Filter: CAST(a AS Timestamp(ns, "America/New_York")) = ANY () + Subquery: + Projection: CAST(CAST(a AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York")) + EmptyRelation: rows=0 + EmptyRelation: rows=0 + "# + ) + } + + #[test] + fn set_comparison_without_session_time_zone() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + None, + set_comparison_plan()?, + @r#" + Filter: CAST(a AS Timestamp(ns, "America/New_York")) = ANY () + Subquery: + Projection: CAST(a AS Timestamp(ns, "America/New_York")) + EmptyRelation: rows=0 + EmptyRelation: rows=0 + "# + ) + } + + fn union_plan() -> Result { + Ok(LogicalPlan::Union(Union::try_new_with_loose_types(vec![ + one_column("a", naive()), + one_column("a", aware()), + ])?)) + } + + #[test] + fn union() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + Some(SESSION), + union_plan()?, + @r#" + Union + Projection: CAST(CAST(a AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York")) AS a + EmptyRelation: rows=0 + Projection: CAST(a AS Timestamp(ns, "America/New_York")) AS a + EmptyRelation: rows=0 + "# + ) + } + + #[test] + fn union_without_session_time_zone() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + None, + union_plan()?, + @r#" + Union + Projection: CAST(a AS Timestamp(ns, "America/New_York")) AS a + EmptyRelation: rows=0 + Projection: CAST(a AS Timestamp(ns, "America/New_York")) AS a + EmptyRelation: rows=0 + "# + ) + } + + fn case_plan() -> Result { + let case = Expr::Case(Case::new( + None, + vec![(Box::new(lit(true)), Box::new(col("ts")))], + Some(Box::new(col("tstz"))), + )); + Ok(LogicalPlan::Projection(Projection::try_new( + vec![case], + timestamps(), + )?)) + } + + #[test] + fn case() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + Some(SESSION), + case_plan()?, + @r#" + Projection: CASE WHEN Boolean(true) THEN CAST(CAST(ts AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York")) ELSE CAST(tstz AS Timestamp(ns, "America/New_York")) END + EmptyRelation: rows=0 + "# + ) + } + + #[test] + fn case_without_session_time_zone() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + None, + case_plan()?, + @r#" + Projection: CASE WHEN Boolean(true) THEN CAST(ts AS Timestamp(ns, "America/New_York")) ELSE CAST(tstz AS Timestamp(ns, "America/New_York")) END + EmptyRelation: rows=0 + "# + ) + } + + /// Stands in for `coalesce`: both arguments are coerced to a common + /// type. `datafusion-optimizer` deliberately does not depend on the + /// function crates, so the real one is out of reach here. + #[derive(Debug, Hash, PartialEq, Eq)] + struct TestCoalesceUDF { + signature: Signature, + } + + impl ScalarUDFImpl for TestCoalesceUDF { + fn name(&self) -> &str { + "test_coalesce" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, args: &[DataType]) -> Result { + Ok(args[0].clone()) + } + + fn invoke_with_args( + &self, + _args: ScalarFunctionArgs, + ) -> Result { + datafusion_common::internal_err!("not executed") + } + } + + fn function_argument_plan() -> Result { + let udf = ScalarUDF::from(TestCoalesceUDF { + signature: Signature::comparable(2, Volatility::Immutable), + }) + .call(vec![col("ts"), col("tstz")]); + Ok(LogicalPlan::Projection(Projection::try_new( + vec![udf], + timestamps(), + )?)) + } + + #[test] + fn function_argument() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + Some(SESSION), + function_argument_plan()?, + @r#" + Projection: test_coalesce(CAST(CAST(ts AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York")), CAST(tstz AS Timestamp(ns, "America/New_York"))) + EmptyRelation: rows=0 + "# + ) + } + + /// A naive timestamp *literal* argument is read in the session time zone + /// too: it must stay an expression cast rather than being folded into a + /// literal at the coerced type, which would read its wall clock in the + /// coerced type's own zone. + fn literal_function_argument_plan() -> Result { + let udf = ScalarUDF::from(TestCoalesceUDF { + signature: Signature::comparable(2, Volatility::Immutable), + }) + .call(vec![ + lit(ScalarValue::TimestampNanosecond(Some(0), None)), + col("tstz"), + ]); + Ok(LogicalPlan::Projection(Projection::try_new( + vec![udf], + timestamps(), + )?)) + } + + #[test] + fn function_argument_literal() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + Some(SESSION), + literal_function_argument_plan()?, + @r#" + Projection: test_coalesce(CAST(CAST(TimestampNanosecond(0, None) AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York")), CAST(tstz AS Timestamp(ns, "America/New_York"))) + EmptyRelation: rows=0 + "# + ) + } + + /// With no session time zone the literal is still folded to the coerced + /// type, exactly as before. + #[test] + fn function_argument_literal_without_session_time_zone() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + None, + literal_function_argument_plan()?, + @r#" + Projection: test_coalesce(TimestampNanosecond(18000000000000, Some("America/New_York")), CAST(tstz AS Timestamp(ns, "America/New_York"))) AS test_coalesce(TimestampNanosecond(0, None),tstz) + EmptyRelation: rows=0 + "# + ) + } + + #[test] + fn function_argument_without_session_time_zone() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + None, + function_argument_plan()?, + @r#" + Projection: test_coalesce(CAST(ts AS Timestamp(ns, "America/New_York")), CAST(tstz AS Timestamp(ns, "America/New_York"))) + EmptyRelation: rows=0 + "# + ) + } + } } diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 44011a125ba96..1ed3fb51eeffb 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -398,9 +398,10 @@ fn build_join( // itself be NULL) otherwise. let mut compensation_exprs = HashMap::new(); if let Some(expr_map) = collected_count_expr_map { - let mut expr_rewrite = TypeCoercionRewriter { - schema: new_plan.schema(), - }; + // No session time zone: this optimizer rule runs over a plan the + // analyzer has already coerced, so any naive -> aware timestamp cast is + // in place. + let mut expr_rewrite = TypeCoercionRewriter::new(new_plan.schema()); let having_arm = pull_up .pull_up_having_expr .as_ref() diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index 5436bd092163e..d5c982ff4c84c 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -243,7 +243,9 @@ impl ExprSimplifier { /// See the [type coercion module](datafusion_expr::type_coercion) /// documentation for more details on type coercion pub fn coerce(&self, expr: Expr, schema: &DFSchema) -> Result { - let mut expr_rewrite = TypeCoercionRewriter { schema }; + let mut expr_rewrite = TypeCoercionRewriter::new(schema).with_session_time_zone( + self.info.config_options().execution.time_zone.as_deref(), + ); expr.rewrite(&mut expr_rewrite).data() } diff --git a/datafusion/optimizer/src/utils.rs b/datafusion/optimizer/src/utils.rs index d4ac31e8a517c..bc5800a81730e 100644 --- a/datafusion/optimizer/src/utils.rs +++ b/datafusion/optimizer/src/utils.rs @@ -244,7 +244,9 @@ fn evaluate_expr_with_null_column<'a>( } fn coerce(expr: Expr, schema: &DFSchema) -> Result { - let mut expr_rewrite = TypeCoercionRewriter { schema }; + // No session time zone: this coerces a predicate that the analyzer has + // already coerced, so any naive -> aware timestamp cast is in place. + let mut expr_rewrite = TypeCoercionRewriter::new(schema); expr.rewrite(&mut expr_rewrite).data() } From 68791d870ad395af6996c4bf549832798bff6708 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:35:17 -0500 Subject: [PATCH 4/6] feat: read naive timestamps in the session time zone for VALUES and INSERT `INSERT`, `UPDATE ... SET` and `VALUES` cast their cells to the target column's type in the SQL planner rather than in type coercion, so they did not pick up the session time zone with the analyzer's casts. Inserting a timezone-naive timestamp into a timezone-aware column read the wall clock in the *column's* zone; PostgreSQL and DuckDB read it in the session zone. `INSERT`/`UPDATE` build each cast one expression at a time, so they call `cast_to_with_session_time_zone` directly. `VALUES` goes through `LogicalPlanBuilder::values`, which casts every cell to the column's common type and takes no configuration; its public API is left alone and the SQL planner splits the top level cast of each cell afterwards instead. Co-Authored-By: Claude Fable 5.1 --- datafusion/sql/src/statement.rs | 54 +++++++++++++++++++++++---------- datafusion/sql/src/values.rs | 25 +++++++++++---- 2 files changed, 57 insertions(+), 22 deletions(-) diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index dde9420a68696..4a566d62eb761 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -46,15 +46,16 @@ use datafusion_expr::dml::{ use datafusion_expr::expr_rewriter::normalize_col_with_schemas_and_ambiguity_check; use datafusion_expr::logical_plan::DdlStatement; use datafusion_expr::logical_plan::builder::project; +use datafusion_expr::type_coercion::session_time_zone::cast_to_with_session_time_zone; use datafusion_expr::utils::expr_to_columns; use datafusion_expr::{ Analyze, CreateCatalog, CreateCatalogSchema, CreateExternalTable as PlanCreateExternalTable, CreateFunction, CreateFunctionBody, CreateIndex as PlanCreateIndex, CreateMemoryTable, CreateView, Deallocate, DescribeTable, DmlStatement, DropCatalogSchema, DropFunction, DropTable, DropView, - EmptyRelation, Execute, Explain, ExplainFormat, Expr, ExprSchemable, Filter, - LogicalPlan, LogicalPlanBuilder, OperateFunctionArg, PlanType, Prepare, - ResetVariable, SetVariable, SortExpr, Statement as PlanStatement, ToStringifiedPlan, + EmptyRelation, Execute, Explain, ExplainFormat, Expr, Filter, LogicalPlan, + LogicalPlanBuilder, OperateFunctionArg, PlanType, Prepare, ResetVariable, + SetVariable, SortExpr, Statement as PlanStatement, ToStringifiedPlan, TransactionAccessMode, TransactionConclusion, TransactionEnd, TransactionIsolationLevel, TransactionStart, Volatility, WriteOp, cast, }; @@ -2390,7 +2391,16 @@ impl SqlToRel<'_, S> { .or_else(|| Some(Arc::clone(field))); } // Cast to target column type, if necessary - expr.cast_to(field.data_type(), source.schema())? + cast_to_with_session_time_zone( + expr, + field.data_type(), + source.schema(), + self.context_provider + .options() + .execution + .time_zone + .as_deref(), + )? } None => { // If the target table has an alias, use it to qualify the column name @@ -2935,25 +2945,37 @@ impl SqlToRel<'_, S> { plan_err!("Column count doesn't match insert query!")?; } + let session_time_zone = self + .context_provider + .options() + .execution + .time_zone + .as_deref(); let exprs = value_indices .into_iter() .enumerate() .map(|(i, value_index)| { let target_field = table_schema.field(i); let expr = match value_index { - Some(v) => { - Expr::Column(Column::from(source.schema().qualified_field(v))) - .cast_to(target_field.data_type(), source.schema())? - } + Some(v) => cast_to_with_session_time_zone( + Expr::Column(Column::from(source.schema().qualified_field(v))), + target_field.data_type(), + source.schema(), + session_time_zone, + )?, // The value is not specified. Fill in the default value for the column. - None => table_source - .get_column_default(target_field.name()) - .cloned() - .unwrap_or_else(|| { - // If there is no default for the column, then the default is NULL - Expr::Literal(ScalarValue::Null, None) - }) - .cast_to(target_field.data_type(), &DFSchema::empty())?, + None => cast_to_with_session_time_zone( + table_source + .get_column_default(target_field.name()) + .cloned() + .unwrap_or_else(|| { + // If there is no default for the column, then the default is NULL + Expr::Literal(ScalarValue::Null, None) + }), + target_field.data_type(), + &DFSchema::empty(), + session_time_zone, + )?, }; Ok(expr.alias(target_field.name())) }) diff --git a/datafusion/sql/src/values.rs b/datafusion/sql/src/values.rs index 2f15785520f77..fa01eef135a47 100644 --- a/datafusion/sql/src/values.rs +++ b/datafusion/sql/src/values.rs @@ -62,11 +62,24 @@ impl SqlToRel<'_, S> { }) .collect::>>()?; - let schema = planner_context.table_schema().unwrap_or(empty_schema); - if schema.fields().is_empty() { - LogicalPlanBuilder::values(values)?.build() - } else { - LogicalPlanBuilder::values_with_schema(values, &schema)?.build() - } + let table_schema = planner_context.table_schema().unwrap_or(empty_schema); + // A VALUES list with no table schema infers its column types from the + // values themselves. + let schema = (!table_schema.fields().is_empty()).then(|| table_schema.as_ref()); + // The builder casts every cell to its column's type itself, so it is + // handed the session time zone to read a timezone-naive cell that it + // makes timezone-aware in. + let session_time_zone = self + .context_provider + .options() + .execution + .time_zone + .as_deref(); + LogicalPlanBuilder::values_with_session_time_zone( + values, + schema, + session_time_zone, + )? + .build() } } From f27628a4480c40a22c86586c663aebbf5218d336 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:44:53 -0500 Subject: [PATCH 5/6] test: cover reading naive timestamps in the session time zone End to end coverage of every site where DataFusion converts a timezone-naive timestamp to a timezone-aware one on the user's behalf: comparisons, `IN`, `BETWEEN`, `CASE`, subtraction, subqueries, joins, `UNION ALL`, `coalesce`, `least`/`greatest`, `nullif`, `nvl2`, the array functions, `VALUES` and `INSERT`. Each expectation is the instant PostgreSQL and DuckDB return. Also pinned: the plans, so that the two casts are visible and provably survive `simplify_expressions` rather than collapsing back into the single cast that reads the wall clock in the wrong zone; the default, where with no session time zone the naive operand is still read in the aware operand's zone; and `AT TIME ZONE` and `arrow_cast`, which keep arrow's semantics whatever the session time zone is. `date_bin('1 day', TIMESTAMPTZ '2022-01-01 20:10:00Z', TIMESTAMP '2020-01-01')` under `+07` was pinned at `2022-01-01T07:00:00+07:00`, which is neither what PostgreSQL nor what DuckDB returns: the naive origin was read as UTC rather than in the session time zone. It now returns `2022-01-02T00:00:00+07:00`, matching both and matching the TIMESTAMPTZ-origin case just above it. Co-Authored-By: Claude Fable 5.1 --- datafusion/core/tests/expr_api/mod.rs | 38 +++ .../optimizer/src/analyzer/type_coercion.rs | 79 ++++- datafusion/sql/tests/sql_integration.rs | 28 ++ .../test_files/datetime/timestamps.slt | 298 +++++++++++++++++- 4 files changed, 440 insertions(+), 3 deletions(-) diff --git a/datafusion/core/tests/expr_api/mod.rs b/datafusion/core/tests/expr_api/mod.rs index 19ff3933193de..5cd420798ec8b 100644 --- a/datafusion/core/tests/expr_api/mod.rs +++ b/datafusion/core/tests/expr_api/mod.rs @@ -17,6 +17,7 @@ use arrow::array::{ ArrayRef, Int64Array, RecordBatch, StringArray, StructArray, + TimestampNanosecondArray, builder::{ListBuilder, StringBuilder}, }; use arrow::datatypes::{DataType, Field}; @@ -364,6 +365,43 @@ async fn test_create_physical_expr_coercion() { ); } +/// `SessionContext::create_physical_expr` coerces the way SQL planning does, so +/// a timezone-naive timestamp is read in `datafusion.execution.time_zone`. +#[test] +fn test_create_physical_expr_timestamp_subtraction_uses_session_timezone() { + // 2024-11-01T04:00:00Z, which is 2024-11-01T00:00:00 in New York. + let timestamp_tz = Arc::new( + TimestampNanosecondArray::from(vec![1_730_433_600_000_000_000]) + .with_timezone("America/New_York"), + ) as ArrayRef; + // The same wall clock reading, with no zone attached. Read in +08:00 it is + // 2024-10-31T16:00:00Z, twelve hours before `timestamp_tz`. + let timestamp = Arc::new(TimestampNanosecondArray::from(vec![ + 1_730_419_200_000_000_000, + ])) as ArrayRef; + let batch = RecordBatch::try_from_iter(vec![ + ("timestamp_tz", timestamp_tz), + ("timestamp", timestamp), + ]) + .unwrap(); + let df_schema = DFSchema::try_from(batch.schema()).unwrap(); + + let mut config = SessionConfig::new(); + config.options_mut().execution.time_zone = Some("+08:00".to_string()); + let ctx = SessionContext::new_with_config(config); + + let physical_expr = ctx + .create_physical_expr(col("timestamp_tz") - col("timestamp"), &df_schema) + .unwrap(); + let result = physical_expr.evaluate(&batch).unwrap(); + let result = ScalarValue::try_from_array(&result.into_array(1).unwrap(), 0).unwrap(); + + assert_eq!( + result, + ScalarValue::DurationNanosecond(Some(12 * 3_600 * 1_000_000_000)) + ); +} + /// Evaluates the specified expr as an aggregate and compares the result to the /// expected result. async fn evaluate_agg_test(expr: Expr, expected_lines: Vec<&str>) { diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index f06e2333a0cea..71871e5c6303f 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -3453,7 +3453,7 @@ mod test { use super::*; use datafusion_expr::expr::{SetComparison, SetQuantifier}; - use datafusion_expr::{Case, Union}; + use datafusion_expr::{Case, Cast, Union}; /// The session time zone: a fixed offset, so that the tests do not /// depend on any time zone database. @@ -3591,6 +3591,48 @@ mod test { ) } + /// A subquery projection that is *already* the common type is not cast + /// at all, so an explicit `AT TIME ZONE` at the top of it keeps arrow's + /// semantics. + fn explicit_set_comparison_plan() -> Result { + let projection = LogicalPlan::Projection(Projection::try_new( + vec![Expr::Cast(Cast::new( + Box::new(col("a")), + DataType::Timestamp(TimeUnit::Nanosecond, Some(NY.into())), + ))], + one_column("a", naive()), + )?); + let set_comparison = Expr::SetComparison(SetComparison::new( + Box::new(col("a")), + Subquery { + subquery: Arc::new(projection), + outer_ref_columns: vec![], + spans: Spans::new(), + }, + Operator::Eq, + SetQuantifier::Any, + )); + Ok(LogicalPlan::Filter(Filter::try_new( + set_comparison, + one_column("a", aware()), + )?)) + } + + #[test] + fn explicit_set_comparison_is_not_split() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + Some(SESSION), + explicit_set_comparison_plan()?, + @r#" + Filter: CAST(a AS Timestamp(ns, "America/New_York")) = ANY () + Subquery: + Projection: CAST(a AS Timestamp(ns, "America/New_York")) + EmptyRelation: rows=0 + EmptyRelation: rows=0 + "# + ) + } + #[test] fn set_comparison_without_session_time_zone() -> Result<()> { assert_analyzed_plan_in_time_zone!( @@ -3628,6 +3670,41 @@ mod test { ) } + /// A `UNION` branch that is *already* the common type is not cast at + /// all, so an explicit `AT TIME ZONE` at the top of it keeps arrow's + /// semantics. + fn explicit_union_plan() -> Result { + let explicit = LogicalPlan::Projection(Projection::try_new( + vec![ + Expr::Cast(Cast::new( + Box::new(col("a")), + DataType::Timestamp(TimeUnit::Nanosecond, Some(NY.into())), + )) + .alias("a"), + ], + one_column("a", naive()), + )?); + Ok(LogicalPlan::Union(Union::try_new_with_loose_types(vec![ + Arc::new(explicit), + one_column("a", naive()), + ])?)) + } + + #[test] + fn explicit_union_branch_is_not_split() -> Result<()> { + assert_analyzed_plan_in_time_zone!( + Some(SESSION), + explicit_union_plan()?, + @r#" + Union + Projection: CAST(a AS Timestamp(ns, "America/New_York")) AS a + EmptyRelation: rows=0 + Projection: CAST(CAST(a AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York")) AS a + EmptyRelation: rows=0 + "# + ) + } + #[test] fn union_without_session_time_zone() -> Result<()> { assert_analyzed_plan_in_time_zone!( diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 2bd06ee4bafbf..65d2564dab751 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -719,6 +719,34 @@ fn plan_insert_no_target_columns() { ); } +/// The builder casts the naive cell to the column type it infers, so that cast +/// reads the wall clock in the session time zone. The `AT TIME ZONE` cell +/// already *has* that type, is therefore not cast at all, and keeps arrow's +/// semantics. +#[test] +fn plan_values_in_session_time_zone() { + let sql = "VALUES ('2024-11-01T00:00:00'::timestamp), \ + ('2024-11-01T00:00:00'::timestamp AT TIME ZONE 'America/New_York')"; + let mut config_options = datafusion_common::config::ConfigOptions::new(); + config_options.execution.time_zone = Some("+08:00".to_string()); + let plan = logical_plan_with_config(sql, config_options).unwrap(); + assert_snapshot!( + plan, + @r#"Values: (CAST(CAST(CAST(Utf8("2024-11-01T00:00:00") AS Timestamp(ns)) AS Timestamp(ns, "+08:00")) AS Timestamp(ns, "America/New_York"))), (CAST(CAST(Utf8("2024-11-01T00:00:00") AS Timestamp(ns)) AS Timestamp(ns, "America/New_York")))"# + ); +} + +#[test] +fn plan_values_without_session_time_zone() { + let sql = "VALUES ('2024-11-01T00:00:00'::timestamp), \ + ('2024-11-01T00:00:00'::timestamp AT TIME ZONE 'America/New_York')"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#"Values: (CAST(CAST(Utf8("2024-11-01T00:00:00") AS Timestamp(ns)) AS Timestamp(ns, "America/New_York"))), (CAST(CAST(Utf8("2024-11-01T00:00:00") AS Timestamp(ns)) AS Timestamp(ns, "America/New_York")))"# + ); +} + #[rstest] #[case::duplicate_columns( "INSERT INTO test_decimal (id, price, price) VALUES (1, 2, 3), (4, 5, 6)", diff --git a/datafusion/sqllogictest/test_files/datetime/timestamps.slt b/datafusion/sqllogictest/test_files/datetime/timestamps.slt index fa4159c1dce23..0e0665122e4c4 100644 --- a/datafusion/sqllogictest/test_files/datetime/timestamps.slt +++ b/datafusion/sqllogictest/test_files/datetime/timestamps.slt @@ -1980,6 +1980,298 @@ SELECT '2000-01-01T00:00:00'::timestamp - '2010-01-01T00:00:00'::timestamp; ---- -3653 days 0 hours 0 mins 0.000000000 secs +########## +## Timezone-naive timestamps and the session time zone +## +## A timezone-naive timestamp is a wall clock reading. Converting one to a +## timezone-aware timestamp has to pick a zone to read that wall clock in, and +## wherever DataFusion inserts the conversion itself it reads it in +## `datafusion.execution.time_zone`, as PostgreSQL and DuckDB do. Explicit +## conversions (`AT TIME ZONE`, `arrow_cast`) are unaffected. +########## + +statement ok +SET TIME ZONE = '+08' + +# https://github.com/apache/datafusion/issues/13212 +# postgresql / duckdb: 08:00:00 +query ? +SELECT '2024-11-01T00:00:00+00:00'::timestamptz - '2024-11-01T00:00:00'::timestamp; +---- +0 days 8 hours 0 mins 0.000000000 secs + +# postgresql / duckdb: -08:00:00 +query ? +SELECT '2024-11-01T00:00:00'::timestamp - '2024-11-01T00:00:00+00:00'::timestamptz; +---- +0 days -8 hours 0 mins 0.000000000 secs + +# the same, over columns rather than literals +query ?? +WITH timestamps(ts_tz, ts) AS ( + VALUES ('2024-11-01T00:00:00+00:00'::timestamptz, '2024-11-01T00:00:00'::timestamp) +) +SELECT ts_tz - ts, ts - ts_tz FROM timestamps; +---- +0 days 8 hours 0 mins 0.000000000 secs 0 days -8 hours 0 mins 0.000000000 secs + +# ... and over scalar subqueries +query ?? +WITH timestamps(ts_tz, ts) AS ( + VALUES ('2024-11-01T00:00:00+00:00'::timestamptz, '2024-11-01T00:00:00'::timestamp) +) +SELECT + (SELECT ts_tz FROM timestamps) - (SELECT ts FROM timestamps), + (SELECT ts FROM timestamps) - (SELECT ts_tz FROM timestamps); +---- +0 days 8 hours 0 mins 0.000000000 secs 0 days -8 hours 0 mins 0.000000000 secs + +# `ts_tz` is 2024-11-01T04:00:00Z. `ts` reads as 2024-10-31T16:00:00Z, twelve +# hours earlier, because the session time zone is +08 and not New York. +statement ok +CREATE TABLE timestamps AS SELECT + arrow_cast('2024-11-01T00:00:00-04:00', 'Timestamp(Nanosecond, Some("America/New_York"))') AS ts_tz, + '2024-11-01T00:00:00'::timestamp AS ts; + +query TT +SELECT arrow_typeof(ts_tz), arrow_typeof(ts) FROM timestamps; +---- +Timestamp(ns, "America/New_York") Timestamp(ns) + +query BBBBBB?? +SELECT + ts_tz = ts, + ts_tz > ts, + ts_tz IN (ts), + ts_tz BETWEEN ts AND ts, + ts_tz BETWEEN ts_tz AND ts, + CASE ts_tz WHEN ts THEN true ELSE false END, + ts_tz - ts, + ts - ts_tz +FROM timestamps; +---- +false true false false false false 0 days 12 hours 0 mins 0.000000000 secs 0 days -12 hours 0 mins 0.000000000 secs + +# nested naive timestamps are read in the session time zone too +query BBB +SELECT + arrow_cast(ts_tz, 'Dictionary(Int32, Timestamp(Nanosecond, Some("America/New_York")))') = ts, + arrow_cast(ts_tz, 'RunEndEncoded("run_ends": non-null Int32, "values": Timestamp(Nanosecond, Some("America/New_York")))') = ts, + make_array(ts_tz) = make_array(ts) +FROM timestamps; +---- +false false false + +query I +SELECT count(*) FROM timestamps WHERE ts_tz IN (SELECT ts FROM timestamps); +---- +0 + +query I +SELECT count(*) FROM timestamps WHERE ts_tz = ANY (SELECT ts FROM timestamps); +---- +0 + +query I +SELECT count(*) FROM timestamps a JOIN timestamps b ON a.ts_tz = b.ts; +---- +0 + +# postgresql / duckdb: 2024-10-31 12:00:00-04, 2024-11-01 00:00:00-04 +query PP +SELECT + CASE WHEN true THEN ts ELSE ts_tz END, + CASE WHEN false THEN ts ELSE ts_tz END +FROM timestamps; +---- +2024-10-31T12:00:00-04:00 2024-11-01T00:00:00-04:00 + +# postgresql / duckdb: 2024-10-31 12:00:00-04 +query P +SELECT coalesce(ts, ts_tz) FROM timestamps; +---- +2024-10-31T12:00:00-04:00 + +query P rowsort +SELECT ts FROM timestamps UNION ALL SELECT ts_tz FROM timestamps; +---- +2024-10-31T12:00:00-04:00 +2024-11-01T00:00:00-04:00 + +query PP +SELECT least(ts, ts_tz), greatest(ts, ts_tz) FROM timestamps; +---- +2024-10-31T12:00:00-04:00 2024-11-01T00:00:00-04:00 + +query B +SELECT nullif(ts, ts_tz) IS NULL FROM timestamps; +---- +false + +query ? +SELECT make_array(ts, ts_tz) FROM timestamps; +---- +[2024-10-31T12:00:00-04:00, 2024-11-01T00:00:00-04:00] + +query B +SELECT array_has(make_array(ts_tz), ts) FROM timestamps; +---- +false + +query I +SELECT array_position(make_array(ts_tz), ts) FROM timestamps; +---- +NULL + +query P +SELECT nvl2(ts, ts, ts_tz) FROM timestamps; +---- +2024-10-31T12:00:00-04:00 + +query P rowsort +VALUES ('2024-11-01T00:00:00'::timestamp), (arrow_cast('2024-11-01T00:00:00-04:00', 'Timestamp(Nanosecond, Some("America/New_York"))')); +---- +2024-10-31T12:00:00-04:00 +2024-11-01T00:00:00-04:00 + +statement ok +CREATE TABLE new_york AS SELECT ts_tz FROM timestamps WHERE false; + +statement ok +INSERT INTO new_york VALUES ('2024-11-01T00:00:00'::timestamp); + +query P +SELECT ts_tz FROM new_york; +---- +2024-10-31T12:00:00-04:00 + +statement ok +set datafusion.explain.logical_plan_only = true; + +# the inner cast reads the wall clock in the session time zone; the outer cast +# only relabels it as the coerced type +query TT +EXPLAIN SELECT ts_tz - ts FROM timestamps; +---- +logical_plan +01)Projection: timestamps.ts_tz - CAST(CAST(timestamps.ts AS Timestamp(ns, "+08")) AS Timestamp(ns, "America/New_York")) +02)--TableScan: timestamps projection=[ts_tz, ts] + +statement ok +set datafusion.explain.logical_plan_only = false; + +# both casts survive the optimizer: `simplify_expressions` must not collapse +# them into the single cast that reads the wall clock in New York +query TT +EXPLAIN FORMAT indent SELECT ts_tz = ts FROM timestamps; +---- +logical_plan +01)Projection: timestamps.ts_tz = CAST(CAST(timestamps.ts AS Timestamp(ns, "+08")) AS Timestamp(ns, "America/New_York")) +02)--TableScan: timestamps projection=[ts_tz, ts] +physical_plan +01)ProjectionExec: expr=[ts_tz@0 = CAST(CAST(ts@1 AS Timestamp(ns, "+08")) AS Timestamp(ns, "America/New_York")) as timestamps.ts_tz = timestamps.ts] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# a named session time zone follows its daylight saving rules +statement ok +SET TIME ZONE = 'America/New_York' + +# postgresql / duckdb: -05:00:00 (EST) +query ? +SELECT '2024-01-15T00:00:00+00:00'::timestamptz - '2024-01-15T00:00:00'::timestamp; +---- +0 days -5 hours 0 mins 0.000000000 secs + +# postgresql / duckdb: -04:00:00 (EDT) +query ? +SELECT '2024-07-15T00:00:00+00:00'::timestamptz - '2024-07-15T00:00:00'::timestamp; +---- +0 days -4 hours 0 mins 0.000000000 secs + +# With no session time zone (the default) the naive operand is read in the zone +# of the aware operand, so the two are the same instant here. +statement ok +RESET datafusion.execution.time_zone + +query ?? +SELECT ts_tz - ts, ts - ts_tz FROM timestamps; +---- +0 days 0 hours 0 mins 0.000000000 secs 0 days 0 hours 0 mins 0.000000000 secs + +statement ok +set datafusion.explain.logical_plan_only = true; + +query TT +EXPLAIN SELECT ts_tz - ts FROM timestamps; +---- +logical_plan +01)Projection: timestamps.ts_tz - CAST(timestamps.ts AS Timestamp(ns, "America/New_York")) +02)--TableScan: timestamps projection=[ts_tz, ts] + +statement ok +set datafusion.explain.logical_plan_only = false; + +# explicit conversions keep arrow's semantics whatever the session time zone is +statement ok +SET TIME ZONE = '+08' + +query P +SELECT '2024-11-01T00:00:00'::timestamp AT TIME ZONE 'America/New_York'; +---- +2024-11-01T00:00:00-04:00 + +query P +SELECT arrow_cast('2024-11-01T00:00:00'::timestamp, 'Timestamp(Nanosecond, Some("America/New_York"))'); +---- +2024-11-01T00:00:00-04:00 + +# An explicit `AT TIME ZONE` at the top of a `UNION` branch, a subquery +# projection, a `VALUES` cell or a view is not rewritten either: it already has +# the type it is being aligned to, so nothing casts it and there is no inserted +# cast to read in the session time zone. +query P rowsort +SELECT ts AT TIME ZONE 'America/New_York' FROM timestamps UNION ALL SELECT ts_tz FROM timestamps; +---- +2024-11-01T00:00:00-04:00 +2024-11-01T00:00:00-04:00 + +query I +SELECT count(*) FROM timestamps WHERE ts_tz IN (SELECT ts AT TIME ZONE 'America/New_York' FROM timestamps); +---- +1 + +query I +SELECT count(*) FROM timestamps WHERE ts_tz = ANY (SELECT ts AT TIME ZONE 'America/New_York' FROM timestamps); +---- +1 + +query P rowsort +VALUES ('2024-11-01T00:00:00'::timestamp AT TIME ZONE 'America/New_York'), (arrow_cast('2024-11-01T00:00:00-04:00', 'Timestamp(Nanosecond, Some("America/New_York"))')); +---- +2024-11-01T00:00:00-04:00 +2024-11-01T00:00:00-04:00 + +statement ok +CREATE VIEW new_york_view AS SELECT ts AT TIME ZONE 'America/New_York' AS x FROM timestamps; + +query P +SELECT x FROM new_york_view; +---- +2024-11-01T00:00:00-04:00 + +statement ok +DROP VIEW new_york_view; + +statement ok +RESET datafusion.execution.time_zone + +statement ok +DROP TABLE timestamps; + +statement ok +DROP TABLE new_york; + + # Interval - Timestamp => error query error Cannot coerce arithmetic expression Interval\(MonthDayNano\) - Timestamp\(ns\) to valid types SELECT i - ts1 from FOO; @@ -2059,11 +2351,13 @@ SELECT date_bin('1 day', TIMESTAMPTZ '2022-01-01 20:10:00Z', TIMESTAMPTZ '2020-0 ---- 2022-01-02T00:00:00+07:00 -# coerce TIMESTAMP to TIMESTAMPTZ +# coerce TIMESTAMP to TIMESTAMPTZ: the naive origin is read in the session time +# zone, as PostgreSQL and DuckDB read it +# postgresql: 2022-01-02 00:00:00+07 query P SELECT date_bin('1 day', TIMESTAMPTZ '2022-01-01 20:10:00Z', TIMESTAMP '2020-01-01') ---- -2022-01-01T07:00:00+07:00 +2022-01-02T00:00:00+07:00 # postgresql: 1 query I From 61cb5dba6204285607d244d81d8fee7343d7addc Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:48:06 -0500 Subject: [PATCH 6/6] docs: document the session time zone for naive timestamps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `datafusion.execution.time_zone`'s description only mentioned `now`, which undersold it: it is now also the zone a timezone-naive timestamp's wall clock is read in wherever DataFusion converts one to a timezone-aware timestamp on the user's behalf. The upgrade guide covers the three behaviour changes a 55.0.0 user can see — the session time zone applying to implicit conversions, equal-unit `aware - naive` subtraction now coercing like every other unit pairing, and the arrow limitation that makes a DST gap or fold error under a named session time zone — plus the new `TypeCoercionRewriter::with_session_time_zone`, which callers who coerce expressions themselves need in order to match SQL planning. Co-Authored-By: Claude Fable 5.1 --- datafusion/common/src/config.rs | 2 +- .../test_files/information_schema.slt | 6 +- .../library-user-guide/upgrading/56.0.0.md | 71 +++++++++++++++++++ docs/source/user-guide/configs.md | 2 +- 4 files changed, 76 insertions(+), 5 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 360586b0e9bae..79171570d6b38 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -940,7 +940,7 @@ config_namespace! { /// The default time zone /// - /// Some functions, e.g. `now` return timestamps in this time zone + /// Some functions, e.g. `now`, return timestamps in this time zone. It is also the time zone a timezone-naive timestamp is read in when it is implicitly converted to a timezone-aware one, for example in comparisons, arithmetic, `UNION`, `CASE`, function arguments, `VALUES` and `INSERT`. Explicit conversions (`AT TIME ZONE`, `arrow_cast`) are unaffected. pub time_zone: Option, default = None /// Parquet options diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index b270eba99d7b0..e9145f08e2592 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -444,7 +444,7 @@ datafusion.execution.sort_spill_reservation_bytes 10485760 Specifies the reserve datafusion.execution.spill_compression uncompressed Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. datafusion.execution.split_file_groups_by_statistics false Attempt to eliminate sorts by packing & sorting files with non-overlapping statistics into the same file groups. Currently experimental datafusion.execution.target_partitions 7 Number of partitions for query execution. Increasing partitions can increase concurrency. Defaults to the number of CPU cores on the system -datafusion.execution.time_zone NULL The default time zone Some functions, e.g. `now` return timestamps in this time zone +datafusion.execution.time_zone NULL The default time zone Some functions, e.g. `now`, return timestamps in this time zone. It is also the time zone a timezone-naive timestamp is read in when it is implicitly converted to a timezone-aware one, for example in comparisons, arithmetic, `UNION`, `CASE`, function arguments, `VALUES` and `INSERT`. Explicit conversions (`AT TIME ZONE`, `arrow_cast`) are unaffected. datafusion.execution.use_row_number_estimates_to_optimize_partitioning false Should DataFusion use row number estimates at the input to decide whether increasing parallelism is beneficial or not. By default, only exact row numbers (not estimates) are used for this decision. Setting this flag to `true` will likely produce better plans. if the source of statistics is accurate. We plan to make this the default in the future. datafusion.explain.analyze_categories all Which metric categories to include in "EXPLAIN ANALYZE" output. Comma-separated list of: "rows", "bytes", "timing", "uncategorized". Use "none" to show plan structure only, or "all" (default) to show everything. Metrics without a declared category are treated as "uncategorized". datafusion.explain.analyze_level dev Verbosity level for "EXPLAIN ANALYZE". Default is "dev" "summary" shows common metrics for high-level insights. "dev" provides deep operator-level introspection for developers. @@ -555,14 +555,14 @@ datafusion.execution.time_zone NULL query TTT SHOW TIME ZONE VERBOSE ---- -datafusion.execution.time_zone NULL The default time zone Some functions, e.g. `now` return timestamps in this time zone +datafusion.execution.time_zone NULL The default time zone Some functions, e.g. `now`, return timestamps in this time zone. It is also the time zone a timezone-naive timestamp is read in when it is implicitly converted to a timezone-aware one, for example in comparisons, arithmetic, `UNION`, `CASE`, function arguments, `VALUES` and `INSERT`. Explicit conversions (`AT TIME ZONE`, `arrow_cast`) are unaffected. # show_timezone_default_utc # https://github.com/apache/datafusion/issues/3255 query TTT SHOW TIMEZONE VERBOSE ---- -datafusion.execution.time_zone NULL The default time zone Some functions, e.g. `now` return timestamps in this time zone +datafusion.execution.time_zone NULL The default time zone Some functions, e.g. `now`, return timestamps in this time zone. It is also the time zone a timezone-naive timestamp is read in when it is implicitly converted to a timezone-aware one, for example in comparisons, arithmetic, `UNION`, `CASE`, function arguments, `VALUES` and `INSERT`. Explicit conversions (`AT TIME ZONE`, `arrow_cast`) are unaffected. # show empty verbose diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 8de2abc5a05fc..8a5af41ef954f 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -25,6 +25,77 @@ in this section pertains to features and changes that have already been merged to the main branch and are awaiting release in this version. +### Timezone-naive timestamps are read in the session time zone + +A `Timestamp(unit, Some(tz))` value is an instant and `tz` only labels it for +display; a `Timestamp(unit, None)` value is a wall clock reading. Converting the +second to the first has to pick a zone to read the wall clock in. + +Wherever DataFusion inserts that conversion on your behalf, it now reads the wall +clock in `datafusion.execution.time_zone` rather than in the zone that happens to +label the other operand. This is what PostgreSQL and DuckDB do. The sites are +comparisons, arithmetic, `IN`, `BETWEEN`, `CASE`, `UNION`, function arguments +(`coalesce`, `greatest`/`least`, `nullif`, the array functions, `date_bin`'s +origin, ...), `VALUES` and `INSERT`. + +```sql +SET TIME ZONE = '+08'; +-- 55.0.0: 0 days 0 hours 0 mins 0.000000000 secs +-- 56.0.0: 0 days 8 hours 0 mins 0.000000000 secs, as PostgreSQL and DuckDB return +SELECT '2024-11-01T00:00:00+00:00'::timestamptz - '2024-11-01T00:00:00'::timestamp; +``` + +`datafusion.execution.time_zone` is unset by default, and with no session time +zone plans are unchanged. + +_Explicit_ conversions are unaffected and keep arrow's semantics, again as in +PostgreSQL and DuckDB: `AT TIME ZONE`, `arrow_cast`, the DataFrame API's +`cast_to`/`cast` and Substrait casts all read the wall clock in the zone they +name. SQL `CAST(x AS TIMESTAMPTZ)` already targeted the session time zone. + +### `Timestamp(u, Some(tz)) - Timestamp(u, None)` reads the naive operand in `tz` + +Subtracting a timezone-naive timestamp from a timezone-aware one _at the same +time unit_ did not coerce at all in `55.0.0`: the naive operand was read as UTC, +which differed both from `=` on the same pair of values and from every other +pairing of time units. It now coerces like the rest, so with no session time zone +the naive operand is read in the aware operand's time zone: + +```sql +-- ts_tz is 2024-11-01T00:00:00-04:00 (America/New_York), ts is the naive +-- 2024-11-01T00:00:00 +-- 55.0.0: 0 days 4 hours 0 mins 0.000000000 secs +-- 56.0.0: 0 days 0 hours 0 mins 0.000000000 secs +SELECT ts_tz - ts FROM timestamps; +``` + +With a session time zone set, the naive operand is read in that zone instead, as +described above. + +### Naive timestamps in a DST gap or fold error under a named session time zone + +With `datafusion.execution.time_zone` set to a _named_ zone, a timezone-naive +wall clock value that falls in a daylight saving gap or fold currently fails with +arrow's `Cannot cast timezone to different timezone` rather than resolving to an +instant. Fixed offsets (`+08`, `-05:00`, `UTC`) are unaffected, as are values +outside a transition. + +This is fixed by and +; see + for the DataFusion side. + +### `TypeCoercionRewriter::with_session_time_zone` + +`TypeCoercionRewriter` has a new `with_session_time_zone` method. Applying the +session time zone is opt in: a rewriter built with `TypeCoercionRewriter::new` +alone coerces exactly as it did before, so if you coerce expressions yourself and +want to match SQL planning, pass `config.execution.time_zone.as_deref()`: + +```rust,ignore +let rewriter = TypeCoercionRewriter::new(&schema) + .with_session_time_zone(config.execution.time_zone.as_deref()); +``` + ### `ForeignSession::create_physical_plan` is unsupported `ForeignSession::create_physical_plan` no longer forwards to the library that diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 16750d79d750d..9d75683b8900f 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -79,7 +79,7 @@ The following configuration settings are available: | datafusion.execution.coalesce_batches | true | When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting | | datafusion.execution.collect_statistics | true | Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Defaults to true. | | datafusion.execution.target_partitions | 0 | Number of partitions for query execution. Increasing partitions can increase concurrency. Defaults to the number of CPU cores on the system | -| datafusion.execution.time_zone | NULL | The default time zone Some functions, e.g. `now` return timestamps in this time zone | +| datafusion.execution.time_zone | NULL | The default time zone Some functions, e.g. `now`, return timestamps in this time zone. It is also the time zone a timezone-naive timestamp is read in when it is implicitly converted to a timezone-aware one, for example in comparisons, arithmetic, `UNION`, `CASE`, function arguments, `VALUES` and `INSERT`. Explicit conversions (`AT TIME ZONE`, `arrow_cast`) are unaffected. | | datafusion.execution.parquet.enable_page_index | true | (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. | | datafusion.execution.parquet.pruning | true | (reading) If true, the parquet reader attempts to skip entire row groups based on the predicate in the query and the metadata (min/max values) stored in the parquet file | | datafusion.execution.parquet.skip_metadata | true | (reading) If true, the parquet reader skip the optional embedded metadata that may be in the file Schema. This setting can help avoid schema conflicts when querying multiple parquet files with schemas containing compatible types but different metadata |