| // 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 std::any::Any; |
| use std::num::NonZeroI64; |
| use std::ops::{Add, Sub}; |
| use std::str::FromStr; |
| use std::sync::Arc; |
| |
| use arrow::array::temporal_conversions::{ |
| as_datetime_with_timezone, timestamp_ns_to_datetime, |
| }; |
| use arrow::array::timezone::Tz; |
| use arrow::array::types::{ |
| ArrowTimestampType, TimestampMicrosecondType, TimestampMillisecondType, |
| TimestampNanosecondType, TimestampSecondType, |
| }; |
| use arrow::array::{Array, ArrayRef, PrimitiveArray}; |
| use arrow::datatypes::DataType::{self, Null, Timestamp, Utf8, Utf8View}; |
| use arrow::datatypes::TimeUnit::{self, Microsecond, Millisecond, Nanosecond, Second}; |
| use datafusion_common::cast::as_primitive_array; |
| use datafusion_common::{ |
| exec_datafusion_err, exec_err, plan_err, DataFusionError, Result, ScalarValue, |
| }; |
| use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; |
| use datafusion_expr::TypeSignature::Exact; |
| use datafusion_expr::{ |
| ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD, |
| }; |
| use datafusion_macros::user_doc; |
| |
| use chrono::{ |
| DateTime, Datelike, Duration, LocalResult, NaiveDateTime, Offset, TimeDelta, Timelike, |
| }; |
| |
| /// Represents the granularity for date truncation operations |
| #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| enum DateTruncGranularity { |
| Microsecond, |
| Millisecond, |
| Second, |
| Minute, |
| Hour, |
| Day, |
| Week, |
| Month, |
| Quarter, |
| Year, |
| } |
| |
| impl DateTruncGranularity { |
| /// List of all supported granularity values |
| /// Cannot use HashMap here as it would require lazy_static or once_cell, |
| /// Rust does not support const HashMap yet. |
| const SUPPORTED_GRANULARITIES: &[&str] = &[ |
| "microsecond", |
| "millisecond", |
| "second", |
| "minute", |
| "hour", |
| "day", |
| "week", |
| "month", |
| "quarter", |
| "year", |
| ]; |
| |
| /// Parse a granularity string into a DateTruncGranularity enum |
| fn from_str(s: &str) -> Result<Self> { |
| // Using match for O(1) lookup - compiler optimizes this into a jump table or perfect hash |
| match s.to_lowercase().as_str() { |
| "microsecond" => Ok(Self::Microsecond), |
| "millisecond" => Ok(Self::Millisecond), |
| "second" => Ok(Self::Second), |
| "minute" => Ok(Self::Minute), |
| "hour" => Ok(Self::Hour), |
| "day" => Ok(Self::Day), |
| "week" => Ok(Self::Week), |
| "month" => Ok(Self::Month), |
| "quarter" => Ok(Self::Quarter), |
| "year" => Ok(Self::Year), |
| _ => { |
| let supported = Self::SUPPORTED_GRANULARITIES.join(", "); |
| exec_err!( |
| "Unsupported date_trunc granularity: '{s}'. Supported values are: {supported}" |
| ) |
| } |
| } |
| } |
| |
| /// Returns true if this granularity can be handled with simple arithmetic |
| /// (fine granularity: second, minute, millisecond, microsecond) |
| fn is_fine_granularity(&self) -> bool { |
| matches!( |
| self, |
| Self::Second | Self::Minute | Self::Millisecond | Self::Microsecond |
| ) |
| } |
| |
| /// Returns true if this granularity can be handled with simple arithmetic in UTC |
| /// (hour and day in addition to fine granularities) |
| fn is_fine_granularity_utc(&self) -> bool { |
| self.is_fine_granularity() || matches!(self, Self::Hour | Self::Day) |
| } |
| } |
| |
| #[user_doc( |
| doc_section(label = "Time and Date Functions"), |
| description = "Truncates a timestamp value to a specified precision.", |
| syntax_example = "date_trunc(precision, expression)", |
| argument( |
| name = "precision", |
| description = r#"Time precision to truncate to. The following precisions are supported: |
| |
| - year / YEAR |
| - quarter / QUARTER |
| - month / MONTH |
| - week / WEEK |
| - day / DAY |
| - hour / HOUR |
| - minute / MINUTE |
| - second / SECOND |
| - millisecond / MILLISECOND |
| - microsecond / MICROSECOND |
| "# |
| ), |
| argument( |
| name = "expression", |
| description = "Time expression to operate on. Can be a constant, column, or function." |
| ) |
| )] |
| #[derive(Debug, PartialEq, Eq, Hash)] |
| pub struct DateTruncFunc { |
| signature: Signature, |
| aliases: Vec<String>, |
| } |
| |
| impl Default for DateTruncFunc { |
| fn default() -> Self { |
| Self::new() |
| } |
| } |
| |
| impl DateTruncFunc { |
| pub fn new() -> Self { |
| Self { |
| signature: Signature::one_of( |
| vec![ |
| Exact(vec![Utf8, Timestamp(Nanosecond, None)]), |
| Exact(vec![Utf8View, Timestamp(Nanosecond, None)]), |
| Exact(vec![ |
| Utf8, |
| Timestamp(Nanosecond, Some(TIMEZONE_WILDCARD.into())), |
| ]), |
| Exact(vec![ |
| Utf8View, |
| Timestamp(Nanosecond, Some(TIMEZONE_WILDCARD.into())), |
| ]), |
| Exact(vec![Utf8, Timestamp(Microsecond, None)]), |
| Exact(vec![Utf8View, Timestamp(Microsecond, None)]), |
| Exact(vec![ |
| Utf8, |
| Timestamp(Microsecond, Some(TIMEZONE_WILDCARD.into())), |
| ]), |
| Exact(vec![ |
| Utf8View, |
| Timestamp(Microsecond, Some(TIMEZONE_WILDCARD.into())), |
| ]), |
| Exact(vec![Utf8, Timestamp(Millisecond, None)]), |
| Exact(vec![Utf8View, Timestamp(Millisecond, None)]), |
| Exact(vec![ |
| Utf8, |
| Timestamp(Millisecond, Some(TIMEZONE_WILDCARD.into())), |
| ]), |
| Exact(vec![ |
| Utf8View, |
| Timestamp(Millisecond, Some(TIMEZONE_WILDCARD.into())), |
| ]), |
| Exact(vec![Utf8, Timestamp(Second, None)]), |
| Exact(vec![Utf8View, Timestamp(Second, None)]), |
| Exact(vec![ |
| Utf8, |
| Timestamp(Second, Some(TIMEZONE_WILDCARD.into())), |
| ]), |
| Exact(vec![ |
| Utf8View, |
| Timestamp(Second, Some(TIMEZONE_WILDCARD.into())), |
| ]), |
| ], |
| Volatility::Immutable, |
| ), |
| aliases: vec![String::from("datetrunc")], |
| } |
| } |
| } |
| |
| impl ScalarUDFImpl for DateTruncFunc { |
| fn as_any(&self) -> &dyn Any { |
| self |
| } |
| |
| fn name(&self) -> &str { |
| "date_trunc" |
| } |
| |
| fn signature(&self) -> &Signature { |
| &self.signature |
| } |
| |
| fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { |
| match &arg_types[1] { |
| Timestamp(Nanosecond, None) | Utf8 | DataType::Date32 | Null => { |
| Ok(Timestamp(Nanosecond, None)) |
| } |
| Timestamp(Nanosecond, tz_opt) => Ok(Timestamp(Nanosecond, tz_opt.clone())), |
| Timestamp(Microsecond, tz_opt) => Ok(Timestamp(Microsecond, tz_opt.clone())), |
| Timestamp(Millisecond, tz_opt) => Ok(Timestamp(Millisecond, tz_opt.clone())), |
| Timestamp(Second, tz_opt) => Ok(Timestamp(Second, tz_opt.clone())), |
| _ => plan_err!( |
| "The date_trunc function can only accept timestamp as the second arg." |
| ), |
| } |
| } |
| |
| fn invoke_with_args( |
| &self, |
| args: datafusion_expr::ScalarFunctionArgs, |
| ) -> Result<ColumnarValue> { |
| let args = args.args; |
| let (granularity, array) = (&args[0], &args[1]); |
| |
| let granularity_str = if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) = |
| granularity |
| { |
| v.to_lowercase() |
| } else if let ColumnarValue::Scalar(ScalarValue::Utf8View(Some(v))) = granularity |
| { |
| v.to_lowercase() |
| } else { |
| return exec_err!("Granularity of `date_trunc` must be non-null scalar Utf8"); |
| }; |
| |
| let granularity = DateTruncGranularity::from_str(&granularity_str)?; |
| |
| fn process_array<T: ArrowTimestampType>( |
| array: &dyn Array, |
| granularity: DateTruncGranularity, |
| tz_opt: &Option<Arc<str>>, |
| ) -> Result<ColumnarValue> { |
| let parsed_tz = parse_tz(tz_opt)?; |
| let array = as_primitive_array::<T>(array)?; |
| |
| // fast path for fine granularity |
| // For modern timezones, it's correct to truncate "minute" in this way. |
| // Both datafusion and arrow are ignoring historical timezone's non-minute granularity |
| // bias (e.g., Asia/Kathmandu before 1919 is UTC+05:41:16). |
| // In UTC, "hour" and "day" have uniform durations and can be truncated with simple arithmetic |
| if granularity.is_fine_granularity() |
| || (parsed_tz.is_none() && granularity.is_fine_granularity_utc()) |
| { |
| let result = general_date_trunc_array_fine_granularity( |
| T::UNIT, |
| array, |
| granularity, |
| )?; |
| return Ok(ColumnarValue::Array(result)); |
| } |
| |
| let array: PrimitiveArray<T> = array |
| .try_unary(|x| general_date_trunc(T::UNIT, x, parsed_tz, granularity))? |
| .with_timezone_opt(tz_opt.clone()); |
| Ok(ColumnarValue::Array(Arc::new(array))) |
| } |
| |
| fn process_scalar<T: ArrowTimestampType>( |
| v: &Option<i64>, |
| granularity: DateTruncGranularity, |
| tz_opt: &Option<Arc<str>>, |
| ) -> Result<ColumnarValue> { |
| let parsed_tz = parse_tz(tz_opt)?; |
| let value = if let Some(v) = v { |
| Some(general_date_trunc(T::UNIT, *v, parsed_tz, granularity)?) |
| } else { |
| None |
| }; |
| let value = ScalarValue::new_timestamp::<T>(value, tz_opt.clone()); |
| Ok(ColumnarValue::Scalar(value)) |
| } |
| |
| Ok(match array { |
| ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(v, tz_opt)) => { |
| process_scalar::<TimestampNanosecondType>(v, granularity, tz_opt)? |
| } |
| ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(v, tz_opt)) => { |
| process_scalar::<TimestampMicrosecondType>(v, granularity, tz_opt)? |
| } |
| ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(v, tz_opt)) => { |
| process_scalar::<TimestampMillisecondType>(v, granularity, tz_opt)? |
| } |
| ColumnarValue::Scalar(ScalarValue::TimestampSecond(v, tz_opt)) => { |
| process_scalar::<TimestampSecondType>(v, granularity, tz_opt)? |
| } |
| ColumnarValue::Array(array) => { |
| let array_type = array.data_type(); |
| if let Timestamp(unit, tz_opt) = array_type { |
| match unit { |
| Second => process_array::<TimestampSecondType>( |
| array, |
| granularity, |
| tz_opt, |
| )?, |
| Millisecond => process_array::<TimestampMillisecondType>( |
| array, |
| granularity, |
| tz_opt, |
| )?, |
| Microsecond => process_array::<TimestampMicrosecondType>( |
| array, |
| granularity, |
| tz_opt, |
| )?, |
| Nanosecond => process_array::<TimestampNanosecondType>( |
| array, |
| granularity, |
| tz_opt, |
| )?, |
| } |
| } else { |
| return exec_err!("second argument of `date_trunc` is an unsupported array type: {array_type}"); |
| } |
| } |
| _ => { |
| return exec_err!( |
| "second argument of `date_trunc` must be timestamp scalar or array" |
| ); |
| } |
| }) |
| } |
| |
| fn aliases(&self) -> &[String] { |
| &self.aliases |
| } |
| |
| fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> { |
| // The DATE_TRUNC function preserves the order of its second argument. |
| let precision = &input[0]; |
| let date_value = &input[1]; |
| |
| if precision.sort_properties.eq(&SortProperties::Singleton) { |
| Ok(date_value.sort_properties) |
| } else { |
| Ok(SortProperties::Unordered) |
| } |
| } |
| fn documentation(&self) -> Option<&Documentation> { |
| self.doc() |
| } |
| } |
| |
| fn _date_trunc_coarse<T>( |
| granularity: DateTruncGranularity, |
| value: Option<T>, |
| ) -> Result<Option<T>> |
| where |
| T: Datelike + Timelike + Sub<Duration, Output = T> + Copy, |
| { |
| let value = match granularity { |
| DateTruncGranularity::Millisecond => value, |
| DateTruncGranularity::Microsecond => value, |
| DateTruncGranularity::Second => value.and_then(|d| d.with_nanosecond(0)), |
| DateTruncGranularity::Minute => value |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)), |
| DateTruncGranularity::Hour => value |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)), |
| DateTruncGranularity::Day => value |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)), |
| DateTruncGranularity::Week => value |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)) |
| .map(|d| { |
| d - TimeDelta::try_seconds(60 * 60 * 24 * d.weekday() as i64).unwrap() |
| }), |
| DateTruncGranularity::Month => value |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)) |
| .and_then(|d| d.with_day0(0)), |
| DateTruncGranularity::Quarter => value |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)) |
| .and_then(|d| d.with_day0(0)) |
| .and_then(|d| d.with_month(quarter_month(&d))), |
| DateTruncGranularity::Year => value |
| .and_then(|d| d.with_nanosecond(0)) |
| .and_then(|d| d.with_second(0)) |
| .and_then(|d| d.with_minute(0)) |
| .and_then(|d| d.with_hour(0)) |
| .and_then(|d| d.with_day0(0)) |
| .and_then(|d| d.with_month0(0)), |
| }; |
| Ok(value) |
| } |
| |
| fn quarter_month<T>(date: &T) -> u32 |
| where |
| T: Datelike, |
| { |
| 1 + 3 * ((date.month() - 1) / 3) |
| } |
| |
| fn _date_trunc_coarse_with_tz( |
| granularity: DateTruncGranularity, |
| value: Option<DateTime<Tz>>, |
| ) -> Result<Option<i64>> { |
| if let Some(value) = value { |
| let local = value.naive_local(); |
| let truncated = _date_trunc_coarse::<NaiveDateTime>(granularity, Some(local))?; |
| let truncated = truncated.and_then(|truncated| { |
| match truncated.and_local_timezone(value.timezone()) { |
| LocalResult::None => { |
| // This can happen if the date_trunc operation moves the time into |
| // an hour that doesn't exist due to daylight savings. On known example where |
| // this can happen is with historic dates in the America/Sao_Paulo time zone. |
| // To account for this adjust the time by a few hours, convert to local time, |
| // and then adjust the time back. |
| truncated |
| .sub(TimeDelta::try_hours(3).unwrap()) |
| .and_local_timezone(value.timezone()) |
| .single() |
| .map(|v| v.add(TimeDelta::try_hours(3).unwrap())) |
| } |
| LocalResult::Single(datetime) => Some(datetime), |
| LocalResult::Ambiguous(datetime1, datetime2) => { |
| // Because we are truncating from an equally or more specific time |
| // the original time must have been within the ambiguous local time |
| // period. Therefore the offset of one of these times should match the |
| // offset of the original time. |
| if datetime1.offset().fix() == value.offset().fix() { |
| Some(datetime1) |
| } else { |
| Some(datetime2) |
| } |
| } |
| } |
| }); |
| Ok(truncated.and_then(|value| value.timestamp_nanos_opt())) |
| } else { |
| _date_trunc_coarse::<NaiveDateTime>(granularity, None)?; |
| Ok(None) |
| } |
| } |
| |
| fn _date_trunc_coarse_without_tz( |
| granularity: DateTruncGranularity, |
| value: Option<NaiveDateTime>, |
| ) -> Result<Option<i64>> { |
| let value = _date_trunc_coarse::<NaiveDateTime>(granularity, value)?; |
| Ok(value.and_then(|value| value.and_utc().timestamp_nanos_opt())) |
| } |
| |
| /// Truncates the single `value`, expressed in nanoseconds since the |
| /// epoch, for granularities greater than 1 second, in taking into |
| /// account that some granularities are not uniform durations of time |
| /// (e.g. months are not always the same lengths, leap seconds, etc) |
| fn date_trunc_coarse( |
| granularity: DateTruncGranularity, |
| value: i64, |
| tz: Option<Tz>, |
| ) -> Result<i64> { |
| let value = match tz { |
| Some(tz) => { |
| // Use chrono DateTime<Tz> to clear the various fields because need to clear per timezone, |
| // and NaiveDateTime (ISO 8601) has no concept of timezones |
| let value = as_datetime_with_timezone::<TimestampNanosecondType>(value, tz) |
| .ok_or(exec_datafusion_err!("Timestamp {value} out of range"))?; |
| _date_trunc_coarse_with_tz(granularity, Some(value)) |
| } |
| None => { |
| // Use chrono NaiveDateTime to clear the various fields, if we don't have a timezone. |
| let value = timestamp_ns_to_datetime(value) |
| .ok_or_else(|| exec_datafusion_err!("Timestamp {value} out of range"))?; |
| _date_trunc_coarse_without_tz(granularity, Some(value)) |
| } |
| }?; |
| |
| // `with_x(0)` are infallible because `0` are always a valid |
| Ok(value.unwrap()) |
| } |
| |
| /// Fast path for fine granularities (hour and smaller) that can be handled |
| /// with simple arithmetic operations without calendar complexity. |
| /// |
| /// This function is timezone-agnostic and should only be used when: |
| /// - No timezone is specified in the input, OR |
| /// - The granularity is less than hour as hour can be affected by DST transitions in some cases |
| fn general_date_trunc_array_fine_granularity<T: ArrowTimestampType>( |
| tu: TimeUnit, |
| array: &PrimitiveArray<T>, |
| granularity: DateTruncGranularity, |
| ) -> Result<ArrayRef> { |
| let unit = match (tu, granularity) { |
| (Second, DateTruncGranularity::Minute) => NonZeroI64::new(60), |
| (Second, DateTruncGranularity::Hour) => NonZeroI64::new(3600), |
| (Second, DateTruncGranularity::Day) => NonZeroI64::new(86400), |
| |
| (Millisecond, DateTruncGranularity::Second) => NonZeroI64::new(1_000), |
| (Millisecond, DateTruncGranularity::Minute) => NonZeroI64::new(60_000), |
| (Millisecond, DateTruncGranularity::Hour) => NonZeroI64::new(3_600_000), |
| (Millisecond, DateTruncGranularity::Day) => NonZeroI64::new(86_400_000), |
| |
| (Microsecond, DateTruncGranularity::Millisecond) => NonZeroI64::new(1_000), |
| (Microsecond, DateTruncGranularity::Second) => NonZeroI64::new(1_000_000), |
| (Microsecond, DateTruncGranularity::Minute) => NonZeroI64::new(60_000_000), |
| (Microsecond, DateTruncGranularity::Hour) => NonZeroI64::new(3_600_000_000), |
| (Microsecond, DateTruncGranularity::Day) => NonZeroI64::new(86_400_000_000), |
| |
| (Nanosecond, DateTruncGranularity::Microsecond) => NonZeroI64::new(1_000), |
| (Nanosecond, DateTruncGranularity::Millisecond) => NonZeroI64::new(1_000_000), |
| (Nanosecond, DateTruncGranularity::Second) => NonZeroI64::new(1_000_000_000), |
| (Nanosecond, DateTruncGranularity::Minute) => NonZeroI64::new(60_000_000_000), |
| (Nanosecond, DateTruncGranularity::Hour) => NonZeroI64::new(3_600_000_000_000), |
| (Nanosecond, DateTruncGranularity::Day) => NonZeroI64::new(86_400_000_000_000), |
| _ => None, |
| }; |
| |
| if let Some(unit) = unit { |
| let unit = unit.get(); |
| let array = PrimitiveArray::<T>::from_iter_values_with_nulls( |
| array |
| .values() |
| .iter() |
| .map(|v| *v - i64::rem_euclid(*v, unit)), |
| array.nulls().cloned(), |
| ); |
| Ok(Arc::new(array)) |
| } else { |
| // truncate to the same or smaller unit |
| Ok(Arc::new(array.clone())) |
| } |
| } |
| |
| // truncates a single value with the given timeunit to the specified granularity |
| fn general_date_trunc( |
| tu: TimeUnit, |
| value: i64, |
| tz: Option<Tz>, |
| granularity: DateTruncGranularity, |
| ) -> Result<i64, DataFusionError> { |
| let scale = match tu { |
| Second => 1_000_000_000, |
| Millisecond => 1_000_000, |
| Microsecond => 1_000, |
| Nanosecond => 1, |
| }; |
| |
| // convert to nanoseconds |
| let nano = date_trunc_coarse(granularity, scale * value, tz)?; |
| |
| let result = match tu { |
| Second => match granularity { |
| DateTruncGranularity::Minute => nano / 1_000_000_000 / 60 * 60, |
| _ => nano / 1_000_000_000, |
| }, |
| Millisecond => match granularity { |
| DateTruncGranularity::Minute => nano / 1_000_000 / 1_000 / 60 * 1_000 * 60, |
| DateTruncGranularity::Second => nano / 1_000_000 / 1_000 * 1_000, |
| _ => nano / 1_000_000, |
| }, |
| Microsecond => match granularity { |
| DateTruncGranularity::Minute => { |
| nano / 1_000 / 1_000_000 / 60 * 60 * 1_000_000 |
| } |
| DateTruncGranularity::Second => nano / 1_000 / 1_000_000 * 1_000_000, |
| DateTruncGranularity::Millisecond => nano / 1_000 / 1_000 * 1_000, |
| _ => nano / 1_000, |
| }, |
| _ => match granularity { |
| DateTruncGranularity::Minute => { |
| nano / 1_000_000_000 / 60 * 1_000_000_000 * 60 |
| } |
| DateTruncGranularity::Second => nano / 1_000_000_000 * 1_000_000_000, |
| DateTruncGranularity::Millisecond => nano / 1_000_000 * 1_000_000, |
| DateTruncGranularity::Microsecond => nano / 1_000 * 1_000, |
| _ => nano, |
| }, |
| }; |
| Ok(result) |
| } |
| |
| fn parse_tz(tz: &Option<Arc<str>>) -> Result<Option<Tz>> { |
| tz.as_ref() |
| .map(|tz| { |
| Tz::from_str(tz) |
| .map_err(|op| exec_datafusion_err!("failed on timezone {tz}: {op:?}")) |
| }) |
| .transpose() |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use std::sync::Arc; |
| |
| use crate::datetime::date_trunc::{ |
| date_trunc_coarse, DateTruncFunc, DateTruncGranularity, |
| }; |
| |
| use arrow::array::cast::as_primitive_array; |
| use arrow::array::types::TimestampNanosecondType; |
| use arrow::array::{Array, TimestampNanosecondArray}; |
| use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos; |
| use arrow::datatypes::{DataType, Field, TimeUnit}; |
| use datafusion_common::config::ConfigOptions; |
| use datafusion_common::ScalarValue; |
| use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; |
| |
| #[test] |
| fn date_trunc_test() { |
| let cases = vec![ |
| ( |
| "2020-09-08T13:42:29.190855Z", |
| "second", |
| "2020-09-08T13:42:29.000000Z", |
| ), |
| ( |
| "2020-09-08T13:42:29.190855Z", |
| "minute", |
| "2020-09-08T13:42:00.000000Z", |
| ), |
| ( |
| "2020-09-08T13:42:29.190855Z", |
| "hour", |
| "2020-09-08T13:00:00.000000Z", |
| ), |
| ( |
| "2020-09-08T13:42:29.190855Z", |
| "day", |
| "2020-09-08T00:00:00.000000Z", |
| ), |
| ( |
| "2020-09-08T13:42:29.190855Z", |
| "week", |
| "2020-09-07T00:00:00.000000Z", |
| ), |
| ( |
| "2020-09-08T13:42:29.190855Z", |
| "month", |
| "2020-09-01T00:00:00.000000Z", |
| ), |
| ( |
| "2020-09-08T13:42:29.190855Z", |
| "year", |
| "2020-01-01T00:00:00.000000Z", |
| ), |
| // week |
| ( |
| "2021-01-01T13:42:29.190855Z", |
| "week", |
| "2020-12-28T00:00:00.000000Z", |
| ), |
| ( |
| "2020-01-01T13:42:29.190855Z", |
| "week", |
| "2019-12-30T00:00:00.000000Z", |
| ), |
| // quarter |
| ( |
| "2020-01-01T13:42:29.190855Z", |
| "quarter", |
| "2020-01-01T00:00:00.000000Z", |
| ), |
| ( |
| "2020-02-01T13:42:29.190855Z", |
| "quarter", |
| "2020-01-01T00:00:00.000000Z", |
| ), |
| ( |
| "2020-03-01T13:42:29.190855Z", |
| "quarter", |
| "2020-01-01T00:00:00.000000Z", |
| ), |
| ( |
| "2020-04-01T13:42:29.190855Z", |
| "quarter", |
| "2020-04-01T00:00:00.000000Z", |
| ), |
| ( |
| "2020-08-01T13:42:29.190855Z", |
| "quarter", |
| "2020-07-01T00:00:00.000000Z", |
| ), |
| ( |
| "2020-11-01T13:42:29.190855Z", |
| "quarter", |
| "2020-10-01T00:00:00.000000Z", |
| ), |
| ( |
| "2020-12-01T13:42:29.190855Z", |
| "quarter", |
| "2020-10-01T00:00:00.000000Z", |
| ), |
| ]; |
| |
| cases.iter().for_each(|(original, granularity, expected)| { |
| let left = string_to_timestamp_nanos(original).unwrap(); |
| let right = string_to_timestamp_nanos(expected).unwrap(); |
| let granularity_enum = DateTruncGranularity::from_str(granularity).unwrap(); |
| let result = date_trunc_coarse(granularity_enum, left, None).unwrap(); |
| assert_eq!(result, right, "{original} = {expected}"); |
| }); |
| } |
| |
| #[test] |
| fn test_date_trunc_timezones() { |
| let cases = [ |
| ( |
| vec![ |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T01:00:00Z", |
| "2020-09-08T02:00:00Z", |
| "2020-09-08T03:00:00Z", |
| "2020-09-08T04:00:00Z", |
| ], |
| Some("+00".into()), |
| vec![ |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T00:00:00Z", |
| ], |
| ), |
| ( |
| vec![ |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T01:00:00Z", |
| "2020-09-08T02:00:00Z", |
| "2020-09-08T03:00:00Z", |
| "2020-09-08T04:00:00Z", |
| ], |
| None, |
| vec![ |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T00:00:00Z", |
| ], |
| ), |
| ( |
| vec![ |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T01:00:00Z", |
| "2020-09-08T02:00:00Z", |
| "2020-09-08T03:00:00Z", |
| "2020-09-08T04:00:00Z", |
| ], |
| Some("-02".into()), |
| vec![ |
| "2020-09-07T02:00:00Z", |
| "2020-09-07T02:00:00Z", |
| "2020-09-08T02:00:00Z", |
| "2020-09-08T02:00:00Z", |
| "2020-09-08T02:00:00Z", |
| ], |
| ), |
| ( |
| vec![ |
| "2020-09-08T00:00:00+05", |
| "2020-09-08T01:00:00+05", |
| "2020-09-08T02:00:00+05", |
| "2020-09-08T03:00:00+05", |
| "2020-09-08T04:00:00+05", |
| ], |
| Some("+05".into()), |
| vec![ |
| "2020-09-08T00:00:00+05", |
| "2020-09-08T00:00:00+05", |
| "2020-09-08T00:00:00+05", |
| "2020-09-08T00:00:00+05", |
| "2020-09-08T00:00:00+05", |
| ], |
| ), |
| ( |
| vec![ |
| "2020-09-08T00:00:00+08", |
| "2020-09-08T01:00:00+08", |
| "2020-09-08T02:00:00+08", |
| "2020-09-08T03:00:00+08", |
| "2020-09-08T04:00:00+08", |
| ], |
| Some("+08".into()), |
| vec![ |
| "2020-09-08T00:00:00+08", |
| "2020-09-08T00:00:00+08", |
| "2020-09-08T00:00:00+08", |
| "2020-09-08T00:00:00+08", |
| "2020-09-08T00:00:00+08", |
| ], |
| ), |
| ( |
| vec![ |
| "2024-10-26T23:00:00Z", |
| "2024-10-27T00:00:00Z", |
| "2024-10-27T01:00:00Z", |
| "2024-10-27T02:00:00Z", |
| ], |
| Some("Europe/Berlin".into()), |
| vec![ |
| "2024-10-27T00:00:00+02", |
| "2024-10-27T00:00:00+02", |
| "2024-10-27T00:00:00+02", |
| "2024-10-27T00:00:00+02", |
| ], |
| ), |
| ( |
| vec![ |
| "2018-02-18T00:00:00Z", |
| "2018-02-18T01:00:00Z", |
| "2018-02-18T02:00:00Z", |
| "2018-02-18T03:00:00Z", |
| "2018-11-04T01:00:00Z", |
| "2018-11-04T02:00:00Z", |
| "2018-11-04T03:00:00Z", |
| "2018-11-04T04:00:00Z", |
| ], |
| Some("America/Sao_Paulo".into()), |
| vec![ |
| "2018-02-17T00:00:00-02", |
| "2018-02-17T00:00:00-02", |
| "2018-02-17T00:00:00-02", |
| "2018-02-18T00:00:00-03", |
| "2018-11-03T00:00:00-03", |
| "2018-11-03T00:00:00-03", |
| "2018-11-04T01:00:00-02", |
| "2018-11-04T01:00:00-02", |
| ], |
| ), |
| ]; |
| |
| cases.iter().for_each(|(original, tz_opt, expected)| { |
| let input = original |
| .iter() |
| .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) |
| .collect::<TimestampNanosecondArray>() |
| .with_timezone_opt(tz_opt.clone()); |
| let right = expected |
| .iter() |
| .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) |
| .collect::<TimestampNanosecondArray>() |
| .with_timezone_opt(tz_opt.clone()); |
| let batch_len = input.len(); |
| let arg_fields = vec![ |
| Field::new("a", DataType::Utf8, false).into(), |
| Field::new("b", input.data_type().clone(), false).into(), |
| ]; |
| let args = datafusion_expr::ScalarFunctionArgs { |
| args: vec![ |
| ColumnarValue::Scalar(ScalarValue::from("day")), |
| ColumnarValue::Array(Arc::new(input)), |
| ], |
| arg_fields, |
| number_rows: batch_len, |
| return_field: Field::new( |
| "f", |
| DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone()), |
| true, |
| ) |
| .into(), |
| config_options: Arc::new(ConfigOptions::default()), |
| }; |
| let result = DateTruncFunc::new().invoke_with_args(args).unwrap(); |
| if let ColumnarValue::Array(result) = result { |
| assert_eq!( |
| result.data_type(), |
| &DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone()) |
| ); |
| let left = as_primitive_array::<TimestampNanosecondType>(&result); |
| assert_eq!(left, &right); |
| } else { |
| panic!("unexpected column type"); |
| } |
| }); |
| } |
| |
| #[test] |
| fn test_date_trunc_hour_timezones() { |
| let cases = [ |
| ( |
| vec![ |
| "2020-09-08T00:30:00Z", |
| "2020-09-08T01:30:00Z", |
| "2020-09-08T02:30:00Z", |
| "2020-09-08T03:30:00Z", |
| "2020-09-08T04:30:00Z", |
| ], |
| Some("+00".into()), |
| vec![ |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T01:00:00Z", |
| "2020-09-08T02:00:00Z", |
| "2020-09-08T03:00:00Z", |
| "2020-09-08T04:00:00Z", |
| ], |
| ), |
| ( |
| vec![ |
| "2020-09-08T00:30:00Z", |
| "2020-09-08T01:30:00Z", |
| "2020-09-08T02:30:00Z", |
| "2020-09-08T03:30:00Z", |
| "2020-09-08T04:30:00Z", |
| ], |
| None, |
| vec![ |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T01:00:00Z", |
| "2020-09-08T02:00:00Z", |
| "2020-09-08T03:00:00Z", |
| "2020-09-08T04:00:00Z", |
| ], |
| ), |
| ( |
| vec![ |
| "2020-09-08T00:30:00Z", |
| "2020-09-08T01:30:00Z", |
| "2020-09-08T02:30:00Z", |
| "2020-09-08T03:30:00Z", |
| "2020-09-08T04:30:00Z", |
| ], |
| Some("-02".into()), |
| vec![ |
| "2020-09-08T00:00:00Z", |
| "2020-09-08T01:00:00Z", |
| "2020-09-08T02:00:00Z", |
| "2020-09-08T03:00:00Z", |
| "2020-09-08T04:00:00Z", |
| ], |
| ), |
| ( |
| vec![ |
| "2020-09-08T00:30:00+05", |
| "2020-09-08T01:30:00+05", |
| "2020-09-08T02:30:00+05", |
| "2020-09-08T03:30:00+05", |
| "2020-09-08T04:30:00+05", |
| ], |
| Some("+05".into()), |
| vec![ |
| "2020-09-08T00:00:00+05", |
| "2020-09-08T01:00:00+05", |
| "2020-09-08T02:00:00+05", |
| "2020-09-08T03:00:00+05", |
| "2020-09-08T04:00:00+05", |
| ], |
| ), |
| ( |
| vec![ |
| "2020-09-08T00:30:00+08", |
| "2020-09-08T01:30:00+08", |
| "2020-09-08T02:30:00+08", |
| "2020-09-08T03:30:00+08", |
| "2020-09-08T04:30:00+08", |
| ], |
| Some("+08".into()), |
| vec![ |
| "2020-09-08T00:00:00+08", |
| "2020-09-08T01:00:00+08", |
| "2020-09-08T02:00:00+08", |
| "2020-09-08T03:00:00+08", |
| "2020-09-08T04:00:00+08", |
| ], |
| ), |
| ( |
| vec![ |
| "2024-10-26T23:30:00Z", |
| "2024-10-27T00:30:00Z", |
| "2024-10-27T01:30:00Z", |
| "2024-10-27T02:30:00Z", |
| ], |
| Some("Europe/Berlin".into()), |
| vec![ |
| "2024-10-27T01:00:00+02", |
| "2024-10-27T02:00:00+02", |
| "2024-10-27T02:00:00+01", |
| "2024-10-27T03:00:00+01", |
| ], |
| ), |
| ( |
| vec![ |
| "2018-02-18T00:30:00Z", |
| "2018-02-18T01:30:00Z", |
| "2018-02-18T02:30:00Z", |
| "2018-02-18T03:30:00Z", |
| "2018-11-04T01:00:00Z", |
| "2018-11-04T02:00:00Z", |
| "2018-11-04T03:00:00Z", |
| "2018-11-04T04:00:00Z", |
| ], |
| Some("America/Sao_Paulo".into()), |
| vec![ |
| "2018-02-17T22:00:00-02", |
| "2018-02-17T23:00:00-02", |
| "2018-02-17T23:00:00-03", |
| "2018-02-18T00:00:00-03", |
| "2018-11-03T22:00:00-03", |
| "2018-11-03T23:00:00-03", |
| "2018-11-04T01:00:00-02", |
| "2018-11-04T02:00:00-02", |
| ], |
| ), |
| ( |
| vec![ |
| "2024-10-26T23:30:00Z", |
| "2024-10-27T00:30:00Z", |
| "2024-10-27T01:30:00Z", |
| "2024-10-27T02:30:00Z", |
| ], |
| Some("Asia/Kathmandu".into()), // UTC+5:45 |
| vec![ |
| "2024-10-27T05:00:00+05:45", |
| "2024-10-27T06:00:00+05:45", |
| "2024-10-27T07:00:00+05:45", |
| "2024-10-27T08:00:00+05:45", |
| ], |
| ), |
| ]; |
| |
| cases.iter().for_each(|(original, tz_opt, expected)| { |
| let input = original |
| .iter() |
| .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) |
| .collect::<TimestampNanosecondArray>() |
| .with_timezone_opt(tz_opt.clone()); |
| let right = expected |
| .iter() |
| .map(|s| Some(string_to_timestamp_nanos(s).unwrap())) |
| .collect::<TimestampNanosecondArray>() |
| .with_timezone_opt(tz_opt.clone()); |
| let batch_len = input.len(); |
| let arg_fields = vec![ |
| Field::new("a", DataType::Utf8, false).into(), |
| Field::new("b", input.data_type().clone(), false).into(), |
| ]; |
| let args = datafusion_expr::ScalarFunctionArgs { |
| args: vec![ |
| ColumnarValue::Scalar(ScalarValue::from("hour")), |
| ColumnarValue::Array(Arc::new(input)), |
| ], |
| arg_fields, |
| number_rows: batch_len, |
| return_field: Field::new( |
| "f", |
| DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone()), |
| true, |
| ) |
| .into(), |
| config_options: Arc::new(ConfigOptions::default()), |
| }; |
| let result = DateTruncFunc::new().invoke_with_args(args).unwrap(); |
| if let ColumnarValue::Array(result) = result { |
| assert_eq!( |
| result.data_type(), |
| &DataType::Timestamp(TimeUnit::Nanosecond, tz_opt.clone()) |
| ); |
| let left = as_primitive_array::<TimestampNanosecondType>(&result); |
| assert_eq!(left, &right); |
| } else { |
| panic!("unexpected column type"); |
| } |
| }); |
| } |
| } |
