use std::str::FromStr;
use chrono::{FixedOffset, LocalResult, NaiveDateTime, Offset, TimeDelta, TimeZone, Utc};
use chrono_tz::Tz;
use crate::arrow::compute::kernels::cast_utils::string_to_datetime;
use crate::expressions::MapToStructOptions;
use crate::{DeltaResult, Error};
#[derive(Clone, Copy, Debug)]
pub(crate) enum TimestampTimezone {
Named(Tz),
Fixed(FixedOffset),
}
impl Default for TimestampTimezone {
fn default() -> Self {
Self::Named(chrono_tz::UTC)
}
}
impl FromStr for TimestampTimezone {
type Err = Error;
fn from_str(value: &str) -> DeltaResult<Self> {
let timezone = if value.starts_with(['+', '-']) {
parse_normalized_fixed_offset(value).map(Self::Fixed)
} else {
value.parse::<Tz>().ok().map(Self::Named)
};
timezone.ok_or_else(|| Error::generic(format!("Invalid timestamp timezone: {value}")))
}
}
impl TimestampTimezone {
pub(crate) fn try_from_options(options: &MapToStructOptions) -> DeltaResult<Self> {
match options.timestamp_timezone() {
Some(value) => value.parse(),
None => Ok(Self::default()),
}
}
pub(crate) fn parse_timestamp(self, raw: &str) -> Option<i64> {
match self {
Self::Named(timezone) => parse_timestamp_in_named_timezone(raw, timezone),
Self::Fixed(timezone) => string_to_datetime(&timezone, raw)
.ok()
.map(|timestamp| timestamp.timestamp_micros()),
}
}
}
fn parse_normalized_fixed_offset(value: &str) -> Option<FixedOffset> {
let sign = match value.as_bytes().first()? {
b'+' => 1,
b'-' => -1,
_ => return None,
};
let mut components = value.get(1..)?.split(':');
let hours = parse_two_digits(components.next()?)?;
let minutes = parse_two_digits(components.next()?)?;
if components.next().is_some() || hours > 18 || minutes > 59 || (hours == 18 && minutes != 0) {
return None;
}
FixedOffset::east_opt(sign * (hours * 3_600 + minutes * 60))
}
fn parse_two_digits(value: &str) -> Option<i32> {
let [tens, ones] = value.as_bytes() else {
return None;
};
(tens.is_ascii_digit() && ones.is_ascii_digit())
.then_some(i32::from(*tens - b'0') * 10 + i32::from(*ones - b'0'))
}
fn parse_timestamp_in_named_timezone(raw: &str, timezone: Tz) -> Option<i64> {
if let Ok(timestamp) = string_to_datetime(&timezone, raw) {
return Some(timestamp.timestamp_micros());
}
let local_datetime = string_to_datetime(&Utc, raw).ok()?.naive_utc();
match timezone.from_local_datetime(&local_datetime) {
LocalResult::Ambiguous(first, second) => Some(first.min(second).timestamp_micros()),
LocalResult::None => resolve_nonexistent_local_timestamp(local_datetime, timezone),
LocalResult::Single(_) => {
None
}
}
}
fn resolve_nonexistent_local_timestamp(local_datetime: NaiveDateTime, timezone: Tz) -> Option<i64> {
let offset = (1..=48).find_map(|hours| {
let before_transition = local_datetime.checked_sub_signed(TimeDelta::hours(hours))?;
timezone
.from_local_datetime(&before_transition)
.earliest()
.map(|timestamp| timestamp.offset().fix())
})?;
local_datetime
.checked_sub_signed(TimeDelta::seconds(i64::from(offset.local_minus_utc())))
.map(|utc_datetime| utc_datetime.and_utc().timestamp_micros())
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::super::expected_timestamp_micros;
use super::*;
fn options(timezone: &str) -> MapToStructOptions {
MapToStructOptions::default().with_timestamp_timezone(timezone)
}
#[rstest]
#[case::named("America/Los_Angeles")]
#[case::minute_offset("+05:30")]
#[case::positive_limit("+18:00")]
#[case::negative_limit("-18:00")]
fn accepts_normalized_reader_timezones(#[case] timezone: &str) {
assert!(TimestampTimezone::try_from_options(&options(timezone)).is_ok());
}
#[rstest]
#[case::empty("")]
#[case::bare_sign("+")]
#[case::compact_hour("+05")]
#[case::compact_hour_minute("+0530")]
#[case::invalid_minutes("+05:60")]
#[case::second_precision("+05:00:01")]
#[case::past_limit("+18:00:01")]
#[case::past_limit_minutes("+18:30")]
#[case::invalid_hours("+19:00")]
#[case::extra_component("+05:00:00:00")]
#[case::unknown_name("Not/AZone")]
fn rejects_noncanonical_reader_timezones(#[case] timezone: &str) {
assert!(TimestampTimezone::try_from_options(&options(timezone)).is_err());
}
#[rstest]
#[case::local("2024-01-15 12:30:45.123456", "2024-01-15T12:30:45.123456Z")]
#[case::compact_offset("2024-01-15T17:30:45+0530", "2024-01-15T12:00:45Z")]
#[case::offset_beyond_reader_limit("2024-01-15T12:30:45+19:00", "2024-01-14T17:30:45Z")]
#[case::negative_offset("2024-01-15T07:00:45-05:30", "2024-01-15T12:30:45Z")]
#[case::spaced_offset("2024-01-15 17:30:45 +05:30", "2024-01-15T12:00:45Z")]
#[case::named("2024-01-15 12:30:45 America/New_York", "2024-01-15T17:30:45Z")]
#[case::lowercase_t("2024-01-15t12:30:45", "2024-01-15T12:30:45Z")]
#[case::compact_clock("2024-01-15 123045", "2024-01-15T12:30:45Z")]
#[case::date_only("2024-01-15", "2024-01-15T00:00:00Z")]
#[case::excess_fraction("2024-01-15 12:30:45.123456789123", "2024-01-15T12:30:45.123456Z")]
fn parses_compatible_partition_timestamps(#[case] raw: &str, #[case] expected: &str) {
assert_eq!(
TimestampTimezone::default().parse_timestamp(raw),
Some(expected_timestamp_micros(expected))
);
}
#[rstest]
#[case::conflicting("2024-01-15 12:30:45+02:00 America/New_York")]
#[case::unknown_embedded("2024-01-15 12:30:45 Foo/Bar")]
#[case::multibyte_prefix("日時 America/New_York")]
#[case::extra_offset_component("2024-01-15T12:30:45+05:30:15:00")]
#[case::bad_compact_offset("2024-01-15T12:30:45+053")]
#[case::ampm("2024-01-15 01:30:45 PM")]
#[case::day_first("15/01/2024 12:30:45")]
#[case::time_only("12:30:45")]
fn rejects_unsupported_partition_timestamps(#[case] raw: &str) {
assert_eq!(TimestampTimezone::default().parse_timestamp(raw), None);
}
}