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, Datelike, Days, Months, TimeZone};
27use std::cmp::Ordering;
28
29/// Shift a date by the given number of months.
30pub(crate) fn shift_months<D>(date: D, months: i32) -> D
31where
32    D: Datelike + std::ops::Add<Months, Output = D> + std::ops::Sub<Months, Output = D>,
33{
34    match months.cmp(&0) {
35        Ordering::Equal => date,
36        Ordering::Greater => date + Months::new(months as u32),
37        Ordering::Less => date - Months::new(months.unsigned_abs()),
38    }
39}
40
41/// Add the given number of months to the given datetime.
42///
43/// Returns `None` when it will result in overflow.
44pub(crate) fn add_months_datetime<Tz: TimeZone>(
45    dt: DateTime<Tz>,
46    months: i32,
47) -> Option<DateTime<Tz>> {
48    match months.cmp(&0) {
49        Ordering::Equal => Some(dt),
50        Ordering::Greater => dt.checked_add_months(Months::new(months as u32)),
51        Ordering::Less => dt.checked_sub_months(Months::new(months.unsigned_abs())),
52    }
53}
54
55/// Add the given number of days to the given datetime.
56///
57/// Returns `None` when it will result in overflow.
58pub(crate) fn add_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
59    match days.cmp(&0) {
60        Ordering::Equal => Some(dt),
61        Ordering::Greater => dt.checked_add_days(Days::new(days as u64)),
62        Ordering::Less => dt.checked_sub_days(Days::new(days.unsigned_abs() as u64)),
63    }
64}
65
66/// Substract the given number of months to the given datetime.
67///
68/// Returns `None` when it will result in overflow.
69pub(crate) fn sub_months_datetime<Tz: TimeZone>(
70    dt: DateTime<Tz>,
71    months: i32,
72) -> Option<DateTime<Tz>> {
73    match months.cmp(&0) {
74        Ordering::Equal => Some(dt),
75        Ordering::Greater => dt.checked_sub_months(Months::new(months as u32)),
76        Ordering::Less => dt.checked_add_months(Months::new(months.unsigned_abs())),
77    }
78}
79
80/// Substract the given number of days to the given datetime.
81///
82/// Returns `None` when it will result in overflow.
83pub(crate) fn sub_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
84    match days.cmp(&0) {
85        Ordering::Equal => Some(dt),
86        Ordering::Greater => dt.checked_sub_days(Days::new(days as u64)),
87        Ordering::Less => dt.checked_add_days(Days::new(days.unsigned_abs() as u64)),
88    }
89}
90
91#[cfg(test)]
92mod tests {
93
94    use chrono::naive::{NaiveDate, NaiveDateTime, NaiveTime};
95
96    use super::*;
97
98    #[test]
99    fn test_shift_months() {
100        let base = NaiveDate::from_ymd_opt(2020, 1, 31).unwrap();
101
102        assert_eq!(
103            shift_months(base, 0),
104            NaiveDate::from_ymd_opt(2020, 1, 31).unwrap()
105        );
106        assert_eq!(
107            shift_months(base, 1),
108            NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
109        );
110        assert_eq!(
111            shift_months(base, 2),
112            NaiveDate::from_ymd_opt(2020, 3, 31).unwrap()
113        );
114        assert_eq!(
115            shift_months(base, 3),
116            NaiveDate::from_ymd_opt(2020, 4, 30).unwrap()
117        );
118        assert_eq!(
119            shift_months(base, 4),
120            NaiveDate::from_ymd_opt(2020, 5, 31).unwrap()
121        );
122        assert_eq!(
123            shift_months(base, 5),
124            NaiveDate::from_ymd_opt(2020, 6, 30).unwrap()
125        );
126        assert_eq!(
127            shift_months(base, 6),
128            NaiveDate::from_ymd_opt(2020, 7, 31).unwrap()
129        );
130        assert_eq!(
131            shift_months(base, 7),
132            NaiveDate::from_ymd_opt(2020, 8, 31).unwrap()
133        );
134        assert_eq!(
135            shift_months(base, 8),
136            NaiveDate::from_ymd_opt(2020, 9, 30).unwrap()
137        );
138        assert_eq!(
139            shift_months(base, 9),
140            NaiveDate::from_ymd_opt(2020, 10, 31).unwrap()
141        );
142        assert_eq!(
143            shift_months(base, 10),
144            NaiveDate::from_ymd_opt(2020, 11, 30).unwrap()
145        );
146        assert_eq!(
147            shift_months(base, 11),
148            NaiveDate::from_ymd_opt(2020, 12, 31).unwrap()
149        );
150        assert_eq!(
151            shift_months(base, 12),
152            NaiveDate::from_ymd_opt(2021, 1, 31).unwrap()
153        );
154        assert_eq!(
155            shift_months(base, 13),
156            NaiveDate::from_ymd_opt(2021, 2, 28).unwrap()
157        );
158
159        assert_eq!(
160            shift_months(base, -1),
161            NaiveDate::from_ymd_opt(2019, 12, 31).unwrap()
162        );
163        assert_eq!(
164            shift_months(base, -2),
165            NaiveDate::from_ymd_opt(2019, 11, 30).unwrap()
166        );
167        assert_eq!(
168            shift_months(base, -3),
169            NaiveDate::from_ymd_opt(2019, 10, 31).unwrap()
170        );
171        assert_eq!(
172            shift_months(base, -4),
173            NaiveDate::from_ymd_opt(2019, 9, 30).unwrap()
174        );
175        assert_eq!(
176            shift_months(base, -5),
177            NaiveDate::from_ymd_opt(2019, 8, 31).unwrap()
178        );
179        assert_eq!(
180            shift_months(base, -6),
181            NaiveDate::from_ymd_opt(2019, 7, 31).unwrap()
182        );
183        assert_eq!(
184            shift_months(base, -7),
185            NaiveDate::from_ymd_opt(2019, 6, 30).unwrap()
186        );
187        assert_eq!(
188            shift_months(base, -8),
189            NaiveDate::from_ymd_opt(2019, 5, 31).unwrap()
190        );
191        assert_eq!(
192            shift_months(base, -9),
193            NaiveDate::from_ymd_opt(2019, 4, 30).unwrap()
194        );
195        assert_eq!(
196            shift_months(base, -10),
197            NaiveDate::from_ymd_opt(2019, 3, 31).unwrap()
198        );
199        assert_eq!(
200            shift_months(base, -11),
201            NaiveDate::from_ymd_opt(2019, 2, 28).unwrap()
202        );
203        assert_eq!(
204            shift_months(base, -12),
205            NaiveDate::from_ymd_opt(2019, 1, 31).unwrap()
206        );
207        assert_eq!(
208            shift_months(base, -13),
209            NaiveDate::from_ymd_opt(2018, 12, 31).unwrap()
210        );
211
212        assert_eq!(
213            shift_months(base, 1265),
214            NaiveDate::from_ymd_opt(2125, 6, 30).unwrap()
215        );
216    }
217
218    #[test]
219    fn test_shift_months_with_overflow() {
220        let base = NaiveDate::from_ymd_opt(2020, 12, 31).unwrap();
221
222        assert_eq!(shift_months(base, 0), base);
223        assert_eq!(
224            shift_months(base, 1),
225            NaiveDate::from_ymd_opt(2021, 1, 31).unwrap()
226        );
227        assert_eq!(
228            shift_months(base, 2),
229            NaiveDate::from_ymd_opt(2021, 2, 28).unwrap()
230        );
231        assert_eq!(
232            shift_months(base, 12),
233            NaiveDate::from_ymd_opt(2021, 12, 31).unwrap()
234        );
235        assert_eq!(
236            shift_months(base, 18),
237            NaiveDate::from_ymd_opt(2022, 6, 30).unwrap()
238        );
239
240        assert_eq!(
241            shift_months(base, -1),
242            NaiveDate::from_ymd_opt(2020, 11, 30).unwrap()
243        );
244        assert_eq!(
245            shift_months(base, -2),
246            NaiveDate::from_ymd_opt(2020, 10, 31).unwrap()
247        );
248        assert_eq!(
249            shift_months(base, -10),
250            NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
251        );
252        assert_eq!(
253            shift_months(base, -12),
254            NaiveDate::from_ymd_opt(2019, 12, 31).unwrap()
255        );
256        assert_eq!(
257            shift_months(base, -18),
258            NaiveDate::from_ymd_opt(2019, 6, 30).unwrap()
259        );
260    }
261
262    #[test]
263    fn test_shift_months_datetime() {
264        let date = NaiveDate::from_ymd_opt(2020, 1, 31).unwrap();
265        let o_clock = NaiveTime::from_hms_opt(1, 2, 3).unwrap();
266
267        let base = NaiveDateTime::new(date, o_clock);
268
269        assert_eq!(
270            shift_months(base, 0).date(),
271            NaiveDate::from_ymd_opt(2020, 1, 31).unwrap()
272        );
273        assert_eq!(
274            shift_months(base, 1).date(),
275            NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
276        );
277        assert_eq!(
278            shift_months(base, 2).date(),
279            NaiveDate::from_ymd_opt(2020, 3, 31).unwrap()
280        );
281        assert_eq!(shift_months(base, 0).time(), o_clock);
282        assert_eq!(shift_months(base, 1).time(), o_clock);
283        assert_eq!(shift_months(base, 2).time(), o_clock);
284    }
285}
arrow_array/delta.rs

arrow_array/
delta.rs
