Skip to content

chore: use enum as date_trunc granularity #18390

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

Merged
merged 2 commits into from
Oct 31, 2025
Merged

chore: use enum as date_trunc granularity #18390

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b6c3582
chore: use enum as `date_trunc` granularity
comphead Oct 30, 2025
a3882ac
address comments
comphead Oct 31, 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
216 changes: 142 additions & 74 deletions datafusion/functions/src/datetime/date_trunc.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,71 @@ 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 {
/// Mapping of string representations to enum variants
const GRANULARITY_MAP: &[(&str, Self)] = &[
("microsecond", Self::Microsecond),
("millisecond", Self::Millisecond),
("second", Self::Second),
("minute", Self::Minute),
("hour", Self::Hour),
("day", Self::Day),
("week", Self::Week),
("month", Self::Month),
("quarter", Self::Quarter),
("year", Self::Year),
];

/// Parse a granularity string into a DateTruncGranularity enum
fn from_str(s: &str) -> Result<Self> {
let s_lower = s.to_lowercase();
Self::GRANULARITY_MAP
.iter()
.find(|(key, _)| *key == s_lower.as_str())
Copy link
Contributor

@alamb alamb Oct 30, 2025

Choose a reason for hiding this comment

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

Perhaps we can make this faster via https://doc.rust-lang.org/std/vec/struct.Vec.html#method.binary_search?

Or maybe just hard coding the strings so the compiler can make a jump table

Copy link
Contributor Author

@comphead comphead Oct 31, 2025

Choose a reason for hiding this comment

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

yeah, Rust still doesn't support const enums :(

.map(|(_, value)| *value)
.ok_or_else(|| {
let supported = Self::GRANULARITY_MAP
.iter()
.map(|(key, _)| *key)
.collect::<Vec<_>>()
.join(", ");
exec_datafusion_err!(
Copy link
Contributor

@alamb alamb Oct 30, 2025

Choose a reason for hiding this comment

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

Can you please add a slt test for this error?

"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.",
Expand Down Expand Up @@ -172,7 +237,7 @@ impl ScalarUDFImpl for DateTruncFunc {
let args = args.args;
let (granularity, array) = (&args[0], &args[1]);

let granularity = if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) =
let granularity_str = if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) =
granularity
{
v.to_lowercase()
Expand All @@ -183,54 +248,46 @@ impl ScalarUDFImpl for DateTruncFunc {
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: String,
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 granularities
if matches!(
granularity.as_str(),
// 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).
"second" | "minute" | "millisecond" | "microsecond"
) ||
// 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
(parsed_tz.is_none() && matches!(granularity.as_str(), "hour" | "day"))
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.as_str(),
granularity,
)?;
return Ok(ColumnarValue::Array(result));
}

let array: PrimitiveArray<T> = array
.try_unary(|x| {
general_date_trunc(T::UNIT, x, parsed_tz, granularity.as_str())
})?
.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: String,
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.as_str(),
)?)
Some(general_date_trunc(T::UNIT, *v, parsed_tz, granularity)?)
} else {
None
};
Expand Down Expand Up @@ -308,57 +365,57 @@ impl ScalarUDFImpl for DateTruncFunc {
}
}

fn _date_trunc_coarse<T>(granularity: &str, value: Option<T>) -> Result<Option<T>>
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 {
"millisecond" => value,
"microsecond" => value,
"second" => value.and_then(|d| d.with_nanosecond(0)),
"minute" => value
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)),
"hour" => value
DateTruncGranularity::Hour => value
.and_then(|d| d.with_nanosecond(0))
.and_then(|d| d.with_second(0))
.and_then(|d| d.with_minute(0)),
"day" => value
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)),
"week" => value
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()
}),
"month" => value
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)),
"quarter" => value
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))),
"year" => value
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)),
unsupported => {
return exec_err!("Unsupported date_trunc granularity: {unsupported}");
}
};
Ok(value)
}
Expand All @@ -371,7 +428,7 @@ where
}

