arrow_cast/cast/

temporal.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements

//! Cast support for temporal arrays: timestamps, dates, durations,
//! intervals, and timezones.

use arrow_array::{cast::*, temporal_conversions::*, timezone::Tz, types::*, *};
use arrow_buffer::IntervalMonthDayNano;
use arrow_schema::{ArrowError, DataType, TimeUnit};
use chrono::{NaiveTime, Offset};
use std::sync::Arc;

use num_traits::AsPrimitive;

use crate::local_time::resolve_local_offset;

use super::{CastOptions, time_unit_multiple, units_per_day};

/// Cast the array from interval year month to month day nano
pub(crate) fn cast_interval_year_month_to_interval_month_day_nano(
    array: &dyn Array,
    _cast_options: &CastOptions,
) -> Result<ArrayRef, ArrowError> {
    let array = array.as_primitive::<IntervalYearMonthType>();

    Ok(Arc::new(array.unary::<_, IntervalMonthDayNanoType>(|v| {
        let months = IntervalYearMonthType::to_months(v);
        IntervalMonthDayNanoType::make_value(months, 0, 0)
    })))
}

/// Cast the array from interval day time to month day nano
pub(crate) fn cast_interval_day_time_to_interval_month_day_nano(
    array: &dyn Array,
    _cast_options: &CastOptions,
) -> Result<ArrayRef, ArrowError> {
    let array = array.as_primitive::<IntervalDayTimeType>();
    let mul = 1_000_000;

    Ok(Arc::new(array.unary::<_, IntervalMonthDayNanoType>(|v| {
        let (days, ms) = IntervalDayTimeType::to_parts(v);
        IntervalMonthDayNanoType::make_value(0, days, ms as i64 * mul)
    })))
}

/// Cast the array from interval to duration
pub(crate) fn cast_month_day_nano_to_duration<D: ArrowTemporalType<Native = i64>>(
    array: &dyn Array,
    cast_options: &CastOptions,
) -> Result<ArrayRef, ArrowError> {
    let array = array.as_primitive::<IntervalMonthDayNanoType>();
    let scale = match D::DATA_TYPE {
        DataType::Duration(TimeUnit::Second) => 1_000_000_000,
        DataType::Duration(TimeUnit::Millisecond) => 1_000_000,
        DataType::Duration(TimeUnit::Microsecond) => 1_000,
        DataType::Duration(TimeUnit::Nanosecond) => 1,
        _ => unreachable!(),
    };

    if cast_options.safe {
        let iter = array.iter().map(|v| {
            let v = v?;
            (v.days == 0 && v.months == 0).then_some(v.nanoseconds / scale)
        });
        Ok(Arc::new(unsafe {
            PrimitiveArray::<D>::from_trusted_len_iter(iter)
        }))
    } else {
        let vec = array
            .iter()
            .map(|v| {
                v.map(|v| match v.days == 0 && v.months == 0 {
                    true => Ok((v.nanoseconds) / scale),
                    _ => Err(ArrowError::ComputeError(
                        "Cannot convert interval containing non-zero months or days to duration"
                            .to_string(),
                    )),
                })
                .transpose()
            })
            .collect::<Result<Vec<_>, _>>()?;
        Ok(Arc::new(unsafe {
            PrimitiveArray::<D>::from_trusted_len_iter(vec.iter())
        }))
    }
}

/// Cast the array from duration and interval
pub(crate) fn cast_duration_to_interval<D: ArrowTemporalType<Native = i64>>(
    array: &dyn Array,
    cast_options: &CastOptions,
) -> Result<ArrayRef, ArrowError> {
    let array = array.as_primitive::<D>();

    let scale = match array.data_type() {
        DataType::Duration(TimeUnit::Second) => 1_000_000_000,
        DataType::Duration(TimeUnit::Millisecond) => 1_000_000,
        DataType::Duration(TimeUnit::Microsecond) => 1_000,
        DataType::Duration(TimeUnit::Nanosecond) => 1,
        _ => unreachable!(),
    };

    if cast_options.safe {
        let iter = array.iter().map(|v| {
            v?.checked_mul(scale)
                .map(|v| IntervalMonthDayNano::new(0, 0, v))
        });
        Ok(Arc::new(unsafe {
            PrimitiveArray::<IntervalMonthDayNanoType>::from_trusted_len_iter(iter)
        }))
    } else {
        let vec = array
            .iter()
            .map(|v| {
                v.map(|v| {
                    if let Ok(v) = v.mul_checked(scale) {
                        Ok(IntervalMonthDayNano::new(0, 0, v))
                    } else {
                        Err(ArrowError::ComputeError(format!(
                            "Cannot cast to {:?}. Overflowing on {:?}",
                            IntervalMonthDayNanoType::DATA_TYPE,
                            v
                        )))
                    }
                })
                .transpose()
            })
            .collect::<Result<Vec<_>, _>>()?;
        Ok(Arc::new(unsafe {
            PrimitiveArray::<IntervalMonthDayNanoType>::from_trusted_len_iter(vec.iter())
        }))
    }
}

