mbutrovich commented on code in PR #26686:
URL: https://github.com/apache/datafusion/pull/26686#discussion_r4367250013
##########
datafusion/functions/src/datetime/date_part.rs:
##########
@@ -564,48 +504,83 @@ fn epoch(array: &dyn Array) -> Result<ArrayRef> {
Ok(Arc::new(f))
}
-/// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the
-/// result to a total number of nanoseconds as an Int64 array.
+/// Return the seconds of `array`, with their fraction, as a number of
+/// milliseconds, microseconds or nanoseconds (`part`): the seconds of the
+/// minute for timestamps, times and intervals, and the total length for
+/// durations.
///
-/// This returns an Int64 rather than Int32 because there 1 billion
+/// `nanosecond` returns an Int64 rather than Int32 because there are 1 billion
/// `nanosecond`s in each second, so representing up to 60 seconds as
/// nanoseconds can be values up to 60 billion, which does not fit in Int32.
-fn seconds_ns(array: &dyn Array) -> Result<ArrayRef> {
- // Fast path for Date32 and Date64 - no nanoseconds
- if array.data_type() == &Date32 || array.data_type() == &Date64 {
- return Ok(Arc::new(Int64Array::from_iter_values_with_nulls(
- repeat_n(0, array.len()),
- array.nulls().cloned(),
- )));
- }
+fn subsecond_part(array: &dyn Array, part: DatePart) -> Result<ArrayRef> {
+ let len = array.len();
+ let nulls = array.nulls().cloned();
+ Ok(match (array.data_type(), part) {
+ // Fast path for Date32 and Date64 - no seconds
+ (Date32 | Date64, DatePart::Nanosecond) => {
+ Arc::new(Int64Array::new(vec![0; len].into(), nulls))
+ }
+ (Date32 | Date64, _) => Arc::new(Int32Array::new(vec![0; len].into(),
nulls)),
+ // arrow's `date_part(Nanosecond)` computes these, but returns Int32,
+ // which overflows past about 2.1 seconds
+ (Interval(ArrowIntervalUnit::DayTime), DatePart::Nanosecond) =>
Arc::new(
+ as_interval_dt_array(array)?
+ .unary::<_, Int64Type>(|x| (x.milliseconds % 60_000) as i64 *
1_000_000),
+ ),
+ (Interval(ArrowIntervalUnit::MonthDayNano), DatePart::Nanosecond) =>
Arc::new(
+ as_interval_mdn_array(array)?
+ .unary::<_, Int64Type>(|x| x.nanoseconds % 60_000_000_000),
+ ),
+ (Duration(Nanosecond), DatePart::Nanosecond) => {
+
Arc::new(as_duration_nanosecond_array(array)?.reinterpret_cast::<Int64Type>())
+ }
+ // arrow's cast returns NULL on overflow, like its `date_part`
+ (Duration(_), DatePart::Nanosecond) => Arc::new(
+ as_duration_nanosecond_array(&cast(array, &Duration(Nanosecond))?)?
+ .reinterpret_cast::<Int64Type>(),
+ ),
+ // arrow's sub-second parts of intervals and durations already include
+ // the whole seconds
+ (
+ Interval(ArrowIntervalUnit::DayTime |
ArrowIntervalUnit::MonthDayNano)
+ | Duration(_),
+ _,
+ ) => date_part(array, part)?,
+ (_, DatePart::Millisecond) => {
+ seconds_with_fraction::<Int32Type, 1_000>(array, part)?
+ }
+ (_, DatePart::Microsecond) => {
+ seconds_with_fraction::<Int32Type, 1_000_000>(array, part)?
+ }
+ (_, DatePart::Nanosecond) => {
+ seconds_with_fraction::<Int64Type, 1_000_000_000>(array, part)?
+ }
+ (_, part) => return internal_err!("{part} is not a sub-second part"),
+ })
+}
+/// Combine arrow's seconds of the minute with its fraction of the second in
+/// `part` units, for types where arrow returns them separately (timestamps
+/// and times).
+fn seconds_with_fraction<O, const PER_SECOND: i64>(
+ array: &dyn Array,
+ part: DatePart,
+) -> Result<ArrayRef>
+where
+ O: ArrowPrimitiveType,
+ i64: AsPrimitive<O::Native>,
+{
let secs = date_part(array, DatePart::Second)?;
// This assumes array is primitive and not a dictionary
let secs = as_int32_array(secs.as_ref())?;
- let subsecs = date_part(array, DatePart::Nanosecond)?;
- let subsecs = as_int32_array(subsecs.as_ref())?;
-
- // Special case where there are no nulls.
- if subsecs.null_count() == 0 {
- let r: Int64Array = binary(secs, subsecs, |secs, subsecs| {
- (secs as i64) * 1_000_000_000 + (subsecs as i64)
- })?;
- Ok(Arc::new(r))
- } else {
- // Nulls in secs are preserved, nulls in subsecs are treated as zero
to account for the case
- // where the number of nanoseconds overflows.
- let r: Int64Array = secs
- .iter()
- .zip(subsecs)
- .map(|(secs, subsecs)| {
- secs.map(|secs| {
- let subsecs = subsecs.unwrap_or(0);
- (secs as i64) * 1_000_000_000 + (subsecs as i64)
- })
- })
- .collect();
- Ok(Arc::new(r))
- }
+ let fraction = date_part(array, part)?;
Review Comment:
> `main` already used two passes, so a follow-up is fine.
I did it in this PR, in
[`date_part.rs`](https://github.com/apache/datafusion/blob/48ff4d8da0d3a990f6ef216e13ef4e1f5cbe758f/datafusion/functions/src/datetime/date_part.rs#L561-L695).
Two things differ from your sketch:
- The timestamp helper takes arrow's converter for the unit
(`timestamp_s_to_datetime` and so on), which arrow's `date_part` calls too.
With `as_datetime::<T>` it measured about 1.5x slower, and `Timestamp(Second)`
came out slower than before.
- Times use arithmetic with the same one-day range check as arrow's
`date_part`, not `as_time`. `time64ns_to_time` computes `(v / NANOSECONDS) as
u32`, so `2^32 + 5` seconds becomes 00:00:05 instead of `None`
(apache/arrow-rs#11893). A [new
query](https://github.com/apache/datafusion/blob/48ff4d8da0d3a990f6ef216e13ef4e1f5cbe758f/datafusion/sqllogictest/test_files/datetime/date_part.slt#L952-L965)
checks that out-of-range times stay NULL.
Over `c920a5b92`, `millisecond` on `Timestamp(ns)` went from 1.07x to 1.12x
before this change to 2.18x to 2.56x, and `microsecond` on `Time64(ns)` from
1.53x to 2.24x to 3.23x to 4.77x, across the four bot runs linked in the
description. The timezone case gained about as much as the others, 2.34x to
2.84x.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
