From c7c3b033c3b9662731c5254e4b5ee796a37cf540 Mon Sep 17 00:00:00 2001 From: Marco Edward Gorelli Date: Sat, 4 Nov 2023 09:52:19 +0000 Subject: [PATCH] refactor(rust): factor out some shared code in `truncate_impl` (#12229) --- crates/polars-time/src/windows/duration.rs | 226 +++++++++++++-------- 1 file changed, 136 insertions(+), 90 deletions(-) diff --git a/crates/polars-time/src/windows/duration.rs b/crates/polars-time/src/windows/duration.rs index 35b7a082fec7..4716c68552a7 100644 --- a/crates/polars-time/src/windows/duration.rs +++ b/crates/polars-time/src/windows/duration.rs @@ -439,6 +439,111 @@ impl Duration { ) } + fn truncate_subweekly( + &self, + t: i64, + tz: Option<&Tz>, + duration: i64, + _timestamp_to_datetime: G, + _datetime_to_timestamp: J, + _ambiguous: Ambiguous, + ) -> PolarsResult + where + G: Fn(i64) -> NaiveDateTime, + J: Fn(NaiveDateTime) -> i64, + { + let t = match tz { + #[cfg(feature = "timezones")] + Some(tz) => _datetime_to_timestamp(unlocalize_datetime(_timestamp_to_datetime(t), tz)), + _ => t, + }; + let mut remainder = t % duration; + if remainder < 0 { + remainder += duration + } + match tz { + #[cfg(feature = "timezones")] + Some(tz) => Ok(_datetime_to_timestamp(localize_datetime( + _timestamp_to_datetime(t - remainder), + tz, + _ambiguous, + )?)), + _ => Ok(t - remainder), + } + } + + fn truncate_weekly( + &self, + t: i64, + tz: Option<&Tz>, + timestamp_to_datetime: G, + datetime_to_timestamp: J, + _ambiguous: Ambiguous, + ) -> PolarsResult + where + G: Fn(i64) -> NaiveDateTime, + J: Fn(NaiveDateTime) -> i64, + { + let dt = match tz { + #[cfg(feature = "timezones")] + Some(tz) => unlocalize_datetime(timestamp_to_datetime(t), tz).date(), + _ => timestamp_to_datetime(t).date(), + }; + let week_timestamp = dt.week(Weekday::Mon); + let first_day_of_week = + week_timestamp.first_day() - chrono::Duration::weeks(self.weeks - 1); + match tz { + #[cfg(feature = "timezones")] + Some(tz) => Ok(datetime_to_timestamp(localize_datetime( + first_day_of_week.and_time(NaiveTime::default()), + tz, + _ambiguous, + )?)), + _ => Ok(datetime_to_timestamp( + first_day_of_week.and_time(NaiveTime::default()), + )), + } + } + fn truncate_monthly( + &self, + t: i64, + tz: Option<&Tz>, + timestamp_to_datetime: G, + datetime_to_timestamp: J, + _ambiguous: Ambiguous, + ) -> PolarsResult + where + G: Fn(i64) -> NaiveDateTime, + J: Fn(NaiveDateTime) -> i64, + { + let ts = match tz { + #[cfg(feature = "timezones")] + Some(tz) => unlocalize_datetime(timestamp_to_datetime(t), tz), + _ => timestamp_to_datetime(t), + }; + let (year, month) = (ts.year(), ts.month()); + + // determine the total number of months and truncate + // the number of months by the duration amount + let mut total = (year * 12) + (month as i32 - 1); + let remainder = total % self.months as i32; + total -= remainder; + + // recreate a new time from the year and month combination + let (year, month) = ((total / 12), ((total % 12) + 1) as u32); + + let dt = 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) => Ok(datetime_to_timestamp(localize_datetime( + dt, tz, _ambiguous, + )?)), + _ => Ok(datetime_to_timestamp(dt)), + } + } + #[inline] pub fn truncate_impl( &self, @@ -458,103 +563,44 @@ impl Duration { (0, 0, 0, 0) => polars_bail!(ComputeError: "duration cannot be zero"), // truncate by ns/us/ms (0, 0, 0, _) => { - let t = match tz { - #[cfg(feature = "timezones")] - Some(tz) => { - datetime_to_timestamp(unlocalize_datetime(timestamp_to_datetime(t), tz)) - }, - _ => t, - }; let duration = nsecs_to_unit(self.nsecs); - let mut remainder = t % duration; - if remainder < 0 { - remainder += duration - } - match tz { - #[cfg(feature = "timezones")] - Some(tz) => Ok(datetime_to_timestamp(localize_datetime( - timestamp_to_datetime(t - remainder), - tz, - _ambiguous, - )?)), - _ => Ok(t - remainder), - } - }, - // truncate by weeks - (0, _, 0, 0) => { - let dt = match tz { - #[cfg(feature = "timezones")] - Some(tz) => unlocalize_datetime(timestamp_to_datetime(t), tz).date(), - _ => timestamp_to_datetime(t).date(), - }; - let week_timestamp = dt.week(Weekday::Mon); - let first_day_of_week = - week_timestamp.first_day() - chrono::Duration::weeks(self.weeks - 1); - match tz { - #[cfg(feature = "timezones")] - Some(tz) => Ok(datetime_to_timestamp(localize_datetime( - first_day_of_week.and_time(NaiveTime::default()), - tz, - _ambiguous, - )?)), - _ => Ok(datetime_to_timestamp( - first_day_of_week.and_time(NaiveTime::default()), - )), - } + self.truncate_subweekly( + t, + tz, + duration, + timestamp_to_datetime, + datetime_to_timestamp, + _ambiguous, + ) }, // truncate by days (0, 0, _, 0) => { - let t = match tz { - #[cfg(feature = "timezones")] - Some(tz) => { - datetime_to_timestamp(unlocalize_datetime(timestamp_to_datetime(t), tz)) - }, - _ => t, - }; let duration = self.days * nsecs_to_unit(NS_DAY); - let mut remainder = t % duration; - if remainder < 0 { - remainder += duration - } - match tz { - #[cfg(feature = "timezones")] - Some(tz) => Ok(datetime_to_timestamp(localize_datetime( - timestamp_to_datetime(t - remainder), - tz, - _ambiguous, - )?)), - _ => Ok(t - remainder), - } + self.truncate_subweekly( + t, + tz, + duration, + timestamp_to_datetime, + datetime_to_timestamp, + _ambiguous, + ) }, + // truncate by weeks + (0, _, 0, 0) => self.truncate_weekly( + t, + tz, + timestamp_to_datetime, + datetime_to_timestamp, + _ambiguous, + ), // truncate by months - (_, 0, 0, 0) => { - let ts = match tz { - #[cfg(feature = "timezones")] - Some(tz) => unlocalize_datetime(timestamp_to_datetime(t), tz), - _ => timestamp_to_datetime(t), - }; - let (year, month) = (ts.year(), ts.month()); - - // determine the total number of months and truncate - // the number of months by the duration amount - let mut total = (year * 12) + (month as i32 - 1); - let remainder = total % self.months as i32; - total -= remainder; - - // recreate a new time from the year and month combination - let (year, month) = ((total / 12), ((total % 12) + 1) as u32); - - let dt = 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) => Ok(datetime_to_timestamp(localize_datetime( - dt, tz, _ambiguous, - )?)), - _ => Ok(datetime_to_timestamp(dt)), - } - }, + (_, 0, 0, 0) => self.truncate_monthly( + t, + tz, + timestamp_to_datetime, + datetime_to_timestamp, + _ambiguous, + ), _ => { polars_bail!(ComputeError: "duration may not mix month, weeks and nanosecond units") },