Skip to content

Commit 48ff4d8

Browse files
committed
refactor: reuse arrow's time constants in date_part sub-second helpers
1 parent 9745f21 commit 48ff4d8

1 file changed

Lines changed: 42 additions & 40 deletions

File tree

‎datafusion/functions/src/datetime/date_part.rs‎

Lines changed: 42 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,9 @@ use std::str::FromStr;
1919
use std::sync::Arc;
2020

2121
use arrow::array::temporal_conversions::{
22-
as_datetime_with_timezone, timestamp_ms_to_datetime, timestamp_ns_to_datetime,
23-
timestamp_s_to_datetime, timestamp_us_to_datetime,
22+
MICROSECONDS, MILLISECONDS, NANOSECONDS, SECONDS_IN_DAY, as_datetime_with_timezone,
23+
timestamp_ms_to_datetime, timestamp_ns_to_datetime, timestamp_s_to_datetime,
24+
timestamp_us_to_datetime,
2425
};
2526
use arrow::array::timezone::Tz;
2627
use arrow::array::{
@@ -545,13 +546,13 @@ fn subsecond_part(array: &dyn Array, part: DatePart) -> Result<ArrayRef> {
545546
// the whole seconds, and it rejects them for `Interval(YearMonth)`
546547
(Interval(_) | Duration(_), _) => date_part(array, part)?,
547548
(_, DatePart::Millisecond) => {
548-
seconds_with_fraction::<Int32Type, 1_000>(array, part)?
549+
seconds_with_fraction::<Int32Type, MILLISECONDS>(array, part)?
549550
}
550551
(_, DatePart::Microsecond) => {
551-
seconds_with_fraction::<Int32Type, 1_000_000>(array, part)?
552+
seconds_with_fraction::<Int32Type, MICROSECONDS>(array, part)?
552553
}
553554
(_, DatePart::Nanosecond) => {
554-
seconds_with_fraction::<Int64Type, 1_000_000_000>(array, part)?
555+
seconds_with_fraction::<Int64Type, NANOSECONDS>(array, part)?
555556
}
556557
(_, part) => return internal_err!("{part} is not a sub-second part"),
557558
})
@@ -596,23 +597,24 @@ where
596597
Time32(Second) => {
597598
time_seconds_with_fraction::<Time32SecondType, O, 1, PER_SECOND>(array)?
598599
}
599-
Time32(Millisecond) => {
600-
time_seconds_with_fraction::<Time32MillisecondType, O, 1_000, PER_SECOND>(
601-
array,
602-
)?
603-
}
604-
Time64(Microsecond) => {
605-
time_seconds_with_fraction::<Time64MicrosecondType, O, 1_000_000, PER_SECOND>(
606-
array,
607-
)?
608-
}
609-
Time64(Nanosecond) => time_seconds_with_fraction::<
610-
Time64NanosecondType,
600+
Time32(Millisecond) => time_seconds_with_fraction::<
601+
Time32MillisecondType,
602+
O,
603+
MILLISECONDS,
604+
PER_SECOND,
605+
>(array)?,
606+
Time64(Microsecond) => time_seconds_with_fraction::<
607+
Time64MicrosecondType,
611608
O,
612-
1_000_000_000,
609+
MICROSECONDS,
613610
PER_SECOND,
614611
>(array)?,
615-
other => return exec_err!("{part} does not support {other}"),
612+
Time64(Nanosecond) => {
613+
time_seconds_with_fraction::<Time64NanosecondType, O, NANOSECONDS, PER_SECOND>(
614+
array,
615+
)?
616+
}
617+
other => return internal_err!("{part} does not support {other}"),
616618
})
617619
}
618620

@@ -628,38 +630,39 @@ where
628630
O: ArrowPrimitiveType,
629631
i64: AsPrimitive<O::Native>,
630632
{
631-
// Like arrow's `date_part`, but one chrono conversion per row instead of
632-
// one each for the seconds and the fraction
633-
fn seconds<O: ArrowPrimitiveType, const PER_SECOND: i64>(
634-
t: &impl Timelike,
635-
) -> O::Native
636-
where
637-
i64: AsPrimitive<O::Native>,
638-
{
639-
// At most 60 seconds, which fits in Int32 for milliseconds and
640-
// microseconds, so the narrowing `as_` never truncates
641-
(t.second() as i64 * PER_SECOND
642-
+ t.nanosecond() as i64 / (1_000_000_000 / PER_SECOND))
643-
.as_()
644-
}
645-
646633
let array = as_primitive_array::<T>(array)?;
647-
// `unary_opt` keeps the input nulls and additionally nulls out any value
634+
// One chrono conversion per row, where arrow's `date_part` needs one each
635+
// for the seconds and the fraction. `unary_opt` keeps the input nulls and additionally nulls out any value
648636
// that cannot be represented as a `DateTime` (out of chrono's range).
649637
let r: PrimitiveArray<O> = match tz {
650638
Some(tz) => array.unary_opt(|v| {
651639
to_datetime(v).map(|naive| {
652-
seconds::<O, PER_SECOND>(
640+
seconds_of_minute::<O, PER_SECOND>(
653641
&Utc.from_utc_datetime(&naive).with_timezone(&tz),
654642
)
655643
})
656644
}),
657-
None => array
658-
.unary_opt(|v| to_datetime(v).map(|naive| seconds::<O, PER_SECOND>(&naive))),
645+
None => array.unary_opt(|v| {
646+
to_datetime(v).map(|naive| seconds_of_minute::<O, PER_SECOND>(&naive))
647+
}),
659648
};
660649
Ok(Arc::new(r))
661650
}
662651

652+
/// The seconds of the minute of `t`, with their fraction, in units of
653+
/// `1 / PER_SECOND` seconds
654+
fn seconds_of_minute<O: ArrowPrimitiveType, const PER_SECOND: i64>(
655+
t: &impl Timelike,
656+
) -> O::Native
657+
where
658+
i64: AsPrimitive<O::Native>,
659+
{
660+
// At most 60 seconds, which fits in Int32 for milliseconds and
661+
// microseconds, so the narrowing `as_` never truncates
662+
(t.second() as i64 * PER_SECOND + t.nanosecond() as i64 / (NANOSECONDS / PER_SECOND))
663+
.as_()
664+
}
665+
663666
/// Like [`timestamp_seconds_with_fraction`] for times, which count
664667
/// `1 / UNITS_PER_SECOND` seconds since midnight.
665668
fn time_seconds_with_fraction<T, O, const UNITS_PER_SECOND: i64, const PER_SECOND: i64>(
@@ -671,13 +674,12 @@ where
671674
O: ArrowPrimitiveType,
672675
i64: AsPrimitive<O::Native>,
673676
{
674-
const SECONDS_PER_DAY: i64 = 86_400;
675677
let array = as_primitive_array::<T>(array)?;
676678
// Values outside one day are NULL, as in arrow's `date_part`. arrow's
677679
// `as_time` would wrap some of them around to a valid time instead.
678680
let r: PrimitiveArray<O> = array.unary_opt(|v| {
679681
let v: i64 = v.as_();
680-
(0..SECONDS_PER_DAY * UNITS_PER_SECOND)
682+
(0..SECONDS_IN_DAY * UNITS_PER_SECOND)
681683
.contains(&v)
682684
.then(|| {
683685
let v = v % (60 * UNITS_PER_SECOND);

0 commit comments

Comments
 (0)