diff --git a/datafusion/functions/src/datetime/date_part.rs b/datafusion/functions/src/datetime/date_part.rs index e3f67db905615..4bfadf7999906 100644 --- a/datafusion/functions/src/datetime/date_part.rs +++ b/datafusion/functions/src/datetime/date_part.rs @@ -21,7 +21,6 @@ use std::sync::Arc; use arrow::array::timezone::Tz; use arrow::array::{Array, ArrayRef, Float64Array, Int32Array, Int64Array}; -use arrow::compute::kernels::cast_utils::IntervalUnit; use arrow::compute::{DatePart, binary, date_part}; use arrow::datatypes::DataType::{ Date32, Date64, Duration, Interval, Time32, Time64, Timestamp, @@ -216,37 +215,14 @@ impl ScalarUDFImpl for DatePartFunc { let part_trim = part_normalization(&part); - // 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 { - 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)?, - IntervalUnit::Day => date_part(array.as_ref(), DatePart::Day)?, - IntervalUnit::Hour => date_part(array.as_ref(), DatePart::Hour)?, - IntervalUnit::Minute => date_part(array.as_ref(), DatePart::Minute)?, - IntervalUnit::Second => seconds_as_i32(array.as_ref(), Second)?, - IntervalUnit::Millisecond => seconds_as_i32(array.as_ref(), Millisecond)?, - IntervalUnit::Microsecond => seconds_as_i32(array.as_ref(), Microsecond)?, - IntervalUnit::Nanosecond => seconds_ns(array.as_ref())?, - // century and decade are not supported by `DatePart`, although they are supported in postgres - _ => return exec_err!("Date part '{part}' not supported"), - } - } else { - // special cases that can be extracted (in postgres) but are not interval units - match part_trim.to_lowercase().as_str() { - "isoyear" => date_part(array.as_ref(), DatePart::YearISO)?, - "qtr" | "quarter" => date_part(array.as_ref(), DatePart::Quarter)?, - "doy" => date_part(array.as_ref(), DatePart::DayOfYear)?, - "dow" => date_part(array.as_ref(), DatePart::DayOfWeekSunday0)?, - "isodow" => { - // Postgres `isodow` is 1..=7 with Mon=1 - date_part(array.as_ref(), DatePart::DayOfWeekMonday1)? - } - "epoch" => epoch(array.as_ref())?, - _ => return exec_err!("Date part '{part}' not supported"), - } + let arr = match DatePart::from_str(part_trim) { + Ok(DatePart::Second) => seconds_as_i32(array.as_ref(), Second)?, + Ok(DatePart::Millisecond) => seconds_as_i32(array.as_ref(), Millisecond)?, + Ok(DatePart::Microsecond) => seconds_as_i32(array.as_ref(), Microsecond)?, + Ok(DatePart::Nanosecond) => seconds_ns(array.as_ref())?, + Ok(part) => date_part(array.as_ref(), part)?, + Err(_) if is_epoch(part_trim) => epoch(array.as_ref())?, + Err(_) => return exec_err!("Date part '{part}' not supported"), }; Ok(if is_scalar { @@ -256,7 +232,7 @@ impl ScalarUDFImpl for DatePartFunc { }) } - // Only casting the year is supported since pruning other IntervalUnit is not possible + // Only casting the year is supported since pruning other date parts is not possible // date_part(col, YEAR) = 2024 => col >= '2024-01-01' and col < '2025-01-01' // But for anything less than YEAR simplifying is not possible without specifying the bigger interval // date_part(col, MONTH) = 1 => col = '2023-01-01' or col = '2024-01-01' or ... or col = '3000-01-01' @@ -268,16 +244,16 @@ impl ScalarUDFImpl for DatePartFunc { ) -> Result { let [part, col_expr] = take_function_args(self.name(), args)?; - // Get the interval unit from the part argument - let interval_unit = part + // Get the date part from the part argument + let date_part = part .as_literal() .and_then(|sv| sv.try_as_str().flatten()) .map(part_normalization) - .and_then(|s| IntervalUnit::from_str(s).ok()); + .and_then(|s| DatePart::from_str(s).ok()); // only support extracting year - match interval_unit { - Some(IntervalUnit::Year) => (), + match date_part { + Some(DatePart::Year) => (), _ => return Ok(PreimageResult::None), } @@ -336,8 +312,8 @@ fn is_epoch(part: &str) -> bool { } fn is_nanosecond(part: &str) -> bool { - IntervalUnit::from_str(part_normalization(part)) - .map(|p| matches!(p, IntervalUnit::Nanosecond)) + DatePart::from_str(part_normalization(part)) + .map(|p| matches!(p, DatePart::Nanosecond)) .unwrap_or(false) } diff --git a/datafusion/functions/src/datetime/date_trunc.rs b/datafusion/functions/src/datetime/date_trunc.rs index 926653c0d1a01..22c8af262a4df 100644 --- a/datafusion/functions/src/datetime/date_trunc.rs +++ b/datafusion/functions/src/datetime/date_trunc.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::fmt; use std::num::NonZeroI64; use std::ops::{Add, Sub}; use std::str::FromStr; @@ -31,6 +30,7 @@ use arrow::array::types::{ TimestampNanosecondType, TimestampSecondType, }; use arrow::array::{Array, ArrayRef, PrimitiveArray}; +use arrow::compute::DatePart; use arrow::datatypes::DataType::{self, Time32, Time64, Timestamp}; use arrow::datatypes::TimeUnit::{self, Microsecond, Millisecond, Nanosecond, Second}; use arrow::datatypes::{Field, FieldRef}; @@ -51,74 +51,58 @@ 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 { - // 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}" - ) - } +fn parse_granularity(value: &str) -> Result { + // DatePart also contains extraction-only fields such as day-of-week, which + // are not valid truncation granularities. + match DatePart::from_str(value) { + Ok(granularity) + if matches!( + granularity, + DatePart::Microsecond + | DatePart::Millisecond + | DatePart::Second + | DatePart::Minute + | DatePart::Hour + | DatePart::Day + | DatePart::Week + | DatePart::Month + | DatePart::Quarter + | DatePart::Year + ) => + { + Ok(granularity) + } + _ => { + exec_err!( + "Unsupported date_trunc granularity: '{value}'. Supported granularities are: microsecond, millisecond, second, minute, hour, day, week, month, quarter, year" + ) } } +} + +trait DateTruncGranularityExt { + fn is_fine_granularity(&self) -> bool; + fn is_fine_granularity_utc(&self) -> bool; + fn valid_for_time(&self) -> bool; +} +impl DateTruncGranularityExt for DatePart { /// 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 + DatePart::Second + | DatePart::Minute + | DatePart::Millisecond + | DatePart::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) + self.is_fine_granularity() || matches!(self, DatePart::Hour | DatePart::Day) } /// Returns true if this granularity is valid for Time types @@ -126,33 +110,15 @@ impl DateTruncGranularity { fn valid_for_time(&self) -> bool { matches!( self, - Self::Hour - | Self::Minute - | Self::Second - | Self::Millisecond - | Self::Microsecond + DatePart::Hour + | DatePart::Minute + | DatePart::Second + | DatePart::Millisecond + | DatePart::Microsecond ) } } -impl fmt::Display for DateTruncGranularity { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let value = match self { - Self::Microsecond => "microsecond", - Self::Millisecond => "millisecond", - Self::Second => "second", - Self::Minute => "minute", - Self::Hour => "hour", - Self::Day => "day", - Self::Week => "week", - Self::Month => "month", - Self::Quarter => "quarter", - Self::Year => "year", - }; - f.write_str(value) - } -} - #[user_doc( doc_section(label = "Time and Date Functions"), description = "Truncates a timestamp or time value to a specified precision.", @@ -286,7 +252,7 @@ impl ScalarUDFImpl for DateTruncFunc { return exec_err!("Granularity of `date_trunc` must be non-null scalar Utf8"); }; - let granularity = DateTruncGranularity::from_str(&granularity_str)?; + let granularity = parse_granularity(&granularity_str)?; // Check upfront if granularity is valid for Time types let is_time_type = matches!(array.data_type(), Time64(_) | Time32(_)); @@ -299,7 +265,7 @@ impl ScalarUDFImpl for DateTruncFunc { fn process_array( array: &dyn Array, - granularity: DateTruncGranularity, + granularity: DatePart, tz_opt: Option<&Arc>, ) -> Result { let parsed_tz = parse_tz(tz_opt)?; @@ -330,7 +296,7 @@ impl ScalarUDFImpl for DateTruncFunc { fn process_scalar( v: Option<&i64>, - granularity: DateTruncGranularity, + granularity: DatePart, tz_opt: Option<&Arc>, ) -> Result { let parsed_tz = parse_tz(tz_opt)?; @@ -500,81 +466,78 @@ const SECS_PER_MINUTE: i32 = 60; const SECS_PER_HOUR: i32 = 60 * SECS_PER_MINUTE; /// Truncate time in nanoseconds to the specified granularity -fn truncate_time_nanos(value: i64, granularity: DateTruncGranularity) -> i64 { +fn truncate_time_nanos(value: i64, granularity: DatePart) -> i64 { match granularity { - DateTruncGranularity::Hour => value - (value % NANOS_PER_HOUR), - DateTruncGranularity::Minute => value - (value % NANOS_PER_MINUTE), - DateTruncGranularity::Second => value - (value % NANOS_PER_SECOND), - DateTruncGranularity::Millisecond => value - (value % NANOS_PER_MILLISECOND), - DateTruncGranularity::Microsecond => value - (value % NANOS_PER_MICROSECOND), + DatePart::Hour => value - (value % NANOS_PER_HOUR), + DatePart::Minute => value - (value % NANOS_PER_MINUTE), + DatePart::Second => value - (value % NANOS_PER_SECOND), + DatePart::Millisecond => value - (value % NANOS_PER_MILLISECOND), + DatePart::Microsecond => value - (value % NANOS_PER_MICROSECOND), // Other granularities are not valid for time - should be caught earlier _ => value, } } /// Truncate time in microseconds to the specified granularity -fn truncate_time_micros(value: i64, granularity: DateTruncGranularity) -> i64 { +fn truncate_time_micros(value: i64, granularity: DatePart) -> i64 { match granularity { - DateTruncGranularity::Hour => value - (value % MICROS_PER_HOUR), - DateTruncGranularity::Minute => value - (value % MICROS_PER_MINUTE), - DateTruncGranularity::Second => value - (value % MICROS_PER_SECOND), - DateTruncGranularity::Millisecond => value - (value % MICROS_PER_MILLISECOND), - DateTruncGranularity::Microsecond => value, // Already at microsecond precision + DatePart::Hour => value - (value % MICROS_PER_HOUR), + DatePart::Minute => value - (value % MICROS_PER_MINUTE), + DatePart::Second => value - (value % MICROS_PER_SECOND), + DatePart::Millisecond => value - (value % MICROS_PER_MILLISECOND), + DatePart::Microsecond => value, // Already at microsecond precision // Other granularities are not valid for time _ => value, } } /// Truncate time in milliseconds to the specified granularity -fn truncate_time_millis(value: i32, granularity: DateTruncGranularity) -> i32 { +fn truncate_time_millis(value: i32, granularity: DatePart) -> i32 { match granularity { - DateTruncGranularity::Hour => value - (value % MILLIS_PER_HOUR), - DateTruncGranularity::Minute => value - (value % MILLIS_PER_MINUTE), - DateTruncGranularity::Second => value - (value % MILLIS_PER_SECOND), - DateTruncGranularity::Millisecond => value, // Already at millisecond precision - DateTruncGranularity::Microsecond => value, // Can't truncate to finer precision + DatePart::Hour => value - (value % MILLIS_PER_HOUR), + DatePart::Minute => value - (value % MILLIS_PER_MINUTE), + DatePart::Second => value - (value % MILLIS_PER_SECOND), + DatePart::Millisecond => value, // Already at millisecond precision + DatePart::Microsecond => value, // Can't truncate to finer precision // Other granularities are not valid for time _ => value, } } /// Truncate time in seconds to the specified granularity -fn truncate_time_secs(value: i32, granularity: DateTruncGranularity) -> i32 { +fn truncate_time_secs(value: i32, granularity: DatePart) -> i32 { match granularity { - DateTruncGranularity::Hour => value - (value % SECS_PER_HOUR), - DateTruncGranularity::Minute => value - (value % SECS_PER_MINUTE), - DateTruncGranularity::Second => value, // Already at second precision - DateTruncGranularity::Millisecond => value, // Can't truncate to finer precision - DateTruncGranularity::Microsecond => value, // Can't truncate to finer precision + DatePart::Hour => value - (value % SECS_PER_HOUR), + DatePart::Minute => value - (value % SECS_PER_MINUTE), + DatePart::Second => value, // Already at second precision + DatePart::Millisecond => value, // Can't truncate to finer precision + DatePart::Microsecond => value, // Can't truncate to finer precision // Other granularities are not valid for time _ => value, } } -fn _date_trunc_coarse( - granularity: DateTruncGranularity, - value: Option, -) -> Result> +fn _date_trunc_coarse(granularity: DatePart, value: Option) -> Result> where T: Datelike + Timelike + Sub + Copy, { let value = match granularity { - DateTruncGranularity::Millisecond => value, - DateTruncGranularity::Microsecond => value, - DateTruncGranularity::Second => value.and_then(|d| d.with_nanosecond(0)), - DateTruncGranularity::Minute => value + DatePart::Millisecond => value, + DatePart::Microsecond => value, + DatePart::Second => value.and_then(|d| d.with_nanosecond(0)), + DatePart::Minute => value .and_then(|d| d.with_nanosecond(0)) .and_then(|d| d.with_second(0)), - DateTruncGranularity::Hour => value + DatePart::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 + DatePart::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 + DatePart::Week => value .and_then(|d| d.with_nanosecond(0)) .and_then(|d| d.with_second(0)) .and_then(|d| d.with_minute(0)) @@ -582,26 +545,27 @@ where .map(|d| { d - TimeDelta::try_seconds(60 * 60 * 24 * d.weekday() as i64).unwrap() }), - DateTruncGranularity::Month => value + DatePart::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 + DatePart::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 + DatePart::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)), + _ => unreachable!("unsupported date_trunc granularity"), }; Ok(value) } @@ -614,7 +578,7 @@ where } fn _date_trunc_coarse_with_tz( - granularity: DateTruncGranularity, + granularity: DatePart, value: DateTime, ) -> Result> { let local = value.naive_local(); @@ -715,10 +679,7 @@ fn days_from_civil(year: i64, month: i64, day: i64) -> i64 { /// Returns `None` when the truncated timestamp is no longer representable as /// nanoseconds since the epoch, which the caller reports as an out of range /// error. -fn _date_trunc_coarse_without_tz( - granularity: DateTruncGranularity, - value: i64, -) -> Option { +fn _date_trunc_coarse_without_tz(granularity: DatePart, value: i64) -> Option { let truncate_to = |unit: i64| value.checked_sub(value.rem_euclid(unit)); let days = || value.div_euclid(NANOS_PER_DAY); let nanos_from_days = |days: i64| days.checked_mul(NANOS_PER_DAY); @@ -726,31 +687,30 @@ fn _date_trunc_coarse_without_tz( match granularity { // Sub-second granularities are applied by the caller, which rescales // the nanoseconds to the time unit of the array. - DateTruncGranularity::Millisecond | DateTruncGranularity::Microsecond => { - Some(value) - } - DateTruncGranularity::Second => truncate_to(NANOS_PER_SECOND), - DateTruncGranularity::Minute => truncate_to(NANOS_PER_MINUTE), - DateTruncGranularity::Hour => truncate_to(NANOS_PER_HOUR), - DateTruncGranularity::Day => nanos_from_days(days()), - DateTruncGranularity::Week => { + DatePart::Millisecond | DatePart::Microsecond => Some(value), + DatePart::Second => truncate_to(NANOS_PER_SECOND), + DatePart::Minute => truncate_to(NANOS_PER_MINUTE), + DatePart::Hour => truncate_to(NANOS_PER_HOUR), + DatePart::Day => nanos_from_days(days()), + DatePart::Week => { let days = days(); // `Weekday::num_days_from_monday` for the epoch (a Thursday) is 3. nanos_from_days(days - (days + 3).rem_euclid(7)) } - DateTruncGranularity::Month => { + DatePart::Month => { let days = days(); let (_, _, day_of_month) = civil_from_days(days); nanos_from_days(days - (day_of_month - 1)) } - DateTruncGranularity::Quarter => { + DatePart::Quarter => { let (year, month, _) = civil_from_days(days()); nanos_from_days(days_from_civil(year, 1 + 3 * ((month - 1) / 3), 1)) } - DateTruncGranularity::Year => { + DatePart::Year => { let (year, _, _) = civil_from_days(days()); nanos_from_days(days_from_civil(year, 1, 1)) } + _ => unreachable!("unsupported date_trunc granularity"), } } @@ -758,11 +718,7 @@ fn _date_trunc_coarse_without_tz( /// 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, -) -> Result { +fn date_trunc_coarse(granularity: DatePart, value: i64, tz: Option) -> Result { let input = value; let value = match tz { Some(tz) => { @@ -776,6 +732,7 @@ fn date_trunc_coarse( }; value.ok_or_else(|| { + let granularity = granularity.to_string().to_lowercase(); exec_datafusion_err!( "Timestamp {input} out of range after truncating to {granularity}" ) @@ -791,31 +748,31 @@ fn date_trunc_coarse( fn general_date_trunc_array_fine_granularity( tu: TimeUnit, array: &PrimitiveArray, - granularity: DateTruncGranularity, + granularity: DatePart, tz_opt: Option>, ) -> Result { 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), + (Second, DatePart::Minute) => NonZeroI64::new(60), + (Second, DatePart::Hour) => NonZeroI64::new(3600), + (Second, DatePart::Day) => NonZeroI64::new(86400), + + (Millisecond, DatePart::Second) => NonZeroI64::new(1_000), + (Millisecond, DatePart::Minute) => NonZeroI64::new(60_000), + (Millisecond, DatePart::Hour) => NonZeroI64::new(3_600_000), + (Millisecond, DatePart::Day) => NonZeroI64::new(86_400_000), + + (Microsecond, DatePart::Millisecond) => NonZeroI64::new(1_000), + (Microsecond, DatePart::Second) => NonZeroI64::new(1_000_000), + (Microsecond, DatePart::Minute) => NonZeroI64::new(60_000_000), + (Microsecond, DatePart::Hour) => NonZeroI64::new(3_600_000_000), + (Microsecond, DatePart::Day) => NonZeroI64::new(86_400_000_000), + + (Nanosecond, DatePart::Microsecond) => NonZeroI64::new(1_000), + (Nanosecond, DatePart::Millisecond) => NonZeroI64::new(1_000_000), + (Nanosecond, DatePart::Second) => NonZeroI64::new(1_000_000_000), + (Nanosecond, DatePart::Minute) => NonZeroI64::new(60_000_000_000), + (Nanosecond, DatePart::Hour) => NonZeroI64::new(3_600_000_000_000), + (Nanosecond, DatePart::Day) => NonZeroI64::new(86_400_000_000_000), _ => None, }; @@ -841,7 +798,7 @@ fn general_date_trunc( tu: TimeUnit, value: i64, tz: Option, - granularity: DateTruncGranularity, + granularity: DatePart, ) -> Result { let scale = match tu { Second => 1_000_000_000, @@ -861,29 +818,25 @@ fn general_date_trunc( let result = match tu { Second => match granularity { - DateTruncGranularity::Minute => nano / 1_000_000_000 / 60 * 60, + DatePart::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, + DatePart::Minute => nano / 1_000_000 / 1_000 / 60 * 1_000 * 60, + DatePart::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, + DatePart::Minute => nano / 1_000 / 1_000_000 / 60 * 60 * 1_000_000, + DatePart::Second => nano / 1_000 / 1_000_000 * 1_000_000, + DatePart::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, + DatePart::Minute => nano / 1_000_000_000 / 60 * 1_000_000_000 * 60, + DatePart::Second => nano / 1_000_000_000 * 1_000_000_000, + DatePart::Millisecond => nano / 1_000_000 * 1_000_000, + DatePart::Microsecond => nano / 1_000 * 1_000, _ => nano, }, }; @@ -903,12 +856,13 @@ mod tests { use std::sync::Arc; use crate::datetime::date_trunc::{ - DateTruncFunc, DateTruncGranularity, date_trunc_coarse, + DateTruncFunc, date_trunc_coarse, parse_granularity, }; use arrow::array::cast::as_primitive_array; use arrow::array::types::TimestampNanosecondType; use arrow::array::{Array, TimestampNanosecondArray}; + use arrow::compute::DatePart; use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos; use arrow::datatypes::{DataType, Field, TimeUnit}; use datafusion_common::ScalarValue; @@ -1005,7 +959,7 @@ mod tests { for (original, granularity, expected) in &cases { 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 granularity_enum = parse_granularity(granularity).unwrap(); let result = date_trunc_coarse(granularity_enum, left, None).unwrap(); assert_eq!(result, right, "{original} = {expected}"); } @@ -1014,7 +968,7 @@ mod tests { #[test] fn date_trunc_out_of_range_lower_bound_returns_error() { let timestamp = string_to_timestamp_nanos("1677-09-22T00:00:00Z").unwrap(); - let err = date_trunc_coarse(DateTruncGranularity::Year, timestamp, None) + let err = date_trunc_coarse(DatePart::Year, timestamp, None) .unwrap_err() .to_string(); diff --git a/datafusion/spark/src/function/datetime/date_part.rs b/datafusion/spark/src/function/datetime/date_part.rs index 91bdb9a55318b..50ef1ece9d1a3 100644 --- a/datafusion/spark/src/function/datetime/date_part.rs +++ b/datafusion/spark/src/function/datetime/date_part.rs @@ -109,10 +109,10 @@ impl ScalarUDFImpl for SparkDatePart { } }; - // Map Spark-specific date part aliases to datafusion ones + // Map Spark-specific date part aliases to DataFusion ones. let part = match part.as_str() { "yearofweek" | "year_iso" => "isoyear", - "dayofweek" => "dow", + "dayofweek" | "dow" => "dow1", "dayofweek_iso" | "dow_iso" => "isodow", other => other, }; @@ -124,15 +124,6 @@ impl ScalarUDFImpl for SparkDatePart { vec![part_expr, date_expr], )); - match part { - // Spark's `dayofweek` is 1..=7 (Sun=1) but df's `dow` is 0..=6 - // (Sun=0); shift by +1. df's `isodow` already returns the - // PG-correct 1..=7 (Mon=1), which matches Spark's - // `dayofweek_iso`/`dow_iso`, so no shift is needed there. - "dow" => Ok(ExprSimplifyResult::Simplified( - date_part_expr + Expr::Literal(ScalarValue::Int32(Some(1)), None), - )), - _ => Ok(ExprSimplifyResult::Simplified(date_part_expr)), - } + Ok(ExprSimplifyResult::Simplified(date_part_expr)) } } diff --git a/datafusion/spark/src/function/datetime/time_trunc.rs b/datafusion/spark/src/function/datetime/time_trunc.rs index a66b8e94685aa..577c19733a71e 100644 --- a/datafusion/spark/src/function/datetime/time_trunc.rs +++ b/datafusion/spark/src/function/datetime/time_trunc.rs @@ -90,24 +90,15 @@ impl ScalarUDFImpl for SparkTimeTrunc { ) -> Result { let fmt_expr = &args[0]; - let fmt = match fmt_expr.as_literal() { - Some(ScalarValue::Utf8(Some(v))) - | Some(ScalarValue::Utf8View(Some(v))) - | Some(ScalarValue::LargeUtf8(Some(v))) => v.to_lowercase(), + match fmt_expr.as_literal() { + Some(ScalarValue::Utf8(Some(_))) + | Some(ScalarValue::Utf8View(Some(_))) + | Some(ScalarValue::LargeUtf8(Some(_))) => {} _ => { return plan_err!( "First argument of `TIME_TRUNC` must be non-null scalar Utf8" ); } - }; - - if !matches!( - fmt.as_str(), - "hour" | "minute" | "second" | "millisecond" | "microsecond" - ) { - return plan_err!( - "The format argument of `TIME_TRUNC` must be one of: hour, minute, second, millisecond, microsecond" - ); } Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction( diff --git a/datafusion/spark/src/function/datetime/trunc.rs b/datafusion/spark/src/function/datetime/trunc.rs index 9d7da5969a525..2c18e3909fffb 100644 --- a/datafusion/spark/src/function/datetime/trunc.rs +++ b/datafusion/spark/src/function/datetime/trunc.rs @@ -15,8 +15,9 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; +use std::{str::FromStr, sync::Arc}; +use arrow::compute::DatePart; use arrow::datatypes::{DataType, Field, FieldRef, TimeUnit}; use datafusion_common::types::{NativeType, logical_date, logical_string}; use datafusion_common::utils::take_function_args; @@ -109,13 +110,22 @@ impl ScalarUDFImpl for SparkTrunc { let fmt = match fmt.as_str() { "yy" | "yyyy" => "year", "mm" | "mon" => "month", - "year" | "month" | "day" | "week" | "quarter" => fmt.as_str(), - _ => { - return plan_err!( - "The format argument of `TRUNC` must be one of: year, yy, yyyy, month, mm, mon, day, week, quarter." - ); - } + other => other, }; + + // Accept shared DatePart aliases for the date parts supported by TRUNC. + if !matches!( + DatePart::from_str(fmt), + Ok(DatePart::Year + | DatePart::Month + | DatePart::Day + | DatePart::Week + | DatePart::Quarter) + ) { + return plan_err!( + "The format argument of `TRUNC` must represent a year, month, day, week, or quarter." + ); + } let return_type = dt_expr.get_type(info.schema())?; let fmt_expr = Expr::Literal(ScalarValue::new_utf8(fmt), None); diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index cd3ac1f3a455b..b33c366d002e3 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -2931,11 +2931,17 @@ mod tests { "MONTH", "EXTRACT(MONTH FROM x)", ), + ( + DateFieldExtractStyle::Extract, + "MONS", + "EXTRACT(MONTH FROM x)", + ), ( DateFieldExtractStyle::Strftime, "MONTH", "strftime('%m', x)", ), + (DateFieldExtractStyle::Strftime, "YRS", "strftime('%Y', x)"), ( DateFieldExtractStyle::DatePart, "DAY", diff --git a/datafusion/sql/src/unparser/utils.rs b/datafusion/sql/src/unparser/utils.rs index 240032d26c845..86f3e23115dcb 100644 --- a/datafusion/sql/src/unparser/utils.rs +++ b/datafusion/sql/src/unparser/utils.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::{cmp::Ordering, sync::Arc, vec}; +use std::{cmp::Ordering, str::FromStr, sync::Arc, vec}; use super::{ Unparser, dialect::CharacterLengthStyle, dialect::DateFieldExtractStyle, @@ -31,6 +31,7 @@ use datafusion_expr::{ Window, expr, utils::grouping_set_to_exprlist, }; +use arrow::compute::DatePart; use indexmap::IndexSet; use sqlparser::ast; use sqlparser::tokenizer::Span; @@ -455,13 +456,13 @@ pub(crate) fn date_part_to_sql( (DateFieldExtractStyle::Extract, 2) => { let date_expr = unparser.expr_to_sql_with_nesting(&date_part_args[1])?; if let Expr::Literal(ScalarValue::Utf8(Some(field)), _) = &date_part_args[0] { - let field = match field.to_lowercase().as_str() { - "year" => ast::DateTimeField::Year, - "month" => ast::DateTimeField::Month, - "day" => ast::DateTimeField::Day, - "hour" => ast::DateTimeField::Hour, - "minute" => ast::DateTimeField::Minute, - "second" => ast::DateTimeField::Second, + let field = match DatePart::from_str(field) { + Ok(DatePart::Year) => ast::DateTimeField::Year, + Ok(DatePart::Month) => ast::DateTimeField::Month, + Ok(DatePart::Day) => ast::DateTimeField::Day, + Ok(DatePart::Hour) => ast::DateTimeField::Hour, + Ok(DatePart::Minute) => ast::DateTimeField::Minute, + Ok(DatePart::Second) => ast::DateTimeField::Second, _ => return Ok(None), }; @@ -476,13 +477,13 @@ pub(crate) fn date_part_to_sql( let column = unparser.expr_to_sql_with_nesting(&date_part_args[1])?; if let Expr::Literal(ScalarValue::Utf8(Some(field)), _) = &date_part_args[0] { - let field = match field.to_lowercase().as_str() { - "year" => "%Y", - "month" => "%m", - "day" => "%d", - "hour" => "%H", - "minute" => "%M", - "second" => "%S", + let field = match DatePart::from_str(field) { + Ok(DatePart::Year) => "%Y", + Ok(DatePart::Month) => "%m", + Ok(DatePart::Day) => "%d", + Ok(DatePart::Hour) => "%H", + Ok(DatePart::Minute) => "%M", + Ok(DatePart::Second) => "%S", _ => return Ok(None), }; diff --git a/datafusion/sqllogictest/test_files/datetime/date_part.slt b/datafusion/sqllogictest/test_files/datetime/date_part.slt index 0a992b2d78a22..c13ee264ef8fd 100644 --- a/datafusion/sqllogictest/test_files/datetime/date_part.slt +++ b/datafusion/sqllogictest/test_files/datetime/date_part.slt @@ -288,6 +288,15 @@ SELECT date_part('QUARTER', CAST('2000-01-01' AS DATE)) ---- 1 +query IIII +SELECT + date_part('QUARTERS', CAST('2000-01-01' AS DATE)), + date_part('ISOWEEK', CAST('2000-01-01' AS DATE)), + date_part('DAYOFYEAR', CAST('2000-01-01' AS DATE)), + date_part('DAYOFWEEK', CAST('2000-01-01' AS DATE)) +---- +1 52 1 6 + query I SELECT EXTRACT(quarter FROM to_timestamp('2020-09-08T12:00:00+00:00')) ---- diff --git a/datafusion/sqllogictest/test_files/datetime/dates.slt b/datafusion/sqllogictest/test_files/datetime/dates.slt index a6a5f480f72e2..6fab426c3c5ae 100644 --- a/datafusion/sqllogictest/test_files/datetime/dates.slt +++ b/datafusion/sqllogictest/test_files/datetime/dates.slt @@ -448,11 +448,11 @@ select to_date('2022-01-23', '%Y-%m-%d'); 2022-01-23 # invalid date_trunc format -query error DataFusion error: Execution error: Unsupported date_trunc granularity: ''. Supported values are: microsecond, millisecond, second, minute, hour, day, week, month, quarter, year +query error DataFusion error: Execution error: Unsupported date_trunc granularity: ''. Supported granularities are: microsecond, millisecond, second, minute, hour, day, week, month, quarter, year SELECT date_trunc('', to_date('2022-02-23', '%Y-%m-%d')) # invalid date_trunc format -query error DataFusion error: Execution error: Unsupported date_trunc granularity: 'invalid'. Supported values are: microsecond, millisecond, second, minute, hour, day, week, month, quarter, year +query error DataFusion error: Execution error: Unsupported date_trunc granularity: 'invalid'. Supported granularities are: microsecond, millisecond, second, minute, hour, day, week, month, quarter, year SELECT date_trunc('invalid', to_date('2022-02-23', '%Y-%m-%d')) query PPPP diff --git a/datafusion/sqllogictest/test_files/datetime/timestamps.slt b/datafusion/sqllogictest/test_files/datetime/timestamps.slt index d73bc6eb06de8..e47a5f0b23d7d 100644 --- a/datafusion/sqllogictest/test_files/datetime/timestamps.slt +++ b/datafusion/sqllogictest/test_files/datetime/timestamps.slt @@ -1407,6 +1407,16 @@ SELECT DATE_TRUNC('MONTH', TIMESTAMP '2022-08-03 14:38:50Z'); ---- 2022-08-01T00:00:00 +# date_trunc uses the same field-name aliases as date_part +query PPPP +SELECT + DATE_TRUNC('mon', TIMESTAMP '2022-08-03 14:38:50Z'), + DATE_TRUNC('months', TIMESTAMP '2022-08-03 14:38:50Z'), + DATE_TRUNC('qtr', TIMESTAMP '2022-08-03 14:38:50Z'), + DATE_TRUNC('yrs', TIMESTAMP '2022-08-03 14:38:50Z'); +---- +2022-08-01T00:00:00 2022-08-01T00:00:00 2022-07-01T00:00:00 2022-01-01T00:00:00 + query P SELECT DATE_TRUNC('month', NULL); ---- diff --git a/datafusion/sqllogictest/test_files/spark/datetime/date_part.slt b/datafusion/sqllogictest/test_files/spark/datetime/date_part.slt index 48216bd551692..106f60bc07bc5 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/date_part.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/date_part.slt @@ -58,6 +58,11 @@ SELECT date_part('QTR'::string, '2000-01-01'::date); ---- 1 +query I +SELECT date_part('QUARTERS'::string, '2000-01-01'::date); +---- +1 + # MONTH query I SELECT date_part('MONTH'::string, '2000-01-01'::date); @@ -139,6 +144,11 @@ SELECT date_part('DOY'::string, '2000-01-01'::date); ---- 1 +query I +SELECT date_part('DAYOFYEAR'::string, '2000-01-01'::date); +---- +1 + # HOUR query I SELECT date_part('HOUR'::string, '2000-01-01 12:30:45'::timestamp); diff --git a/datafusion/sqllogictest/test_files/spark/datetime/date_trunc.slt b/datafusion/sqllogictest/test_files/spark/datetime/date_trunc.slt index 7fc1583bb9310..499ac75f76d8c 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/date_trunc.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/date_trunc.slt @@ -53,6 +53,11 @@ SELECT date_trunc('MON', '2015-03-05T09:32:05.123456'::timestamp); ---- 2015-03-01T00:00:00 +query P +SELECT date_trunc('MONTHS', '2015-03-05T09:32:05.123456'::timestamp); +---- +2015-03-01T00:00:00 + # WEEK - truncate to Monday of the week, time zeroed query P SELECT date_trunc('WEEK', '2015-03-05T09:32:05.123456'::timestamp); @@ -129,7 +134,7 @@ SELECT date_trunc('YEAR', NULL::timestamp); NULL # incorrect format -query error DataFusion error: Execution error: Unsupported date_trunc granularity: 'test'. Supported values are: microsecond, millisecond, second, minute, hour, day, week, month, quarter, year +query error DataFusion error: Execution error: Unsupported date_trunc granularity: 'test'. Supported granularities are: microsecond, millisecond, second, minute, hour, day, week, month, quarter, year SELECT date_trunc('test', '2015-03-05T09:32:05.123456'); # Timezone handling - Spark-compatible behavior diff --git a/datafusion/sqllogictest/test_files/spark/datetime/time_trunc.slt b/datafusion/sqllogictest/test_files/spark/datetime/time_trunc.slt index 35ffa483bb068..11f4704520f5e 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/time_trunc.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/time_trunc.slt @@ -27,6 +27,11 @@ SELECT time_trunc('MINUTE', '09:32:05.123456'::time); ---- 09:32:00 +query D +SELECT time_trunc('MINS', '09:32:05.123456'::time); +---- +09:32:00 + # SECOND - zero out fraction query D SELECT time_trunc('SECOND', '09:32:05.123456'::time); @@ -69,5 +74,5 @@ SELECT time_trunc('HOUR', NULL::time); NULL # incorrect format -query error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaused by\nError during planning: The format argument of `TIME_TRUNC` must be one of: hour, minute, second, millisecond, microsecond +query error DataFusion error: Execution error: Unsupported date_trunc granularity: 'test'. Supported granularities are: microsecond, millisecond, second, minute, hour, day, week, month, quarter, year SELECT time_trunc('test', '09:32:05.123456'::time); diff --git a/datafusion/sqllogictest/test_files/spark/datetime/trunc.slt b/datafusion/sqllogictest/test_files/spark/datetime/trunc.slt index aa26d7bd0ef06..f79bfffc7d321 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/trunc.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/trunc.slt @@ -53,6 +53,11 @@ SELECT trunc('2009-02-12'::date, 'MON'::string); ---- 2009-02-01 +query D +SELECT trunc('2009-02-12'::date, 'MONTHS'::string); +---- +2009-02-01 + # WEEK - truncate to Monday of the week query D SELECT trunc('2009-02-12'::date, 'WEEK'::string); @@ -88,5 +93,5 @@ query error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaus SELECT trunc('2009-02-12'::date, NULL::string); # incorrect format -query error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaused by\nError during planning: The format argument of `TRUNC` must be one of: year, yy, yyyy, month, mm, mon, day, week, quarter. +query error DataFusion error: Optimizer rule 'simplify_expressions' failed\ncaused by\nError during planning: The format argument of `TRUNC` must represent a year, month, day, week, or quarter. SELECT trunc('2009-02-12'::date, 'test'::string);