Skip to content

Commit

Permalink
refactor(rust): factor out some shared code in truncate_impl (pola-…
Browse files Browse the repository at this point in the history 
…rs#12229)
  • Loading branch information
@MarcoGorelli
MarcoGorelli authored Nov 4, 2023
1 parent 019dacd commit c7c3b03
Showing 1 changed file with 136 additions and 90 deletions.
226 changes: 136 additions & 90 deletions crates/polars-time/src/windows/duration.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -439,6 +439,111 @@ impl Duration {
)
}

fn truncate_subweekly<G, J>(
&self,
t: i64,
tz: Option<&Tz>,
duration: i64,
_timestamp_to_datetime: G,
_datetime_to_timestamp: J,
_ambiguous: Ambiguous,
) -> PolarsResult<i64>
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<G, J>(
&self,
t: i64,
tz: Option<&Tz>,
timestamp_to_datetime: G,
datetime_to_timestamp: J,
_ambiguous: Ambiguous,
) -> PolarsResult<i64>
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<G, J>(
&self,
t: i64,
tz: Option<&Tz>,
timestamp_to_datetime: G,
datetime_to_timestamp: J,
_ambiguous: Ambiguous,
) -> PolarsResult<i64>
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<F, G, J>(
&self,
Expand All @@ -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")
},
Expand Down

0 comments on commit c7c3b03

Please sign in to comment.