mbutrovich commented on code in PR #11199:
URL: https://github.com/apache/arrow-rs/pull/11199#discussion_r4106546454
##########
arrow-cast/src/cast/mod.rs:
##########
@@ -608,48 +608,171 @@ fn make_duration_array(array:
&PrimitiveArray<Int64Type>, unit: TimeUnit) -> Arr
}
}
-fn as_time_res_with_timezone<T: ArrowPrimitiveType>(
- v: i64,
- tz: Option<Tz>,
-) -> Result<NaiveTime, ArrowError> {
- let time = match tz {
- Some(tz) => as_datetime_with_timezone::<T>(v, tz).map(|d| d.time()),
- None => as_datetime::<T>(v).map(|d| d.time()),
- };
+/// Casts timestamps to the time of day in the unit of `to_type`.
+fn cast_timestamp_to_time<T: ArrowTimestampType>(
+ array: &dyn Array,
+ to_type: &DataType,
+ cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError> {
+ let array = array.as_primitive::<T>();
+ match to_type {
+ DataType::Time32(TimeUnit::Second) => {
+ timestamp_to_time::<T, Time32SecondType>(array, cast_options)
+ }
+ DataType::Time32(TimeUnit::Millisecond) => {
+ timestamp_to_time::<T, Time32MillisecondType>(array, cast_options)
+ }
+ DataType::Time64(TimeUnit::Microsecond) => {
+ timestamp_to_time::<T, Time64MicrosecondType>(array, cast_options)
+ }
+ DataType::Time64(TimeUnit::Nanosecond) => {
+ timestamp_to_time::<T, Time64NanosecondType>(array, cast_options)
+ }
+ _ => Err(ArrowError::CastError(format!(
+ "Casting from {} to {to_type} not supported",
+ array.data_type()
+ ))),
+ }
+}
- time.ok_or_else(|| {
- ArrowError::CastError(format!(
- "Failed to create naive time with {} {}",
- std::any::type_name::<T>(),
- v
- ))
- })
+/// A `Time32` or `Time64` type with a fixed unit and a Chrono conversion.
+trait TimeType: ArrowTemporalType {
+ /// Number of units in one second.
+ const UNIT_MULTIPLE: i64;
+
+ fn from_naive_time(time: NaiveTime) -> Self::Native;
+}
+
+impl TimeType for Time32SecondType {
+ const UNIT_MULTIPLE: i64 = 1;
+
+ fn from_naive_time(time: NaiveTime) -> Self::Native {
+ time_to_time32s(time)
+ }
+}
+
+impl TimeType for Time32MillisecondType {
+ const UNIT_MULTIPLE: i64 = MILLISECONDS;
+
+ fn from_naive_time(time: NaiveTime) -> Self::Native {
+ time_to_time32ms(time)
+ }
+}
+
+impl TimeType for Time64MicrosecondType {
+ const UNIT_MULTIPLE: i64 = MICROSECONDS;
+
+ fn from_naive_time(time: NaiveTime) -> Self::Native {
+ time_to_time64us(time)
+ }
+}
+
+impl TimeType for Time64NanosecondType {
+ const UNIT_MULTIPLE: i64 = NANOSECONDS;
+
+ fn from_naive_time(time: NaiveTime) -> Self::Native {
+ time_to_time64ns(time)
+ }
+}
+
+/// Casts timestamps to the time of day in the unit of `O`.
+fn timestamp_to_time<T, O>(
+ array: &PrimitiveArray<T>,
+ cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError>
+where
+ T: ArrowTimestampType,
+ O: TimeType,
+ i64: AsPrimitive<O::Native>,
+{
+ // A time within one day fits its Time32 or Time64 representation.
+ let array = match array.timezone() {
+ Some(tz) => {
+ let tz: Tz = tz.parse()?;
+ let time = |v: i64| {
+ as_datetime_with_timezone::<T>(v, tz).map(|d|
O::from_naive_time(d.time()))
+ };
+ if cast_options.safe {
+ array.unary_opt::<_, O>(time)
+ } else {
+ array.try_unary::<_, O, _>(|v| {
+ time(v).ok_or_else(|| {
+ ArrowError::CastError(format!(
+ "Failed to create naive time with {} {}",
+ std::any::type_name::<T>(),
+ v
+ ))
+ })
+ })?
+ }
+ }
+ None => array.unary::<_, O>(|v| {
+ // The remainder within a day is the time of day; `rem_euclid`
keeps it
+ // nonnegative for timestamps before the epoch. The units are
constants,
+ // so the branch and the divisions fold at compile time.
+ let from = time_unit_multiple(&T::UNIT);
+ let time = v.rem_euclid(SECONDS_IN_DAY * from);
+ let time = if from >= O::UNIT_MULTIPLE {
+ time / (from / O::UNIT_MULTIPLE)
+ } else {
+ time * (O::UNIT_MULTIPLE / from)
+ };
+ time.as_()
+ }),
+ };
+ Ok(Arc::new(array))
}
fn timestamp_to_date32<T: ArrowTimestampType>(
array: &PrimitiveArray<T>,
+ cast_options: &CastOptions,
) -> Result<ArrayRef, ArrowError> {
let err = |x: i64| {
ArrowError::CastError(format!(
- "Cannot convert {} {x} to datetime",
+ "Cannot convert {} {x} to Date32",
std::any::type_name::<T>()
))
};
let array: Date32Array = match array.timezone() {
Some(tz) => {
let tz: Tz = tz.parse()?;
- array.try_unary(|x| {
+ let date = |x: i64| {
as_datetime_with_timezone::<T>(x, tz)
- .ok_or_else(|| err(x))
.map(|d| Date32Type::from_naive_date(d.date_naive()))
- })?
+ };
+ if cast_options.safe {
+ array.unary_opt(date)
+ } else {
+ array.try_unary(|x| date(x).ok_or_else(|| err(x)))?
+ }
+ }
+ None => {
+ // Date32 stores days since the epoch. Round down so that a
timestamp
+ // just before the epoch belongs to the preceding day. The unit is
a
+ // constant, so the divisor folds at compile time.
+ let days = |x: i64| x.div_euclid(SECONDS_IN_DAY *
time_unit_multiple(&T::UNIT));
+ let all_in_range = match T::UNIT {
+ // Every microsecond or nanosecond timestamp lies within the
Date32 range.
+ TimeUnit::Microsecond | TimeUnit::Nanosecond => true,
+ // A branch-free scan lets the common case skip the per-value
check.
+ _ => {
+ let day = SECONDS_IN_DAY * time_unit_multiple(&T::UNIT);
+ let (lo, hi) = (i32::MIN as i64 * day, (i32::MAX as i64 +
1) * day);
Review Comment:
The range check is correct only as long as `day` here matches the divisor in
`days`, and the two are computed separately. Could we compute `day` once and
use it in both? I tried this locally. The timestamp tests pass, and the
`TimestampSecondType` listing on aarch64 differs only in register allocation,
with no `sdiv` in either version, so the comment about the divisor folding at
compile time still holds.
```suggestion
let day = SECONDS_IN_DAY * time_unit_multiple(&T::UNIT);
let days = |x: i64| x.div_euclid(day);
let all_in_range = match T::UNIT {
// Every microsecond or nanosecond timestamp lies within the
Date32 range.
TimeUnit::Microsecond | TimeUnit::Nanosecond => true,
// A branch-free scan lets the common case skip the
per-value check.
_ => {
let (lo, hi) = (i32::MIN as i64 * day, (i32::MAX as i64
+ 1) * day);
```
--
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]
