From c42a082cddfa9027811cae9fd0485bde65d785cb Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Wed, 9 Sep 2026 09:35:41 +0530 Subject: [PATCH 1/3] fix: use session timezone for timestamp subtraction --- datafusion/core/tests/expr_api/mod.rs | 38 +++++ .../optimizer/src/analyzer/type_coercion.rs | 140 ++++++++++++++++-- .../optimizer/src/scalar_subquery_to_join.rs | 4 +- .../simplify_expressions/expr_simplifier.rs | 4 +- datafusion/optimizer/src/utils.rs | 2 +- .../test_files/datetime/timestamps.slt | 43 ++++++ 6 files changed, 216 insertions(+), 15 deletions(-) diff --git a/datafusion/core/tests/expr_api/mod.rs b/datafusion/core/tests/expr_api/mod.rs index 19ff3933193de..b4583f41e2dee 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() { ); } +#[test] +fn test_create_physical_expr_timestamp_subtraction_uses_session_timezone() { + const NANOS_PER_SECOND: i64 = 1_000_000_000; + // 2024-11-01 00:00:00 in America/New_York (04:00 UTC). + let timestamp_tz: ArrayRef = Arc::new( + TimestampNanosecondArray::from(vec![1_730_433_600 * NANOS_PER_SECOND]) + .with_timezone("America/New_York"), + ); + // The same wall-clock time without a timezone. + let timestamp: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![ + 1_730_419_200 * NANOS_PER_SECOND, + ])); + 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".into()); + 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() + .into_array(batch.num_rows()) + .unwrap(); + + assert_eq!( + ScalarValue::try_from_array(&result, 0).unwrap(), + ScalarValue::DurationNanosecond(Some(12 * 60 * 60 * NANOS_PER_SECOND)), + ); +} + /// 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 07d3173d63342..510c784139280 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -102,7 +102,13 @@ impl AnalyzerRule for TypeCoercion { // 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, + config.execution.time_zone.as_deref(), + ) + })? .data; // finish @@ -116,6 +122,7 @@ impl AnalyzerRule for TypeCoercion { 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 @@ -157,7 +164,8 @@ fn analyze_internal( 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 +183,25 @@ fn analyze_internal( /// Rewrite expressions to apply type coercion. pub struct TypeCoercionRewriter<'a> { pub(crate) schema: &'a DFSchema, + 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. pub fn new(schema: &'a DFSchema) -> Self { - Self { schema } + Self { + schema, + session_time_zone: None, + } + } + + pub(crate) fn with_session_time_zone( + mut self, + session_time_zone: Option<&'a str>, + ) -> Self { + self.session_time_zone = session_time_zone; + self } /// Coerce the [`LogicalPlan`]. @@ -398,9 +418,14 @@ impl<'a> TypeCoercionRewriter<'a> { ) -> Result<(Expr, Expr)> { let left_data_type = left.get_type(left_schema)?; let right_data_type = right.get_type(right_schema)?; - let (left_type, right_type) = + let (left_type, right_type) = if let Some(types) = + self.timestamp_subtraction_input_types(&left_data_type, &op, &right_data_type) + { + types + } else { BinaryTypeCoercer::new(&left_data_type, &op, &right_data_type) - .get_input_types()?; + .get_input_types()? + }; let left_cast_ok = can_cast_types(&left_data_type, &left_type); let right_cast_ok = can_cast_types(&right_data_type, &right_type); @@ -434,6 +459,47 @@ impl<'a> TypeCoercionRewriter<'a> { Ok((left_expr, right_expr)) } + /// Coerces the timezone-naive side of timestamp subtraction using the session + /// timezone, matching PostgreSQL and DuckDB. The timezone-aware side keeps its + /// timezone while both operands are widened to the same precision. + fn timestamp_subtraction_input_types( + &self, + left_type: &DataType, + op: &Operator, + right_type: &DataType, + ) -> Option<(DataType, DataType)> { + if op != &Operator::Minus { + return None; + } + let session_time_zone = self.session_time_zone?; + let (left_time_zone, right_time_zone) = match (left_type, right_type) { + ( + DataType::Timestamp(_, Some(left_time_zone)), + DataType::Timestamp(_, None), + ) => ( + Some(Arc::clone(left_time_zone)), + Some(Arc::from(session_time_zone)), + ), + ( + DataType::Timestamp(_, None), + DataType::Timestamp(_, Some(right_time_zone)), + ) => ( + Some(Arc::from(session_time_zone)), + Some(Arc::clone(right_time_zone)), + ), + _ => return None, + }; + let DataType::Timestamp(unit, _) = comparison_coercion(left_type, right_type)? + else { + return None; + }; + + Some(( + DataType::Timestamp(unit, left_time_zone), + DataType::Timestamp(unit, right_time_zone), + )) + } + fn coerce_date_time_math_op( expr: Expr, op: &Operator, @@ -588,8 +654,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 +670,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 +690,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)?; @@ -648,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)?; @@ -2721,7 +2794,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 +2805,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 +2816,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()?; @@ -3201,6 +3274,53 @@ mod test { ) } + #[test] + fn timestamp_subtraction_uses_session_timezone() -> Result<()> { + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new( + "tstz", + DataType::Timestamp( + TimeUnit::Second, + Some("America/New_York".into()), + ), + true, + ), + Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + ] + .into(), + std::collections::HashMap::new(), + )?; + let input = Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(schema), + })); + let subtract = |left, right| { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col(left)), + Operator::Minus, + Box::new(col(right)), + )) + }; + let plan = LogicalPlan::Projection(Projection::try_new( + vec![subtract("tstz", "ts"), subtract("ts", "tstz")], + input, + )?); + + let mut options = ConfigOptions::default(); + options.execution.time_zone = Some("+08:00".to_string()); + let rule = Arc::new(TypeCoercion::new()); + assert_analyzed_plan_with_config_eq_snapshot!( + options, + rule, + plan, + @r#" + Projection: CAST(tstz AS Timestamp(ns, "America/New_York")) - CAST(ts AS Timestamp(ns, "+08:00")), CAST(ts AS Timestamp(ns, "+08:00")) - CAST(tstz AS Timestamp(ns, "America/New_York")) + EmptyRelation: rows=0 + "#, + ) + } + #[test] fn in_subquery_cast_subquery() -> Result<()> { let empty_int32 = empty_with_type(DataType::Int32); diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 44011a125ba96..ebb037471f1d8 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -398,9 +398,7 @@ 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(), - }; + 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..7654aa19b484c 100644 --- a/datafusion/optimizer/src/utils.rs +++ b/datafusion/optimizer/src/utils.rs @@ -244,7 +244,7 @@ fn evaluate_expr_with_null_column<'a>( } fn coerce(expr: Expr, schema: &DFSchema) -> Result { - let mut expr_rewrite = TypeCoercionRewriter { schema }; + let mut expr_rewrite = TypeCoercionRewriter::new(schema); expr.rewrite(&mut expr_rewrite).data() } diff --git a/datafusion/sqllogictest/test_files/datetime/timestamps.slt b/datafusion/sqllogictest/test_files/datetime/timestamps.slt index e47a5f0b23d7d..66d453404e041 100644 --- a/datafusion/sqllogictest/test_files/datetime/timestamps.slt +++ b/datafusion/sqllogictest/test_files/datetime/timestamps.slt @@ -1980,6 +1980,49 @@ SELECT '2000-01-01T00:00:00'::timestamp - '2010-01-01T00:00:00'::timestamp; ---- -3653 days 0 hours 0 mins 0.000000000 secs +# Regression test for https://github.com/apache/datafusion/issues/13212 +statement ok +SET TIME ZONE = '+08' + +query ??? +WITH timestamps(ts_tz, ts) AS ( + VALUES ('2024-11-01T00:00:00+00:00'::timestamptz, '2024-11-01T00:00:00'::timestamp) +) +SELECT + '2024-11-01T00:00:00+00:00'::timestamptz - '2024-11-01T00:00:00'::timestamp, + 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 0 days -8 hours 0 mins 0.000000000 secs + +# The session timezone, not the aware operand's timezone, controls the cast. +query ?? +WITH timestamps(ts_tz, ts) AS ( + VALUES ( + arrow_cast('2024-11-01T00:00:00-04:00', 'Timestamp(Nanosecond, Some("America/New_York"))'), + '2024-11-01T00:00:00'::timestamp + ) +) +SELECT ts_tz - ts, ts - ts_tz +FROM timestamps; +---- +0 days 12 hours 0 mins 0.000000000 secs 0 days -12 hours 0 mins 0.000000000 secs + +statement ok +SET TIME ZONE = 'America/New_York' + +# The timezone-naive operand is interpreted in EST and EDT, respectively. +query ?? +SELECT + '2024-01-15T00:00:00+00:00'::timestamptz - '2024-01-15T00:00:00'::timestamp, + '2024-07-15T00:00:00+00:00'::timestamptz - '2024-07-15T00:00:00'::timestamp; +---- +0 days -5 hours 0 mins 0.000000000 secs 0 days -4 hours 0 mins 0.000000000 secs + +statement ok +RESET datafusion.execution.time_zone + # Interval - Timestamp => error query error Cannot coerce arithmetic expression Interval\(MonthDayNano\) - Timestamp\(ns\) to valid types SELECT i - ts1 from FOO; From 4f209e52232a2886e650bafb1719efc425366bca Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Wed, 9 Sep 2026 10:25:35 +0530 Subject: [PATCH 2/3] test: cover timezone coercion in scalar subquery --- .../sqllogictest/test_files/datetime/timestamps.slt | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/datafusion/sqllogictest/test_files/datetime/timestamps.slt b/datafusion/sqllogictest/test_files/datetime/timestamps.slt index 66d453404e041..55d941125e878 100644 --- a/datafusion/sqllogictest/test_files/datetime/timestamps.slt +++ b/datafusion/sqllogictest/test_files/datetime/timestamps.slt @@ -1996,6 +1996,17 @@ FROM timestamps; ---- 0 days 8 hours 0 mins 0.000000000 secs 0 days 8 hours 0 mins 0.000000000 secs 0 days -8 hours 0 mins 0.000000000 secs +# The session timezone is propagated when coercing a scalar subquery. +query ? +SELECT ( + SELECT ts_tz - ts + FROM ( + VALUES ('2024-11-01T00:00:00+00:00'::timestamptz, '2024-11-01T00:00:00'::timestamp) + ) AS timestamps(ts_tz, ts) +); +---- +0 days 8 hours 0 mins 0.000000000 secs + # The session timezone, not the aware operand's timezone, controls the cast. query ?? WITH timestamps(ts_tz, ts) AS ( From 0390171d92356a8f2b87542f464fefeb8c590812 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Thu, 10 Sep 2026 09:46:02 +0530 Subject: [PATCH 3/3] fix: apply session timezone to timestamp comparisons --- .../expr-common/src/type_coercion/binary.rs | 67 +++++- .../type_coercion/binary/tests/arithmetic.rs | 21 ++ .../optimizer/src/analyzer/type_coercion.rs | 210 ++++++++++-------- .../test_files/datetime/timestamps.slt | 57 ++++- 4 files changed, 255 insertions(+), 100 deletions(-) diff --git a/datafusion/expr-common/src/type_coercion/binary.rs b/datafusion/expr-common/src/type_coercion/binary.rs index 381897ae86fdc..9d20f016f3e1e 100644 --- a/datafusion/expr-common/src/type_coercion/binary.rs +++ b/datafusion/expr-common/src/type_coercion/binary.rs @@ -79,6 +79,7 @@ pub struct BinaryTypeCoercer<'a> { lhs: &'a DataType, op: &'a Operator, rhs: &'a DataType, + session_time_zone: Option<&'a str>, lhs_spans: Spans, op_spans: Spans, @@ -93,12 +94,20 @@ impl<'a> BinaryTypeCoercer<'a> { lhs, op, rhs, + session_time_zone: None, lhs_spans: Spans::new(), op_spans: Spans::new(), rhs_spans: Spans::new(), } } + /// Sets the session timezone used to coerce mixed timezone-aware and + /// timezone-naive timestamps for comparisons and subtraction. + pub fn with_session_time_zone(mut self, session_time_zone: Option<&'a str>) -> Self { + self.session_time_zone = session_time_zone; + self + } + /// Sets the spans information for the left side of the binary expression, /// so better diagnostics can be provided in case of errors. pub fn set_lhs_spans(&mut self, spans: Spans) { @@ -197,14 +206,24 @@ impl<'a> BinaryTypeCoercer<'a> { GtEq | IsDistinctFrom | IsNotDistinctFrom => { - comparison_coercion(lhs, rhs).map(Signature::comparison).ok_or_else(|| { - plan_datafusion_err!( - "Cannot infer common argument type for comparison operation {} {} {}", - self.lhs, - self.op, - self.rhs - ) - }) + if let Some((lhs, rhs)) = + self.timestamp_types_with_session_timezone(lhs, rhs) + { + Ok(Signature { + lhs, + rhs, + ret: Boolean, + }) + } else { + comparison_coercion(lhs, rhs).map(Signature::comparison).ok_or_else(|| { + plan_datafusion_err!( + "Cannot infer common argument type for comparison operation {} {} {}", + self.lhs, + self.op, + self.rhs + ) + }) + } } And | Or => if matches!((lhs, rhs), (Boolean | Null, Boolean | Null)) { // Logical binary boolean operators can only be evaluated for @@ -285,6 +304,17 @@ impl<'a> BinaryTypeCoercer<'a> { return Ok(Signature { lhs, rhs, ret }); } Plus | Minus | Multiply | Divide | Modulo => { + if self.op == &Minus + && let Some((lhs, rhs)) = + self.timestamp_types_with_session_timezone(lhs, rhs) + { + let ret = self.get_result(&lhs, &rhs).map_err(|e| { + plan_datafusion_err!( + "Cannot get result type for temporal operation {} {} {}: {e}", self.lhs, self.op, self.rhs + ) + })?; + return Ok(Signature { lhs, rhs, ret }); + } if let Ok(ret) = self.get_result(lhs, rhs) { // Temporal arithmetic, e.g. Date32 + Interval @@ -357,6 +387,27 @@ impl<'a> BinaryTypeCoercer<'a> { }) } + /// Coerces a mixed timezone-aware and timezone-naive timestamp pair to the + /// session timezone and widens both operands to the same precision. + fn timestamp_types_with_session_timezone( + &self, + lhs: &DataType, + rhs: &DataType, + ) -> Option<(DataType, DataType)> { + use DataType::Timestamp; + + let session_time_zone = self.session_time_zone?; + let ((Timestamp(lhs_unit, Some(_)), Timestamp(rhs_unit, None)) + | (Timestamp(lhs_unit, None), Timestamp(rhs_unit, Some(_)))) = (lhs, rhs) + else { + return None; + }; + let unit = timeunit_coercion(lhs_unit, rhs_unit); + let data_type = Timestamp(unit, Some(Arc::from(session_time_zone))); + + Some((data_type.clone(), data_type)) + } + /// Returns the resulting type of a binary expression evaluating the `op` with the left and right hand types pub fn get_result_type(&'a self) -> Result { self.signature().map(|sig| sig.ret) 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..dab1ae08b4812 100644 --- a/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs +++ b/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs @@ -57,6 +57,27 @@ fn test_date_timestamp_arithmetic_error() -> Result<()> { Ok(()) } +#[test] +fn test_timestamp_session_timezone_coercion() -> Result<()> { + let aware = DataType::Timestamp(Millisecond, Some("America/New_York".into())); + let naive = DataType::Timestamp(Nanosecond, None); + let expected = DataType::Timestamp(Nanosecond, Some("+08:00".into())); + + for op in [Operator::Minus, Operator::Eq, Operator::Gt] { + let (lhs, rhs) = BinaryTypeCoercer::new(&aware, &op, &naive) + .with_session_time_zone(Some("+08:00")) + .get_input_types()?; + assert_eq!((lhs, rhs), (expected.clone(), expected.clone())); + + let (lhs, rhs) = BinaryTypeCoercer::new(&naive, &op, &aware) + .with_session_time_zone(Some("+08:00")) + .get_input_types()?; + assert_eq!((lhs, rhs), (expected.clone(), expected.clone())); + } + + Ok(()) +} + #[test] fn test_decimal_mathematics_op_type() { // Decimal32 diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 510c784139280..f2e0eae5deede 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -43,15 +43,12 @@ use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema; use datafusion_expr::expr_schema::cast_subquery; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::type_coercion::binary::{ - comparison_coercion, like_coercion, regex_coercion, type_union_coercion, + like_coercion, regex_coercion, type_union_coercion, }; use datafusion_expr::type_coercion::functions::{ UDFCoercionExt, fields_with_udf, value_fields_with_higher_order_udf_and_lambdas, }; -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::other::get_coerce_type_for_case_expression; use datafusion_expr::type_coercion::{ is_datetime, is_interval, is_signed_numeric, is_timestamp, }; @@ -418,14 +415,10 @@ impl<'a> TypeCoercionRewriter<'a> { ) -> Result<(Expr, Expr)> { let left_data_type = left.get_type(left_schema)?; let right_data_type = right.get_type(right_schema)?; - let (left_type, right_type) = if let Some(types) = - self.timestamp_subtraction_input_types(&left_data_type, &op, &right_data_type) - { - types - } else { + let (left_type, right_type) = BinaryTypeCoercer::new(&left_data_type, &op, &right_data_type) - .get_input_types()? - }; + .with_session_time_zone(self.session_time_zone) + .get_input_types()?; let left_cast_ok = can_cast_types(&left_data_type, &left_type); let right_cast_ok = can_cast_types(&right_data_type, &right_type); @@ -459,47 +452,6 @@ impl<'a> TypeCoercionRewriter<'a> { Ok((left_expr, right_expr)) } - /// Coerces the timezone-naive side of timestamp subtraction using the session - /// timezone, matching PostgreSQL and DuckDB. The timezone-aware side keeps its - /// timezone while both operands are widened to the same precision. - fn timestamp_subtraction_input_types( - &self, - left_type: &DataType, - op: &Operator, - right_type: &DataType, - ) -> Option<(DataType, DataType)> { - if op != &Operator::Minus { - return None; - } - let session_time_zone = self.session_time_zone?; - let (left_time_zone, right_time_zone) = match (left_type, right_type) { - ( - DataType::Timestamp(_, Some(left_time_zone)), - DataType::Timestamp(_, None), - ) => ( - Some(Arc::clone(left_time_zone)), - Some(Arc::from(session_time_zone)), - ), - ( - DataType::Timestamp(_, None), - DataType::Timestamp(_, Some(right_time_zone)), - ) => ( - Some(Arc::from(session_time_zone)), - Some(Arc::clone(right_time_zone)), - ), - _ => return None, - }; - let DataType::Timestamp(unit, _) = comparison_coercion(left_type, right_type)? - else { - return None; - }; - - Some(( - DataType::Timestamp(unit, left_time_zone), - DataType::Timestamp(unit, right_time_zone), - )) - } - fn coerce_date_time_math_op( expr: Expr, op: &Operator, @@ -695,11 +647,14 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { .data; let expr_type = expr.get_type(self.schema)?; let subquery_type = new_plan.schema().field(0).data_type(); - let common_type = comparison_coercion(&expr_type, subquery_type).ok_or( - plan_datafusion_err!( + let common_type = comparison_coercion_with_session_timezone( + &expr_type, + subquery_type, + self.session_time_zone, + ) + .ok_or(plan_datafusion_err!( "expr type {expr_type} can't cast to {subquery_type} in InSubquery" - ), - )?; + ))?; let new_subquery = Subquery { subquery: Arc::new(new_plan), outer_ref_columns: subquery.outer_ref_columns, @@ -732,11 +687,14 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { "expr type {expr_type} can't cast to {subquery_type} in SetComparison" ); } - let common_type = comparison_coercion(&expr_type, subquery_type).ok_or( - plan_datafusion_err!( - "expr type {expr_type} can't cast to {subquery_type} in SetComparison" - ), - )?; + let common_type = comparison_coercion_with_session_timezone( + &expr_type, + subquery_type, + self.session_time_zone, + ) + .ok_or(plan_datafusion_err!( + "expr type {expr_type} can't cast to {subquery_type} in SetComparison" + ))?; let new_subquery = Subquery { subquery: Arc::new(new_plan), outer_ref_columns: subquery.outer_ref_columns, @@ -851,26 +809,27 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { }) => { let expr_type = expr.get_type(self.schema)?; let low_type = low.get_type(self.schema)?; - let low_coerced_type = comparison_coercion(&expr_type, &low_type) - .ok_or_else(|| { - internal_datafusion_err!( - "Failed to coerce types {expr_type} and {low_type} in BETWEEN expression" - ) - })?; + let low_coerced_type = comparison_coercion_with_session_timezone( + &expr_type, + &low_type, + self.session_time_zone, + ) + .ok_or_else(|| { + internal_datafusion_err!( + "Failed to coerce types {expr_type} and {low_type} in BETWEEN expression" + ) + })?; let high_type = high.get_type(self.schema)?; - let high_coerced_type = comparison_coercion(&expr_type, &high_type) + let coercion_type = comparison_coercion_with_session_timezone( + &low_coerced_type, + &high_type, + self.session_time_zone, + ) .ok_or_else(|| { internal_datafusion_err!( "Failed to coerce types {expr_type} and {high_type} in BETWEEN expression" ) })?; - let coercion_type = - comparison_coercion(&low_coerced_type, &high_coerced_type) - .ok_or_else(|| { - internal_datafusion_err!( - "Failed to coerce types {expr_type} and {high_type} in BETWEEN expression" - ) - })?; Ok(Transformed::yes(Expr::Between(Between::new( Box::new(expr.cast_to(&coercion_type, self.schema)?), negated, @@ -888,8 +847,16 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { .iter() .map(|list_expr| list_expr.get_type(self.schema)) .collect::>>()?; - let result_type = - get_coerce_type_for_list(&expr_data_type, &list_data_types); + let result_type = list_data_types.iter().try_fold( + expr_data_type.clone(), + |coerced_type, list_data_type| { + comparison_coercion_with_session_timezone( + &coerced_type, + list_data_type, + self.session_time_zone, + ) + }, + ); match result_type { None => plan_err!( "Can not find compatible types to compare {expr_data_type} with [{}]", @@ -913,7 +880,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 }) => { @@ -1334,7 +1302,26 @@ fn coerce_scalar_function_argument( Expr::Literal(value, metadata).cast_to(data_type, schema) } -fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { +/// Returns the common type for non-binary comparison expressions using the +/// same coercion rules as binary equality. +fn comparison_coercion_with_session_timezone( + lhs_type: &DataType, + rhs_type: &DataType, + session_time_zone: Option<&str>, +) -> Option { + let (lhs_type, rhs_type) = BinaryTypeCoercer::new(lhs_type, &Operator::Eq, rhs_type) + .with_session_time_zone(session_time_zone) + .get_input_types() + .ok()?; + + (lhs_type == rhs_type).then_some(lhs_type) +} + +fn coerce_case_expression( + case: Case, + schema: &DFSchema, + session_time_zone: Option<&str>, +) -> Result { // Given expressions like: // // CASE a1 @@ -1391,7 +1378,16 @@ fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result { .iter() .map(|(when, _then)| when.get_type(schema)) .collect::>>()?; - let coerced_type = get_coerce_type_for_case_when(&when_types, case_type); + let coerced_type = when_types.iter().try_fold( + case_type.clone(), + |coerced_type, when_type| { + comparison_coercion_with_session_timezone( + &coerced_type, + when_type, + session_time_zone, + ) + }, + ); coerced_type.ok_or_else(|| { plan_datafusion_err!( "Failed to coerce case ({case_type}) and when ({}) \ @@ -1625,7 +1621,9 @@ mod test { use datafusion_common::{ DFSchema, DFSchemaRef, Result, ScalarValue, Spans, TableReference, }; - use datafusion_expr::expr::{self, InSubquery, Like, ScalarFunction}; + use datafusion_expr::expr::{ + self, InSubquery, Like, ScalarFunction, SetComparison, SetQuantifier, + }; use datafusion_expr::logical_plan::{EmptyRelation, Projection, Sort}; use datafusion_expr::test::function_stub::avg_udaf; use datafusion_expr::{ @@ -2933,7 +2931,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 @@ -2956,7 +2954,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 { @@ -2968,7 +2966,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" @@ -2983,7 +2981,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" @@ -3003,7 +3001,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); }; } @@ -3315,7 +3313,7 @@ mod test { rule, plan, @r#" - Projection: CAST(tstz AS Timestamp(ns, "America/New_York")) - CAST(ts AS Timestamp(ns, "+08:00")), CAST(ts AS Timestamp(ns, "+08:00")) - CAST(tstz AS Timestamp(ns, "America/New_York")) + Projection: CAST(tstz AS Timestamp(ns, "+08:00")) - CAST(ts AS Timestamp(ns, "+08:00")), CAST(ts AS Timestamp(ns, "+08:00")) - CAST(tstz AS Timestamp(ns, "+08:00")) EmptyRelation: rows=0 "#, ) @@ -3406,4 +3404,40 @@ mod test { " ) } + + #[test] + fn timestamp_set_comparison_uses_session_timezone() -> Result<()> { + let outer = empty_with_type(DataType::Timestamp( + TimeUnit::Second, + Some("America/New_York".into()), + )); + let subquery = empty_with_type(DataType::Timestamp(TimeUnit::Nanosecond, None)); + let set_comparison = Expr::SetComparison(SetComparison::new( + Box::new(col("a")), + Subquery { + subquery, + outer_ref_columns: vec![], + spans: Spans::new(), + }, + Operator::Eq, + SetQuantifier::Any, + )); + let plan = LogicalPlan::Filter(Filter::try_new(set_comparison, outer)?); + + let mut options = ConfigOptions::default(); + options.execution.time_zone = Some("+08:00".to_string()); + let rule = Arc::new(TypeCoercion::new()); + assert_analyzed_plan_with_config_eq_snapshot!( + options, + rule, + plan, + @r#" + Filter: CAST(a AS Timestamp(ns, "+08:00")) = ANY () + Subquery: + Projection: CAST(a AS Timestamp(ns, "+08:00")) + EmptyRelation: rows=0 + EmptyRelation: rows=0 + "#, + ) + } } diff --git a/datafusion/sqllogictest/test_files/datetime/timestamps.slt b/datafusion/sqllogictest/test_files/datetime/timestamps.slt index 55d941125e878..7646945d04667 100644 --- a/datafusion/sqllogictest/test_files/datetime/timestamps.slt +++ b/datafusion/sqllogictest/test_files/datetime/timestamps.slt @@ -2007,18 +2007,40 @@ SELECT ( ---- 0 days 8 hours 0 mins 0.000000000 secs -# The session timezone, not the aware operand's timezone, controls the cast. -query ?? +# The session timezone, not the aware operand's timezone, controls coercion. +query BBBBBB?? WITH timestamps(ts_tz, ts) AS ( VALUES ( arrow_cast('2024-11-01T00:00:00-04:00', 'Timestamp(Nanosecond, Some("America/New_York"))'), '2024-11-01T00:00:00'::timestamp ) ) -SELECT ts_tz - ts, ts - ts_tz +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; ---- -0 days 12 hours 0 mins 0.000000000 secs 0 days -12 hours 0 mins 0.000000000 secs +false true false false false false 0 days 12 hours 0 mins 0.000000000 secs 0 days -12 hours 0 mins 0.000000000 secs + +# IN subqueries use the same session-timezone comparison coercion. +query I +WITH timestamps(ts_tz, ts) AS ( + VALUES ( + arrow_cast('2024-11-01T00:00:00-04:00', 'Timestamp(Nanosecond, Some("America/New_York"))'), + '2024-11-01T00:00:00'::timestamp + ) +) +SELECT count(*) +FROM timestamps +WHERE ts_tz IN (SELECT ts FROM timestamps); +---- +0 statement ok SET TIME ZONE = 'America/New_York' @@ -2034,6 +2056,33 @@ SELECT statement ok RESET datafusion.execution.time_zone +# With no session timezone, the naive operand is still read as UTC. +statement ok +CREATE TABLE no_session_tz 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 ?? +SELECT ts_tz - ts, ts - ts_tz FROM no_session_tz; +---- +0 days 4 hours 0 mins 0.000000000 secs 0 days -4 hours 0 mins 0.000000000 secs + +statement ok +SET datafusion.explain.logical_plan_only = true + +query TT +EXPLAIN SELECT ts_tz - ts FROM no_session_tz; +---- +logical_plan +01)Projection: no_session_tz.ts_tz - no_session_tz.ts +02)--TableScan: no_session_tz projection=[ts_tz, ts] + +statement ok +SET datafusion.explain.logical_plan_only = false + +statement ok +DROP TABLE no_session_tz + # Interval - Timestamp => error query error Cannot coerce arithmetic expression Interval\(MonthDayNano\) - Timestamp\(ns\) to valid types SELECT i - ts1 from FOO;