pub(crate) fn make_timestamp_array(
    array: &PrimitiveArray<Int64Type>,
    unit: TimeUnit,
    tz: Option<Arc<str>>,
) -> ArrayRef {
    match unit {
        TimeUnit::Second => Arc::new(
            array
                .reinterpret_cast::<TimestampSecondType>()
                .with_timezone_opt(tz),
        ),
        TimeUnit::Millisecond => Arc::new(
            array
                .reinterpret_cast::<TimestampMillisecondType>()
                .with_timezone_opt(tz),
        ),
        TimeUnit::Microsecond => Arc::new(
            array
                .reinterpret_cast::<TimestampMicrosecondType>()
                .with_timezone_opt(tz),
        ),
        TimeUnit::Nanosecond => Arc::new(
            array
                .reinterpret_cast::<TimestampNanosecondType>()
                .with_timezone_opt(tz),
        ),
    }
}

pub(crate) fn make_duration_array(array: &PrimitiveArray<Int64Type>, unit: TimeUnit) -> ArrayRef {
    match unit {
        TimeUnit::Second => Arc::new(array.reinterpret_cast::<DurationSecondType>()),
        TimeUnit::Millisecond => Arc::new(array.reinterpret_cast::<DurationMillisecondType>()),
        TimeUnit::Microsecond => Arc::new(array.reinterpret_cast::<DurationMicrosecondType>()),
        TimeUnit::Nanosecond => Arc::new(array.reinterpret_cast::<DurationNanosecondType>()),
    }
}

pub(crate) 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::InvalidArgumentError(format!(
            "Casting from {} to {to_type} not supported",
            array.data_type()
        ))),
    }
}

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)
    }
}

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))
}

pub(crate) 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 Date32",
            std::any::type_name::<T>()
        ))
    };

    let array: Date32Array = match array.timezone() {
        Some(tz) => {
            let tz: Tz = tz.parse()?;
            let date = |x: i64| {
                as_datetime_with_timezone::<T>(x, tz)
                    .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)))?
            }
        }
    };
    Ok(Arc::new(array))
}

/// Casts a timestamp array to `Date64`.
///
/// `Date64` values must be whole days, so this truncates to the start of the
/// day the same way [`timestamp_to_date32`] does: in the array's timezone when it
/// has one, otherwise rounding down so that a timestamp just before the epoch
/// belongs to the preceding day.
pub(crate) fn timestamp_to_date64<T: ArrowTimestampType>(
    array: &PrimitiveArray<T>,
    cast_options: &CastOptions,
) -> Result<ArrayRef, ArrowError> {
    let err = |x: i64| {
        ArrowError::CastError(format!(
            "Cannot convert {} {x} to Date64",
            std::any::type_name::<T>()
        ))
    };

    let array: Date64Array = match array.timezone() {
        Some(tz) => {
            let tz: Tz = tz.parse()?;
            // Chrono's date range is small enough that the product cannot overflow.
            let date = |x: i64| {
                as_datetime_with_timezone::<T>(x, tz).map(|d| {
                    Date32Type::from_naive_date(d.date_naive()) as i64 * MILLISECONDS_IN_DAY
                })
            };
            if cast_options.safe {
                array.unary_opt(date)
            } else {
                array.try_unary(|x| date(x).ok_or_else(|| err(x)))?
            }
        }
        None => {
            // Rescaling whole days to milliseconds can overflow for second and
            // millisecond timestamps near the ends of the i64 range.
            let date = |x: i64| {
                x.div_euclid(units_per_day::<T>())
                    .checked_mul(MILLISECONDS_IN_DAY)
            };
            if cast_options.safe {
                array.unary_opt(date)
            } else {
                array.try_unary(|x| date(x).ok_or_else(|| err(x)))?
            }
        }
    };
    Ok(Arc::new(array))
}

pub(crate) fn adjust_timestamp_to_timezone<T: ArrowTimestampType>(
    array: PrimitiveArray<Int64Type>,
    to_tz: &Tz,
    cast_options: &CastOptions,
) -> Result<PrimitiveArray<Int64Type>, ArrowError> {
    let adjust = |o| {
        let local = as_datetime::<T>(o)?;
        let offset = resolve_local_offset(to_tz, &local)?.fix();
        T::from_naive_datetime(local - offset, None)
    };
    let adjusted = if cast_options.safe {
        array.unary_opt::<_, Int64Type>(adjust)
    } else {
        array.try_unary::<_, Int64Type, _>(|o| {
            adjust(o).ok_or_else(|| {
                ArrowError::CastError("Cannot cast timezone to different timezone".to_string())
            })
        })?
    };
    Ok(adjusted)
}