Skip to content

Timezone aware extract SQL expression #18990

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
+333 −78
Open

Timezone aware extract SQL expression #18990

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ce9be4d
resolve conflicts with main
codetyri0n Nov 10, 2025
edcdc63
make extract timezone aware
codetyri0n Nov 29, 2025
9d04c99
cargo fmt
codetyri0n Nov 29, 2025
85d7c0d
refactoring and setting volatility to stable
codetyri0n Nov 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 109 additions & 20 deletions datafusion/functions/src/datetime/date_part.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,16 +19,23 @@ use std::any::Any;
use std::str::FromStr;
use std::sync::Arc;

use arrow::array::{Array, ArrayRef, Float64Array, Int32Array};
use arrow::array::timezone::Tz;
use arrow::array::{Array, ArrayRef, Float64Array, Int32Array, PrimitiveBuilder};
use arrow::compute::kernels::cast_utils::IntervalUnit;
use arrow::compute::{binary, date_part, DatePart};
use arrow::datatypes::DataType::{
Date32, Date64, Duration, Interval, Time32, Time64, Timestamp,
};
use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
use arrow::datatypes::{DataType, Field, FieldRef, TimeUnit};
use arrow::datatypes::{
ArrowTimestampType, DataType, Field, FieldRef, TimeUnit, TimestampMicrosecondType,
TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
};

use datafusion_common::cast::as_primitive_array;
use datafusion_common::types::{logical_date, NativeType};

