diff --git a/datafusion/functions/src/datetime/to_local_time.rs b/datafusion/functions/src/datetime/to_local_time.rs index 973afa549e0b2..1e98f2d654033 100644 --- a/datafusion/functions/src/datetime/to_local_time.rs +++ b/datafusion/functions/src/datetime/to_local_time.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::ops::Add; use std::sync::Arc; use arrow::array::timezone::Tz; @@ -30,8 +29,8 @@ use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc}; use datafusion_common::cast::as_primitive_array; use datafusion_common::{ - Result, ScalarValue, exec_err, internal_datafusion_err, internal_err, - utils::take_function_args, + Result, ScalarValue, exec_datafusion_err, exec_err, internal_datafusion_err, + internal_err, utils::take_function_args, }; use datafusion_expr::{ Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -307,7 +306,10 @@ fn to_local_time(time_value: &ColumnarValue) -> Result { /// ``` /// /// See `test_adjust_to_local_time()` for example -pub fn adjust_to_local_time(ts: i64, tz: Tz) -> Result { +pub fn adjust_to_local_time( + ts: i64, + tz: impl TimeZone + Copy, +) -> Result { fn convert_timestamp(ts: i64, converter: F) -> Result> where F: Fn(i64) -> MappedLocalTime>, @@ -337,12 +339,18 @@ pub fn adjust_to_local_time(ts: i64, tz: Tz) -> Result RecordBatch { + let mut fields = vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + Field::new("timezone", DataType::Utf8, false), + ]; + let mut arrays: Vec = vec![ + Arc::new(TimestampMicrosecondArray::from_iter((0..rows).map(|row| { + (!half_null || row % 2 != 0).then_some(row as i64 * 86_400_000_000) + }))), + Arc::new(StringArray::from_iter_values( + (0..rows).map(|row| if (row / 2) % 2 == 0 { "PST" } else { "EST" }), + )), + ]; + + for column in 2..columns { + fields.push(Field::new( + format!("unused_{column}"), + DataType::Int64, + false, + )); + arrays.push(Arc::new(Int64Array::from_iter_values( + (0..rows).map(|row| (column * rows + row) as i64), + ))); + } + + RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap() +} + +fn criterion_benchmark(c: &mut Criterion) { + let rows = 8192; + let expr = SparkFromUtcTimestampExpr::new( + Arc::new(Column::new("timestamp", 0)), + Arc::new(Column::new("timezone", 1)), + ); + let mut group = c.benchmark_group("from_utc_timestamp_column"); + group.throughput(Throughput::Elements(rows as u64)); + + for columns in [2, 64, 256] { + for (half_null, name) in [(false, "no_nulls"), (true, "half_nulls")] { + let batch = batch(rows, columns, half_null); + let result = expr.evaluate(&batch).unwrap().into_array(rows).unwrap(); + assert_eq!(result.len(), rows); + assert_eq!(result.null_count(), if half_null { rows / 2 } else { 0 }); + + group.bench_with_input( + BenchmarkId::new(format!("{columns}_columns"), name), + &batch, + |b, batch| { + b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())); + }, + ); + } + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/spark/src/function/datetime/from_utc_timestamp.rs b/datafusion/spark/src/function/datetime/from_utc_timestamp.rs index bfca677c1dcce..2e18be931fac4 100644 --- a/datafusion/spark/src/function/datetime/from_utc_timestamp.rs +++ b/datafusion/spark/src/function/datetime/from_utc_timestamp.rs @@ -15,24 +15,30 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{Display, Formatter}; +use std::hash::{Hash, Hasher}; use std::sync::Arc; -use arrow::array::timezone::Tz; -use arrow::array::{Array, ArrayRef, AsArray, PrimitiveBuilder, StringArrayType}; +use super::timezone::SparkTimeZone; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, PrimitiveBuilder, StringArrayType, + new_empty_array, +}; use arrow::datatypes::TimeUnit; use arrow::datatypes::{ - ArrowTimestampType, DataType, Field, FieldRef, TimestampMicrosecondType, + ArrowTimestampType, DataType, Field, FieldRef, Schema, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; +use arrow::record_batch::RecordBatch; use datafusion_common::types::{NativeType, logical_string}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, exec_datafusion_err, exec_err, internal_err}; use datafusion_expr::{ - Coercion, ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, - Signature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, ExpressionPlacement, ReturnFieldArgs, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; -use datafusion_functions::datetime::to_local_time::adjust_to_local_time; use datafusion_functions::utils::make_scalar_function; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; /// Apache Spark `from_utc_timestamp` function. /// @@ -42,6 +48,21 @@ use datafusion_functions::utils::make_scalar_function; /// timezone-agnostic. So in Apache Spark this function just shift the timestamp value from UTC timezone to /// the given timezone. /// +/// # Compatibility +/// +/// This function accepts Spark's additional offset spellings, Java short IDs, and +/// legacy SystemV zone IDs. They are parsed only when both arguments are non-null. +/// Spark's generated execution may validate a foldable timezone before processing +/// rows, so this implementation can return nulls where Spark raises an error. +/// +/// Timezone conversion uses Chrono's calendar range, which is narrower than Spark's +/// implemented timestamp range, although it includes Spark's documented range. +/// Region timezone rules from `chrono-tz` may also differ after 2099 because it does +/// not extrapolate later daylight saving time transitions. +/// +/// This UDF operates on already-evaluated arguments. Callers that need conditional +/// evaluation of computed timezone expressions should construct [`SparkFromUtcTimestampExpr`]. +/// /// See #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkFromUtcTimestamp { @@ -169,16 +190,21 @@ where { let ts_primitive = ts_array.as_primitive::(); let mut builder = PrimitiveBuilder::::with_capacity(ts_array.len()); + // Parse a scalar zone once per batch and avoid re-parsing consecutive column values + let mut last_timezone = None; for (ts_opt, tz_opt) in ts_primitive.iter().zip(tz_array.iter()) { match (ts_opt, tz_opt) { (Some(ts), Some(tz_str)) => { - let tz: Tz = tz_str.parse().map_err(|e| { - exec_datafusion_err!( - "`from_utc_timestamp`: invalid timezone '{tz_str}': {e}" - ) - })?; - let val = adjust_to_local_time::(ts, tz)?; + let timezone = match last_timezone { + Some((previous, timezone)) if previous == tz_str => timezone, + _ => { + let timezone = parse_timezone(tz_str)?; + last_timezone = Some((tz_str, timezone)); + timezone + } + }; + let val = timezone.adjust_to_local_time::(ts)?; builder.append_value(val); } _ => builder.append_null(), @@ -188,3 +214,529 @@ where builder = builder.with_timezone_opt(return_tz_opt); Ok(Arc::new(builder.finish())) } + +fn parse_timezone(timezone: &str) -> Result { + SparkTimeZone::parse(timezone).map_err(|e| { + exec_datafusion_err!("`from_utc_timestamp`: invalid timezone '{timezone}': {e}") + }) +} + +/// Physical `from_utc_timestamp` expression with conditional timezone evaluation. +/// +/// The timestamp child is evaluated once for each non-empty batch. If every timestamp is null, +/// the timezone child is skipped. Otherwise, literal and column timezones are read directly, and +/// their values are ignored on null timestamp rows. Computed timezone expressions are evaluated +/// only for rows with non-null timestamps. Empty batches evaluate neither child. +/// +/// This intentionally differs from Spark's generated evaluation of foldable timezones, which +/// evaluates the timezone and validates a non-null value before processing timestamp rows. This +/// expression can therefore skip an invalid timezone or an error in the timezone expression when +/// no timestamp is non-null. It also evaluates the timestamp first when the timezone is null, so a +/// timestamp error can be raised even though Spark's generated path would return null. Planners +/// using this expression should expose these differences as an incompatibility. +#[derive(Debug, Clone)] +pub struct SparkFromUtcTimestampExpr { + timestamp: Arc, + timezone: Arc, +} + +impl SparkFromUtcTimestampExpr { + /// Construct an expression that skips computed timezone evaluation for null timestamps. + pub fn new( + timestamp: Arc, + timezone: Arc, + ) -> Self { + Self { + timestamp, + timezone, + } + } +} + +impl Display for SparkFromUtcTimestampExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "from_utc_timestamp({}, {})", + self.timestamp, self.timezone + ) + } +} + +impl PartialEq for SparkFromUtcTimestampExpr { + fn eq(&self, other: &Self) -> bool { + self.timestamp.eq(&other.timestamp) && self.timezone.eq(&other.timezone) + } +} + +impl Eq for SparkFromUtcTimestampExpr {} + +impl Hash for SparkFromUtcTimestampExpr { + fn hash(&self, state: &mut H) { + self.timestamp.hash(state); + self.timezone.hash(state); + } +} + +impl PhysicalExpr for SparkFromUtcTimestampExpr { + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } + + fn return_field(&self, input_schema: &Schema) -> Result { + let timestamp = self.timestamp.return_field(input_schema)?; + Ok(Arc::new(Field::new( + "from_utc_timestamp", + timestamp.data_type().clone(), + timestamp.is_nullable() || self.timezone.nullable(input_schema)?, + ))) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + // There are no rows on which either child should run. + if batch.num_rows() == 0 { + return Ok(ColumnarValue::Array(new_empty_array( + &self.data_type(batch.schema_ref())?, + ))); + } + + let timestamp = self.timestamp.evaluate(batch)?; + let timezone = match ×tamp { + ColumnarValue::Scalar(value) if value.is_null() => { + return Ok(timestamp); + } + ColumnarValue::Scalar(_) => self.timezone.evaluate(batch)?, + ColumnarValue::Array(array) if array.null_count() == array.len() => { + return Ok(timestamp); + } + ColumnarValue::Array(_) + if matches!( + self.timezone.placement(), + ExpressionPlacement::Literal | ExpressionPlacement::Column + ) => + { + // Literals and columns only read existing values. The conversion loop skips null + // timestamps, so avoid filtering and copying unrelated batch columns. + self.timezone.evaluate(batch)? + } + ColumnarValue::Array(array) => match array.nulls() { + Some(nulls) => { + let selection = BooleanArray::new(nulls.inner().clone(), None); + self.timezone.evaluate_selection(batch, &selection)? + } + None => self.timezone.evaluate(batch)?, + }, + }; + make_scalar_function(spark_from_utc_timestamp, vec![])(&[timestamp, timezone]) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.timestamp, &self.timezone] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [timestamp, timezone] = children.as_slice() else { + return internal_err!( + "SparkFromUtcTimestampExpr expected 2 children, got {}", + children.len() + ); + }; + Ok(Arc::new(Self::new( + Arc::clone(timestamp), + Arc::clone(timezone), + ))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + LargeStringArray, StringArray, StringViewArray, TimestampMicrosecondArray, + }; + use arrow::compute::CastOptions; + use chrono::{DateTime, Utc}; + use datafusion::physical_expr::expressions::{ + BinaryExpr, CaseExpr, CastExpr, Column, Literal, + }; + use datafusion_common::{ScalarValue, config::ConfigOptions}; + use datafusion_expr::Operator; + + fn expression_batch(timestamps: Vec>, values: Vec<&str>) -> RecordBatch { + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))), + true, + ), + Field::new("value", DataType::Utf8, false), + ])), + vec![ + Arc::new( + TimestampMicrosecondArray::from(timestamps).with_timezone("UTC"), + ), + Arc::new(StringArray::from(values)), + ], + ) + .unwrap() + } + + fn throwing_timestamp() -> Arc { + Arc::new(CastExpr::new( + Arc::new(Literal::new(ScalarValue::Utf8(Some("invalid".to_owned())))), + DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))), + Some(CastOptions { + safe: false, + ..Default::default() + }), + )) + } + + fn computed_timezone() -> Arc { + let value = Arc::new(CastExpr::new( + Arc::new(Column::new("value", 1)), + DataType::Int32, + Some(CastOptions { + safe: false, + ..Default::default() + }), + )); + let is_zero = Arc::new(BinaryExpr::new( + value, + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(0)))), + )); + Arc::new( + CaseExpr::try_new( + None, + vec![( + is_zero, + Arc::new(Literal::new(ScalarValue::Utf8(Some("UTC".to_owned())))), + )], + Some(Arc::new(Literal::new(ScalarValue::Utf8(Some( + "PST".to_owned(), + ))))), + ) + .unwrap(), + ) + } + + fn invoke( + timestamp: ColumnarValue, + timezone: ColumnarValue, + ) -> Result { + let return_field = Arc::new(Field::new("result", timestamp.data_type(), true)); + SparkFromUtcTimestamp::new().invoke_with_args(ScalarFunctionArgs { + args: vec![timestamp, timezone], + arg_fields: vec![], + number_rows: 0, + return_field, + config_options: Arc::new(ConfigOptions::default()), + }) + } + + fn timestamps(values: Vec>) -> ColumnarValue { + ColumnarValue::Array(Arc::new( + TimestampMicrosecondArray::from(values).with_timezone("UTC"), + )) + } + + fn scalar_zone(value: Option<&str>) -> ColumnarValue { + ColumnarValue::Scalar(ScalarValue::Utf8(value.map(str::to_owned))) + } + + fn assert_timestamps(result: ColumnarValue, expected: &[Option]) { + let ColumnarValue::Array(array) = result else { + panic!("expected timestamp array"); + }; + let array = array.as_primitive::(); + assert_eq!(array.iter().collect::>(), expected); + assert_eq!(array.timezone(), Some("UTC")); + } + + #[test] + fn column_zones_preserve_values_nulls_and_timestamp_metadata() { + let summer = 1_719_792_000_123_456; + let zones = vec![ + Some("+00:00:01"), + Some("GMT+1"), + Some("PST"), + Some("invalid"), + None, + ]; + let zone_arrays: Vec = vec![ + Arc::new(StringArray::from(zones.clone())), + Arc::new(LargeStringArray::from(zones.clone())), + Arc::new(StringViewArray::from(zones)), + ]; + for zones in zone_arrays { + let result = invoke( + timestamps(vec![Some(-1), Some(0), Some(summer), None, Some(summer)]), + ColumnarValue::Array(zones), + ) + .unwrap(); + assert_timestamps( + result, + &[ + Some(999_999), + Some(3_600_000_000), + Some(summer - 25_200_000_000), + None, + None, + ], + ); + } + } + + #[test] + fn scalar_and_array_arguments_work_across_batches() { + for values in [vec![Some(-1), None], vec![Some(123_456), Some(-1_000_001)]] { + let expected = values + .iter() + .map(|value| value.map(|value| value + 3_723_000_000)) + .collect::>(); + let result = + invoke(timestamps(values), scalar_zone(Some("UTC+1:02:03"))).unwrap(); + assert_timestamps(result, &expected); + } + + let timestamp = 1_719_792_000_123_456; + let result = invoke( + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond( + Some(timestamp), + Some(Arc::from("UTC")), + )), + ColumnarValue::Array(Arc::new(StringArray::from(vec![ + Some("GMT+1"), + Some("PST"), + None, + Some("GMT+1"), + ]))), + ) + .unwrap(); + assert_timestamps( + result, + &[ + Some(timestamp + 3_600_000_000), + Some(timestamp - 25_200_000_000), + None, + Some(timestamp + 3_600_000_000), + ], + ); + } + + #[test] + fn fixed_offset_seconds_preserve_each_timestamp_unit() { + for (input, expected) in [ + ( + ScalarValue::TimestampSecond(Some(-1), None), + ScalarValue::TimestampSecond(Some(0), None), + ), + ( + ScalarValue::TimestampMillisecond(Some(-1), None), + ScalarValue::TimestampMillisecond(Some(999), None), + ), + ( + ScalarValue::TimestampMicrosecond(Some(-1), None), + ScalarValue::TimestampMicrosecond(Some(999_999), None), + ), + ( + ScalarValue::TimestampNanosecond(Some(-1), None), + ScalarValue::TimestampNanosecond(Some(999_999_999), None), + ), + ] { + let result = invoke( + ColumnarValue::Scalar(input.clone()), + scalar_zone(Some("+00:00:01")), + ) + .unwrap(); + let ColumnarValue::Scalar(actual) = result else { + panic!("expected scalar timestamp"); + }; + assert_eq!(actual, expected); + + let expr = SparkFromUtcTimestampExpr::new( + Arc::new(Literal::new(input)), + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "+00:00:01".to_owned(), + )))), + ); + let result = expr + .evaluate(&expression_batch(vec![Some(0)], vec!["0"])) + .unwrap(); + let ColumnarValue::Scalar(actual) = result else { + panic!("expected scalar timestamp"); + }; + assert_eq!(actual, expected); + } + } + + #[test] + fn invalid_literal_timezone_is_validated_lazily() { + let expr = SparkFromUtcTimestampExpr::new( + Arc::new(Column::new("timestamp", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some("+19:00".to_owned())))), + ); + assert_timestamps( + expr.evaluate(&expression_batch(vec![None, None], vec!["0", "0"])) + .unwrap(), + &[None, None], + ); + let error = expr + .evaluate(&expression_batch(vec![Some(0)], vec!["0"])) + .unwrap_err(); + assert!(error.to_string().contains("invalid timezone '+19:00'")); + } + + #[test] + fn null_timezone_does_not_skip_timestamp_evaluation() { + let expr = SparkFromUtcTimestampExpr::new( + throwing_timestamp(), + Arc::new(Literal::new(ScalarValue::Utf8(None))), + ); + let error = expr + .evaluate(&expression_batch(vec![Some(0)], vec!["0"])) + .unwrap_err(); + assert!(error.to_string().contains("invalid")); + } + + #[test] + fn timezone_evaluates_only_non_null_timestamp_rows() { + let summer = 1_719_792_000_123_456; + let expr = SparkFromUtcTimestampExpr::new( + Arc::new(Column::new("timestamp", 0)), + computed_timezone(), + ); + let batch = expression_batch( + vec![None, Some(0), None, Some(summer)], + vec!["invalid", "0", "invalid", "1"], + ); + assert_timestamps( + expr.evaluate(&batch).unwrap(), + &[None, Some(0), None, Some(summer - 25_200_000_000)], + ); + + // The same timezone child must still throw when its timestamp is non-null. + let error = expr + .evaluate(&expression_batch(vec![Some(0)], vec!["invalid"])) + .unwrap_err(); + assert!(error.to_string().contains("invalid")); + } + + #[test] + fn timezone_column_ignores_values_on_null_timestamp_rows() { + let summer = 1_719_792_000_123_456; + let expr = SparkFromUtcTimestampExpr::new( + Arc::new(Column::new("timestamp", 0)), + Arc::new(Column::new("value", 1)), + ); + let batch = expression_batch( + vec![None, Some(0), None, Some(summer)], + vec!["invalid", "UTC", "+19:00", "PST"], + ); + assert_timestamps( + expr.evaluate(&batch).unwrap(), + &[None, Some(0), None, Some(summer - 25_200_000_000)], + ); + + let error = expr + .evaluate(&expression_batch(vec![Some(0)], vec!["+19:00"])) + .unwrap_err(); + assert!(error.to_string().contains("invalid timezone '+19:00'")); + } + + #[test] + fn timezone_skips_all_null_and_scalar_null_timestamps() { + let batch = expression_batch(vec![None, None], vec!["invalid", "invalid"]); + let expr = SparkFromUtcTimestampExpr::new( + Arc::new(Column::new("timestamp", 0)), + computed_timezone(), + ); + assert_timestamps(expr.evaluate(&batch).unwrap(), &[None, None]); + + let expr = SparkFromUtcTimestampExpr::new( + Arc::new(Literal::new(ScalarValue::TimestampMicrosecond(None, None))), + computed_timezone(), + ); + assert!(matches!( + expr.evaluate(&batch).unwrap(), + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(None, None)) + )); + } + + #[test] + fn empty_batch_skips_children() { + let batch = expression_batch(vec![], vec![]); + let expr = + SparkFromUtcTimestampExpr::new(throwing_timestamp(), computed_timezone()); + assert_timestamps(expr.evaluate(&batch).unwrap(), &[]); + } + + #[test] + fn fixed_offset_outside_calendar_range_returns_error() { + for (timestamp, timezone) in [ + (DateTime::::MAX_UTC.timestamp_micros(), "+00:00:01"), + (DateTime::::MIN_UTC.timestamp_micros(), "-00:00:01"), + ] { + let error = invoke( + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond( + Some(timestamp), + None, + )), + scalar_zone(Some(timezone)), + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("exceeds Chrono's calendar range") + ); + } + } + + #[test] + fn invalid_zones_are_parsed_only_for_non_null_pairs() { + let zones = + ColumnarValue::Array(Arc::new(StringArray::from(vec!["invalid", "+19:00"]))); + assert_timestamps( + invoke(timestamps(vec![None, None]), zones).unwrap(), + &[None, None], + ); + + let error = invoke(timestamps(vec![None, Some(0)]), scalar_zone(Some("+19:00"))) + .unwrap_err(); + assert!(error.to_string().contains("invalid timezone '+19:00'")); + + let result = invoke( + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(None, None)), + scalar_zone(Some("invalid")), + ) + .unwrap(); + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(None, None)) + )); + + let result = invoke( + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond( + Some(i64::MAX), + None, + )), + scalar_zone(None), + ) + .unwrap(); + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(None, None)) + )); + + assert_timestamps( + invoke(timestamps(vec![]), scalar_zone(Some("invalid"))).unwrap(), + &[], + ); + } +} diff --git a/datafusion/spark/src/function/datetime/mod.rs b/datafusion/spark/src/function/datetime/mod.rs index 70ab024c329aa..a365fbe4a971a 100644 --- a/datafusion/spark/src/function/datetime/mod.rs +++ b/datafusion/spark/src/function/datetime/mod.rs @@ -29,6 +29,7 @@ pub mod make_interval; pub mod monthname; pub mod next_day; pub mod time_trunc; +mod timezone; pub mod to_utc_timestamp; pub mod trunc; pub mod unix; diff --git a/datafusion/spark/src/function/datetime/timezone.rs b/datafusion/spark/src/function/datetime/timezone.rs new file mode 100644 index 0000000000000..17971c85ea4ef --- /dev/null +++ b/datafusion/spark/src/function/datetime/timezone.rs @@ -0,0 +1,442 @@ +// 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. + +use arrow::array::timezone::Tz; +use arrow::datatypes::{ArrowTimestampType, TimeUnit}; +use chrono::{DateTime, Datelike, FixedOffset, NaiveDate, Utc}; +use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err}; +use datafusion_functions::datetime::to_local_time::adjust_to_local_time; + +/// A timezone resolved using Spark's `getZoneId` spelling rules. +/// +/// Fixed offsets are kept as seconds to preserve Spark's second-resolution offsets. +/// Region IDs use Arrow's Chrono-backed timezone database. +/// Java's legacy SystemV zones require their own rules because that database omits them. +#[derive(Debug, Clone, Copy)] +pub(super) enum SparkTimeZone { + FixedOffset(i32), + Region(Tz), + SystemV(i32), +} + +impl SparkTimeZone { + pub(super) fn parse(timezone: &str) -> Result { + // These are the exact aliases in java.time.ZoneId.SHORT_IDS. The regional + // aliases retain their daylight saving rules; EST, HST, and MST are fixed. + let timezone = match timezone { + "Z" => return Ok(Self::FixedOffset(0)), + // The JDK explicitly excludes this legacy IANA alias. + "ROC" => return Err(exec_datafusion_err!("Unknown timezone '{timezone}'")), + "EST" => return Ok(Self::FixedOffset(-5 * 3600)), + "HST" => return Ok(Self::FixedOffset(-10 * 3600)), + "MST" => return Ok(Self::FixedOffset(-7 * 3600)), + "ACT" => "Australia/Darwin", + "AET" => "Australia/Sydney", + "AGT" => "America/Argentina/Buenos_Aires", + "ART" => "Africa/Cairo", + "AST" => "America/Anchorage", + "BET" => "America/Sao_Paulo", + "BST" => "Asia/Dhaka", + "CAT" => "Africa/Harare", + "CNT" => "America/St_Johns", + "CST" => "America/Chicago", + "CTT" => "Asia/Shanghai", + "EAT" => "Africa/Addis_Ababa", + "ECT" => "Europe/Paris", + "IET" => "America/Indiana/Indianapolis", + "IST" => "Asia/Kolkata", + "JST" => "Asia/Tokyo", + "MIT" => "Pacific/Apia", + "NET" => "Asia/Yerevan", + "NST" => "Pacific/Auckland", + "PLT" => "Asia/Karachi", + "PNT" => "America/Phoenix", + "PRT" => "America/Puerto_Rico", + "PST" => "America/Los_Angeles", + "SST" => "Pacific/Guadalcanal", + "VST" => "Asia/Ho_Chi_Minh", + "SystemV/AST4" => return Ok(Self::FixedOffset(-4 * 3600)), + "SystemV/EST5" => return Ok(Self::FixedOffset(-5 * 3600)), + "SystemV/CST6" => return Ok(Self::FixedOffset(-6 * 3600)), + "SystemV/MST7" => return Ok(Self::FixedOffset(-7 * 3600)), + "SystemV/PST8" => return Ok(Self::FixedOffset(-8 * 3600)), + "SystemV/YST9" => return Ok(Self::FixedOffset(-9 * 3600)), + "SystemV/HST10" => return Ok(Self::FixedOffset(-10 * 3600)), + "SystemV/AST4ADT" => return Ok(Self::SystemV(-4 * 3600)), + "SystemV/EST5EDT" => return Ok(Self::SystemV(-5 * 3600)), + "SystemV/CST6CDT" => return Ok(Self::SystemV(-6 * 3600)), + "SystemV/MST7MDT" => return Ok(Self::SystemV(-7 * 3600)), + "SystemV/PST8PDT" => return Ok(Self::SystemV(-8 * 3600)), + "SystemV/YST9YDT" => return Ok(Self::SystemV(-9 * 3600)), + _ => timezone, + }; + + let offset = if timezone.starts_with(['+', '-']) { + Some(timezone) + } else { + ["UTC", "GMT", "UT"].into_iter().find_map(|prefix| { + timezone + .strip_prefix(prefix) + .filter(|suffix| suffix.is_empty() || suffix.starts_with(['+', '-'])) + }) + }; + if let Some(offset) = offset { + if offset.is_empty() { + return Ok(Self::FixedOffset(0)); + } + return parse_offset(offset).map(Self::FixedOffset).ok_or_else(|| { + exec_datafusion_err!("Invalid timezone offset '{offset}'") + }); + } + + Ok(Self::Region(timezone.parse()?)) + } + + pub(super) fn adjust_to_local_time( + self, + ts: i64, + ) -> Result { + match self { + Self::Region(timezone) => adjust_to_local_time::(ts, timezone), + Self::FixedOffset(seconds) => adjust_by_offset::(ts, seconds), + Self::SystemV(standard_offset) => { + adjust_by_offset::(ts, system_v_offset::(ts, standard_offset)?) + } + } + } +} + +/// Parses Java's signed offsets after Spark's single-digit hour/minute normalization. +/// +/// Spark permits one-digit hours before a colon, and one-digit minutes only when +/// minutes are the final component. Java also accepts compact HHMM and HHMMSS forms. +fn parse_offset(offset: &str) -> Option { + let sign = match offset.as_bytes().first()? { + b'+' => 1, + b'-' => -1, + _ => return None, + }; + let value = &offset[1..]; + let (hours, minutes, seconds) = if let Some((hours, rest)) = value.split_once(':') { + let hours = parse_digits(hours, 1, 2)?; + if let Some((minutes, seconds)) = rest.split_once(':') { + ( + hours, + parse_digits(minutes, 2, 2)?, + parse_digits(seconds, 2, 2)?, + ) + } else { + (hours, parse_digits(rest, 1, 2)?, 0) + } + } else { + // Check ASCII before byte slicing so malformed Unicode inputs cannot panic. + if !value.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + match value.len() { + 1 | 2 => (parse_digits(value, 1, 2)?, 0, 0), + 4 => ( + parse_digits(&value[..2], 2, 2)?, + parse_digits(&value[2..], 2, 2)?, + 0, + ), + 6 => ( + parse_digits(&value[..2], 2, 2)?, + parse_digits(&value[2..4], 2, 2)?, + parse_digits(&value[4..], 2, 2)?, + ), + _ => return None, + } + }; + if hours > 18 + || minutes > 59 + || seconds > 59 + || (hours == 18 && (minutes != 0 || seconds != 0)) + { + return None; + } + Some(sign * (hours * 3600 + minutes * 60 + seconds)) +} + +fn parse_digits(value: &str, min_len: usize, max_len: usize) -> Option { + if !(min_len..=max_len).contains(&value.len()) + || !value.bytes().all(|byte| byte.is_ascii_digit()) + { + return None; + } + value.parse().ok() +} + +fn units_per_second() -> i64 { + match T::UNIT { + TimeUnit::Second => 1, + TimeUnit::Millisecond => 1_000, + TimeUnit::Microsecond => 1_000_000, + TimeUnit::Nanosecond => 1_000_000_000, + } +} + +fn adjust_by_offset(ts: i64, offset_seconds: i32) -> Result { + let timezone = FixedOffset::east_opt(offset_seconds).ok_or_else(|| { + internal_datafusion_err!("Spark timezone offsets must be within 18 hours") + })?; + adjust_to_local_time::(ts, timezone) +} + +/// Java's SystemV daylight zones use the same calendar with different standard offsets. +/// The JDK starts these transitions in 1900. Most years change on the last Sundays of +/// April and October; 1974 and 1975 have exceptional start/end dates. +/// The rule definitions are in OpenJDK's `make/data/tzdata/jdk11_backward`. +fn system_v_offset(ts: i64, standard_offset: i32) -> Result { + let seconds = ts.div_euclid(units_per_second::()); + let year = DateTime::::from_timestamp(seconds, 0) + .ok_or_else(|| { + exec_datafusion_err!("Timestamp is outside Chrono's calendar range") + })? + .year(); + if year < 1900 { + return Ok(standard_offset); + } + + let start = match year { + 1974 => NaiveDate::from_ymd_opt(year, 1, 6), + 1975 => NaiveDate::from_ymd_opt(year, 2, 23), + _ => last_sunday(year, 4, 30), + }; + let end = if year == 1974 { + last_sunday(year, 11, 30) + } else { + last_sunday(year, 10, 31) + }; + let transition = |date: Option, offset_before: i32| { + date.and_then(|date| date.and_hms_opt(2, 0, 0)) + .map(|local| local.and_utc().timestamp() - i64::from(offset_before)) + .ok_or_else(|| { + exec_datafusion_err!("Timestamp is outside Chrono's calendar range") + }) + }; + let start = transition(start, standard_offset)?; + let end = transition(end, standard_offset + 3600)?; + Ok(if (start..end).contains(&seconds) { + standard_offset + 3600 + } else { + standard_offset + }) +} + +fn last_sunday(year: i32, month: u32, last_day: u32) -> Option { + let date = NaiveDate::from_ymd_opt(year, month, last_day)?; + date.with_day(last_day - date.weekday().num_days_from_sunday()) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::TimestampMicrosecondType; + + fn shift(timezone: &str, timestamp: i64) -> i64 { + SparkTimeZone::parse(timezone) + .unwrap() + .adjust_to_local_time::(timestamp) + .unwrap() + } + + #[test] + fn spark_offset_spellings_preserve_sign_and_seconds() { + let offsets = [ + ("+0", 0), + ("-00", 0), + ("+1", 3600), + ("-01", -3600), + ("+0102", 3720), + ("-01:02", -3720), + ("+1:02", 3720), + ("+01:2", 3720), + ("+1:2", 3720), + ("+010203", 3723), + ("-01:02:03", -3723), + ("+1:02:03", 3723), + ("+00:00:01", 1), + ("-000001", -1), + ("+18:00", 64800), + ("-180000", -64800), + ]; + for prefix in ["", "UTC", "GMT", "UT"] { + for (offset, seconds) in offsets { + let timezone = format!("{prefix}{offset}"); + assert_eq!( + shift(&timezone, 123_456), + 123_456 + i64::from(seconds) * 1_000_000, + "{timezone}" + ); + } + } + for timezone in ["Z", "UTC", "GMT", "UT", "GMT0", "Etc/UTC"] { + assert_eq!(shift(timezone, -123_456), -123_456, "{timezone}"); + } + assert_eq!(shift("GMT+1", 0), 3_600_000_000); + assert_eq!(shift("Etc/GMT+1", 0), -3_600_000_000); + } + + #[test] + fn malformed_offsets_and_unknown_ids_are_rejected() { + for prefix in ["", "UTC", "GMT", "UT"] { + for offset in [ + "+", + "-", + "+19", + "-18:00:01", + "+18:01", + "+01:60", + "+01:99", + "+01:00:60", + "+1:2:03", + "+01:02:3", + "+123", + "+12345", + "+001:00", + "+01:", + "+:01", + "+01::00", + "+01:00:00:00", + "+01:00 ", + "+０1", + "+01:٠0", + "++01", + "--01", + ] { + let timezone = format!("{prefix}{offset}"); + assert!(SparkTimeZone::parse(&timezone).is_err(), "{timezone}"); + } + } + for timezone in [ + "", + " ", + "utc", + "pst", + "PDT", + "ROC", + "America/Unknown", + "SystemV/PST9PDT", + " UTC", + ] { + assert!(SparkTimeZone::parse(timezone).is_err(), "{timezone}"); + } + } + + #[test] + fn java_short_ids_match_their_regions_in_winter_and_summer() { + let aliases = [ + ("ACT", "Australia/Darwin"), + ("AET", "Australia/Sydney"), + ("AGT", "America/Argentina/Buenos_Aires"), + ("ART", "Africa/Cairo"), + ("AST", "America/Anchorage"), + ("BET", "America/Sao_Paulo"), + ("BST", "Asia/Dhaka"), + ("CAT", "Africa/Harare"), + ("CNT", "America/St_Johns"), + ("CST", "America/Chicago"), + ("CTT", "Asia/Shanghai"), + ("EAT", "Africa/Addis_Ababa"), + ("ECT", "Europe/Paris"), + ("IET", "America/Indiana/Indianapolis"), + ("IST", "Asia/Kolkata"), + ("JST", "Asia/Tokyo"), + ("MIT", "Pacific/Apia"), + ("NET", "Asia/Yerevan"), + ("NST", "Pacific/Auckland"), + ("PLT", "Asia/Karachi"), + ("PNT", "America/Phoenix"), + ("PRT", "America/Puerto_Rico"), + ("PST", "America/Los_Angeles"), + ("SST", "Pacific/Guadalcanal"), + ("VST", "Asia/Ho_Chi_Minh"), + ("EST", "-05:00"), + ("HST", "-10:00"), + ("MST", "-07:00"), + ]; + for (alias, canonical) in aliases { + for timestamp in [1_704_067_200_123_456, 1_719_792_000_123_456] { + assert_eq!( + shift(alias, timestamp), + shift(canonical, timestamp), + "{alias}" + ); + } + } + } + + #[test] + fn system_v_zones_match_java_transition_boundaries() { + let zones = [ + ("SystemV/AST4ADT", -4), + ("SystemV/EST5EDT", -5), + ("SystemV/CST6CDT", -6), + ("SystemV/MST7MDT", -7), + ("SystemV/PST8PDT", -8), + ("SystemV/YST9YDT", -9), + ]; + for (timezone, hours) in zones { + let standard = hours * 3600; + for (local, offset_before, offset_after) in [ + ("1900-04-29T02:00:00Z", standard, standard + 3600), + ("1974-01-06T02:00:00Z", standard, standard + 3600), + ("1974-11-24T02:00:00Z", standard + 3600, standard), + ("1975-02-23T02:00:00Z", standard, standard + 3600), + ("2024-04-28T02:00:00Z", standard, standard + 3600), + ("2024-10-27T02:00:00Z", standard + 3600, standard), + ] { + let transition = DateTime::parse_from_rfc3339(local) + .unwrap() + .timestamp_micros() + - i64::from(offset_before) * 1_000_000; + assert_eq!( + shift(timezone, transition - 1), + transition - 1 + i64::from(offset_before) * 1_000_000, + "{timezone} before {local}" + ); + assert_eq!( + shift(timezone, transition), + transition + i64::from(offset_after) * 1_000_000, + "{timezone} at {local}" + ); + } + let before_transitions = DateTime::parse_from_rfc3339("1899-07-01T00:00:00Z") + .unwrap() + .timestamp_micros(); + assert_eq!( + shift(timezone, before_transitions), + before_transitions + i64::from(standard) * 1_000_000 + ); + } + for (timezone, hours) in [ + ("SystemV/AST4", -4), + ("SystemV/EST5", -5), + ("SystemV/CST6", -6), + ("SystemV/MST7", -7), + ("SystemV/PST8", -8), + ("SystemV/YST9", -9), + ("SystemV/HST10", -10), + ] { + assert_eq!( + shift(timezone, 1_719_792_000_000_000), + 1_719_792_000_000_000 + i64::from(hours) * 3_600_000_000 + ); + } + } +} diff --git a/datafusion/sqllogictest/test_files/spark/datetime/from_utc_timestamp.slt b/datafusion/sqllogictest/test_files/spark/datetime/from_utc_timestamp.slt index 8200fccfc1984..8b04956cb309e 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/from_utc_timestamp.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/from_utc_timestamp.slt @@ -154,3 +154,28 @@ query P SELECT from_utc_timestamp('2020-11-04T14:06:40'::timestamp, 'America/New_York'::string); ---- 2020-11-04T09:06:40 + +# Spark timezone spellings supplied as column values, including second-resolution offsets. +query P +SELECT from_utc_timestamp(column1, column2) +FROM VALUES +('2024-07-01T00:00:00.123456'::timestamp, 'PST'::string), +('2024-07-01T00:00:00.123456'::timestamp, 'GMT+1'::string), +('2024-07-01T00:00:00.123456'::timestamp, 'UTC+1:2'::string), +('2024-07-01T00:00:00.123456'::timestamp, 'UT-010203'::string), +('2024-04-15T00:00:00.123456'::timestamp, 'SystemV/PST8PDT'::string), +('2024-07-01T00:00:00.123456'::timestamp, 'SystemV/HST10'::string), +(NULL::timestamp, '+19:00'::string); +---- +2024-06-30T17:00:00.123456 +2024-07-01T01:00:00.123456 +2024-07-01T01:02:00.123456 +2024-06-30T22:57:57.123456 +2024-04-14T16:00:00.123456 +2024-06-30T14:00:00.123456 +NULL + +# Arrow accepts this offset, but Spark limits offsets to 18 hours. +query error invalid timezone '\+19:00' +SELECT from_utc_timestamp(column1, column2) +FROM VALUES ('2024-07-01'::timestamp, '+19:00'::string); diff --git a/typos.toml b/typos.toml index 196766f12fbc0..6a958616a9f17 100644 --- a/typos.toml +++ b/typos.toml @@ -17,6 +17,7 @@ Ois = "Ois" alo = "alo" # abbreviations, common words, etc. +IST = "IST" typ = "typ" datas = "datas" YOUY = "YOUY"