Skip to main content

ferray_ufunc/ops/
datetime.rs

1// ferray-ufunc: datetime64 / timedelta64 predicates and arithmetic
2//
3// Provides:
4//   - isnat() for the time element types — the "is Not a Time"
5//     predicate that returns true where a value equals the NaT sentinel
6//     (i64::MIN).
7//   - Element-wise arithmetic kernels:
8//       datetime - datetime → timedelta
9//       datetime + timedelta → datetime
10//       datetime - timedelta → datetime
11//       timedelta + timedelta → timedelta
12//       timedelta - timedelta → timedelta
13//
14// Same-unit operands required (NumPy promotes to the finer unit; we
15// surface a ShapeMismatch with a clear message and let callers cast
16// explicitly — that matches ferray's "no implicit precision changes"
17// stance from the rest of the workspace).
18//
19// NaT propagation: any operand that is NaT yields NaT in the output
20// (matching NumPy).
21
22use ferray_core::Array;
23use ferray_core::dimension::Dimension;
24use ferray_core::dtype::{DateTime64, NAT, TimeUnit, Timedelta64};
25use ferray_core::error::{FerrayError, FerrayResult};
26
27/// Element-wise: true where the input is the NaT ("Not a Time") sentinel.
28///
29/// Equivalent to `numpy.isnat` on a `datetime64` array. NumPy raises
30/// `TypeError` for any other dtype; in ferray that constraint is
31/// statically enforced — call sites that try to pass a non-time array
32/// fail to type-check.
33///
34/// # Errors
35/// Returns an error only if internal array construction fails.
36pub fn isnat_datetime<D: Dimension>(input: &Array<DateTime64, D>) -> FerrayResult<Array<bool, D>> {
37    let data: Vec<bool> = input.iter().map(|v| v.is_nat()).collect();
38    Array::from_vec(input.dim().clone(), data)
39}
40
41/// Element-wise: true where the input is the NaT sentinel.
42///
43/// Equivalent to `numpy.isnat` on a `timedelta64` array.
44///
45/// # Errors
46/// Returns an error only if internal array construction fails.
47pub fn isnat_timedelta<D: Dimension>(
48    input: &Array<Timedelta64, D>,
49) -> FerrayResult<Array<bool, D>> {
50    let data: Vec<bool> = input.iter().map(|v| v.is_nat()).collect();
51    Array::from_vec(input.dim().clone(), data)
52}
53
54// ===========================================================================
55// Arithmetic kernels
56// ===========================================================================
57
58/// Helper: combine two i64 ticks with NaT propagation. If either side is
59/// NaT, the result is NaT; otherwise the closure runs.
60#[inline]
61fn nat_propagate(a: i64, b: i64, op: impl FnOnce(i64, i64) -> i64) -> i64 {
62    if a == NAT || b == NAT { NAT } else { op(a, b) }
63}
64
65fn check_same_shape<A, B, DA, DB>(
66    a: &Array<A, DA>,
67    b: &Array<B, DB>,
68    name: &str,
69) -> FerrayResult<()>
70where
71    A: ferray_core::Element,
72    B: ferray_core::Element,
73    DA: Dimension,
74    DB: Dimension,
75{
76    if a.shape() != b.shape() {
77        return Err(FerrayError::shape_mismatch(format!(
78            "{name}: shapes {:?} and {:?} differ",
79            a.shape(),
80            b.shape()
81        )));
82    }
83    Ok(())
84}
85
86/// Element-wise `datetime - datetime → timedelta` (same unit assumed).
87///
88/// NaT propagates: if either side is NaT at position `i`, the output is
89/// NaT at that position. Mirrors NumPy's `datetime64 - datetime64`
90/// semantics.
91///
92/// # Errors
93/// Returns `FerrayError::ShapeMismatch` if the shapes differ.
94pub fn sub_datetime<D: Dimension>(
95    a: &Array<DateTime64, D>,
96    b: &Array<DateTime64, D>,
97) -> FerrayResult<Array<Timedelta64, D>> {
98    check_same_shape(a, b, "sub_datetime")?;
99    let data: Vec<Timedelta64> = a
100        .iter()
101        .zip(b.iter())
102        .map(|(x, y)| Timedelta64(nat_propagate(x.0, y.0, |p, q| p.wrapping_sub(q))))
103        .collect();
104    Array::from_vec(a.dim().clone(), data)
105}
106
107/// Element-wise `datetime + timedelta → datetime`.
108///
109/// # Errors
110/// Returns `FerrayError::ShapeMismatch` if the shapes differ.
111pub fn add_datetime_timedelta<D: Dimension>(
112    a: &Array<DateTime64, D>,
113    b: &Array<Timedelta64, D>,
114) -> FerrayResult<Array<DateTime64, D>> {
115    check_same_shape(a, b, "add_datetime_timedelta")?;
116    let data: Vec<DateTime64> = a
117        .iter()
118        .zip(b.iter())
119        .map(|(x, y)| DateTime64(nat_propagate(x.0, y.0, |p, q| p.wrapping_add(q))))
120        .collect();
121    Array::from_vec(a.dim().clone(), data)
122}
123
124/// Element-wise `datetime - timedelta → datetime`.
125///
126/// # Errors
127/// Returns `FerrayError::ShapeMismatch` if the shapes differ.
128pub fn sub_datetime_timedelta<D: Dimension>(
129    a: &Array<DateTime64, D>,
130    b: &Array<Timedelta64, D>,
131) -> FerrayResult<Array<DateTime64, D>> {
132    check_same_shape(a, b, "sub_datetime_timedelta")?;
133    let data: Vec<DateTime64> = a
134        .iter()
135        .zip(b.iter())
136        .map(|(x, y)| DateTime64(nat_propagate(x.0, y.0, |p, q| p.wrapping_sub(q))))
137        .collect();
138    Array::from_vec(a.dim().clone(), data)
139}
140
141/// Element-wise `timedelta + timedelta → timedelta`.
142///
143/// # Errors
144/// Returns `FerrayError::ShapeMismatch` if the shapes differ.
145pub fn add_timedelta<D: Dimension>(
146    a: &Array<Timedelta64, D>,
147    b: &Array<Timedelta64, D>,
148) -> FerrayResult<Array<Timedelta64, D>> {
149    check_same_shape(a, b, "add_timedelta")?;
150    let data: Vec<Timedelta64> = a
151        .iter()
152        .zip(b.iter())
153        .map(|(x, y)| Timedelta64(nat_propagate(x.0, y.0, |p, q| p.wrapping_add(q))))
154        .collect();
155    Array::from_vec(a.dim().clone(), data)
156}
157
158/// Element-wise `timedelta - timedelta → timedelta`.
159///
160/// # Errors
161/// Returns `FerrayError::ShapeMismatch` if the shapes differ.
162pub fn sub_timedelta<D: Dimension>(
163    a: &Array<Timedelta64, D>,
164    b: &Array<Timedelta64, D>,
165) -> FerrayResult<Array<Timedelta64, D>> {
166    check_same_shape(a, b, "sub_timedelta")?;
167    let data: Vec<Timedelta64> = a
168        .iter()
169        .zip(b.iter())
170        .map(|(x, y)| Timedelta64(nat_propagate(x.0, y.0, |p, q| p.wrapping_sub(q))))
171        .collect();
172    Array::from_vec(a.dim().clone(), data)
173}
174
175// ===========================================================================
176// Implicit unit promotion
177// ===========================================================================
178//
179// The promoted-* variants below take both operands together with their
180// units, rescale both to the finer unit using `TimeUnit::finer` /
181// `TimeUnit::scale_to`, and then call the same-unit kernel above. They
182// return `(result, target_unit)` so callers (e.g. DynArray dispatch in
183// ferray-numpy-interop) know which unit to tag the output array with.
184//
185// NaT propagation is preserved: rescaling treats `NAT` as "still NaT"
186// regardless of the multiplicative factor.
187
188#[inline]
189fn rescale_datetime(arr: &[DateTime64], factor: i64) -> Vec<DateTime64> {
190    arr.iter()
191        .map(|v| {
192            if v.is_nat() {
193                DateTime64(NAT)
194            } else {
195                DateTime64(v.0.wrapping_mul(factor))
196            }
197        })
198        .collect()
199}
200
201#[inline]
202fn rescale_timedelta(arr: &[Timedelta64], factor: i64) -> Vec<Timedelta64> {
203    arr.iter()
204        .map(|v| {
205            if v.is_nat() {
206                Timedelta64(NAT)
207            } else {
208                Timedelta64(v.0.wrapping_mul(factor))
209            }
210        })
211        .collect()
212}
213
214/// Element-wise `datetime[unit_a] - datetime[unit_b] → timedelta[finer]`.
215///
216/// Both operands are rescaled to `TimeUnit::finer(unit_a, unit_b)` before
217/// the per-element subtraction. Returns `(result, target_unit)`.
218///
219/// # Errors
220/// Returns `FerrayError::ShapeMismatch` if the shapes differ.
221pub fn sub_datetime_promoted<D: Dimension>(
222    a: &Array<DateTime64, D>,
223    unit_a: TimeUnit,
224    b: &Array<DateTime64, D>,
225    unit_b: TimeUnit,
226) -> FerrayResult<(Array<Timedelta64, D>, TimeUnit)> {
227    check_same_shape(a, b, "sub_datetime_promoted")?;
228    let target = unit_a.finer(unit_b);
229    let scale_a = unit_a.scale_to(target).ok_or_else(|| {
230        FerrayError::invalid_value(format!(
231            "sub_datetime_promoted: cannot rescale {unit_a:?} → {target:?} (non-divisible)"
232        ))
233    })?;
234    let scale_b = unit_b.scale_to(target).ok_or_else(|| {
235        FerrayError::invalid_value(format!(
236            "sub_datetime_promoted: cannot rescale {unit_b:?} → {target:?} (non-divisible)"
237        ))
238    })?;
239    let a_data: Vec<DateTime64> = a.iter().copied().collect();
240    let b_data: Vec<DateTime64> = b.iter().copied().collect();
241    let a_rescaled = if scale_a == 1 {
242        a_data
243    } else {
244        rescale_datetime(&a_data, scale_a)
245    };
246    let b_rescaled = if scale_b == 1 {
247        b_data
248    } else {
249        rescale_datetime(&b_data, scale_b)
250    };
251    let data: Vec<Timedelta64> = a_rescaled
252        .iter()
253        .zip(b_rescaled.iter())
254        .map(|(x, y)| Timedelta64(nat_propagate(x.0, y.0, |p, q| p.wrapping_sub(q))))
255        .collect();
256    let arr = Array::from_vec(a.dim().clone(), data)?;
257    Ok((arr, target))
258}
259
260/// Element-wise `datetime[unit_a] + timedelta[unit_b] → datetime[finer]`.
261///
262/// # Errors
263/// Returns `FerrayError::ShapeMismatch` if the shapes differ.
264pub fn add_datetime_timedelta_promoted<D: Dimension>(
265    a: &Array<DateTime64, D>,
266    unit_a: TimeUnit,
267    b: &Array<Timedelta64, D>,
268    unit_b: TimeUnit,
269) -> FerrayResult<(Array<DateTime64, D>, TimeUnit)> {
270    check_same_shape(a, b, "add_datetime_timedelta_promoted")?;
271    let target = unit_a.finer(unit_b);
272    let scale_a = unit_a.scale_to(target).ok_or_else(|| {
273        FerrayError::invalid_value(format!(
274            "add_datetime_timedelta_promoted: cannot rescale {unit_a:?} → {target:?}"
275        ))
276    })?;
277    let scale_b = unit_b.scale_to(target).ok_or_else(|| {
278        FerrayError::invalid_value(format!(
279            "add_datetime_timedelta_promoted: cannot rescale {unit_b:?} → {target:?}"
280        ))
281    })?;
282    let a_data: Vec<DateTime64> = a.iter().copied().collect();
283    let b_data: Vec<Timedelta64> = b.iter().copied().collect();
284    let a_rescaled = if scale_a == 1 {
285        a_data
286    } else {
287        rescale_datetime(&a_data, scale_a)
288    };
289    let b_rescaled = if scale_b == 1 {
290        b_data
291    } else {
292        rescale_timedelta(&b_data, scale_b)
293    };
294    let data: Vec<DateTime64> = a_rescaled
295        .iter()
296        .zip(b_rescaled.iter())
297        .map(|(x, y)| DateTime64(nat_propagate(x.0, y.0, |p, q| p.wrapping_add(q))))
298        .collect();
299    let arr = Array::from_vec(a.dim().clone(), data)?;
300    Ok((arr, target))
301}
302
303/// Element-wise `timedelta[unit_a] + timedelta[unit_b] → timedelta[finer]`.
304///
305/// # Errors
306/// Returns `FerrayError::ShapeMismatch` if the shapes differ.
307pub fn add_timedelta_promoted<D: Dimension>(
308    a: &Array<Timedelta64, D>,
309    unit_a: TimeUnit,
310    b: &Array<Timedelta64, D>,
311    unit_b: TimeUnit,
312) -> FerrayResult<(Array<Timedelta64, D>, TimeUnit)> {
313    check_same_shape(a, b, "add_timedelta_promoted")?;
314    let target = unit_a.finer(unit_b);
315    let scale_a = unit_a.scale_to(target).ok_or_else(|| {
316        FerrayError::invalid_value(format!(
317            "add_timedelta_promoted: cannot rescale {unit_a:?} → {target:?}"
318        ))
319    })?;
320    let scale_b = unit_b.scale_to(target).ok_or_else(|| {
321        FerrayError::invalid_value(format!(
322            "add_timedelta_promoted: cannot rescale {unit_b:?} → {target:?}"
323        ))
324    })?;
325    let a_data: Vec<Timedelta64> = a.iter().copied().collect();
326    let b_data: Vec<Timedelta64> = b.iter().copied().collect();
327    let a_rescaled = if scale_a == 1 {
328        a_data
329    } else {
330        rescale_timedelta(&a_data, scale_a)
331    };
332    let b_rescaled = if scale_b == 1 {
333        b_data
334    } else {
335        rescale_timedelta(&b_data, scale_b)
336    };
337    let data: Vec<Timedelta64> = a_rescaled
338        .iter()
339        .zip(b_rescaled.iter())
340        .map(|(x, y)| Timedelta64(nat_propagate(x.0, y.0, |p, q| p.wrapping_add(q))))
341        .collect();
342    let arr = Array::from_vec(a.dim().clone(), data)?;
343    Ok((arr, target))
344}
345
346#[cfg(test)]
347mod tests {
348    use super::*;
349    use ferray_core::dimension::Ix1;
350
351    #[test]
352    fn isnat_datetime_flags_min_int() {
353        let data = vec![
354            DateTime64(0),
355            DateTime64::nat(),
356            DateTime64(1_700_000_000_000_000_000),
357            DateTime64::nat(),
358        ];
359        let arr = Array::<DateTime64, Ix1>::from_vec(Ix1::new([4]), data).unwrap();
360        let r = isnat_datetime(&arr).unwrap();
361        let v: Vec<bool> = r.iter().copied().collect();
362        assert_eq!(v, vec![false, true, false, true]);
363    }
364
365    #[test]
366    fn isnat_timedelta_flags_min_int() {
367        let data = vec![Timedelta64(1_000), Timedelta64::nat(), Timedelta64(0)];
368        let arr = Array::<Timedelta64, Ix1>::from_vec(Ix1::new([3]), data).unwrap();
369        let r = isnat_timedelta(&arr).unwrap();
370        let v: Vec<bool> = r.iter().copied().collect();
371        assert_eq!(v, vec![false, true, false]);
372    }
373
374    #[test]
375    fn isnat_datetime_all_finite() {
376        let data = vec![DateTime64(1), DateTime64(2), DateTime64(3)];
377        let arr = Array::<DateTime64, Ix1>::from_vec(Ix1::new([3]), data).unwrap();
378        let r = isnat_datetime(&arr).unwrap();
379        let v: Vec<bool> = r.iter().copied().collect();
380        assert_eq!(v, vec![false, false, false]);
381    }
382
383    #[test]
384    fn isnat_datetime_all_nat() {
385        let data = vec![DateTime64::nat(); 5];
386        let arr = Array::<DateTime64, Ix1>::from_vec(Ix1::new([5]), data).unwrap();
387        let r = isnat_datetime(&arr).unwrap();
388        let v: Vec<bool> = r.iter().copied().collect();
389        assert_eq!(v, vec![true; 5]);
390    }
391
392    // ---- Arithmetic kernels ----
393
394    #[test]
395    fn sub_datetime_basic() {
396        let a = Array::<DateTime64, Ix1>::from_vec(
397            Ix1::new([3]),
398            vec![DateTime64(100), DateTime64(200), DateTime64(50)],
399        )
400        .unwrap();
401        let b = Array::<DateTime64, Ix1>::from_vec(
402            Ix1::new([3]),
403            vec![DateTime64(40), DateTime64(150), DateTime64(50)],
404        )
405        .unwrap();
406        let r = sub_datetime(&a, &b).unwrap();
407        let v: Vec<i64> = r.iter().map(|x| x.0).collect();
408        assert_eq!(v, vec![60, 50, 0]);
409    }
410
411    #[test]
412    fn sub_datetime_propagates_nat() {
413        let a = Array::<DateTime64, Ix1>::from_vec(
414            Ix1::new([3]),
415            vec![DateTime64(100), DateTime64::nat(), DateTime64(50)],
416        )
417        .unwrap();
418        let b = Array::<DateTime64, Ix1>::from_vec(
419            Ix1::new([3]),
420            vec![DateTime64(40), DateTime64(150), DateTime64::nat()],
421        )
422        .unwrap();
423        let r = sub_datetime(&a, &b).unwrap();
424        let v: Vec<i64> = r.iter().map(|x| x.0).collect();
425        assert_eq!(v[0], 60);
426        assert_eq!(v[1], NAT);
427        assert_eq!(v[2], NAT);
428    }
429
430    #[test]
431    fn add_datetime_timedelta_basic() {
432        let a = Array::<DateTime64, Ix1>::from_vec(
433            Ix1::new([2]),
434            vec![DateTime64(1000), DateTime64(2000)],
435        )
436        .unwrap();
437        let b = Array::<Timedelta64, Ix1>::from_vec(
438            Ix1::new([2]),
439            vec![Timedelta64(50), Timedelta64::nat()],
440        )
441        .unwrap();
442        let r = add_datetime_timedelta(&a, &b).unwrap();
443        let v: Vec<i64> = r.iter().map(|x| x.0).collect();
444        assert_eq!(v[0], 1050);
445        assert_eq!(v[1], NAT); // NaT propagates
446    }
447
448    #[test]
449    fn sub_datetime_timedelta_basic() {
450        let a = Array::<DateTime64, Ix1>::from_vec(
451            Ix1::new([2]),
452            vec![DateTime64(1000), DateTime64(500)],
453        )
454        .unwrap();
455        let b = Array::<Timedelta64, Ix1>::from_vec(
456            Ix1::new([2]),
457            vec![Timedelta64(100), Timedelta64(50)],
458        )
459        .unwrap();
460        let r = sub_datetime_timedelta(&a, &b).unwrap();
461        let v: Vec<i64> = r.iter().map(|x| x.0).collect();
462        assert_eq!(v, vec![900, 450]);
463    }
464
465    #[test]
466    fn add_sub_timedelta_basic() {
467        let a = Array::<Timedelta64, Ix1>::from_vec(
468            Ix1::new([2]),
469            vec![Timedelta64(10), Timedelta64(20)],
470        )
471        .unwrap();
472        let b = Array::<Timedelta64, Ix1>::from_vec(
473            Ix1::new([2]),
474            vec![Timedelta64(3), Timedelta64(7)],
475        )
476        .unwrap();
477        let s = add_timedelta(&a, &b).unwrap();
478        let sv: Vec<i64> = s.iter().map(|x| x.0).collect();
479        assert_eq!(sv, vec![13, 27]);
480        let d = sub_timedelta(&a, &b).unwrap();
481        let dv: Vec<i64> = d.iter().map(|x| x.0).collect();
482        assert_eq!(dv, vec![7, 13]);
483    }
484
485    #[test]
486    fn arithmetic_shape_mismatch_errors() {
487        let a =
488            Array::<DateTime64, Ix1>::from_vec(Ix1::new([2]), vec![DateTime64(1), DateTime64(2)])
489                .unwrap();
490        let b = Array::<DateTime64, Ix1>::from_vec(
491            Ix1::new([3]),
492            vec![DateTime64(1), DateTime64(2), DateTime64(3)],
493        )
494        .unwrap();
495        assert!(sub_datetime(&a, &b).is_err());
496    }
497
498    // ---- Promoted (mixed-unit) arithmetic ----
499
500    #[test]
501    fn timeunit_finer_picks_finer() {
502        assert_eq!(TimeUnit::Ns.finer(TimeUnit::Us), TimeUnit::Ns);
503        assert_eq!(TimeUnit::Us.finer(TimeUnit::Ns), TimeUnit::Ns);
504        assert_eq!(TimeUnit::S.finer(TimeUnit::Ms), TimeUnit::Ms);
505        assert_eq!(TimeUnit::Ns.finer(TimeUnit::Ns), TimeUnit::Ns);
506    }
507
508    #[test]
509    fn timeunit_scale_to_correct_factors() {
510        assert_eq!(TimeUnit::Us.scale_to(TimeUnit::Ns), Some(1_000));
511        assert_eq!(TimeUnit::S.scale_to(TimeUnit::Us), Some(1_000_000));
512        assert_eq!(TimeUnit::Ns.scale_to(TimeUnit::S), None);
513        assert_eq!(TimeUnit::Ns.scale_to(TimeUnit::Ns), Some(1));
514    }
515
516    #[test]
517    fn sub_datetime_promoted_us_minus_ns() {
518        let a =
519            Array::<DateTime64, Ix1>::from_vec(Ix1::new([2]), vec![DateTime64(5), DateTime64(10)])
520                .unwrap();
521        let b = Array::<DateTime64, Ix1>::from_vec(
522            Ix1::new([2]),
523            vec![DateTime64(1234), DateTime64(7777)],
524        )
525        .unwrap();
526        let (result, unit) = sub_datetime_promoted(&a, TimeUnit::Us, &b, TimeUnit::Ns).unwrap();
527        assert_eq!(unit, TimeUnit::Ns);
528        let v: Vec<i64> = result.iter().map(|x| x.0).collect();
529        assert_eq!(v, vec![5_000 - 1_234, 10_000 - 7_777]);
530    }
531
532    #[test]
533    fn sub_datetime_promoted_propagates_nat() {
534        let a = Array::<DateTime64, Ix1>::from_vec(
535            Ix1::new([2]),
536            vec![DateTime64(5), DateTime64::nat()],
537        )
538        .unwrap();
539        let b = Array::<DateTime64, Ix1>::from_vec(
540            Ix1::new([2]),
541            vec![DateTime64(1000), DateTime64(2000)],
542        )
543        .unwrap();
544        let (result, _) = sub_datetime_promoted(&a, TimeUnit::Us, &b, TimeUnit::Ns).unwrap();
545        let v: Vec<i64> = result.iter().map(|x| x.0).collect();
546        assert_eq!(v[0], 5_000 - 1_000);
547        assert_eq!(v[1], NAT);
548    }
549
550    #[test]
551    fn add_datetime_timedelta_promoted_basic() {
552        let a = Array::<DateTime64, Ix1>::from_vec(Ix1::new([1]), vec![DateTime64(2)]).unwrap();
553        let b = Array::<Timedelta64, Ix1>::from_vec(Ix1::new([1]), vec![Timedelta64(500)]).unwrap();
554        let (result, unit) =
555            add_datetime_timedelta_promoted(&a, TimeUnit::S, &b, TimeUnit::Ms).unwrap();
556        assert_eq!(unit, TimeUnit::Ms);
557        // 2 s = 2000 ms + 500 ms = 2500 ms
558        assert_eq!(result.iter().next().unwrap().0, 2_500);
559    }
560
561    #[test]
562    fn add_timedelta_promoted_basic() {
563        let a = Array::<Timedelta64, Ix1>::from_vec(Ix1::new([1]), vec![Timedelta64(3)]).unwrap();
564        let b = Array::<Timedelta64, Ix1>::from_vec(Ix1::new([1]), vec![Timedelta64(250)]).unwrap();
565        let (result, unit) = add_timedelta_promoted(&a, TimeUnit::S, &b, TimeUnit::Ms).unwrap();
566        assert_eq!(unit, TimeUnit::Ms);
567        assert_eq!(result.iter().next().unwrap().0, 3_250);
568    }
569
570    // ---- Array-level operator overloads ----
571
572    #[test]
573    fn array_sub_datetime_via_operator() {
574        // The operator overload below routes to sub_datetime on `&Array - &Array`.
575        let a = Array::<DateTime64, Ix1>::from_vec(
576            Ix1::new([3]),
577            vec![DateTime64(100), DateTime64(200), DateTime64(50)],
578        )
579        .unwrap();
580        let b = Array::<DateTime64, Ix1>::from_vec(
581            Ix1::new([3]),
582            vec![DateTime64(40), DateTime64(150), DateTime64(50)],
583        )
584        .unwrap();
585        let result = (&a - &b).unwrap();
586        let v: Vec<i64> = result.iter().map(|x| x.0).collect();
587        assert_eq!(v, vec![60, 50, 0]);
588    }
589
590    #[test]
591    fn array_add_datetime_timedelta_via_operator() {
592        let a = Array::<DateTime64, Ix1>::from_vec(
593            Ix1::new([2]),
594            vec![DateTime64(1000), DateTime64(2000)],
595        )
596        .unwrap();
597        let b = Array::<Timedelta64, Ix1>::from_vec(
598            Ix1::new([2]),
599            vec![Timedelta64(50), Timedelta64(75)],
600        )
601        .unwrap();
602        let result = (&a + &b).unwrap();
603        let v: Vec<i64> = result.iter().map(|x| x.0).collect();
604        assert_eq!(v, vec![1050, 2075]);
605    }
606
607    #[test]
608    fn array_timedelta_arith_via_operators() {
609        let a = Array::<Timedelta64, Ix1>::from_vec(
610            Ix1::new([2]),
611            vec![Timedelta64(10), Timedelta64(20)],
612        )
613        .unwrap();
614        let b = Array::<Timedelta64, Ix1>::from_vec(
615            Ix1::new([2]),
616            vec![Timedelta64(3), Timedelta64(7)],
617        )
618        .unwrap();
619        let s = (&a + &b).unwrap();
620        let sv: Vec<i64> = s.iter().map(|x| x.0).collect();
621        assert_eq!(sv, vec![13, 27]);
622        let d = (&a - &b).unwrap();
623        let dv: Vec<i64> = d.iter().map(|x| x.0).collect();
624        assert_eq!(dv, vec![7, 13]);
625    }
626
627    #[test]
628    fn promoted_same_unit_passes_through() {
629        let a = Array::<DateTime64, Ix1>::from_vec(
630            Ix1::new([3]),
631            vec![DateTime64(100), DateTime64(200), DateTime64(50)],
632        )
633        .unwrap();
634        let b = Array::<DateTime64, Ix1>::from_vec(
635            Ix1::new([3]),
636            vec![DateTime64(40), DateTime64(150), DateTime64(50)],
637        )
638        .unwrap();
639        let basic = sub_datetime(&a, &b).unwrap();
640        let (promoted, unit) = sub_datetime_promoted(&a, TimeUnit::Ns, &b, TimeUnit::Ns).unwrap();
641        assert_eq!(unit, TimeUnit::Ns);
642        let bv: Vec<i64> = basic.iter().map(|x| x.0).collect();
643        let pv: Vec<i64> = promoted.iter().map(|x| x.0).collect();
644        assert_eq!(bv, pv);
645    }
646}