use super::adjust_to_local_time;
use datafusion_common::{
cast::{
as_date32_array, as_date64_array, as_int32_array, as_time32_millisecond_array,
Expand Down Expand Up @@ -56,7 +63,7 @@ use datafusion_macros::user_doc;
argument(
name = "part",
description = r#"Part of the date to return. The following date parts are supported:

- year
- quarter (emits value in inclusive range [1, 4] based on which quartile of the year the date is in)
- month
Expand Down Expand Up @@ -122,9 +129,9 @@ impl DatePartFunc {
Coercion::new_exact(TypeSignatureClass::Duration),
]),
],
Volatility::Immutable,
Volatility::Stable,
),
aliases: vec![String::from("datepart")],
aliases: vec![String::from("datepart"), String::from("extract")],
}
}
}
Expand Down Expand Up @@ -173,6 +180,7 @@ impl ScalarUDFImpl for DatePartFunc {
&self,
args: datafusion_expr::ScalarFunctionArgs,
) -> Result<ColumnarValue> {
let config = &args.config_options;
let args = args.args;
let [part, array] = take_function_args(self.name(), args)?;

Expand All @@ -193,12 +201,77 @@ impl ScalarUDFImpl for DatePartFunc {
ColumnarValue::Scalar(scalar) => scalar.to_array()?,
};

let (is_timezone_aware, tz_str_opt) = match array.data_type() {
Timestamp(_, Some(tz_str)) => (true, Some(Arc::clone(tz_str))),
_ => (false, None),
};

let part_trim = part_normalization(&part);
let is_epoch = is_epoch(&part);

// Epoch is timezone-independent - it always returns seconds since 1970-01-01 UTC
let array = if is_epoch {
array
} else if is_timezone_aware {
// For timezone-aware timestamps, extract in their own timezone
match tz_str_opt.as_ref() {
Some(tz_str) => {
let tz = interpret_session_timezone(tz_str)?;
match array.data_type() {
Timestamp(time_unit, _) => match time_unit {
Nanosecond => adjust_timestamp_array::<
TimestampNanosecondType,
>(&array, tz)?,
Microsecond => adjust_timestamp_array::<
TimestampMicrosecondType,
>(&array, tz)?,
Millisecond => adjust_timestamp_array::<
TimestampMillisecondType,
>(&array, tz)?,
Second => {
adjust_timestamp_array::<TimestampSecondType>(&array, tz)?
}
},
_ => array,
}
}
None => array,
}
} else if let Timestamp(time_unit, None) = array.data_type() {
// For naive timestamps, interpret in session timezone if available
match config.execution.time_zone.as_ref() {
Some(tz_str) => {
let tz = interpret_session_timezone(tz_str)?;

match time_unit {
Nanosecond => {
adjust_timestamp_array::<TimestampNanosecondType>(&array, tz)?
}
Microsecond => {
adjust_timestamp_array::<TimestampMicrosecondType>(
&array, tz,
)?
}
Millisecond => {
adjust_timestamp_array::<TimestampMillisecondType>(
&array, tz,
)?
}
Second => {
adjust_timestamp_array::<TimestampSecondType>(&array, tz)?
}
}
}
None => array,
}
} else {
array
};

// using IntervalUnit here means we hand off all the work of supporting plurals (like "seconds")
// and synonyms ( like "ms,msec,msecond,millisecond") to Arrow
let arr = if let Ok(interval_unit) = IntervalUnit::from_str(part_trim) {
match interval_unit {
let extracted = match interval_unit {
IntervalUnit::Year => date_part(array.as_ref(), DatePart::Year)?,
IntervalUnit::Month => date_part(array.as_ref(), DatePart::Month)?,
IntervalUnit::Week => date_part(array.as_ref(), DatePart::Week)?,
Expand All @@ -209,9 +282,10 @@ impl ScalarUDFImpl for DatePartFunc {
IntervalUnit::Millisecond => seconds_as_i32(array.as_ref(), Millisecond)?,
IntervalUnit::Microsecond => seconds_as_i32(array.as_ref(), Microsecond)?,
IntervalUnit::Nanosecond => seconds_as_i32(array.as_ref(), Nanosecond)?,
// century and decade are not supported by `DatePart`, although they are supported in postgres
_ => return exec_err!("Date part '{part}' not supported"),
}
};

extracted
} else {
// special cases that can be extracted (in postgres) but are not interval units
match part_trim.to_lowercase().as_str() {
Expand Down Expand Up @@ -240,23 +314,43 @@ impl ScalarUDFImpl for DatePartFunc {
}
}

fn adjust_timestamp_array<T: ArrowTimestampType>(
array: &ArrayRef,
tz: Tz,
) -> Result<ArrayRef> {
let mut builder = PrimitiveBuilder::<T>::new();
let primitive_array = as_primitive_array::<T>(array)?;
for ts_opt in primitive_array.iter() {
match ts_opt {
None => builder.append_null(),
Some(ts) => {
let adjusted_ts = adjust_to_local_time::<T>(ts, tz)?;
builder.append_value(adjusted_ts);
}
}
}
Ok(Arc::new(builder.finish()))
}

fn is_epoch(part: &str) -> bool {
let part = part_normalization(part);
matches!(part.to_lowercase().as_str(), "epoch")
}

// Try to remove quote if exist, if the quote is invalid, return original string and let the downstream function handle the error
fn part_normalization(part: &str) -> &str {
part.strip_prefix(|c| c == '\'' || c == '\"')
.and_then(|s| s.strip_suffix(|c| c == '\'' || c == '\"'))
.unwrap_or(part)
}

/// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
/// result to a total number of seconds, milliseconds, microseconds or
/// nanoseconds
fn interpret_session_timezone(tz_str: &str) -> Result<Tz> {
match tz_str.parse::<Tz>() {
Ok(tz) => Ok(tz),
Err(err) => exec_err!("Invalid timezone '{tz_str}': {err}"),
}
}

fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
// Nanosecond is neither supported in Postgres nor DuckDB, to avoid dealing
// with overflow and precision issue we don't support nanosecond
if unit == Nanosecond {
return not_impl_err!("Date part {unit:?} not supported");
Expand All @@ -277,7 +371,6 @@ fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
};

let secs = date_part(array, DatePart::Second)?;
// This assumes array is primitive and not a dictionary
let secs = as_int32_array(secs.as_ref())?;
let subsecs = date_part(array, DatePart::Nanosecond)?;
let subsecs = as_int32_array(subsecs.as_ref())?;
Expand Down Expand Up @@ -305,11 +398,8 @@ fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
}
}

/// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
/// result to a total number of seconds, milliseconds, microseconds or
/// nanoseconds
///
/// Given epoch return f64, this is a duplicated function to optimize for f64 type
// Converts seconds to f64 with the specified time unit.
// Used for Interval and Duration types that need floating-point precision.
fn seconds(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
let sf = match unit {
Second => 1_f64,
Expand All @@ -318,7 +408,6 @@ fn seconds(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
Nanosecond => 1_000_000_000_f64,
};
let secs = date_part(array, DatePart::Second)?;
// This assumes array is primitive and not a dictionary
let secs = as_int32_array(secs.as_ref())?;
let subsecs = date_part(array, DatePart::Nanosecond)?;
let subsecs = as_int32_array(subsecs.as_ref())?;
Expand Down
56 changes: 56 additions & 0 deletions datafusion/functions/src/datetime/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,13 @@
use std::sync::Arc;

use arrow::array::timezone::Tz;
use arrow::datatypes::ArrowTimestampType;
use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
use datafusion_common::{exec_err, internal_datafusion_err, Result};
use std::ops::Add;

use datafusion_expr::ScalarUDF;

pub mod common;
Expand All @@ -37,6 +44,55 @@ pub mod to_local_time;
pub mod to_timestamp;
pub mod to_unixtime;

// Adjusts a timestamp to local time by applying the timezone offset.
Copy link
Contributor Author

@codetyri0n codetyri0n Nov 29, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

adjust_to_local_time refactored to mod.rs as a common helper function for both date_part and to_local_time

CC: @martin-g

pub fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
where
F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
{
match converter(ts) {
MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
"Ambiguous timestamp. Do you mean {:?} or {:?}",
earliest,
latest
),
MappedLocalTime::None => exec_err!(
"The local time does not exist because there is a gap in the local time."
),
MappedLocalTime::Single(date_time) => Ok(date_time),
}
}

let date_time = match T::UNIT {
Nanosecond => Utc.timestamp_nanos(ts),
Microsecond => convert_timestamp(ts, |ts| Utc.timestamp_micros(ts))?,
Millisecond => convert_timestamp(ts, |ts| Utc.timestamp_millis_opt(ts))?,
Second => convert_timestamp(ts, |ts| Utc.timestamp_opt(ts, 0))?,
};

let offset_seconds: i64 = tz
.offset_from_utc_datetime(&date_time.naive_utc())
.fix()
.local_minus_utc() as i64;

let adjusted_date_time = date_time.add(
TimeDelta::try_seconds(offset_seconds)
.ok_or_else(|| internal_datafusion_err!("Offset seconds should be less than i64::MAX / 1_000 or greater than -i64::MAX / 1_000"))?,
);

// convert back to i64
match T::UNIT {
Nanosecond => adjusted_date_time.timestamp_nanos_opt().ok_or_else(|| {
internal_datafusion_err!(
"Failed to convert DateTime to timestamp in nanosecond. This error may occur if the date is out of range. The supported date ranges are between 1677-09-21T00:12:43.145224192 and 2262-04-11T23:47:16.854775807"
)
}),
Microsecond => Ok(adjusted_date_time.timestamp_micros()),
Millisecond => Ok(adjusted_date_time.timestamp_millis()),
Second => Ok(adjusted_date_time.timestamp()),
}
}

// create UDFs
make_udf_function!(current_date::CurrentDateFunc, current_date);
make_udf_function!(current_time::CurrentTimeFunc, current_time);
Expand Down
62 changes: 4 additions & 58 deletions datafusion/functions/src/datetime/to_local_time.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
// under the License.

use std::any::Any;
use std::ops::Add;
use std::sync::Arc;

use arrow::array::timezone::Tz;
Expand All @@ -27,13 +26,10 @@ use arrow::datatypes::{
ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
TimestampNanosecondType, TimestampSecondType,
};
use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};

