arrow_cast/cast/
temporal.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Cast support for temporal arrays: timestamps, dates, durations,
19//! intervals, and timezones.
20
21use arrow_array::{cast::*, temporal_conversions::*, timezone::Tz, types::*, *};
22use arrow_buffer::IntervalMonthDayNano;
23use arrow_schema::{ArrowError, DataType, TimeUnit};
24use chrono::{FixedOffset, LocalResult, NaiveDateTime, NaiveTime, Offset, TimeDelta, TimeZone};
25use std::sync::Arc;
26
27use num_traits::AsPrimitive;
28
29use super::{CastOptions, time_unit_multiple, units_per_day};
30
31/// Cast the array from interval year month to month day nano
32pub(crate) fn cast_interval_year_month_to_interval_month_day_nano(
33    array: &dyn Array,
34    _cast_options: &CastOptions,
35) -> Result<ArrayRef, ArrowError> {
36    let array = array.as_primitive::<IntervalYearMonthType>();
37
38    Ok(Arc::new(array.unary::<_, IntervalMonthDayNanoType>(|v| {
39        let months = IntervalYearMonthType::to_months(v);
40        IntervalMonthDayNanoType::make_value(months, 0, 0)
41    })))
42}
43
44/// Cast the array from interval day time to month day nano
45pub(crate) fn cast_interval_day_time_to_interval_month_day_nano(
46    array: &dyn Array,
47    _cast_options: &CastOptions,
48) -> Result<ArrayRef, ArrowError> {
49    let array = array.as_primitive::<IntervalDayTimeType>();
50    let mul = 1_000_000;
51
52    Ok(Arc::new(array.unary::<_, IntervalMonthDayNanoType>(|v| {
53        let (days, ms) = IntervalDayTimeType::to_parts(v);
54        IntervalMonthDayNanoType::make_value(0, days, ms as i64 * mul)
55    })))
56}
57
58/// Cast the array from interval to duration
59pub(crate) fn cast_month_day_nano_to_duration<D: ArrowTemporalType<Native = i64>>(
60    array: &dyn Array,
61    cast_options: &CastOptions,
62) -> Result<ArrayRef, ArrowError> {
63    let array = array.as_primitive::<IntervalMonthDayNanoType>();
64    let scale = match D::DATA_TYPE {
65        DataType::Duration(TimeUnit::Second) => 1_000_000_000,
66        DataType::Duration(TimeUnit::Millisecond) => 1_000_000,
67        DataType::Duration(TimeUnit::Microsecond) => 1_000,
68        DataType::Duration(TimeUnit::Nanosecond) => 1,
69        _ => unreachable!(),
70    };
71
72    if cast_options.safe {
73        let iter = array.iter().map(|v| {
74            let v = v?;
75            (v.days == 0 && v.months == 0).then_some(v.nanoseconds / scale)
76        });
77        Ok(Arc::new(unsafe {
78            PrimitiveArray::<D>::from_trusted_len_iter(iter)
79        }))
80    } else {
81        let vec = array
82            .iter()
83            .map(|v| {
84                v.map(|v| match v.days == 0 && v.months == 0 {
85                    true => Ok((v.nanoseconds) / scale),
86                    _ => Err(ArrowError::ComputeError(
87                        "Cannot convert interval containing non-zero months or days to duration"
88                            .to_string(),
89                    )),
90                })
91                .transpose()
92            })
93            .collect::<Result<Vec<_>, _>>()?;
94        Ok(Arc::new(unsafe {
95            PrimitiveArray::<D>::from_trusted_len_iter(vec.iter())
96        }))
97    }
98}
99
100/// Cast the array from duration and interval
101pub(crate) fn cast_duration_to_interval<D: ArrowTemporalType<Native = i64>>(
102    array: &dyn Array,
103    cast_options: &CastOptions,
104) -> Result<ArrayRef, ArrowError> {
105    let array = array
106        .as_any()
107        .downcast_ref::<PrimitiveArray<D>>()
108        .ok_or_else(|| {
109            ArrowError::ComputeError(
110                "Internal Error: Cannot cast duration to DurationArray of expected type"
111                    .to_string(),
112            )
113        })?;
114
115    let scale = match array.data_type() {
116        DataType::Duration(TimeUnit::Second) => 1_000_000_000,
117        DataType::Duration(TimeUnit::Millisecond) => 1_000_000,
118        DataType::Duration(TimeUnit::Microsecond) => 1_000,
119        DataType::Duration(TimeUnit::Nanosecond) => 1,
120        _ => unreachable!(),
121    };
122
123    if cast_options.safe {
124        let iter = array.iter().map(|v| {
125            v?.checked_mul(scale)
126                .map(|v| IntervalMonthDayNano::new(0, 0, v))
127        });
128        Ok(Arc::new(unsafe {
129            PrimitiveArray::<IntervalMonthDayNanoType>::from_trusted_len_iter(iter)
130        }))
131    } else {
132        let vec = array
133            .iter()
134            .map(|v| {
135                v.map(|v| {
136                    if let Ok(v) = v.mul_checked(scale) {
137                        Ok(IntervalMonthDayNano::new(0, 0, v))
138                    } else {
139                        Err(ArrowError::ComputeError(format!(
140                            "Cannot cast to {:?}. Overflowing on {:?}",
141                            IntervalMonthDayNanoType::DATA_TYPE,
142                            v
143                        )))
144                    }
145                })
146                .transpose()
147            })
148            .collect::<Result<Vec<_>, _>>()?;
149        Ok(Arc::new(unsafe {
150            PrimitiveArray::<IntervalMonthDayNanoType>::from_trusted_len_iter(vec.iter())
151        }))
152    }
153}
154
155pub(crate) fn make_timestamp_array(
156    array: &PrimitiveArray<Int64Type>,
157    unit: TimeUnit,
158    tz: Option<Arc<str>>,
159) -> ArrayRef {
160    match unit {
161        TimeUnit::Second => Arc::new(
162            array
163                .reinterpret_cast::<TimestampSecondType>()
164                .with_timezone_opt(tz),
165        ),
166        TimeUnit::Millisecond => Arc::new(
167            array
168                .reinterpret_cast::<TimestampMillisecondType>()
169                .with_timezone_opt(tz),
170        ),
171        TimeUnit::Microsecond => Arc::new(
172            array
173                .reinterpret_cast::<TimestampMicrosecondType>()
174                .with_timezone_opt(tz),
175        ),
176        TimeUnit::Nanosecond => Arc::new(
177            array
178                .reinterpret_cast::<TimestampNanosecondType>()
179                .with_timezone_opt(tz),
180        ),
181    }
182}
183
184pub(crate) fn make_duration_array(array: &PrimitiveArray<Int64Type>, unit: TimeUnit) -> ArrayRef {
185    match unit {
186        TimeUnit::Second => Arc::new(array.reinterpret_cast::<DurationSecondType>()),
187        TimeUnit::Millisecond => Arc::new(array.reinterpret_cast::<DurationMillisecondType>()),
188        TimeUnit::Microsecond => Arc::new(array.reinterpret_cast::<DurationMicrosecondType>()),
189        TimeUnit::Nanosecond => Arc::new(array.reinterpret_cast::<DurationNanosecondType>()),
190    }
191}
192
193pub(crate) fn cast_timestamp_to_time<T: ArrowTimestampType>(
194    array: &dyn Array,
195    to_type: &DataType,
196    cast_options: &CastOptions,
197) -> Result<ArrayRef, ArrowError> {
198    let array = array.as_primitive::<T>();
199    match to_type {
200        DataType::Time32(TimeUnit::Second) => {
201            timestamp_to_time::<T, Time32SecondType>(array, cast_options)
202        }
203        DataType::Time32(TimeUnit::Millisecond) => {
204            timestamp_to_time::<T, Time32MillisecondType>(array, cast_options)
205        }
206        DataType::Time64(TimeUnit::Microsecond) => {
207            timestamp_to_time::<T, Time64MicrosecondType>(array, cast_options)
208        }
209        DataType::Time64(TimeUnit::Nanosecond) => {
210            timestamp_to_time::<T, Time64NanosecondType>(array, cast_options)
211        }
212        _ => Err(ArrowError::CastError(format!(
213            "Casting from {} to {to_type} not supported",
214            array.data_type()
215        ))),
216    }
217}
218
219trait TimeType: ArrowTemporalType {
220    /// Number of units in one second.
221    const UNIT_MULTIPLE: i64;
222
223    fn from_naive_time(time: NaiveTime) -> Self::Native;
224}
225
226impl TimeType for Time32SecondType {
227    const UNIT_MULTIPLE: i64 = 1;
228
229    fn from_naive_time(time: NaiveTime) -> Self::Native {
230        time_to_time32s(time)
231    }
232}
233
234impl TimeType for Time32MillisecondType {
235    const UNIT_MULTIPLE: i64 = MILLISECONDS;
236
237    fn from_naive_time(time: NaiveTime) -> Self::Native {
238        time_to_time32ms(time)
239    }
240}
241
242impl TimeType for Time64MicrosecondType {
243    const UNIT_MULTIPLE: i64 = MICROSECONDS;
244
245    fn from_naive_time(time: NaiveTime) -> Self::Native {
246        time_to_time64us(time)
247    }
248}
249
250impl TimeType for Time64NanosecondType {
251    const UNIT_MULTIPLE: i64 = NANOSECONDS;
252
253    fn from_naive_time(time: NaiveTime) -> Self::Native {
254        time_to_time64ns(time)
255    }
256}
257
258fn timestamp_to_time<T, O>(
259    array: &PrimitiveArray<T>,
260    cast_options: &CastOptions,
261) -> Result<ArrayRef, ArrowError>
262where
263    T: ArrowTimestampType,
264    O: TimeType,
265    i64: AsPrimitive<O::Native>,
266{
267    // A time within one day fits its Time32 or Time64 representation.
268    let array = match array.timezone() {
269        Some(tz) => {
270            let tz: Tz = tz.parse()?;
271            let time = |v: i64| {
272                as_datetime_with_timezone::<T>(v, tz).map(|d| O::from_naive_time(d.time()))
273            };
274            if cast_options.safe {
275                array.unary_opt::<_, O>(time)
276            } else {
277                array.try_unary::<_, O, _>(|v| {
278                    time(v).ok_or_else(|| {
279                        ArrowError::CastError(format!(
280                            "Failed to create naive time with {} {}",
281                            std::any::type_name::<T>(),
282                            v
283                        ))
284                    })
285                })?
286            }
287        }
288        None => array.unary::<_, O>(|v| {
289            // The remainder within a day is the time of day; `rem_euclid` keeps it
290            // nonnegative for timestamps before the epoch. The units are constants,
291            // so the branch and the divisions fold at compile time.
292            let from = time_unit_multiple(&T::UNIT);
293            let time = v.rem_euclid(SECONDS_IN_DAY * from);
294            let time = if from >= O::UNIT_MULTIPLE {
295                time / (from / O::UNIT_MULTIPLE)
296            } else {
297                time * (O::UNIT_MULTIPLE / from)
298            };
299            time.as_()
300        }),
301    };
302    Ok(Arc::new(array))
303}
304
305pub(crate) fn timestamp_to_date32<T: ArrowTimestampType>(
306    array: &PrimitiveArray<T>,
307    cast_options: &CastOptions,
308) -> Result<ArrayRef, ArrowError> {
309    let err = |x: i64| {
310        ArrowError::CastError(format!(
311            "Cannot convert {} {x} to Date32",
312            std::any::type_name::<T>()
313        ))
314    };
315
316    let array: Date32Array = match array.timezone() {
317        Some(tz) => {
318            let tz: Tz = tz.parse()?;
319            let date = |x: i64| {
320                as_datetime_with_timezone::<T>(x, tz)
321                    .map(|d| Date32Type::from_naive_date(d.date_naive()))
322            };
323            if cast_options.safe {
324                array.unary_opt(date)
325            } else {
326                array.try_unary(|x| date(x).ok_or_else(|| err(x)))?
327            }
328        }
329        None => {
330            // Date32 stores days since the epoch. Round down so that a timestamp
331            // just before the epoch belongs to the preceding day.
332            let days = |x: i64| x.div_euclid(units_per_day::<T>());
333            let all_in_range = match T::UNIT {
334                // Every microsecond or nanosecond timestamp lies within the Date32 range.
335                TimeUnit::Microsecond | TimeUnit::Nanosecond => true,
336                // A branch-free scan lets the common case skip the per-value check.
337                _ => {
338                    let day = units_per_day::<T>();
339                    let (lo, hi) = (i32::MIN as i64 * day, (i32::MAX as i64 + 1) * day);
340                    array
341                        .values()
342                        .iter()
343                        .fold(true, |ok, &x| ok & (lo <= x) & (x < hi))
344                }
345            };
346            if all_in_range {
347                array.unary(|x| days(x) as i32)
348            } else if cast_options.safe {
349                array.unary_opt(|x| i32::try_from(days(x)).ok())
350            } else {
351                array.try_unary(|x| i32::try_from(days(x)).map_err(|_| err(x)))?
352            }
353        }
354    };
355    Ok(Arc::new(array))
356}
357
358/// Returns the offset to use when interpreting `local` as a wall clock reading
359/// in `tz`, or `None` if it cannot be resolved.
360///
361/// `None` is not expected in practice. With the current timezone database no
362/// reading reaches it, because every ambiguous or nonexistent reading resolves
363/// as described below. The `None` path is a safeguard against a future
364/// timezone database that breaks the assumptions of the gap handling. Callers
365/// then apply their usual error or null handling.
366///
367/// In an IANA timezone a wall clock reading does not always identify a unique
368/// instant, and this function picks one following the same rules as PostgreSQL
369/// and DuckDB:
370///
371/// * **Ambiguous** -- when the clocks go back ("fall back") the same reading
372///   occurs twice. The *later* instant is chosen, i.e. the offset in effect
373///   after the transition. For example `2024-11-03T01:30:00` in
374///   `America/New_York` is read as `-05:00` (EST), not `-04:00` (EDT).
375/// * **Nonexistent** -- when the clocks go forward ("spring forward") the
376///   reading never occurs. It is shifted forward by the length of the gap,
377///   which is the same as reading it with the offset in effect *before* the
378///   transition. For example `2024-03-10T02:30:00` in `America/New_York` is
379///   read as `-05:00` (EST) and therefore denotes `2024-03-10T03:30:00-04:00`.
380///
381/// Timezones with a fixed offset are never ambiguous and have no gaps.
382///
383/// See <https://github.com/apache/arrow-rs/issues/11037> for the PostgreSQL and
384/// ICU (DuckDB) sources these rules are taken from.
385fn resolve_local_offset(tz: &Tz, local: &NaiveDateTime) -> Option<FixedOffset> {
386    match tz.offset_from_local_datetime(local) {
387        LocalResult::Single(offset) => Some(offset.fix()),
388        LocalResult::Ambiguous(_earlier, later) => Some(later.fix()),
389        LocalResult::None => {
390            // The reading falls in a gap. Recover the offset in effect before
391            // the transition by probing 24 hours earlier.
392            //
393            // Two separate properties of the timezone database make this sound:
394            //
395            // 1. No local gap is longer than 24 hours, so the probe lands
396            //    outside this gap and is itself resolvable. Seven zones have a
397            //    gap of exactly 24 hours -- the dateline changes, such as
398            //    `Pacific/Apia` in 2011 and `Pacific/Kiritimati` in 1994. At the
399            //    last second of one of those the probe lands one second before
400            //    the gap starts, so the true margin here is one second, not a
401            //    comfortable one.
402            // 2. No two transitions are closer together than 24 hours, so the
403            //    offset the probe finds is the one in effect immediately before
404            //    this transition, and not some older offset. The smallest
405            //    observed interval is 167 hours (`America/Boa_Vista`, 2000).
406            //
407            // Property 1 is what makes the probe resolvable; property 2 is what
408            // makes the answer correct. If the probe is still unresolvable, give
409            // up and let the caller apply the usual error / null handling.
410            tz.offset_from_local_datetime(&(*local - TimeDelta::hours(24)))
411                .earliest()
412                .map(|offset| offset.fix())
413        }
414    }
415}
416
417pub(crate) fn adjust_timestamp_to_timezone<T: ArrowTimestampType>(
418    array: PrimitiveArray<Int64Type>,
419    to_tz: &Tz,
420    cast_options: &CastOptions,
421) -> Result<PrimitiveArray<Int64Type>, ArrowError> {
422    let adjust = |o| {
423        let local = as_datetime::<T>(o)?;
424        let offset = resolve_local_offset(to_tz, &local)?;
425        T::from_naive_datetime(local - offset, None)
426    };
427    let adjusted = if cast_options.safe {
428        array.unary_opt::<_, Int64Type>(adjust)
429    } else {
430        array.try_unary::<_, Int64Type, _>(|o| {
431            adjust(o).ok_or_else(|| {
432                ArrowError::CastError("Cannot cast timezone to different timezone".to_string())
433            })
434        })?
435    };
436    Ok(adjusted)
437}
arrow_cast/cast/temporal.rs

arrow_cast/cast/
temporal.rs
