use std::cmp::Ordering;
use std::ops::Mul;
#[cfg(feature = "timezones")]
use arrow::legacy::kernels::Ambiguous;
use arrow::legacy::time_zone::Tz;
use arrow::temporal_conversions::{
timestamp_ms_to_datetime, timestamp_ns_to_datetime, timestamp_us_to_datetime, MILLISECONDS,
};
use chrono::{Datelike, NaiveDate, NaiveDateTime, NaiveTime, Timelike, Weekday};
use polars_core::export::arrow::temporal_conversions::MICROSECONDS;
use polars_core::prelude::{
datetime_to_timestamp_ms, datetime_to_timestamp_ns, datetime_to_timestamp_us, polars_bail,
PolarsResult,
};
use polars_core::utils::arrow::temporal_conversions::NANOSECONDS;
use polars_error::*;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use super::calendar::{
NS_DAY, NS_HOUR, NS_MICROSECOND, NS_MILLISECOND, NS_MINUTE, NS_SECOND, NS_WEEK,
};
#[cfg(feature = "timezones")]
use crate::utils::{localize_datetime_opt, try_localize_datetime, unlocalize_datetime};
use crate::windows::calendar::{is_leap_year, last_day_of_month};
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Duration {
months: i64,
weeks: i64,
days: i64,
nsecs: i64,
pub(crate) negative: bool,
pub parsed_int: bool,
}
impl PartialOrd<Self> for Duration {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for Duration {
fn cmp(&self, other: &Self) -> Ordering {
self.duration_ns().cmp(&other.duration_ns())
}
}
impl Duration {
pub fn new(fixed_slots: i64) -> Self {
Duration {
months: 0,
weeks: 0,
days: 0,
nsecs: fixed_slots.abs(),
negative: fixed_slots < 0,
parsed_int: true,
}
}
pub fn parse(duration: &str) -> Self {
let num_minus_signs = duration.matches('-').count();
if num_minus_signs > 1 {
panic!("a Duration string can only have a single minus sign")
}
if (num_minus_signs > 0) & !duration.starts_with('-') {
panic!("only a single minus sign is allowed, at the front of the string")
}
let mut nsecs = 0;
let mut weeks = 0;
let mut days = 0;
let mut months = 0;
let negative = duration.starts_with('-');
let mut iter = duration.char_indices();
let mut start = 0;
if negative {
start += 1;
iter.next().unwrap();
}
let mut parsed_int = false;
let mut unit = String::with_capacity(2);
while let Some((i, mut ch)) = iter.next() {
if !ch.is_ascii_digit() {
let n = duration[start..i]
.parse::<i64>()
.expect("expected an integer in the duration string");
loop {
if ch.is_ascii_alphabetic() {
unit.push(ch)
} else {
break;
}
match iter.next() {
Some((i, ch_)) => {
ch = ch_;
start = i
},
None => {
break;
},
}
}
if unit.is_empty() {
panic!("expected a unit in the duration string")
}
match &*unit {
"ns" => nsecs += n,
"us" => nsecs += n * NS_MICROSECOND,
"ms" => nsecs += n * NS_MILLISECOND,
"s" => nsecs += n * NS_SECOND,
"m" => nsecs += n * NS_MINUTE,
"h" => nsecs += n * NS_HOUR,
"d" => days += n,
"w" => weeks += n,
"mo" => {
months += n
}
"q" => months += n * 3,
"y" => months += n * 12,
"i" => {
nsecs += n;
parsed_int = true;
}
unit => panic!("unit: '{unit}' not supported. Available units are: 'ns', 'us', 'ms', 's', 'm', 'h', 'd', 'w', 'q', 'mo', 'y', 'i'"),
}
unit.clear();
}
}
Duration {
nsecs: nsecs.abs(),
days: days.abs(),
weeks: weeks.abs(),
months: months.abs(),
negative,
parsed_int,
}
}
fn to_positive(v: i64) -> (bool, i64) {
if v < 0 {
(true, -v)
} else {
(false, v)
}
}
#[allow(dead_code)]
pub(crate) fn normalize(&self, interval: &Duration) -> Self {
if self.months_only() && interval.months_only() {
let mut months = self.months() % interval.months();
match (self.negative, interval.negative) {
(true, true) | (true, false) => months = -months + interval.months(),
_ => {},
}
Duration::from_months(months)
} else if self.weeks_only() && interval.weeks_only() {
let mut weeks = self.weeks() % interval.weeks();
match (self.negative, interval.negative) {
(true, true) | (true, false) => weeks = -weeks + interval.weeks(),
_ => {},
}
Duration::from_weeks(weeks)
} else if self.days_only() && interval.days_only() {
let mut days = self.days() % interval.days();
match (self.negative, interval.negative) {
(true, true) | (true, false) => days = -days + interval.days(),
_ => {},
}
Duration::from_days(days)
} else {
let mut offset = self.duration_ns();
if offset == 0 {
return *self;
}
let every = interval.duration_ns();
if offset < 0 {
offset += every * ((offset / -every) + 1)
} else {
offset -= every * (offset / every)
}
Duration::from_nsecs(offset)
}
}
pub(crate) fn from_nsecs(v: i64) -> Self {
let (negative, nsecs) = Self::to_positive(v);
Self {
months: 0,
weeks: 0,
days: 0,
nsecs,
negative,
parsed_int: false,
}
}
pub(crate) fn from_months(v: i64) -> Self {
let (negative, months) = Self::to_positive(v);
Self {
months,
weeks: 0,
days: 0,
nsecs: 0,
negative,
parsed_int: false,
}
}
pub(crate) fn from_weeks(v: i64) -> Self {
let (negative, weeks) = Self::to_positive(v);
Self {
months: 0,
weeks,
days: 0,
nsecs: 0,
negative,
parsed_int: false,
}
}
pub(crate) fn from_days(v: i64) -> Self {
let (negative, days) = Self::to_positive(v);
Self {
months: 0,
weeks: 0,
days,
nsecs: 0,
negative,
parsed_int: false,
}
}
pub fn is_zero(&self) -> bool {
self.months == 0 && self.weeks == 0 && self.days == 0 && self.nsecs == 0
}
pub fn months_only(&self) -> bool {
self.months != 0 && self.weeks == 0 && self.days == 0 && self.nsecs == 0
}
pub fn months(&self) -> i64 {
self.months
}
pub fn weeks_only(&self) -> bool {
self.months == 0 && self.weeks != 0 && self.days == 0 && self.nsecs == 0
}
pub fn weeks(&self) -> i64 {
self.weeks
}
pub fn days_only(&self) -> bool {
self.months == 0 && self.weeks == 0 && self.days != 0 && self.nsecs == 0
}
pub fn days(&self) -> i64 {
self.days
}
pub fn is_full_days(&self) -> bool {
self.nsecs == 0
}
pub fn is_constant_duration(&self) -> bool {
self.months == 0 && self.weeks == 0 && self.days == 0
}
pub fn nanoseconds(&self) -> i64 {
self.nsecs
}
#[doc(hidden)]
pub const fn duration_ns(&self) -> i64 {
self.months * 28 * 24 * 3600 * NANOSECONDS
+ self.weeks * NS_WEEK
+ self.days * NS_DAY
+ self.nsecs
}
#[doc(hidden)]
pub const fn duration_us(&self) -> i64 {
self.months * 28 * 24 * 3600 * MICROSECONDS
+ (self.weeks * NS_WEEK / 1000 + self.nsecs / 1000 + self.days * NS_DAY / 1000)
}
#[doc(hidden)]
pub const fn duration_ms(&self) -> i64 {
self.months * 28 * 24 * 3600 * MILLISECONDS
+ (self.weeks * NS_WEEK / 1_000_000
+ self.nsecs / 1_000_000
+ self.days * NS_DAY / 1_000_000)
}
#[doc(hidden)]
fn add_month(ts: NaiveDateTime, n_months: i64, negative: bool) -> NaiveDateTime {
let mut months = n_months;
if negative {
months = -months;
}
let mut year = ts.year();
let mut month = ts.month() as i32;
let mut day = ts.day();
year += (months / 12) as i32;
month += (months % 12) as i32;
if month > 12 {
year += 1;
month -= 12;
} else if month <= 0 {
year -= 1;
month += 12;
}
let mut last_day_of_month = last_day_of_month(month);
if month == (chrono::Month::February.number_from_month() as i32) && is_leap_year(year) {
last_day_of_month += 1;
}
if day > last_day_of_month {
day = last_day_of_month
}
let hour = ts.hour();
let minute = ts.minute();
let sec = ts.second();
let nsec = ts.nanosecond();
new_datetime(year, month as u32, day, hour, minute, sec, nsec).expect(
"Expected valid datetime, please open an issue at https://github.com/pola-rs/polars/issues"
)
}
#[cfg(feature = "timezones")]
fn localize_result(
&self,
original_dt_naive: NaiveDateTime,
original_dt_utc: NaiveDateTime,
result_dt_naive: NaiveDateTime,
tz: &Tz,
) -> NaiveDateTime {
match localize_datetime_opt(result_dt_naive, tz, Ambiguous::Raise) {
Some(dt) => dt,
None => {
if try_localize_datetime(original_dt_naive, tz, Ambiguous::Earliest).unwrap()
== original_dt_utc
{
try_localize_datetime(result_dt_naive, tz, Ambiguous::Earliest).unwrap()
} else if try_localize_datetime(original_dt_naive, tz, Ambiguous::Latest).unwrap()
== original_dt_utc
{
try_localize_datetime(result_dt_naive, tz, Ambiguous::Latest).unwrap()
} else {
unreachable!()
}
},
}
}
fn truncate_subweekly<G, J>(
&self,
t: i64,
tz: Option<&Tz>,
duration: i64,
_timestamp_to_datetime: G,
_datetime_to_timestamp: J,
) -> PolarsResult<i64>
where
G: Fn(i64) -> NaiveDateTime,
J: Fn(NaiveDateTime) -> i64,
{
match tz {
#[cfg(feature = "timezones")]
Some(tz) => {
let original_dt_utc = _timestamp_to_datetime(t);
let original_dt_naive = unlocalize_datetime(original_dt_utc, tz);
let t = _datetime_to_timestamp(original_dt_naive);
let mut remainder = t % duration;
if remainder < 0 {
remainder += duration
}
let result_timestamp = t - remainder;
let result_dt_naive = _timestamp_to_datetime(result_timestamp);
let result_dt_utc =
self.localize_result(original_dt_naive, original_dt_utc, result_dt_naive, tz);
Ok(_datetime_to_timestamp(result_dt_utc))
},
_ => {
let mut remainder = t % duration;
if remainder < 0 {
remainder += duration
}
Ok(t - remainder)
},
}
}
fn truncate_weekly<G, J>(
&self,
t: i64,
tz: Option<&Tz>,
timestamp_to_datetime: G,
datetime_to_timestamp: J,
) -> PolarsResult<i64>
where
G: Fn(i64) -> NaiveDateTime,
J: Fn(NaiveDateTime) -> i64,
{
let _original_dt_utc;
let original_dt_naive;
let dt = match tz {
#[cfg(feature = "timezones")]
Some(tz) => {
_original_dt_utc = timestamp_to_datetime(t);
original_dt_naive = unlocalize_datetime(_original_dt_utc, tz);
original_dt_naive.date()
},
_ => {
_original_dt_utc = timestamp_to_datetime(t);
original_dt_naive = _original_dt_utc;
original_dt_naive.date()
},
};
let week_timestamp = dt.week(Weekday::Mon);
let first_day_of_week =
week_timestamp.first_day() - chrono::Duration::weeks(self.weeks - 1);
let result_dt_naive = first_day_of_week.and_time(NaiveTime::default());
match tz {
#[cfg(feature = "timezones")]
Some(tz) => {
let result_dt_utc =
self.localize_result(original_dt_naive, _original_dt_utc, result_dt_naive, tz);
Ok(datetime_to_timestamp(result_dt_utc))
},
_ => Ok(datetime_to_timestamp(result_dt_naive)),
}
}
fn truncate_monthly<G, J>(
&self,
t: i64,
tz: Option<&Tz>,
timestamp_to_datetime: G,
datetime_to_timestamp: J,
) -> PolarsResult<i64>
where
G: Fn(i64) -> NaiveDateTime,
J: Fn(NaiveDateTime) -> i64,
{
let original_dt_utc;
let original_dt_naive;
match tz {
#[cfg(feature = "timezones")]
Some(tz) => {
original_dt_utc = timestamp_to_datetime(t);
original_dt_naive = unlocalize_datetime(original_dt_utc, tz);
},
_ => {
original_dt_utc = timestamp_to_datetime(t);
original_dt_naive = original_dt_utc;
},
};
let (year, month) = (original_dt_naive.year(), original_dt_naive.month());
let mut total = (year * 12) + (month as i32 - 1);
let remainder = total % self.months as i32;
total -= remainder;
let (year, month) = ((total / 12), ((total % 12) + 1) as u32);
let result_dt_naive = new_datetime(year, month, 1, 0, 0, 0, 0).ok_or(polars_err!(
ComputeError: format!("date '{}-{}-1' does not exist", year, month)
))?;
match tz {
#[cfg(feature = "timezones")]
Some(tz) => {
let result_dt_utc =
self.localize_result(original_dt_naive, original_dt_utc, result_dt_naive, tz);
Ok(datetime_to_timestamp(result_dt_utc))
},
_ => Ok(datetime_to_timestamp(result_dt_naive)),
}
}
#[inline]
pub fn truncate_impl<F, G, J>(
&self,
t: i64,
tz: Option<&Tz>,
nsecs_to_unit: F,
timestamp_to_datetime: G,
datetime_to_timestamp: J,
) -> PolarsResult<i64>
where
F: Fn(i64) -> i64,
G: Fn(i64) -> NaiveDateTime,
J: Fn(NaiveDateTime) -> i64,
{
match (self.months, self.weeks, self.days, self.nsecs) {
(0, 0, 0, 0) => polars_bail!(ComputeError: "duration cannot be zero"),
(0, 0, 0, _) => {
let duration = nsecs_to_unit(self.nsecs);
self.truncate_subweekly(
t,
tz,
duration,
timestamp_to_datetime,
datetime_to_timestamp,
)
},
(0, 0, _, 0) => {
let duration = self.days * nsecs_to_unit(NS_DAY);
self.truncate_subweekly(
t,
tz,
duration,
timestamp_to_datetime,
datetime_to_timestamp,
)
},
(0, _, 0, 0) => {
self.truncate_weekly(t, tz, timestamp_to_datetime, datetime_to_timestamp)
},
(_, 0, 0, 0) => {
self.truncate_monthly(t, tz, timestamp_to_datetime, datetime_to_timestamp)
},
_ => {
polars_bail!(ComputeError: "duration may not mix month, weeks and nanosecond units")
},
}
}
#[inline]
pub fn truncate_ns(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
self.truncate_impl(
t,
tz,
|nsecs| nsecs,
timestamp_ns_to_datetime,
datetime_to_timestamp_ns,
)
}
#[inline]
pub fn truncate_us(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
self.truncate_impl(
t,
tz,
|nsecs| nsecs / 1000,
timestamp_us_to_datetime,
datetime_to_timestamp_us,
)
}
#[inline]
pub fn truncate_ms(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
self.truncate_impl(
t,
tz,
|nsecs| nsecs / 1_000_000,
timestamp_ms_to_datetime,
datetime_to_timestamp_ms,
)
}
fn add_impl_month_week_or_day<F, G, J>(
&self,
t: i64,
tz: Option<&Tz>,
nsecs_to_unit: F,
timestamp_to_datetime: G,
datetime_to_timestamp: J,
) -> PolarsResult<i64>
where
F: Fn(i64) -> i64,
G: Fn(i64) -> NaiveDateTime,
J: Fn(NaiveDateTime) -> i64,
{
let d = self;
let mut new_t = t;
if d.months > 0 {
let ts = match tz {
#[cfg(feature = "timezones")]
Some(tz) => unlocalize_datetime(timestamp_to_datetime(t), tz),
_ => timestamp_to_datetime(t),
};
let dt = Self::add_month(ts, d.months, d.negative);
new_t = match tz {
#[cfg(feature = "timezones")]
Some(tz) => datetime_to_timestamp(try_localize_datetime(dt, tz, Ambiguous::Raise)?),
_ => datetime_to_timestamp(dt),
};
}
if d.weeks > 0 {
let t_weeks = nsecs_to_unit(NS_WEEK) * self.weeks;
match tz {
#[cfg(feature = "timezones")]
Some(tz) => {
new_t =
datetime_to_timestamp(unlocalize_datetime(timestamp_to_datetime(t), tz));
new_t += if d.negative { -t_weeks } else { t_weeks };
new_t = datetime_to_timestamp(try_localize_datetime(
timestamp_to_datetime(new_t),
tz,
Ambiguous::Raise,
)?);
},
_ => new_t += if d.negative { -t_weeks } else { t_weeks },
};
}
if d.days > 0 {
let t_days = nsecs_to_unit(NS_DAY) * self.days;
match tz {
#[cfg(feature = "timezones")]
Some(tz) => {
new_t =
datetime_to_timestamp(unlocalize_datetime(timestamp_to_datetime(t), tz));
new_t += if d.negative { -t_days } else { t_days };
new_t = datetime_to_timestamp(try_localize_datetime(
timestamp_to_datetime(new_t),
tz,
Ambiguous::Raise,
)?);
},
_ => new_t += if d.negative { -t_days } else { t_days },
};
}
Ok(new_t)
}
pub fn add_ns(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
let d = self;
let new_t = self.add_impl_month_week_or_day(
t,
tz,
|nsecs| nsecs,
timestamp_ns_to_datetime,
datetime_to_timestamp_ns,
);
let nsecs = if d.negative { -d.nsecs } else { d.nsecs };
Ok(new_t? + nsecs)
}
pub fn add_us(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
let d = self;
let new_t = self.add_impl_month_week_or_day(
t,
tz,
|nsecs| nsecs / 1000,
timestamp_us_to_datetime,
datetime_to_timestamp_us,
);
let nsecs = if d.negative { -d.nsecs } else { d.nsecs };
Ok(new_t? + nsecs / 1_000)
}
pub fn add_ms(&self, t: i64, tz: Option<&Tz>) -> PolarsResult<i64> {
let d = self;
let new_t = self.add_impl_month_week_or_day(
t,
tz,
|nsecs| nsecs / 1_000_000,
timestamp_ms_to_datetime,
datetime_to_timestamp_ms,
);
let nsecs = if d.negative { -d.nsecs } else { d.nsecs };
Ok(new_t? + nsecs / 1_000_000)
}
}
impl Mul<i64> for Duration {
type Output = Self;
fn mul(mut self, mut rhs: i64) -> Self {
if rhs < 0 {
rhs = -rhs;
self.negative = !self.negative
}
self.months *= rhs;
self.weeks *= rhs;
self.days *= rhs;
self.nsecs *= rhs;
self
}
}
fn new_datetime(
year: i32,
month: u32,
days: u32,
hour: u32,
min: u32,
sec: u32,
nano: u32,
) -> Option<NaiveDateTime> {
let date = NaiveDate::from_ymd_opt(year, month, days)?;
let time = NaiveTime::from_hms_nano_opt(hour, min, sec, nano)?;
Some(NaiveDateTime::new(date, time))
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_parse() {
let out = Duration::parse("1ns");
assert_eq!(out.nsecs, 1);
let out = Duration::parse("1ns1ms");
assert_eq!(out.nsecs, NS_MILLISECOND + 1);
let out = Duration::parse("123ns40ms");
assert_eq!(out.nsecs, 40 * NS_MILLISECOND + 123);
let out = Duration::parse("123ns40ms1w");
assert_eq!(out.nsecs, 40 * NS_MILLISECOND + 123);
assert_eq!(out.duration_ns(), 40 * NS_MILLISECOND + 123 + NS_WEEK);
let out = Duration::parse("-123ns40ms1w");
assert!(out.negative);
let out = Duration::parse("5w");
assert_eq!(out.weeks(), 5);
}
#[test]
fn test_add_ns() {
let t = 1;
let seven_days = Duration::parse("7d");
let one_week = Duration::parse("1w");
assert_eq!(
seven_days.add_ns(t, None).unwrap(),
one_week.add_ns(t, None).unwrap()
);
let seven_days_negative = Duration::parse("-7d");
let one_week_negative = Duration::parse("-1w");
assert_eq!(
seven_days_negative.add_ns(t, None).unwrap(),
one_week_negative.add_ns(t, None).unwrap()
);
}
}