1use ferray_core::Array;
23use ferray_core::dimension::Dimension;
24use ferray_core::dtype::{DateTime64, NAT, TimeUnit, Timedelta64};
25use ferray_core::error::{FerrayError, FerrayResult};
26
27pub 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
41pub 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#[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
86pub 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
107pub 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
124pub 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
141pub 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
158pub 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#[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
214pub 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
260pub 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
303pub 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 #[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); }
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 #[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 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 #[test]
573 fn array_sub_datetime_via_operator() {
574 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}