Timezone aware extract SQL expression

datafusion/functions/src/datetime/common.rs

@@ -17,26 +17,78 @@
 
 use std::sync::Arc;
 
+use arrow::array::timezone::Tz;
 use arrow::array::{
     Array, ArrowPrimitiveType, AsArray, GenericStringArray, PrimitiveArray,
     StringArrayType, StringViewArray,
 };
 use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
-use arrow::datatypes::DataType;
+use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
+use arrow::datatypes::{ArrowTimestampType, DataType};
 use chrono::format::{parse, Parsed, StrftimeItems};
 use chrono::LocalResult::Single;
-use chrono::{DateTime, TimeZone, Utc};
+use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
+use std::ops::Add;
 
 use datafusion_common::cast::as_generic_string_array;
 use datafusion_common::{
-    exec_datafusion_err, exec_err, unwrap_or_internal_err, DataFusionError, Result,
-    ScalarType, ScalarValue,
+    exec_datafusion_err, exec_err, internal_datafusion_err, unwrap_or_internal_err,
+    DataFusionError, Result, ScalarType, ScalarValue,
 };
 use datafusion_expr::ColumnarValue;
 
 /// Error message if nanosecond conversion request beyond supported interval
 const ERR_NANOSECONDS_NOT_SUPPORTED: &str = "The dates that can be represented as nanoseconds have to be between 1677-09-21T00:12:44.0 and 2262-04-11T23:47:16.854775804";
 
+/// Adjusts a timestamp to local time by applying the timezone offset.
+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()),
+    }
+}
+
 /// Calls string_to_timestamp_nanos and converts the error type
 pub(crate) fn string_to_timestamp_nanos_shim(s: &str) -> Result<i64> {
     string_to_timestamp_nanos(s).map_err(|e| e.into())

datafusion/functions/src/datetime/date_part.rs

@@ -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,
@@ -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
@@ -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")],
         }
     }
 }
@@ -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)?;
 
@@ -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_trim);
+
+        // 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)?,
@@ -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() {
@@ -240,21 +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
+// 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
@@ -277,7 +373,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())?;
@@ -305,11 +400,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,
@@ -318,7 +410,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())?;

datafusion/functions/src/datetime/mod.rs

@@ -22,6 +22,7 @@ use std::sync::Arc;
 use datafusion_expr::ScalarUDF;
 
 pub mod common;
+pub use common::adjust_to_local_time;
 pub mod current_date;
 pub mod current_time;
 pub mod date_bin;