This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11199-6d44c7c5ec45eb54ba2517179bcda5f786aceb85 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 61c69cd16dcc9dc15a109adbacba479b8d08d9f5 Author: Neil Conway <[email protected]> AuthorDate: Sat Oct 3 16:05:49 2026 +0000 perf: Optimize casting timestamp -> `Date32`, `Time32`, `Time64` (#11199) # Which issue does this PR close? - N/A # Rationale for this change Casting from timestamp to another datetime type currently goes through Chrono. When casting timestamp values without a timezone to `Date32`, `Time32`, or `Time64`, we can avoid going through Chrono and do the cast directly from the stored integer timestamp value. This is similar to the approach taken in #11187; it improves performance by about 10-20x. Along the way, adjust casts from timestamp to `Date32`, `Time32` and `Time64` to correctly respect "safe" mode; previously, such casts returned an error for out-of-range values. Benchmarks: (Arm64) - timestamp_s/date32/no_nulls: 88.338 → 3.155 µs, −96.4% - timestamp_s/date32/mixed_nulls: 73.260 → 3.155 µs, −95.7% - timestamp_s/date32/timezone_control: 145.930 → 145.460 µs, −0.3% - timestamp_s/time32_s/no_nulls: 66.287 → 2.293 µs, −96.5% - timestamp_s/time32_s/mixed_nulls: 54.928 → 2.297 µs, −95.8% - timestamp_s/time32_s/timezone_control: 76.130 → 73.408 µs, −3.6% - timestamp_s/time32_ms/no_nulls: 68.243 → 2.367 µs, −96.5% - timestamp_s/time32_ms/mixed_nulls: 56.229 → 2.369 µs, −95.8% - timestamp_s/time32_ms/timezone_control: 79.243 → 77.735 µs, −1.9% - timestamp_s/time64_us/no_nulls: 68.461 → 4.648 µs, −93.2% - timestamp_s/time64_us/mixed_nulls: 58.283 → 4.657 µs, −92.0% - timestamp_s/time64_us/timezone_control: 79.850 → 79.258 µs, −0.7% - timestamp_s/time64_ns/no_nulls: 67.501 → 4.510 µs, −93.3% - timestamp_s/time64_ns/mixed_nulls: 57.150 → 4.570 µs, −92.0% - timestamp_s/time64_ns/timezone_control: 78.252 → 77.134 µs, −1.4% - timestamp_ms/date32/no_nulls: 102.210 → 3.176 µs, −96.9% - timestamp_ms/date32/mixed_nulls: 84.143 → 3.150 µs, −96.3% - timestamp_ms/date32/timezone_control: 165.530 → 165.770 µs, +0.1% - timestamp_ms/time32_s/no_nulls: 78.368 → 2.614 µs, −96.7% - timestamp_ms/time32_s/mixed_nulls: 63.744 → 2.612 µs, −95.9% - timestamp_ms/time32_s/timezone_control: 88.934 → 88.578 µs, −0.4% - timestamp_ms/time32_ms/no_nulls: 80.194 → 2.295 µs, −97.1% - timestamp_ms/time32_ms/mixed_nulls: 65.715 → 2.289 µs, −96.5% - timestamp_ms/time32_ms/timezone_control: 91.717 → 91.280 µs, −0.5% - timestamp_ms/time64_us/no_nulls: 80.959 → 4.274 µs, −94.7% - timestamp_ms/time64_us/mixed_nulls: 67.123 → 4.290 µs, −93.6% - timestamp_ms/time64_us/timezone_control: 93.714 → 93.021 µs, −0.7% - timestamp_ms/time64_ns/no_nulls: 79.801 → 4.688 µs, −94.1% - timestamp_ms/time64_ns/mixed_nulls: 66.226 → 4.687 µs, −92.9% - timestamp_ms/time64_ns/timezone_control: 91.657 → 90.890 µs, −0.8% - timestamp_us/date32/no_nulls: 100.840 → 2.410 µs, −97.6% - timestamp_us/date32/mixed_nulls: 82.909 → 2.415 µs, −97.1% - timestamp_us/date32/timezone_control: 160.590 → 160.570 µs, −0.0% - timestamp_us/time32_s/no_nulls: 78.063 → 4.334 µs, −94.4% - timestamp_us/time32_s/mixed_nulls: 63.884 → 4.385 µs, −93.1% - timestamp_us/time32_s/timezone_control: 86.076 → 86.178 µs, +0.1% - timestamp_us/time32_ms/no_nulls: 78.962 → 4.263 µs, −94.6% - timestamp_us/time32_ms/mixed_nulls: 65.085 → 4.306 µs, −93.4% - timestamp_us/time32_ms/timezone_control: 87.605 → 89.651 µs, +2.3% - timestamp_us/time64_us/no_nulls: 80.197 → 2.359 µs, −97.1% - timestamp_us/time64_us/mixed_nulls: 66.402 → 2.370 µs, −96.4% - timestamp_us/time64_us/timezone_control: 88.763 → 90.953 µs, +2.5% - timestamp_us/time64_ns/no_nulls: 79.202 → 4.476 µs, −94.3% - timestamp_us/time64_ns/mixed_nulls: 66.193 → 4.479 µs, −93.2% - timestamp_us/time64_ns/timezone_control: 87.896 → 89.171 µs, +1.5% - timestamp_ns/date32/no_nulls: 100.430 → 2.394 µs, −97.6% - timestamp_ns/date32/mixed_nulls: 82.835 → 2.397 µs, −97.1% - timestamp_ns/date32/timezone_control: 160.090 → 160.820 µs, +0.5% - timestamp_ns/time32_s/no_nulls: 77.068 → 4.346 µs, −94.4% - timestamp_ns/time32_s/mixed_nulls: 63.343 → 4.560 µs, −92.8% - timestamp_ns/time32_s/timezone_control: 84.787 → 85.908 µs, +1.3% - timestamp_ns/time32_ms/no_nulls: 78.946 → 4.114 µs, −94.8% - timestamp_ns/time32_ms/mixed_nulls: 64.964 → 4.258 µs, −93.4% - timestamp_ns/time32_ms/timezone_control: 87.032 → 89.097 µs, +2.4% - timestamp_ns/time64_us/no_nulls: 79.828 → 5.607 µs, −93.0% - timestamp_ns/time64_us/mixed_nulls: 65.953 → 5.651 µs, −91.4% - timestamp_ns/time64_us/timezone_control: 88.268 → 90.504 µs, +2.5% - timestamp_ns/time64_ns/no_nulls: 78.514 → 2.372 µs, −97.0% - timestamp_ns/time64_ns/mixed_nulls: 65.090 → 2.371 µs, −96.4% - timestamp_ns/time64_ns/timezone_control: 87.227 → 88.661 µs, +1.6% # What changes are included in this PR? * Optimize casting timestamp without timezone to time-of-day / days-since-epoch * Add tests ensuring that optimized code path preserves the same behavior as the Chrono code path * Add benchmarks # Are these changes tested? Yes; existing tests pass, new tests added. # Are there any user-facing changes? Previously, attempting to cast a timestamp value that was outside Chrono's supported calendar range produced an error. There are two behavior changes here: * When taking the optimized code path (the timestamp does not have a timezone and the target of the cast is `Date32`, `Time32`, or `Time64`), the cast will succeed and produce a correct result value * Otherwise, "safe" mode will be respected. --------- Co-authored-by: Jeffrey Vo <[email protected]> --- arrow-cast/Cargo.toml | 4 + arrow-cast/benches/cast_timestamp.rs | 75 +++++ arrow-cast/src/cast/mod.rs | 632 ++++++++++++++++++++++------------- 3 files changed, 471 insertions(+), 240 deletions(-) diff --git a/arrow-cast/Cargo.toml b/arrow-cast/Cargo.toml index 94a17fa2a8..dfe10803c6 100644 --- a/arrow-cast/Cargo.toml +++ b/arrow-cast/Cargo.toml @@ -87,6 +87,10 @@ harness = false name = "format_temporal" harness = false +[[bench]] +name = "cast_timestamp" +harness = false + [lints] workspace = true diff --git a/arrow-cast/benches/cast_timestamp.rs b/arrow-cast/benches/cast_timestamp.rs new file mode 100644 index 0000000000..5e516aae8b --- /dev/null +++ b/arrow-cast/benches/cast_timestamp.rs @@ -0,0 +1,75 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; + +use arrow_array::{PrimitiveArray, types::*}; +use arrow_cast::cast; +use arrow_schema::{DataType, TimeUnit}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; + +const SIZE: usize = 8192; + +fn benchmark<T: ArrowTimestampType>(c: &mut Criterion, unit: &str, units_per_second: i64) { + // Timestamps spread over roughly 63 years either side of the epoch. + let mut rng = StdRng::seed_from_u64(42); + let range = -2_000_000_000 * units_per_second..2_000_000_000 * units_per_second; + let values: Vec<i64> = (0..SIZE).map(|_| rng.random_range(range.clone())).collect(); + let mixed_nulls = PrimitiveArray::<T>::from_iter( + values + .iter() + .enumerate() + .map(|(i, &v)| (i % 5 != 0).then_some(v)), + ); + let no_nulls = PrimitiveArray::<T>::new(values.into(), None); + // Timezone-aware casts go through Chrono and serve as a control. + let zoned = no_nulls.clone().with_timezone("+05:45"); + let inputs = [ + ("no_nulls", &no_nulls), + ("mixed_nulls", &mixed_nulls), + ("timezone_control", &zoned), + ]; + let targets = [ + ("date32", DataType::Date32), + ("time32_s", DataType::Time32(TimeUnit::Second)), + ("time32_ms", DataType::Time32(TimeUnit::Millisecond)), + ("time64_us", DataType::Time64(TimeUnit::Microsecond)), + ("time64_ns", DataType::Time64(TimeUnit::Nanosecond)), + ]; + let mut group = c.benchmark_group(format!("timestamp_{unit}")); + group.throughput(Throughput::Elements(SIZE as u64)); + for (name, target) in &targets { + for (input, array) in inputs { + group.bench_function(format!("{name}/{input}"), |b| { + b.iter(|| black_box(cast(black_box(array), target).unwrap())) + }); + } + } + group.finish(); +} + +fn cast_timestamp(c: &mut Criterion) { + benchmark::<TimestampSecondType>(c, "s", 1); + benchmark::<TimestampMillisecondType>(c, "ms", 1_000); + benchmark::<TimestampMicrosecondType>(c, "us", 1_000_000); + benchmark::<TimestampNanosecondType>(c, "ns", 1_000_000_000); +} + +criterion_group!(benches, cast_timestamp); +criterion_main!(benches); diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs index 6cf6fcc5cd..cb347d4e98 100644 --- a/arrow-cast/src/cast/mod.rs +++ b/arrow-cast/src/cast/mod.rs @@ -611,30 +611,128 @@ 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>() )) }; @@ -642,17 +740,41 @@ fn timestamp_to_date32<T: ArrowTimestampType>( 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. + let days = |x: i64| x.div_euclid(units_per_day::<T>()); + 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 = units_per_day::<T>(); + let (lo, hi) = (i32::MIN as i64 * day, (i32::MAX as i64 + 1) * day); + array + .values() + .iter() + .fold(true, |ok, &x| ok & (lo <= x) & (x < hi)) + } + }; + if all_in_range { + array.unary(|x| days(x) as i32) + } else if cast_options.safe { + array.unary_opt(|x| i32::try_from(days(x)).ok()) + } else { + array.try_unary(|x| i32::try_from(days(x)).map_err(|_| err(x)))? + } } - None => array.try_unary(|x| { - as_datetime::<T>(x) - .ok_or_else(|| err(x)) - .map(|d| Date32Type::from_naive_date(d.date())) - })?, }; Ok(Arc::new(array)) } @@ -677,6 +799,10 @@ fn timestamp_to_date32<T: ArrowTimestampType>( /// * `Date32` and `Date64`: precision lost when going to higher interval /// * `Time32` and `Time64`: precision lost when going to higher interval /// * `Timestamp` and `Date{32|64}`: precision lost when going to higher interval +/// * `Timestamp` to `Date32`, `Time32`, or `Time64`: timestamps without a timezone +/// support the full `i64` range; timestamps with a timezone must lie within Chrono's +/// date range. If safe is true, a timestamp outside that range or a `Date32` day +/// count that does not fit in `i32` becomes NULL, otherwise an error is returned /// * Temporal to/from backing Primitive: zero-copy with data type change /// * `Float16/Float32/Float64` to `Decimal(precision, scale)` rounds to the `scale` decimals /// (i.e. casting `6.4999` to `Decimal(10, 1)` becomes `6.5`). @@ -1958,18 +2084,21 @@ pub fn cast_with_options( }; Ok(make_timestamp_array(&adjusted, *to_unit, to_tz.clone())) } - (Timestamp(TimeUnit::Microsecond, _), Date32) => { - timestamp_to_date32(array.as_primitive::<TimestampMicrosecondType>()) - } - (Timestamp(TimeUnit::Millisecond, _), Date32) => { - timestamp_to_date32(array.as_primitive::<TimestampMillisecondType>()) - } + (Timestamp(TimeUnit::Microsecond, _), Date32) => timestamp_to_date32( + array.as_primitive::<TimestampMicrosecondType>(), + cast_options, + ), + (Timestamp(TimeUnit::Millisecond, _), Date32) => timestamp_to_date32( + array.as_primitive::<TimestampMillisecondType>(), + cast_options, + ), (Timestamp(TimeUnit::Second, _), Date32) => { - timestamp_to_date32(array.as_primitive::<TimestampSecondType>()) - } - (Timestamp(TimeUnit::Nanosecond, _), Date32) => { - timestamp_to_date32(array.as_primitive::<TimestampNanosecondType>()) + timestamp_to_date32(array.as_primitive::<TimestampSecondType>(), cast_options) } + (Timestamp(TimeUnit::Nanosecond, _), Date32) => timestamp_to_date32( + array.as_primitive::<TimestampNanosecondType>(), + cast_options, + ), (Timestamp(TimeUnit::Second, _), Date64) => Ok(Arc::new(match cast_options.safe { true => { // change error to None @@ -1994,197 +2123,17 @@ pub fn cast_with_options( .as_primitive::<TimestampNanosecondType>() .unary::<_, Date64Type>(|x| x / (NANOSECONDS / MILLISECONDS)), )), - (Timestamp(TimeUnit::Second, tz), Time64(TimeUnit::Microsecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampSecondType>() - .try_unary::<_, Time64MicrosecondType, ArrowError>(|x| { - Ok(time_to_time64us(as_time_res_with_timezone::< - TimestampSecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Second, tz), Time64(TimeUnit::Nanosecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampSecondType>() - .try_unary::<_, Time64NanosecondType, ArrowError>(|x| { - Ok(time_to_time64ns(as_time_res_with_timezone::< - TimestampSecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Millisecond, tz), Time64(TimeUnit::Microsecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampMillisecondType>() - .try_unary::<_, Time64MicrosecondType, ArrowError>(|x| { - Ok(time_to_time64us(as_time_res_with_timezone::< - TimestampMillisecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Millisecond, tz), Time64(TimeUnit::Nanosecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampMillisecondType>() - .try_unary::<_, Time64NanosecondType, ArrowError>(|x| { - Ok(time_to_time64ns(as_time_res_with_timezone::< - TimestampMillisecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Microsecond, tz), Time64(TimeUnit::Microsecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampMicrosecondType>() - .try_unary::<_, Time64MicrosecondType, ArrowError>(|x| { - Ok(time_to_time64us(as_time_res_with_timezone::< - TimestampMicrosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Microsecond, tz), Time64(TimeUnit::Nanosecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampMicrosecondType>() - .try_unary::<_, Time64NanosecondType, ArrowError>(|x| { - Ok(time_to_time64ns(as_time_res_with_timezone::< - TimestampMicrosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Nanosecond, tz), Time64(TimeUnit::Microsecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampNanosecondType>() - .try_unary::<_, Time64MicrosecondType, ArrowError>(|x| { - Ok(time_to_time64us(as_time_res_with_timezone::< - TimestampNanosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Nanosecond, tz), Time64(TimeUnit::Nanosecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampNanosecondType>() - .try_unary::<_, Time64NanosecondType, ArrowError>(|x| { - Ok(time_to_time64ns(as_time_res_with_timezone::< - TimestampNanosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Second, tz), Time32(TimeUnit::Second)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampSecondType>() - .try_unary::<_, Time32SecondType, ArrowError>(|x| { - Ok(time_to_time32s(as_time_res_with_timezone::< - TimestampSecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Second, tz), Time32(TimeUnit::Millisecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampSecondType>() - .try_unary::<_, Time32MillisecondType, ArrowError>(|x| { - Ok(time_to_time32ms(as_time_res_with_timezone::< - TimestampSecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Millisecond, tz), Time32(TimeUnit::Second)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampMillisecondType>() - .try_unary::<_, Time32SecondType, ArrowError>(|x| { - Ok(time_to_time32s(as_time_res_with_timezone::< - TimestampMillisecondType, - >(x, tz)?)) - })?, - )) + (Timestamp(TimeUnit::Second, _), Time32(_) | Time64(_)) => { + cast_timestamp_to_time::<TimestampSecondType>(array, to_type, cast_options) } - (Timestamp(TimeUnit::Millisecond, tz), Time32(TimeUnit::Millisecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampMillisecondType>() - .try_unary::<_, Time32MillisecondType, ArrowError>(|x| { - Ok(time_to_time32ms(as_time_res_with_timezone::< - TimestampMillisecondType, - >(x, tz)?)) - })?, - )) + (Timestamp(TimeUnit::Millisecond, _), Time32(_) | Time64(_)) => { + cast_timestamp_to_time::<TimestampMillisecondType>(array, to_type, cast_options) } - (Timestamp(TimeUnit::Microsecond, tz), Time32(TimeUnit::Second)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampMicrosecondType>() - .try_unary::<_, Time32SecondType, ArrowError>(|x| { - Ok(time_to_time32s(as_time_res_with_timezone::< - TimestampMicrosecondType, - >(x, tz)?)) - })?, - )) + (Timestamp(TimeUnit::Microsecond, _), Time32(_) | Time64(_)) => { + cast_timestamp_to_time::<TimestampMicrosecondType>(array, to_type, cast_options) } - (Timestamp(TimeUnit::Microsecond, tz), Time32(TimeUnit::Millisecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampMicrosecondType>() - .try_unary::<_, Time32MillisecondType, ArrowError>(|x| { - Ok(time_to_time32ms(as_time_res_with_timezone::< - TimestampMicrosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Nanosecond, tz), Time32(TimeUnit::Second)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampNanosecondType>() - .try_unary::<_, Time32SecondType, ArrowError>(|x| { - Ok(time_to_time32s(as_time_res_with_timezone::< - TimestampNanosecondType, - >(x, tz)?)) - })?, - )) - } - (Timestamp(TimeUnit::Nanosecond, tz), Time32(TimeUnit::Millisecond)) => { - let tz = tz.as_ref().map(|tz| tz.parse()).transpose()?; - Ok(Arc::new( - array - .as_primitive::<TimestampNanosecondType>() - .try_unary::<_, Time32MillisecondType, ArrowError>(|x| { - Ok(time_to_time32ms(as_time_res_with_timezone::< - TimestampNanosecondType, - >(x, tz)?)) - })?, - )) + (Timestamp(TimeUnit::Nanosecond, _), Time32(_) | Time64(_)) => { + cast_timestamp_to_time::<TimestampNanosecondType>(array, to_type, cast_options) } (Date64, Timestamp(TimeUnit::Second, _)) => { let array = array @@ -2491,6 +2440,11 @@ where } } +/// Number of `T` units in one day +const fn units_per_day<T: ArrowTimestampType>() -> i64 { + SECONDS_IN_DAY * time_unit_multiple(&T::UNIT) +} + /// Get the time unit as a multiple of a second const fn time_unit_multiple(unit: &TimeUnit) -> i64 { match unit { @@ -2798,7 +2752,7 @@ mod tests { use arrow_buffer::{BooleanBuffer, Buffer, IntervalDayTime, NullBuffer}; use arrow_buffer::{ScalarBuffer, i256}; use arrow_schema::{DataType, Field}; - use chrono::NaiveDate; + use chrono::{NaiveDate, NaiveTime}; use half::f16; use std::sync::Arc; @@ -6665,6 +6619,185 @@ mod tests { assert_eq!(17890, c.value(1)); } + /// Compares arithmetic casts with Chrono for representable dates, including + /// negative and fractional boundaries, nulls, and nonzero-offset slices. + fn check_timestamp_casts<T: ArrowTimestampType>() { + use rand::rngs::StdRng; + use rand::{RngExt, SeedableRng}; + + let units = time_unit_multiple(&T::UNIT); + let around = |x: i64| [x - 1, x, x + 1, -x - 1, -x, -x + 1]; + let mut values = vec![i64::MIN, i64::MAX]; + for base in [0, 60, 3_600, 86_400] { + for seconds in around(base) { + for fraction in [-1, 0, 1, units - 1] { + values.push(seconds * units + fraction); + } + } + } + for boundary in [1_000, 1_000_000] { + values.extend(around(boundary)); + } + let mut rng = StdRng::seed_from_u64(42); + for _ in 0..256 { + values.push(rng.random()); + let seconds = rng.random_range(-2_000_000_000..2_000_000_000); + values.push(seconds * units + rng.random_range(0..units)); + } + values.retain(|&v| as_datetime::<T>(v).is_some()); + let no_nulls = PrimitiveArray::<T>::new(values.clone().into(), None); + let mixed = PrimitiveArray::<T>::from_iter( + values + .iter() + .enumerate() + .map(|(i, &v)| (i % 5 != 0).then_some(v)), + ); + for array in [no_nulls, mixed.slice(1, mixed.len() - 2)] { + let dates: Vec<_> = array + .iter() + .map(|v| v.map(|v| as_datetime::<T>(v).unwrap())) + .collect(); + let expected: [ArrayRef; 5] = [ + Arc::new(Date32Array::from_iter( + dates + .iter() + .map(|d| d.map(|d| Date32Type::from_naive_date(d.date()))), + )), + Arc::new(Time32SecondArray::from_iter( + dates.iter().map(|d| d.map(|d| time_to_time32s(d.time()))), + )), + Arc::new(Time32MillisecondArray::from_iter( + dates.iter().map(|d| d.map(|d| time_to_time32ms(d.time()))), + )), + Arc::new(Time64MicrosecondArray::from_iter( + dates.iter().map(|d| d.map(|d| time_to_time64us(d.time()))), + )), + Arc::new(Time64NanosecondArray::from_iter( + dates.iter().map(|d| d.map(|d| time_to_time64ns(d.time()))), + )), + ]; + for expected in expected { + let actual = cast(&array, expected.data_type()).unwrap(); + assert_eq!(actual.as_ref(), expected.as_ref(), "{}", T::DATA_TYPE); + } + } + } + + #[test] + fn test_cast_timestamp_without_timezone_matches_chrono() { + check_timestamp_casts::<TimestampSecondType>(); + check_timestamp_casts::<TimestampMillisecondType>(); + check_timestamp_casts::<TimestampMicrosecondType>(); + check_timestamp_casts::<TimestampNanosecondType>(); + } + + #[test] + fn test_cast_timestamp_to_time_extremes() { + fn check<T: ArrowTimestampType>(min: (u32, u32, u32, u32), max: (u32, u32, u32, u32)) { + let time = |(h, m, s, ns)| NaiveTime::from_hms_nano_opt(h, m, s, ns).unwrap(); + let times = [Some(time(min)), None, Some(time(max))]; + let array = PrimitiveArray::<T>::from_iter([Some(i64::MIN), None, Some(i64::MAX)]); + let expected: [ArrayRef; 4] = [ + Arc::new(Time32SecondArray::from_iter( + times.map(|v| v.map(time_to_time32s)), + )), + Arc::new(Time32MillisecondArray::from_iter( + times.map(|v| v.map(time_to_time32ms)), + )), + Arc::new(Time64MicrosecondArray::from_iter( + times.map(|v| v.map(time_to_time64us)), + )), + Arc::new(Time64NanosecondArray::from_iter( + times.map(|v| v.map(time_to_time64ns)), + )), + ]; + for expected in expected { + assert_eq!( + cast(&array, expected.data_type()).unwrap().as_ref(), + expected.as_ref() + ); + } + } + check::<TimestampSecondType>((8, 29, 52, 0), (15, 30, 7, 0)); + check::<TimestampMillisecondType>((16, 47, 4, 192_000_000), (7, 12, 55, 807_000_000)); + check::<TimestampMicrosecondType>((19, 59, 5, 224_192_000), (4, 0, 54, 775_807_000)); + // Nanosecond extremes are representable in Chrono and covered above. + } + + #[test] + fn test_cast_timestamp_to_date32_range() { + fn check<T: ArrowTimestampType>() { + let day = SECONDS_IN_DAY * time_unit_multiple(&T::UNIT); + let first = i32::MIN as i64 * day; + let last = (i32::MAX as i64 + 1) * day - 1; + let array = PrimitiveArray::<T>::new( + vec![first, i64::MAX, last].into(), + Some(NullBuffer::from(vec![true, false, true])), + ); + let expected = Date32Array::from(vec![Some(i32::MIN), None, Some(i32::MAX)]); + for safe in [true, false] { + let options = CastOptions { + safe, + ..Default::default() + }; + assert_eq!( + cast_with_options(&array, &Date32, &options) + .unwrap() + .as_ref(), + &expected + ); + } + // Day counts that do not fit in i32 become null when safe, otherwise an error. + let array = PrimitiveArray::<T>::from_iter_values([ + first - 1, + last, + last + 1, + i64::MIN, + i64::MAX, + ]); + let expected = Date32Array::from(vec![None, Some(i32::MAX), None, None, None]); + assert_eq!(cast(&array, &Date32).unwrap().as_ref(), &expected); + let options = CastOptions { + safe: false, + ..Default::default() + }; + assert!(cast_with_options(&array, &Date32, &options).is_err()); + } + check::<TimestampSecondType>(); + check::<TimestampMillisecondType>(); + // The entire microsecond domain fits in Date32, including dates outside Chrono's range. + let array = TimestampMicrosecondArray::from(vec![Some(i64::MIN), None, Some(i64::MAX)]); + assert_eq!( + cast(&array, &Date32).unwrap().as_ref(), + &Date32Array::from(vec![Some(-106_751_992), None, Some(106_751_991)]), + ); + } + + #[test] + fn test_cast_timestamp_date_time_timezone_validation() { + let invalid = TimestampSecondArray::from(vec![0]).with_timezone("invalid timezone"); + let extreme = TimestampSecondArray::from(vec![0, i64::MAX]).with_timezone("+00:00"); + let options = CastOptions { + safe: false, + ..Default::default() + }; + for target in [ + Date32, + Time32(TimeUnit::Second), + Time32(TimeUnit::Millisecond), + Time64(TimeUnit::Microsecond), + Time64(TimeUnit::Nanosecond), + ] { + // A timezone that fails to parse is an error even when the cast is safe. + assert!(cast(&invalid, &target).is_err()); + // A timestamp outside Chrono's date range is null when safe, otherwise an error. + let actual = cast(&extreme, &target).unwrap(); + assert!(actual.is_valid(0)); + assert!(actual.is_null(1)); + assert!(cast_with_options(&extreme, &target, &options).is_err()); + } + } + #[test] fn test_cast_timestamp_to_date32() { let array = @@ -6772,29 +6905,39 @@ mod tests { assert_eq!(3601000000000, c.value(1)); assert!(c.is_null(2)); - // test timestamp nanoseconds - let a = TimestampNanosecondArray::from(vec![Some(86405000000000), Some(1000000000), None]) + // Test fractional timestamps on both sides of the epoch. + let a = TimestampNanosecondArray::from(vec![Some(86_405_123_456_789), Some(-1), None]) .with_timezone("+01:00".to_string()); let array = Arc::new(a) as ArrayRef; let b = cast(&array, &DataType::Time64(TimeUnit::Microsecond)).unwrap(); let c = b.as_primitive::<Time64MicrosecondType>(); - assert_eq!(3605000000, c.value(0)); - assert_eq!(3601000000, c.value(1)); + assert_eq!(3_605_123_456, c.value(0)); + assert_eq!(3_599_999_999, c.value(1)); assert!(c.is_null(2)); let b = cast(&array, &DataType::Time64(TimeUnit::Nanosecond)).unwrap(); let c = b.as_primitive::<Time64NanosecondType>(); - assert_eq!(3605000000000, c.value(0)); - assert_eq!(3601000000000, c.value(1)); + assert_eq!(3_605_123_456_789, c.value(0)); + assert_eq!(3_599_999_999_999, c.value(1)); assert!(c.is_null(2)); - // test overflow + // test overflow, safe cast let a = TimestampSecondArray::from(vec![Some(i64::MAX)]).with_timezone("+01:00".to_string()); let array = Arc::new(a) as ArrayRef; - let b = cast(&array, &DataType::Time64(TimeUnit::Microsecond)); + let b = cast(&array, &DataType::Time64(TimeUnit::Microsecond)).unwrap(); + assert!(b.is_null(0)); + let b = cast(&array, &DataType::Time64(TimeUnit::Nanosecond)).unwrap(); + assert!(b.is_null(0)); + // test overflow, unsafe cast + let options = CastOptions { + safe: false, + ..Default::default() + }; + let b = cast_with_options(&array, &DataType::Time64(TimeUnit::Microsecond), &options); assert!(b.is_err()); - let b = cast(&array, &DataType::Time64(TimeUnit::Nanosecond)); + let b = cast_with_options(&array, &DataType::Time64(TimeUnit::Nanosecond), &options); assert!(b.is_err()); + // Time64 only supports microseconds and nanoseconds let b = cast(&array, &DataType::Time64(TimeUnit::Millisecond)); assert!(b.is_err()); } @@ -6846,28 +6989,37 @@ mod tests { assert_eq!(3601000, c.value(1)); assert!(c.is_null(2)); - // test timestamp nanoseconds - let a = TimestampNanosecondArray::from(vec![Some(86405000000000), Some(1000000000), None]) + // Test fractional timestamps on both sides of the epoch. + let a = TimestampNanosecondArray::from(vec![Some(86_405_123_456_789), Some(-1), None]) .with_timezone("+01:00".to_string()); let array = Arc::new(a) as ArrayRef; let b = cast(&array, &DataType::Time32(TimeUnit::Second)).unwrap(); let c = b.as_primitive::<Time32SecondType>(); assert_eq!(3605, c.value(0)); - assert_eq!(3601, c.value(1)); + assert_eq!(3599, c.value(1)); assert!(c.is_null(2)); let b = cast(&array, &DataType::Time32(TimeUnit::Millisecond)).unwrap(); let c = b.as_primitive::<Time32MillisecondType>(); - assert_eq!(3605000, c.value(0)); - assert_eq!(3601000, c.value(1)); + assert_eq!(3_605_123, c.value(0)); + assert_eq!(3_599_999, c.value(1)); assert!(c.is_null(2)); - // test overflow + // test overflow, safe cast let a = TimestampSecondArray::from(vec![Some(i64::MAX)]).with_timezone("+01:00".to_string()); let array = Arc::new(a) as ArrayRef; - let b = cast(&array, &DataType::Time32(TimeUnit::Second)); + let b = cast(&array, &DataType::Time32(TimeUnit::Second)).unwrap(); + assert!(b.is_null(0)); + let b = cast(&array, &DataType::Time32(TimeUnit::Millisecond)).unwrap(); + assert!(b.is_null(0)); + // test overflow, unsafe cast + let options = CastOptions { + safe: false, + ..Default::default() + }; + let b = cast_with_options(&array, &DataType::Time32(TimeUnit::Second), &options); assert!(b.is_err()); - let b = cast(&array, &DataType::Time32(TimeUnit::Millisecond)); + let b = cast_with_options(&array, &DataType::Time32(TimeUnit::Millisecond), &options); assert!(b.is_err()); }