use crate::datetime::adjust_to_local_time;
use datafusion_common::cast::as_primitive_array;
use datafusion_common::{
exec_err, internal_datafusion_err, internal_err, utils::take_function_args, Result,
ScalarValue,
};
use datafusion_common::{internal_err, utils::take_function_args, Result, ScalarValue};
use datafusion_expr::{
Coercion, ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignatureClass,
Volatility,
Expand Down Expand Up @@ -324,60 +320,12 @@ fn to_local_time(time_value: &ColumnarValue) -> Result<ColumnarValue> {
/// ```
///
/// See `test_adjust_to_local_time()` for example
fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
where
F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
{
match converter(ts) {
MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
"Ambiguous timestamp. Do you mean {:?} or {:?}",
earliest,
latest
),
MappedLocalTime::None => exec_err!(
"The local time does not exist because there is a gap in the local time."
),
MappedLocalTime::Single(date_time) => Ok(date_time),
}
}

let date_time = match T::UNIT {
Nanosecond => Utc.timestamp_nanos(ts),
Microsecond => convert_timestamp(ts, |ts| Utc.timestamp_micros(ts))?,
Millisecond => convert_timestamp(ts, |ts| Utc.timestamp_millis_opt(ts))?,
Second => convert_timestamp(ts, |ts| Utc.timestamp_opt(ts, 0))?,
};

let offset_seconds: i64 = tz
.offset_from_utc_datetime(&date_time.naive_utc())
.fix()
.local_minus_utc() as i64;

let adjusted_date_time = date_time.add(
// This should not fail under normal circumstances as the
// maximum possible offset is 26 hours (93,600 seconds)
TimeDelta::try_seconds(offset_seconds)
.ok_or_else(|| internal_datafusion_err!("Offset seconds should be less than i64::MAX / 1_000 or greater than -i64::MAX / 1_000"))?,
);

// convert the naive datetime back to i64
match T::UNIT {
Nanosecond => adjusted_date_time.timestamp_nanos_opt().ok_or_else(||
internal_datafusion_err!(
"Failed to convert DateTime to timestamp in nanosecond. This error may occur if the date is out of range. The supported date ranges are between 1677-09-21T00:12:43.145224192 and 2262-04-11T23:47:16.854775807"
)
),
Microsecond => Ok(adjusted_date_time.timestamp_micros()),
Millisecond => Ok(adjusted_date_time.timestamp_millis()),
Second => Ok(adjusted_date_time.timestamp()),
}
}

#[cfg(test)]
mod tests {
use std::sync::Arc;

use super::ToLocalTimeFunc;
use crate::datetime::adjust_to_local_time;
use arrow::array::{types::TimestampNanosecondType, Array, TimestampNanosecondArray};
use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
use arrow::datatypes::{DataType, Field, TimeUnit};
Expand All @@ -386,8 +334,6 @@ mod tests {
use datafusion_common::ScalarValue;
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};

use super::{adjust_to_local_time, ToLocalTimeFunc};

#[test]
fn test_adjust_to_local_time() {
let timestamp_str = "2020-03-31T13:40:00";
Expand Down
Loading