align formulae between datetime and duration

pola-rs · Jul 15, 2024 · 1554a57 · 1554a57
1 parent c72d3c4
commit 1554a57
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 20 deletions.
diff --git a/crates/polars-time/src/round.rs b/crates/polars-time/src/round.rs
@@ -12,6 +12,11 @@ pub trait PolarsRound {
         Self: Sized;
 }
 
+fn simple_round(t: i64, every: i64) -> i64 {
+    let half_away = t.signum() * every / 2;
+    t + half_away - (t + half_away) % every
+}
+
 impl PolarsRound for DatetimeChunked {
     fn round(&self, every: &StringChunked, tz: Option<&Tz>) -> PolarsResult<Self> {
         let time_zone = self.time_zone();
@@ -33,11 +38,7 @@ impl PolarsRound for DatetimeChunked {
                         TimeUnit::Nanoseconds => every_parsed.duration_ns(),
                     };
                     return Ok(self
-                        .apply_values(|t| {
-                            // Round half-way values away from zero
-                            let half_away = t.signum() * every / 2;
-                            t + half_away - (t + half_away) % every
-                        })
+                        .apply_values(|t| simple_round(t, every))
                         .into_datetime(self.time_unit(), time_zone.clone()));
                 } else {
                     let w = Window::new(every_parsed, every_parsed, offset);
@@ -143,11 +144,7 @@ impl PolarsRound for DurationChunked {
                     TimeUnit::Nanoseconds => every_parsed.duration_ns(),
                 };
                 return Ok(self
-                    .apply_values(|t| {
-                        // Round half-way values away from zero
-                        let half_away = t.signum() * every / 2;
-                        t + half_away - (t + half_away) % every
-                    })
+                    .apply_values(|t| simple_round(t, every))
                     .into_duration(self.time_unit()));
             } else {
                 return Ok(Int64Chunked::full_null(self.name(), self.len())
@@ -172,9 +169,7 @@ impl PolarsRound for DurationChunked {
                     TimeUnit::Microseconds => every_parsed.duration_us(),
                     TimeUnit::Nanoseconds => every_parsed.duration_ns(),
                 };
-                // Round half-way values away from zero
-                let half_away = t.signum() * every / 2;
-                Ok(Some(t + half_away - (t + half_away) % every))
+                Ok(Some(simple_round(t, every)))
             },
             _ => Ok(None),
         });

diff --git a/crates/polars-time/src/truncate.rs b/crates/polars-time/src/truncate.rs
@@ -12,6 +12,11 @@ pub trait PolarsTruncate {
         Self: Sized;
 }
 
+fn simple_truncate(t: i64, every: i64) -> i64 {
+    let remainder = t % every;
+    t - (remainder + every * (remainder < 0) as i64)
+}
+
 impl PolarsTruncate for DatetimeChunked {
     fn truncate(&self, tz: Option<&Tz>, every: &StringChunked) -> PolarsResult<Self> {
         let time_zone = self.time_zone();
@@ -33,10 +38,7 @@ impl PolarsTruncate for DatetimeChunked {
                         TimeUnit::Nanoseconds => every_parsed.duration_ns(),
                     };
                     return Ok(self
-                        .apply_values(|t| {
-                            let remainder = t % every;
-                            t - (remainder + every * (remainder < 0) as i64)
-                        })
+                        .apply_values(|t| simple_truncate(t, every))
                         .into_datetime(self.time_unit(), time_zone.clone()));
                 } else {
                     let w = Window::new(every_parsed, every_parsed, offset);
@@ -140,7 +142,7 @@ impl PolarsTruncate for DurationChunked {
                     TimeUnit::Nanoseconds => every_parsed.duration_ns(),
                 };
                 return Ok(self
-                    .apply_values(|t| t - t % every)
+                    .apply_values(|t: i64| simple_truncate(t, every))
                     .into_duration(self.time_unit()));
             } else {
                 return Ok(Int64Chunked::full_null(self.name(), self.len())
@@ -165,7 +167,7 @@ impl PolarsTruncate for DurationChunked {
                     TimeUnit::Microseconds => every_parsed.duration_us(),
                     TimeUnit::Nanoseconds => every_parsed.duration_ns(),
                 };
-                Ok(Some(t - t % every))
+                Ok(Some(simple_truncate(t, every)))
             },
             _ => Ok(None),
         });

diff --git a/py-polars/tests/unit/operations/namespaces/temporal/test_truncate.py b/py-polars/tests/unit/operations/namespaces/temporal/test_truncate.py
@@ -174,7 +174,7 @@ def test_truncate_duration(time_unit: TimeUnit) -> None:
             timedelta(seconds=30),
             timedelta(seconds=50),
             None,
-            timedelta(seconds=-30),
+            timedelta(seconds=-40),
         ]
     ).dt.cast_time_unit(time_unit)