fn _date_trunc_coarse_with_tz(
granularity: &str,
granularity: DateTruncGranularity,
value: Option<DateTime<Tz>>,
) -> Result<Option<i64>> {
if let Some(value) = value {
Expand Down Expand Up @@ -413,7 +470,7 @@ fn _date_trunc_coarse_with_tz(
}

fn _date_trunc_coarse_without_tz(
granularity: &str,
granularity: DateTruncGranularity,
value: Option<NaiveDateTime>,
) -> Result<Option<i64>> {
let value = _date_trunc_coarse::<NaiveDateTime>(granularity, value)?;
Expand All @@ -424,7 +481,11 @@ 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: &str, value: i64, tz: Option<Tz>) -> Result<i64> {
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,
Expand Down Expand Up @@ -454,30 +515,30 @@ fn date_trunc_coarse(granularity: &str, value: i64, tz: Option<Tz>) -> Result<i6
fn general_date_trunc_array_fine_granularity<T: ArrowTimestampType>(
tu: TimeUnit,
array: &PrimitiveArray<T>,
granularity: &str,
granularity: DateTruncGranularity,
) -> Result<ArrayRef> {
let unit = match (tu, granularity) {
(Second, "minute") => NonZeroI64::new(60),
(Second, "hour") => NonZeroI64::new(3600),
(Second, "day") => NonZeroI64::new(86400),

(Millisecond, "second") => NonZeroI64::new(1_000),
(Millisecond, "minute") => NonZeroI64::new(60_000),
(Millisecond, "hour") => NonZeroI64::new(3_600_000),
(Millisecond, "day") => NonZeroI64::new(86_400_000),

(Microsecond, "millisecond") => NonZeroI64::new(1_000),
(Microsecond, "second") => NonZeroI64::new(1_000_000),
(Microsecond, "minute") => NonZeroI64::new(60_000_000),
(Microsecond, "hour") => NonZeroI64::new(3_600_000_000),
(Microsecond, "day") => NonZeroI64::new(86_400_000_000),

(Nanosecond, "microsecond") => NonZeroI64::new(1_000),
(Nanosecond, "millisecond") => NonZeroI64::new(1_000_000),
(Nanosecond, "second") => NonZeroI64::new(1_000_000_000),
(Nanosecond, "minute") => NonZeroI64::new(60_000_000_000),
(Nanosecond, "hour") => NonZeroI64::new(3_600_000_000_000),
(Nanosecond, "day") => NonZeroI64::new(86_400_000_000_000),
(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,
};

Expand All @@ -502,7 +563,7 @@ fn general_date_trunc(
tu: TimeUnit,
value: i64,
tz: Option<Tz>,
granularity: &str,
granularity: DateTruncGranularity,
) -> Result<i64, DataFusionError> {
let scale = match tu {
Second => 1_000_000_000,
Expand All @@ -516,25 +577,29 @@ fn general_date_trunc(

let result = match tu {
Second => match granularity {
"minute" => nano / 1_000_000_000 / 60 * 60,
DateTruncGranularity::Minute => nano / 1_000_000_000 / 60 * 60,
_ => nano / 1_000_000_000,
},
Millisecond => match granularity {
"minute" => nano / 1_000_000 / 1_000 / 60 * 1_000 * 60,
"second" => nano / 1_000_000 / 1_000 * 1_000,
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 {
"minute" => nano / 1_000 / 1_000_000 / 60 * 60 * 1_000_000,
"second" => nano / 1_000 / 1_000_000 * 1_000_000,
"millisecond" => nano / 1_000 / 1_000 * 1_000,
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 {
"minute" => nano / 1_000_000_000 / 60 * 1_000_000_000 * 60,
"second" => nano / 1_000_000_000 * 1_000_000_000,
"millisecond" => nano / 1_000_000 * 1_000_000,
"microsecond" => nano / 1_000 * 1_000,
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,
},
};
Expand All @@ -554,7 +619,9 @@ fn parse_tz(tz: &Option<Arc<str>>) -> Result<Option<Tz>> {
mod tests {
use std::sync::Arc;

use crate::datetime::date_trunc::{date_trunc_coarse, DateTruncFunc};
use crate::datetime::date_trunc::{
date_trunc_coarse, DateTruncFunc, DateTruncGranularity,
};

use arrow::array::cast::as_primitive_array;
use arrow::array::types::TimestampNanosecondType;
Expand Down Expand Up @@ -655,7 +722,8 @@ mod tests {
cases.iter().for_each(|(original, granularity, expected)| {
let left = string_to_timestamp_nanos(original).unwrap();
let right = string_to_timestamp_nanos(expected).unwrap();
let result = date_trunc_coarse(granularity, left, None).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}");
});
}
Expand Down