Implement diff_months kernels

jhorstmann · Oct 25, 2023 · 38e71b5 · 38e71b5
1 parent cb4378c
commit 38e71b5
Show file tree

Hide file tree

Showing 3 changed files with 172 additions and 1 deletion.
diff --git a/benches/bench_date_add.rs b/benches/bench_date_add.rs
@@ -2,12 +2,13 @@ use chrono::NaiveDateTime;
 use chronoutil::shift_months;
 use criterion::{criterion_group, criterion_main, Criterion, Throughput};
 
-use packedtime_rs::date_add_month_timestamp_millis;
+use packedtime_rs::{date_add_month_timestamp_millis, date_add_month_timestamp_millis_float, date_diff_month_timestamp_millis, date_diff_month_timestamp_millis_float};
 use rand::rngs::StdRng;
 use rand::{Rng, SeedableRng};
 
 #[inline(never)]
 fn bench_date_add_month(input: &[i64], output: &mut [i64], months: i32) {
+    assert_eq!(input.len(), output.len());
     output
         .iter_mut()
         .zip(input.iter().copied())
@@ -16,8 +17,46 @@ fn bench_date_add_month(input: &[i64], output: &mut [i64], months: i32) {
         });
 }
 
+#[inline(never)]
+fn bench_date_add_month_float(input: &[f64], output: &mut [f64], months: i32) {
+    assert_eq!(input.len(), output.len());
+    output
+        .iter_mut()
+        .zip(input.iter().copied())
+        .for_each(|(output, input)| {
+            *output = date_add_month_timestamp_millis_float(input, months);
+        });
+}
+
+#[inline(never)]
+fn bench_date_diff_month(start: &[i64], end: &[i64], output: &mut [i32]) {
+    assert_eq!(start.len(), end.len());
+    assert_eq!(start.len(), output.len());
+    output
+        .iter_mut()
+        .zip(start.iter().copied().zip(end.iter().copied()))
+        .for_each(|(output, (start, end))| {
+            *output = date_diff_month_timestamp_millis(start, end);
+        });
+}
+
+#[inline(never)]
+fn bench_date_diff_month_float(start: &[f64], end: &[f64], output: &mut [i32]) {
+    assert_eq!(start.len(), end.len());
+    assert_eq!(start.len(), output.len());
+    output
+        .iter_mut()
+        .zip(start.iter().copied().zip(end.iter().copied()))
+        .for_each(|(output, (start, end))| {
+            *output = date_diff_month_timestamp_millis_float(start, end);
+        });
+}
+
+
+
 #[inline(never)]
 fn bench_date_add_month_chronoutil(input: &[i64], output: &mut [i64], months: i32) {
+    assert_eq!(input.len(), output.len());
     output
         .iter_mut()
         .zip(input.iter().copied())
@@ -38,7 +77,21 @@ pub fn bench_date_add(c: &mut Criterion) {
         .map(|_| rng.gen_range(0..4102444800_000_i64))
         .collect::<Vec<_>>();
 
+    let input2 = (0..BATCH_SIZE)
+        .map(|_| rng.gen_range(0..4102444800_000_i64))
+        .collect::<Vec<_>>();
+
+    let input_float = (0..BATCH_SIZE)
+        .map(|_| rng.gen_range(0..4102444800_000_i64) as f64)
+        .collect::<Vec<_>>();
+
+    let input_float2 = (0..BATCH_SIZE)
+        .map(|_| rng.gen_range(0..4102444800_000_i64) as f64)
+        .collect::<Vec<_>>();
+
     let mut output = vec![0_i64; BATCH_SIZE];
+    let mut output_float = vec![0_f64; BATCH_SIZE];
+    let mut output_diff = vec![0_i32; BATCH_SIZE];
 
     c.benchmark_group("date_add_month")
         .throughput(Throughput::Bytes(
@@ -47,9 +100,25 @@ pub fn bench_date_add(c: &mut Criterion) {
         .bench_function("date_add_month", |b| {
             b.iter(|| bench_date_add_month(&input, &mut output, 1))
         })
+        .bench_function("date_add_month_float", |b| {
+            b.iter(|| bench_date_add_month_float(&input_float, &mut output_float, 1))
+        })
         .bench_function("date_add_month_chronoutil", |b| {
             b.iter(|| bench_date_add_month_chronoutil(&input, &mut output, 1))
         });
+
+    c.benchmark_group("date_diff_month")
+        .throughput(Throughput::Bytes(
+            (BATCH_SIZE * 2 * std::mem::size_of::<i64>() + BATCH_SIZE*std::mem::size_of::<i32>()) as u64,
+        ))
+        .bench_function("date_diff_month", |b| {
+            b.iter(|| bench_date_diff_month(&input, &input2, &mut output_diff))
+        })
+        .bench_function("date_diff_month_float", |b| {
+            b.iter(|| bench_date_diff_month_float(&input_float, &input_float2, &mut output_diff))
+        })
+    ;
+
 }
 
 criterion_group!(benches, bench_date_add);

diff --git a/src/epoch_days.rs b/src/epoch_days.rs
@@ -183,6 +183,28 @@ impl EpochDays {
         Self::from_ymd(y, m, d)
     }
 
+    /// Adds the given number of `years` to `epoch_days`.
+    /// If the day would be out of range for the resulting month
+    /// then the date will be clamped to the end of the month.
+    ///
+    /// For example: 2020-02-29 + 1year => 2021-02-28
+    #[inline]
+    pub fn add_years(&self, years: i32) -> Self {
+        let (mut y, m, mut d) = self.to_ymd();
+        y += years;
+        d = d.min(days_per_month(y, m-1));
+        Self::from_ymd(y, m, d)
+    }
+
+    #[inline]
+    pub fn diff_months(&self, other: EpochDays) -> i32 {
+        let (y0, m0, d0) = self.to_ymd();
+        let (y1, m1, d1) = other.to_ymd();
+
+        // TODO: Special-case the end of month? IOW, should diff(2023-10-31, 2023-11-30) return 1?
+        (y1 * 12 + m1) - (y0 * 12 + m0) - (d1 < d0) as i32
+    }
+
     #[inline]
     pub fn date_trunc_month(&self) -> Self {
         let (y, m, d) = self.to_ymd();
@@ -269,6 +291,47 @@ mod tests {
         assert_eq!(18993, EpochDays::new(19198).date_trunc_year().days());
     }
 
+    #[test]
+    fn test_date_trunc_month_epoch_days() {
+        assert_eq!(19174, EpochDays::new(19198).date_trunc_month().days());
+    }
+
+    #[test]
+    fn test_date_diff_month_epoch_days() {
+        assert_eq!(
+            EpochDays::from_ymd(2023, 10, 1).diff_months(EpochDays::from_ymd(2023, 11, 1)),
+            1
+        );
+        assert_eq!(
+            EpochDays::from_ymd(2023, 10, 1).diff_months(EpochDays::from_ymd(2023, 12, 1)),
+            2
+        );
+        assert_eq!(
+            EpochDays::from_ymd(2023, 10, 1).diff_months(EpochDays::from_ymd(2023, 12, 31)),
+            2
+        );
+        assert_eq!(
+            EpochDays::from_ymd(2023, 10, 22).diff_months(EpochDays::from_ymd(2023, 11, 22)),
+            1
+        );
+        assert_eq!(
+            EpochDays::from_ymd(2023, 10, 22).diff_months(EpochDays::from_ymd(2023, 11, 21)),
+            0
+        );
+        assert_eq!(
+            EpochDays::from_ymd(2023, 10, 31).diff_months(EpochDays::from_ymd(2023, 11, 30)),
+            0
+        );
+    }
+
+    #[test]
+    fn test_date_diff_month_epoch_days_negative() {
+        assert_eq!(
+            EpochDays::from_ymd(2023, 11, 1).diff_months(EpochDays::from_ymd(2023, 10, 1)),
+            -1
+        );
+    }
+
     #[test]
     fn test_extract_year() {
         assert_eq!(2022, EpochDays::from_ymd(2022, 1, 1).extract_year());

diff --git a/src/kernels.rs b/src/kernels.rs
@@ -102,6 +102,38 @@ pub fn date_part_month_timestamp_millis(ts: i64) -> i32 {
     epoch_days.extract_month()
 }
 
+#[inline]
+fn timestamp_to_year_month_millis_of_month(ts: i64) -> (i32, i32, i64) {
+    let (days, millis) = (ts.div_euclid(MILLIS_PER_DAY), ts.rem_euclid(MILLIS_PER_DAY));
+    let (year, month, d0) = EpochDays::new(days as i32).to_ymd();
+    let millis_of_month = (d0 as i64)*MILLIS_PER_DAY + millis;
+    (year, month, millis_of_month)
+}
+
+#[inline]
+fn timestamp_to_year_month_millis_of_month_float(ts: f64) -> (i32, i32, f64) {
+    let days = (ts * (1.0 / MILLIS_PER_DAY as f64)).floor();
+    let epoch_days = EpochDays::new(unsafe { days.to_int_unchecked() });
+    let millis = ts - days * (MILLIS_PER_DAY as f64);
+    let (year, month, d0) = epoch_days.to_ymd();
+    let millis_of_month = (d0 as f64)*(MILLIS_PER_DAY as f64) + millis;
+    (year, month, millis_of_month)
+}
+
+#[inline]
+pub fn date_diff_month_timestamp_millis(t0: i64, t1: i64) -> i32 {
+    let (y0, m0, ms0) = timestamp_to_year_month_millis_of_month(t0);
+    let (y1, m1, ms1) = timestamp_to_year_month_millis_of_month(t1);
+    (y1*12 + m1) - (y0*12 + m0) - ((ms1 < ms0) as i32)
+}
+
+#[inline]
+pub fn date_diff_month_timestamp_millis_float(t0: f64, t1: f64) -> i32 {
+    let (y0, m0, ms0) = timestamp_to_year_month_millis_of_month_float(t0);
+    let (y1, m1, ms1) = timestamp_to_year_month_millis_of_month_float(t1);
+    (y1*12 + m1) - (y0*12 + m0) - ((ms1 < ms0) as i32)
+}
+
 #[cfg(test)]
 mod tests {
     use crate::epoch_days::EpochDays;
@@ -209,6 +241,7 @@ mod tests {
         assert_eq!(epoch_day.add_months(-4), EpochDays::from_ymd(2022, 3, 31));
         assert_eq!(epoch_day.add_months(-5), EpochDays::from_ymd(2022, 2, 28));
         assert_eq!(epoch_day.add_months(-6), EpochDays::from_ymd(2022, 1, 31));
+        assert_eq!(epoch_day.add_months(-7), EpochDays::from_ymd(2021, 12, 31));
     }
 
     #[test]
@@ -229,6 +262,12 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_date_diff_months() {
+        // assert_eq!(epoch_day.add_months(-1), EpochDays::from_ymd(2022, 6, 30));
+
+    }
+
     #[test]
     #[cfg_attr(any(miri, not(feature = "expensive_tests")), ignore)]
     fn test_date_trunc_year_exhaustive() {