arrow_array/
delta.rs

1// MIT License
2//
3// Copyright (c) 2020-2022 Oliver Margetts
4//
5// Permission is hereby granted, free of charge, to any person obtaining a copy
6// of this software and associated documentation files (the "Software"), to deal
7// in the Software without restriction, including without limitation the rights
8// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9// copies of the Software, and to permit persons to whom the Software is
10// furnished to do so, subject to the following conditions:
11//
12// The above copyright notice and this permission notice shall be included in all
13// copies or substantial portions of the Software.
14//
15// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21// SOFTWARE.
22
23// Copied from chronoutil crate
24
25//! Contains utility functions for shifting Date objects.
26use chrono::{DateTime, Days, Months, NaiveDate, TimeZone};
27use std::cmp::Ordering;
28
29/// Add the given number of months to the given date.
30///
31/// Returns `None` when it will result in overflow.
32pub(crate) fn add_months_date(date: NaiveDate, months: i32) -> Option<NaiveDate> {
33    match months.cmp(&0) {
34        Ordering::Equal => Some(date),
35        Ordering::Greater => date.checked_add_months(Months::new(months as u32)),
36        Ordering::Less => date.checked_sub_months(Months::new(months.unsigned_abs())),
37    }
38}
39
40/// Add the given number of months to the given datetime.
41///
42/// Returns `None` when it will result in overflow.
43pub(crate) fn add_months_datetime<Tz: TimeZone>(
44    dt: DateTime<Tz>,
45    months: i32,
46) -> Option<DateTime<Tz>> {
47    match months.cmp(&0) {
48        Ordering::Equal => Some(dt),
49        Ordering::Greater => dt.checked_add_months(Months::new(months as u32)),
50        Ordering::Less => dt.checked_sub_months(Months::new(months.unsigned_abs())),
51    }
52}
53
54/// Add the given number of days to the given datetime.
55///
56/// Returns `None` when it will result in overflow.
57pub(crate) fn add_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
58    match days.cmp(&0) {
59        Ordering::Equal => Some(dt),
60        Ordering::Greater => dt.checked_add_days(Days::new(days as u64)),
61        Ordering::Less => dt.checked_sub_days(Days::new(days.unsigned_abs() as u64)),
62    }
63}
64
65/// Substract the given number of months to the given datetime.
66///
67/// Returns `None` when it will result in overflow.
68pub(crate) fn sub_months_datetime<Tz: TimeZone>(
69    dt: DateTime<Tz>,
70    months: i32,
71) -> Option<DateTime<Tz>> {
72    match months.cmp(&0) {
73        Ordering::Equal => Some(dt),
74        Ordering::Greater => dt.checked_sub_months(Months::new(months as u32)),
75        Ordering::Less => dt.checked_add_months(Months::new(months.unsigned_abs())),
76    }
77}
78
79/// Substract the given number of days to the given datetime.
80///
81/// Returns `None` when it will result in overflow.
82pub(crate) fn sub_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
83    match days.cmp(&0) {
84        Ordering::Equal => Some(dt),
85        Ordering::Greater => dt.checked_sub_days(Days::new(days as u64)),
86        Ordering::Less => dt.checked_add_days(Days::new(days.unsigned_abs() as u64)),
87    }
88}
89
90#[cfg(test)]
91mod tests {
92
93    use chrono::naive::NaiveDate;
94
95    use super::*;
96
97    #[test]
98    fn test_add_monts_months() {
99        let base = NaiveDate::from_ymd_opt(2020, 1, 31).unwrap();
100
101        assert_eq!(
102            add_months_date(base, 0),
103            NaiveDate::from_ymd_opt(2020, 1, 31)
104        );
105        assert_eq!(
106            add_months_date(base, 1),
107            NaiveDate::from_ymd_opt(2020, 2, 29)
108        );
109        assert_eq!(
110            add_months_date(base, 2),
111            NaiveDate::from_ymd_opt(2020, 3, 31)
112        );
113        assert_eq!(
114            add_months_date(base, 3),
115            NaiveDate::from_ymd_opt(2020, 4, 30)
116        );
117        assert_eq!(
118            add_months_date(base, 4),
119            NaiveDate::from_ymd_opt(2020, 5, 31)
120        );
121        assert_eq!(
122            add_months_date(base, 5),
123            NaiveDate::from_ymd_opt(2020, 6, 30)
124        );
125        assert_eq!(
126            add_months_date(base, 6),
127            NaiveDate::from_ymd_opt(2020, 7, 31)
128        );
129        assert_eq!(
130            add_months_date(base, 7),
131            NaiveDate::from_ymd_opt(2020, 8, 31)
132        );
133        assert_eq!(
134            add_months_date(base, 8),
135            NaiveDate::from_ymd_opt(2020, 9, 30)
136        );
137        assert_eq!(
138            add_months_date(base, 9),
139            NaiveDate::from_ymd_opt(2020, 10, 31)
140        );
141        assert_eq!(
142            add_months_date(base, 10),
143            NaiveDate::from_ymd_opt(2020, 11, 30)
144        );
145        assert_eq!(
146            add_months_date(base, 11),
147            NaiveDate::from_ymd_opt(2020, 12, 31)
148        );
149        assert_eq!(
150            add_months_date(base, 12),
151            NaiveDate::from_ymd_opt(2021, 1, 31)
152        );
153        assert_eq!(
154            add_months_date(base, 13),
155            NaiveDate::from_ymd_opt(2021, 2, 28)
156        );
157
158        assert_eq!(
159            add_months_date(base, -1),
160            NaiveDate::from_ymd_opt(2019, 12, 31)
161        );
162        assert_eq!(
163            add_months_date(base, -2),
164            NaiveDate::from_ymd_opt(2019, 11, 30)
165        );
166        assert_eq!(
167            add_months_date(base, -3),
168            NaiveDate::from_ymd_opt(2019, 10, 31)
169        );
170        assert_eq!(
171            add_months_date(base, -4),
172            NaiveDate::from_ymd_opt(2019, 9, 30)
173        );
174        assert_eq!(
175            add_months_date(base, -5),
176            NaiveDate::from_ymd_opt(2019, 8, 31)
177        );
178        assert_eq!(
179            add_months_date(base, -6),
180            NaiveDate::from_ymd_opt(2019, 7, 31)
181        );
182        assert_eq!(
183            add_months_date(base, -7),
184            NaiveDate::from_ymd_opt(2019, 6, 30)
185        );
186        assert_eq!(
187            add_months_date(base, -8),
188            NaiveDate::from_ymd_opt(2019, 5, 31)
189        );
190        assert_eq!(
191            add_months_date(base, -9),
192            NaiveDate::from_ymd_opt(2019, 4, 30)
193        );
194        assert_eq!(
195            add_months_date(base, -10),
196            NaiveDate::from_ymd_opt(2019, 3, 31)
197        );
198        assert_eq!(
199            add_months_date(base, -11),
200            NaiveDate::from_ymd_opt(2019, 2, 28)
201        );
202        assert_eq!(
203            add_months_date(base, -12),
204            NaiveDate::from_ymd_opt(2019, 1, 31)
205        );
206        assert_eq!(
207            add_months_date(base, -13),
208            NaiveDate::from_ymd_opt(2018, 12, 31)
209        );
210
211        assert_eq!(
212            add_months_date(base, 1265),
213            NaiveDate::from_ymd_opt(2125, 6, 30)
214        );
215
216        // overflow handling
217        assert_eq!(add_months_date(base, i32::MAX), None);
218        assert_eq!(add_months_date(base, i32::MIN), None);
219    }
220
221    #[test]
222    fn test_add_months_date_with_overflow() {
223        let base = NaiveDate::from_ymd_opt(2020, 12, 31).unwrap();
224
225        assert_eq!(add_months_date(base, 0), Some(base));
226        assert_eq!(
227            add_months_date(base, 1),
228            NaiveDate::from_ymd_opt(2021, 1, 31)
229        );
230        assert_eq!(
231            add_months_date(base, 2),
232            NaiveDate::from_ymd_opt(2021, 2, 28)
233        );
234        assert_eq!(
235            add_months_date(base, 12),
236            NaiveDate::from_ymd_opt(2021, 12, 31)
237        );
238        assert_eq!(
239            add_months_date(base, 18),
240            NaiveDate::from_ymd_opt(2022, 6, 30)
241        );
242
243        assert_eq!(
244            add_months_date(base, -1),
245            NaiveDate::from_ymd_opt(2020, 11, 30)
246        );
247        assert_eq!(
248            add_months_date(base, -2),
249            NaiveDate::from_ymd_opt(2020, 10, 31)
250        );
251        assert_eq!(
252            add_months_date(base, -10),
253            NaiveDate::from_ymd_opt(2020, 2, 29)
254        );
255        assert_eq!(
256            add_months_date(base, -12),
257            NaiveDate::from_ymd_opt(2019, 12, 31)
258        );
259        assert_eq!(
260            add_months_date(base, -18),
261            NaiveDate::from_ymd_opt(2019, 6, 30)
262        );
263    }
264}
arrow_array/delta.rs

arrow_array/
delta.rs
