1include!(concat!(env!("OUT_DIR"), "/simd_lanes.rs"));
20
21use std::marker::PhantomData;
22
23use minarrow::{
24 Bitmask, BooleanAVT, BooleanArray, FloatArray, Integer, IntegerArray, Length, MaskedArray,
25 Offset, StringArray, Vec64,
26 aliases::{FloatAVT, IntegerAVT},
27 enums::error::KernelError,
28 vec64,
29};
30use num_traits::{Float, Num, NumCast, One, Zero};
31
32use minarrow::StringAVT;
33use minarrow::utils::confirm_mask_capacity;
34
35#[inline(always)]
37fn new_null_mask(len: usize) -> Bitmask {
38 Bitmask::new_set_all(len, false)
39}
40
41#[inline(always)]
42fn prealloc_vec<T: Copy>(len: usize) -> Vec64<T> {
43 let mut v = Vec64::<T>::with_capacity(len);
44 unsafe { v.set_len(len) };
45 v
46}
47
48#[inline(always)]
59fn rolling_push_pop_to<T, FAdd, FRem>(
60 data: &[T],
61 mask: Option<&Bitmask>,
62 subwindow: usize,
63 mut add: FAdd,
64 mut remove: FRem,
65 zero: T,
66 out: &mut [T],
67 out_mask: &mut Bitmask,
68) where
69 T: Copy,
70 FAdd: FnMut(T, T) -> T,
71 FRem: FnMut(T, T) -> T,
72{
73 let n = data.len();
74 assert_eq!(
75 n,
76 out.len(),
77 "rolling_push_pop_to: input/output length mismatch"
78 );
79
80 if subwindow == 0 {
81 for slot in out.iter_mut() {
82 *slot = zero;
83 }
84 return;
85 }
86
87 let mut agg = zero;
88 let mut invalids = 0usize;
89 for i in 0..n {
90 if mask.map_or(true, |m| unsafe { m.get_unchecked(i) }) {
91 agg = add(agg, data[i]);
92 } else {
93 invalids += 1;
94 }
95 if i + 1 > subwindow {
96 let j = i + 1 - subwindow - 1;
97 if mask.map_or(true, |m| unsafe { m.get_unchecked(j) }) {
98 agg = remove(agg, data[j]);
99 } else {
100 invalids -= 1;
101 }
102 }
103 if i + 1 < subwindow {
104 unsafe { out_mask.set_unchecked(i, false) };
105 out[i] = zero;
106 } else {
107 let ok = invalids == 0;
108 unsafe { out_mask.set_unchecked(i, ok) };
109 out[i] = agg;
110 }
111 }
112}
113
114#[inline(always)]
116pub fn rolling_push_pop<T, FAdd, FRem>(
117 data: &[T],
118 mask: Option<&Bitmask>,
119 subwindow: usize,
120 add: FAdd,
121 remove: FRem,
122 zero: T,
123) -> (Vec64<T>, Bitmask)
124where
125 T: Copy,
126 FAdd: FnMut(T, T) -> T,
127 FRem: FnMut(T, T) -> T,
128{
129 let n = data.len();
130 let mut out = prealloc_vec::<T>(n);
131 let mut out_mask = new_null_mask(n);
132 rolling_push_pop_to(
133 data,
134 mask,
135 subwindow,
136 add,
137 remove,
138 zero,
139 &mut out,
140 &mut out_mask,
141 );
142 (out, out_mask)
143}
144
145#[inline(always)]
151pub fn rolling_extreme_to<T, F>(
152 data: &[T],
153 mask: Option<&Bitmask>,
154 subwindow: usize,
155 mut better: F,
156 zero: T,
157 out: &mut [T],
158 out_mask: &mut Bitmask,
159) where
160 T: Copy,
161 F: FnMut(&T, &T) -> bool,
162{
163 let n = data.len();
164 assert_eq!(
165 n,
166 out.len(),
167 "rolling_extreme_to: input/output length mismatch"
168 );
169
170 if subwindow == 0 {
171 return;
172 }
173
174 for i in 0..n {
175 if i + 1 < subwindow {
176 unsafe { out_mask.set_unchecked(i, false) };
177 out[i] = zero;
178 continue;
179 }
180 let start = i + 1 - subwindow;
181 let mut found = false;
182 let mut extreme = zero;
183 for j in start..=i {
184 if mask.map_or(true, |m| unsafe { m.get_unchecked(j) }) {
185 if !found {
186 extreme = data[j];
187 found = true;
188 } else if better(&data[j], &extreme) {
189 extreme = data[j];
190 }
191 } else {
192 found = false;
193 break;
194 }
195 }
196 unsafe { out_mask.set_unchecked(i, found) };
197 out[i] = if found { extreme } else { zero };
198 }
199}
200
201#[inline(always)]
203pub fn rolling_extreme<T, F>(
204 data: &[T],
205 mask: Option<&Bitmask>,
206 subwindow: usize,
207 better: F,
208 zero: T,
209) -> (Vec64<T>, Bitmask)
210where
211 T: Copy,
212 F: FnMut(&T, &T) -> bool,
213{
214 let n = data.len();
215 let mut out = prealloc_vec::<T>(n);
216 let mut out_mask = new_null_mask(n);
217 rolling_extreme_to(data, mask, subwindow, better, zero, &mut out, &mut out_mask);
218 (out, out_mask)
219}
220
221#[inline]
226pub fn rolling_sum_int_to<T: Num + Copy + Zero>(
227 data: &[T],
228 mask: Option<&Bitmask>,
229 subwindow: usize,
230 out: &mut [T],
231 out_mask: &mut Bitmask,
232) {
233 rolling_push_pop_to(
234 data,
235 mask,
236 subwindow,
237 |a, b| a + b,
238 |a, b| a - b,
239 T::zero(),
240 out,
241 out_mask,
242 );
243 if mask.is_some() && subwindow > 0 && subwindow - 1 < out.len() {
244 unsafe { out_mask.set_unchecked(subwindow - 1, false) };
245 out[subwindow - 1] = T::zero();
246 }
247}
248
249#[inline]
265pub fn rolling_sum_int<T: Num + Copy + Zero>(
266 window: IntegerAVT<'_, T>,
267 subwindow: usize,
268) -> IntegerArray<T> {
269 let (arr, offset, len) = window;
270 let data = &arr.data[offset..offset + len];
271 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
272 let mut out = prealloc_vec::<T>(len);
273 let mut out_mask = new_null_mask(len);
274 rolling_sum_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
275 IntegerArray {
276 data: out.into(),
277 null_mask: Some(out_mask),
278 }
279}
280
281#[inline]
286pub fn rolling_sum_float_to<T: Float + Copy + Zero>(
287 data: &[T],
288 mask: Option<&Bitmask>,
289 subwindow: usize,
290 out: &mut [T],
291 out_mask: &mut Bitmask,
292) {
293 rolling_push_pop_to(
294 data,
295 mask,
296 subwindow,
297 |a, b| a + b,
298 |a, b| a - b,
299 T::zero(),
300 out,
301 out_mask,
302 );
303 if subwindow > 0 && subwindow - 1 < out.len() {
304 out_mask.set(subwindow - 1, false);
305 out[subwindow - 1] = T::zero();
306 }
307}
308
309#[inline]
325pub fn rolling_sum_float<T: Float + Copy + Zero>(
326 window: FloatAVT<'_, T>,
327 subwindow: usize,
328) -> FloatArray<T> {
329 let (arr, offset, len) = window;
330 let data = &arr.data[offset..offset + len];
331 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
332 let mut out = prealloc_vec::<T>(len);
333 let mut out_mask = new_null_mask(len);
334 rolling_sum_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
335 FloatArray {
336 data: out.into(),
337 null_mask: Some(out_mask),
338 }
339}
340
341#[inline]
346pub fn rolling_sum_bool_to(
347 data: &[i32],
348 mask: Option<&Bitmask>,
349 subwindow: usize,
350 out: &mut [i32],
351 out_mask: &mut Bitmask,
352) {
353 rolling_push_pop_to(
354 data,
355 mask,
356 subwindow,
357 |a, b| a + b,
358 |a, b| a - b,
359 0,
360 out,
361 out_mask,
362 );
363 if subwindow > 0 && subwindow - 1 < out.len() {
364 out_mask.set(subwindow - 1, false);
365 out[subwindow - 1] = 0;
366 }
367}
368
369#[inline]
385pub fn rolling_sum_bool(window: BooleanAVT<'_, ()>, subwindow: usize) -> IntegerArray<i32> {
386 let (arr, offset, len) = window;
387 let bools: Vec<i32> = arr.iter_range(offset, len).map(|b| b as i32).collect();
388 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
389 let mut out = prealloc_vec::<i32>(len);
390 let mut out_mask = new_null_mask(len);
391 rolling_sum_bool_to(&bools, mask.as_ref(), subwindow, &mut out, &mut out_mask);
392 IntegerArray {
393 data: out.into(),
394 null_mask: Some(out_mask),
395 }
396}
397
398#[inline]
400pub fn rolling_product_int_to<T: Num + Copy + One + Zero>(
401 data: &[T],
402 mask: Option<&Bitmask>,
403 subwindow: usize,
404 out: &mut [T],
405 out_mask: &mut Bitmask,
406) {
407 rolling_push_pop_to(
408 data,
409 mask,
410 subwindow,
411 |a, b| a * b,
412 |a, b| a / b,
413 T::one(),
414 out,
415 out_mask,
416 );
417}
418
419#[inline]
425pub fn rolling_product_int<T: Num + Copy + One + Zero>(
426 window: IntegerAVT<'_, T>,
427 subwindow: usize,
428) -> IntegerArray<T> {
429 let (arr, offset, len) = window;
430 let data = &arr.data[offset..offset + len];
431 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
432 let mut out = prealloc_vec::<T>(len);
433 let mut out_mask = new_null_mask(len);
434 rolling_product_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
435 IntegerArray {
436 data: out.into(),
437 null_mask: Some(out_mask),
438 }
439}
440
441#[inline]
443pub fn rolling_product_float_to<T: Float + Copy + One + Zero>(
444 data: &[T],
445 mask: Option<&Bitmask>,
446 subwindow: usize,
447 out: &mut [T],
448 out_mask: &mut Bitmask,
449) {
450 rolling_push_pop_to(
451 data,
452 mask,
453 subwindow,
454 |a, b| a * b,
455 |a, b| a / b,
456 T::one(),
457 out,
458 out_mask,
459 );
460}
461
462#[inline]
468pub fn rolling_product_float<T: Float + Copy + One + Zero>(
469 window: FloatAVT<'_, T>,
470 subwindow: usize,
471) -> FloatArray<T> {
472 let (arr, offset, len) = window;
473 let data = &arr.data[offset..offset + len];
474 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
475 let mut out = prealloc_vec::<T>(len);
476 let mut out_mask = new_null_mask(len);
477 rolling_product_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
478 FloatArray {
479 data: out.into(),
480 null_mask: Some(out_mask),
481 }
482}
483
484#[inline]
488pub fn rolling_product_bool_to(
489 data: &[i32],
490 mask: Option<&Bitmask>,
491 subwindow: usize,
492 out: &mut Bitmask,
493 out_mask: &mut Bitmask,
494) {
495 let n = data.len();
496 for i in 0..n {
497 let start = if i + 1 >= subwindow {
498 i + 1 - subwindow
499 } else {
500 0
501 };
502 let mut acc = true;
503 let mut valid = subwindow > 0 && i + 1 >= subwindow;
504 for j in start..=i {
505 let is_valid = mask.map_or(true, |m| unsafe { m.get_unchecked(j) });
506 if is_valid {
507 acc &= data[j] != 0;
508 } else {
509 valid = false;
510 break;
511 }
512 }
513 unsafe { out_mask.set_unchecked(i, valid) };
514 out.set(i, valid && acc);
515 }
516}
517
518#[inline]
524pub fn rolling_product_bool(window: BooleanAVT<'_, ()>, subwindow: usize) -> BooleanArray<()> {
525 let (arr, offset, len) = window;
526 let n = len;
527 let mut out_mask = new_null_mask(n);
528 let mut out = Bitmask::new_set_all(n, false);
529
530 for i in 0..n {
531 let start = if i + 1 >= subwindow {
532 i + 1 - subwindow
533 } else {
534 0
535 };
536 let mut acc = true;
537 let mut valid = subwindow > 0 && i + 1 >= subwindow;
538 for j in start..=i {
539 match unsafe { arr.get_unchecked(offset + j) } {
540 Some(val) => acc &= val,
541 None => {
542 valid = false;
543 break;
544 }
545 }
546 }
547 unsafe { out_mask.set_unchecked(i, valid) };
548 out.set(i, valid && acc);
549 }
550
551 BooleanArray {
552 data: out.into(),
553 null_mask: Some(out_mask),
554 len: n,
555 _phantom: PhantomData,
556 }
557}
558
559#[inline]
561pub fn rolling_mean_int_to<T: NumCast + Copy + Zero>(
562 data: &[T],
563 mask: Option<&Bitmask>,
564 subwindow: usize,
565 out: &mut [f64],
566 out_mask: &mut Bitmask,
567) {
568 let n = data.len();
569 if subwindow == 0 {
570 return;
571 }
572 for i in 0..n {
573 if i + 1 < subwindow {
574 unsafe { out_mask.set_unchecked(i, false) };
575 out[i] = 0.0;
576 continue;
577 }
578 let start = i + 1 - subwindow;
579 let mut sum = 0.0;
580 let mut valid = true;
581 for j in start..=i {
582 if mask.map_or(true, |m| unsafe { m.get_unchecked(j) }) {
583 sum += num_traits::cast(data[j]).unwrap_or(0.0);
584 } else {
585 valid = false;
586 break;
587 }
588 }
589 unsafe { out_mask.set_unchecked(i, valid) };
590 out[i] = if valid { sum / subwindow as f64 } else { 0.0 };
591 }
592 if subwindow > 0 && subwindow - 1 < out.len() {
593 unsafe { out_mask.set_unchecked(subwindow - 1, false) };
594 out[subwindow - 1] = 0.0;
595 }
596}
597
598#[inline]
600pub fn rolling_mean_int<T: NumCast + Copy + Zero>(
601 window: IntegerAVT<'_, T>,
602 subwindow: usize,
603) -> FloatArray<f64> {
604 let (arr, offset, len) = window;
605 let data = &arr.data[offset..offset + len];
606 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
607 let mut out = prealloc_vec::<f64>(len);
608 let mut out_mask = new_null_mask(len);
609 rolling_mean_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
610 FloatArray {
611 data: out.into(),
612 null_mask: Some(out_mask),
613 }
614}
615
616#[inline]
618pub fn rolling_mean_float_to<T: Float + Copy + Zero>(
619 data: &[T],
620 mask: Option<&Bitmask>,
621 subwindow: usize,
622 out: &mut [T],
623 out_mask: &mut Bitmask,
624) {
625 let n = data.len();
626 if subwindow == 0 {
627 return;
628 }
629 for i in 0..n {
630 if i + 1 < subwindow {
631 unsafe { out_mask.set_unchecked(i, false) };
632 out[i] = T::zero();
633 continue;
634 }
635 let start = i + 1 - subwindow;
636 let mut sum = T::zero();
637 let mut valid = true;
638 for j in start..=i {
639 if mask.map_or(true, |m| unsafe { m.get_unchecked(j) }) {
640 sum = sum + data[j];
641 } else {
642 valid = false;
643 break;
644 }
645 }
646 unsafe { out_mask.set_unchecked(i, valid) };
647 out[i] = if valid {
648 sum / T::from(subwindow as u32).unwrap()
649 } else {
650 T::zero()
651 };
652 }
653 if subwindow > 0 && subwindow - 1 < out.len() {
654 unsafe { out_mask.set_unchecked(subwindow - 1, false) };
655 out[subwindow - 1] = T::zero();
656 }
657}
658
659#[inline]
661pub fn rolling_mean_float<T: Float + Copy + Zero>(
662 window: FloatAVT<'_, T>,
663 subwindow: usize,
664) -> FloatArray<T> {
665 let (arr, offset, len) = window;
666 let data = &arr.data[offset..offset + len];
667 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
668 let mut out = prealloc_vec::<T>(len);
669 let mut out_mask = new_null_mask(len);
670 rolling_mean_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
671 FloatArray {
672 data: out.into(),
673 null_mask: Some(out_mask),
674 }
675}
676
677#[inline]
679pub fn rolling_min_int_to<T: Ord + Copy + Zero>(
680 data: &[T],
681 mask: Option<&Bitmask>,
682 subwindow: usize,
683 out: &mut [T],
684 out_mask: &mut Bitmask,
685) {
686 rolling_extreme_to(
687 data,
688 mask,
689 subwindow,
690 |a, b| a < b,
691 T::zero(),
692 out,
693 out_mask,
694 );
695 if subwindow > 0 && subwindow - 1 < out.len() {
696 out_mask.set(subwindow - 1, false);
697 out[subwindow - 1] = T::zero();
698 }
699}
700
701#[inline]
703pub fn rolling_min_int<T: Ord + Copy + Zero>(
704 window: IntegerAVT<'_, T>,
705 subwindow: usize,
706) -> IntegerArray<T> {
707 let (arr, offset, len) = window;
708 let data = &arr.data[offset..offset + len];
709 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
710 let mut out = prealloc_vec::<T>(len);
711 let mut out_mask = new_null_mask(len);
712 rolling_min_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
713 IntegerArray {
714 data: out.into(),
715 null_mask: Some(out_mask),
716 }
717}
718
719#[inline]
721pub fn rolling_max_int_to<T: Ord + Copy + Zero>(
722 data: &[T],
723 mask: Option<&Bitmask>,
724 subwindow: usize,
725 out: &mut [T],
726 out_mask: &mut Bitmask,
727) {
728 rolling_extreme_to(
729 data,
730 mask,
731 subwindow,
732 |a, b| a > b,
733 T::zero(),
734 out,
735 out_mask,
736 );
737}
738
739#[inline]
741pub fn rolling_max_int<T: Ord + Copy + Zero>(
742 window: IntegerAVT<'_, T>,
743 subwindow: usize,
744) -> IntegerArray<T> {
745 let (arr, offset, len) = window;
746 let data = &arr.data[offset..offset + len];
747 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
748 let mut out = prealloc_vec::<T>(len);
749 let mut out_mask = new_null_mask(len);
750 rolling_max_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
751 IntegerArray {
752 data: out.into(),
753 null_mask: Some(out_mask),
754 }
755}
756
757#[inline]
759pub fn rolling_min_float_to<T: Float + Copy + Zero>(
760 data: &[T],
761 mask: Option<&Bitmask>,
762 subwindow: usize,
763 out: &mut [T],
764 out_mask: &mut Bitmask,
765) {
766 rolling_extreme_to(
767 data,
768 mask,
769 subwindow,
770 |a, b| a < b,
771 T::zero(),
772 out,
773 out_mask,
774 );
775 if subwindow > 0 && subwindow - 1 < out.len() {
776 out_mask.set(subwindow - 1, false);
777 out[subwindow - 1] = T::zero();
778 }
779}
780
781#[inline]
783pub fn rolling_min_float<T: Float + Copy + Zero>(
784 window: FloatAVT<'_, T>,
785 subwindow: usize,
786) -> FloatArray<T> {
787 let (arr, offset, len) = window;
788 let data = &arr.data[offset..offset + len];
789 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
790 let mut out = prealloc_vec::<T>(len);
791 let mut out_mask = new_null_mask(len);
792 rolling_min_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
793 FloatArray {
794 data: out.into(),
795 null_mask: Some(out_mask),
796 }
797}
798
799#[inline]
801pub fn rolling_max_float_to<T: Float + Copy + Zero>(
802 data: &[T],
803 mask: Option<&Bitmask>,
804 subwindow: usize,
805 out: &mut [T],
806 out_mask: &mut Bitmask,
807) {
808 rolling_extreme_to(
809 data,
810 mask,
811 subwindow,
812 |a, b| a > b,
813 T::zero(),
814 out,
815 out_mask,
816 );
817}
818
819#[inline]
821pub fn rolling_max_float<T: Float + Copy + Zero>(
822 window: FloatAVT<'_, T>,
823 subwindow: usize,
824) -> FloatArray<T> {
825 let (arr, offset, len) = window;
826 let data = &arr.data[offset..offset + len];
827 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
828 let mut out = prealloc_vec::<T>(len);
829 let mut out_mask = new_null_mask(len);
830 rolling_max_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
831 FloatArray {
832 data: out.into(),
833 null_mask: Some(out_mask),
834 }
835}
836
837#[inline]
860pub fn rolling_count(window: (Offset, Length), subwindow: usize) -> IntegerArray<i32> {
861 let (_offset, len) = window;
862 let mut out = prealloc_vec::<i32>(len);
863 let mut out_mask = new_null_mask(len);
864 for i in 0..len {
865 let start = if i + 1 >= subwindow {
866 i + 1 - subwindow
867 } else {
868 0
869 };
870 let count = (i - start + 1) as i32;
871 let valid_row = subwindow > 0 && i + 1 >= subwindow;
872 unsafe { out_mask.set_unchecked(i, valid_row) };
873 out[i] = if valid_row { count } else { 0 };
874 }
875 IntegerArray {
876 data: out.into(),
877 null_mask: Some(out_mask),
878 }
879}
880
881#[inline(always)]
884fn rank_numeric<T, F>(data: &[T], mask: Option<&Bitmask>, mut cmp: F) -> IntegerArray<i32>
885where
886 T: Copy,
887 F: FnMut(&T, &T) -> std::cmp::Ordering,
888{
889 let n = data.len();
890 let mut indices: Vec<usize> = (0..n).collect();
891 indices.sort_by(|&i, &j| cmp(&data[i], &data[j]));
892
893 let mut out = vec64![0i32; n];
894 let mut out_mask = Bitmask::new_set_all(n, false);
895
896 for (rank, &i) in indices.iter().enumerate() {
897 if mask.map_or(true, |m| unsafe { m.get_unchecked(i) }) {
898 out[i] = (rank + 1) as i32;
899 unsafe { out_mask.set_unchecked(i, true) };
900 }
901 }
902
903 IntegerArray {
904 data: out.into(),
905 null_mask: Some(out_mask),
906 }
907}
908
909#[inline(always)]
946pub fn rank_int<T: Ord + Copy>(window: IntegerAVT<T>) -> IntegerArray<i32> {
947 let (arr, offset, len) = window;
948 let data = &arr.data[offset..offset + len];
949 let null_mask = if len != arr.data.len() {
950 &arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
951 } else {
952 &arr.null_mask
953 };
954 rank_numeric(data, null_mask.as_ref(), |a, b| a.cmp(b))
955}
956
957#[inline(always)]
1000pub fn rank_float<T: Float + Copy>(window: FloatAVT<T>) -> IntegerArray<i32> {
1001 let (arr, offset, len) = window;
1002 let data = &arr.data[offset..offset + len];
1003 let null_mask = if len != arr.data.len() {
1004 &arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
1005 } else {
1006 &arr.null_mask
1007 };
1008 rank_numeric(data, null_mask.as_ref(), |a, b| a.partial_cmp(b).unwrap())
1009}
1010
1011#[inline(always)]
1053pub fn rank_str<T: Integer>(arr: StringAVT<T>) -> Result<IntegerArray<i32>, KernelError> {
1054 let (array, offset, len) = arr;
1055 let mask = array.null_mask.as_ref();
1056 let _ = confirm_mask_capacity(array.len(), mask)?;
1057
1058 let mut tuples: Vec<(usize, &str)> = (0..len)
1060 .filter(|&i| mask.map_or(true, |m| unsafe { m.get_unchecked(offset + i) }))
1061 .map(|i| (i, unsafe { array.get_unchecked(offset + i) }.unwrap_or("")))
1062 .collect();
1063
1064 tuples.sort_by(|a, b| a.1.cmp(&b.1));
1066
1067 let mut out = vec64![0i32; len];
1068 let mut out_mask = new_null_mask(len);
1069
1070 for (rank, (i, _)) in tuples.iter().enumerate() {
1072 out[*i] = (rank + 1) as i32;
1073 unsafe { out_mask.set_unchecked(*i, true) };
1074 }
1075
1076 Ok(IntegerArray {
1077 data: out.into(),
1078 null_mask: Some(out_mask),
1079 })
1080}
1081
1082#[inline(always)]
1083fn dense_rank_numeric<T, F, G>(
1084 data: &[T],
1085 mask: Option<&Bitmask>,
1086 mut sort: F,
1087 mut eq: G,
1088) -> Result<IntegerArray<i32>, KernelError>
1089where
1090 T: Copy,
1091 F: FnMut(&T, &T) -> std::cmp::Ordering,
1092 G: FnMut(&T, &T) -> bool,
1093{
1094 let n = data.len();
1095 let _ = confirm_mask_capacity(n, mask)?;
1096 let mut uniqs: Vec<T> = (0..n)
1097 .filter(|&i| mask.map_or(true, |m| unsafe { m.get_unchecked(i) }))
1098 .map(|i| data[i])
1099 .collect();
1100
1101 uniqs.sort_by(&mut sort);
1102 uniqs.dedup_by(|a, b| eq(&*a, &*b));
1103
1104 let mut out = prealloc_vec::<i32>(n);
1105 let mut out_mask = Bitmask::new_set_all(n, false);
1106
1107 for i in 0..n {
1108 if mask.map_or(true, |m| unsafe { m.get_unchecked(i) }) {
1109 let rank = uniqs.binary_search_by(|x| sort(x, &data[i])).unwrap() + 1;
1110 out[i] = rank as i32;
1111 unsafe { out_mask.set_unchecked(i, true) };
1112 } else {
1113 out[i] = 0;
1114 }
1115 }
1116
1117 Ok(IntegerArray {
1118 data: out.into(),
1119 null_mask: Some(out_mask),
1120 })
1121}
1122
1123#[inline(always)]
1165pub fn dense_rank_int<T: Ord + Copy>(
1166 window: IntegerAVT<T>,
1167) -> Result<IntegerArray<i32>, KernelError> {
1168 let (arr, offset, len) = window;
1169 let data = &arr.data[offset..offset + len];
1170 let null_mask = if len != arr.data.len() {
1171 &arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
1172 } else {
1173 &arr.null_mask
1174 };
1175 dense_rank_numeric(data, null_mask.as_ref(), |a, b| a.cmp(b), |a, b| a == b)
1176}
1177
1178#[inline(always)]
1210pub fn dense_rank_float<T: Float + Copy>(
1211 window: FloatAVT<T>,
1212) -> Result<IntegerArray<i32>, KernelError> {
1213 let (arr, offset, len) = window;
1214 let data = &arr.data[offset..offset + len];
1215 let null_mask = if len != arr.data.len() {
1216 &arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
1217 } else {
1218 &arr.null_mask
1219 };
1220 dense_rank_numeric(
1221 data,
1222 null_mask.as_ref(),
1223 |a, b| a.partial_cmp(b).unwrap(),
1224 |a, b| a == b,
1225 )
1226}
1227
1228#[inline(always)]
1266pub fn dense_rank_str<T: Integer>(arr: StringAVT<T>) -> Result<IntegerArray<i32>, KernelError> {
1267 let (array, offset, len) = arr;
1268 let mask = array.null_mask.as_ref();
1269 let _ = confirm_mask_capacity(array.len(), mask)?;
1270
1271 let mut vals: Vec<&str> = (0..len)
1273 .filter(|&i| mask.map_or(true, |m| unsafe { m.get_unchecked(offset + i) }))
1274 .map(|i| unsafe { array.get_unchecked(offset + i) }.unwrap_or(""))
1275 .collect();
1276 vals.sort();
1277 vals.dedup();
1278
1279 let mut out = prealloc_vec::<i32>(len);
1280 let mut out_mask = Bitmask::new_set_all(len, false);
1281
1282 for i in 0..len {
1283 let valid = mask.map_or(true, |m| unsafe { m.get_unchecked(offset + i) });
1284 if valid {
1285 let rank = vals
1286 .binary_search(&unsafe { array.get_unchecked(offset + i) }.unwrap_or(""))
1287 .unwrap()
1288 + 1;
1289 out[i] = rank as i32;
1290 unsafe { out_mask.set_unchecked(i, true) };
1291 } else {
1292 out[i] = 0;
1293 }
1294 }
1295
1296 Ok(IntegerArray {
1297 data: out.into(),
1298 null_mask: Some(out_mask),
1299 })
1300}
1301
1302#[inline(always)]
1308fn shift_with_bounds_to<T: Copy>(
1309 data: &[T],
1310 mask: Option<&Bitmask>,
1311 len: usize,
1312 offset_fn: impl Fn(usize) -> Option<usize>,
1313 default: T,
1314 out: &mut [T],
1315 out_mask: &mut Bitmask,
1316) {
1317 assert_eq!(
1318 len,
1319 out.len(),
1320 "shift_with_bounds_to: input/output length mismatch"
1321 );
1322 for i in 0..len {
1323 if let Some(j) = offset_fn(i) {
1324 out[i] = data[j];
1325 let is_valid = mask.map_or(true, |m| unsafe { m.get_unchecked(j) });
1326 unsafe { out_mask.set_unchecked(i, is_valid) };
1327 } else {
1328 out[i] = default;
1329 }
1330 }
1331}
1332
1333#[inline(always)]
1335fn shift_with_bounds<T: Copy>(
1336 data: &[T],
1337 mask: Option<&Bitmask>,
1338 len: usize,
1339 offset_fn: impl Fn(usize) -> Option<usize>,
1340 default: T,
1341) -> (Vec64<T>, Bitmask) {
1342 let mut out = prealloc_vec::<T>(len);
1343 let mut out_mask = Bitmask::new_set_all(len, false);
1344 shift_with_bounds_to(data, mask, len, offset_fn, default, &mut out, &mut out_mask);
1345 (out, out_mask)
1346}
1347
1348#[inline(always)]
1349fn shift_str_with_bounds<T: Integer>(
1350 arr: StringAVT<T>,
1351 offset_fn: impl Fn(usize) -> Option<usize>,
1352) -> Result<StringArray<T>, KernelError> {
1353 let (array, offset, len) = arr;
1354 let src_mask = array.null_mask.as_ref();
1355 let _ = confirm_mask_capacity(array.len(), src_mask)?;
1356
1357 let mut offsets = Vec64::<T>::with_capacity(len + 1);
1359 unsafe {
1360 offsets.set_len(len + 1);
1361 }
1362 offsets[0] = T::zero();
1363
1364 let mut total_bytes = 0;
1365 let mut string_lengths = vec![0usize; len];
1366
1367 for i in 0..len {
1368 let byte_len = if let Some(j) = offset_fn(i) {
1369 let src_idx = offset + j;
1370 let valid = src_mask.map_or(true, |m| unsafe { m.get_unchecked(src_idx) });
1371 if valid {
1372 unsafe { array.get_unchecked(src_idx).unwrap_or("").len() }
1373 } else {
1374 0
1375 }
1376 } else {
1377 0
1378 };
1379 total_bytes += byte_len;
1380 string_lengths[i] = byte_len;
1381 offsets[i + 1] = T::from_usize(total_bytes);
1382 }
1383
1384 let mut data = Vec64::<u8>::with_capacity(total_bytes);
1386 let mut out_mask = Bitmask::new_set_all(len, false);
1387
1388 for i in 0..len {
1390 if let Some(j) = offset_fn(i) {
1391 let src_idx = offset + j;
1392 let valid = src_mask.map_or(true, |m| unsafe { m.get_unchecked(src_idx) });
1393 if valid {
1394 let s = unsafe { array.get_unchecked(src_idx).unwrap_or("") };
1395 data.extend_from_slice(s.as_bytes());
1396 unsafe { out_mask.set_unchecked(i, true) };
1397 continue;
1398 }
1399 }
1400 }
1402
1403 Ok(StringArray {
1404 offsets: offsets.into(),
1405 data: data.into(),
1406 null_mask: Some(out_mask),
1407 })
1408}
1409
1410#[inline]
1438pub fn lag_int_to<T: Copy + Default>(
1439 data: &[T],
1440 mask: Option<&Bitmask>,
1441 n: usize,
1442 out: &mut [T],
1443 out_mask: &mut Bitmask,
1444) {
1445 let len = data.len();
1446 shift_with_bounds_to(
1447 data,
1448 mask,
1449 len,
1450 |i| if i >= n { Some(i - n) } else { None },
1451 T::default(),
1452 out,
1453 out_mask,
1454 );
1455}
1456
1457#[inline]
1458pub fn lag_int<T: Copy + Default>(window: IntegerAVT<T>, n: usize) -> IntegerArray<T> {
1459 let (arr, offset, len) = window;
1460 let data_window = &arr.data[offset..offset + len];
1461 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
1462 let mut out = prealloc_vec::<T>(len);
1463 let mut out_mask = Bitmask::new_set_all(len, false);
1464 lag_int_to(data_window, mask.as_ref(), n, &mut out, &mut out_mask);
1465 IntegerArray {
1466 data: out.into(),
1467 null_mask: Some(out_mask),
1468 }
1469}
1470
1471#[inline]
1497pub fn lead_int_to<T: Copy + Default>(
1498 data: &[T],
1499 mask: Option<&Bitmask>,
1500 n: usize,
1501 out: &mut [T],
1502 out_mask: &mut Bitmask,
1503) {
1504 let len = data.len();
1505 shift_with_bounds_to(
1506 data,
1507 mask,
1508 len,
1509 |i| if i + n < len { Some(i + n) } else { None },
1510 T::default(),
1511 out,
1512 out_mask,
1513 );
1514}
1515
1516#[inline]
1517pub fn lead_int<T: Copy + Default>(window: IntegerAVT<T>, n: usize) -> IntegerArray<T> {
1518 let (arr, offset, len) = window;
1519 let data_window = &arr.data[offset..offset + len];
1520 let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
1521 let mut out = prealloc_vec::<T>(len);
1522 let mut out_mask = Bitmask::new_set_all(len, false);
1523 lead_int_to(data_window, mask.as_ref(), n, &mut out, &mut out_mask);
1524 IntegerArray {
1525 data: out.into(),
1526 null_mask: Some(out_mask),
1527 }
1528}
1529
1530#[inline]
1532pub fn lag_float_to<T: Copy + num_traits::Zero>(
1533 data: &[T],
1534 mask: Option<&Bitmask>,
1535 n: usize,
1536 out: &mut [T],
1537 out_mask: &mut Bitmask,
1538) {
1539 let len = data.len();
1540 shift_with_bounds_to(
1541 data,
1542 mask,
1543 len,
1544 |i| if i >= n { Some(i - n) } else { None },
1545 T::zero(),
1546 out,
1547 out_mask,
1548 );
1549}
1550
1551#[inline]
1553pub fn lag_float<T: Copy + num_traits::Zero>(window: FloatAVT<T>, n: usize) -> FloatArray<T> {
1554 let (arr, offset, len) = window;
1555 let data_window = &arr.data[offset..offset + len];
1556 let mask_window = if len != arr.data.len() {
1557 arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
1558 } else {
1559 arr.null_mask.clone()
1560 };
1561 let (data, null_mask) = shift_with_bounds(
1562 data_window,
1563 mask_window.as_ref(),
1564 len,
1565 |i| if i >= n { Some(i - n) } else { None },
1566 T::zero(),
1567 );
1568 FloatArray {
1569 data: data.into(),
1570 null_mask: Some(null_mask),
1571 }
1572}
1573
1574#[inline]
1606pub fn lead_float<T: Copy + num_traits::Zero>(window: FloatAVT<T>, n: usize) -> FloatArray<T> {
1607 let (arr, offset, len) = window;
1608 let data_window = &arr.data[offset..offset + len];
1609 let mask_window = if len != arr.data.len() {
1610 arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
1611 } else {
1612 arr.null_mask.clone()
1613 };
1614 let (data, null_mask) = shift_with_bounds(
1615 data_window,
1616 mask_window.as_ref(),
1617 len,
1618 |i| if i + n < len { Some(i + n) } else { None },
1619 T::zero(),
1620 );
1621 FloatArray {
1622 data: data.into(),
1623 null_mask: Some(null_mask),
1624 }
1625}
1626
1627#[inline]
1668pub fn lag_str<T: Integer>(arr: StringAVT<T>, n: usize) -> Result<StringArray<T>, KernelError> {
1669 shift_str_with_bounds(arr, |i| if i >= n { Some(i - n) } else { None })
1670}
1671
1672#[inline]
1698pub fn lead_str<T: Integer>(arr: StringAVT<T>, n: usize) -> Result<StringArray<T>, KernelError> {
1699 let (_array, _offset, len) = arr;
1700 shift_str_with_bounds(arr, |i| if i + n < len { Some(i + n) } else { None })
1701}
1702
1703#[inline(always)]
1744pub fn shift_int<T: Copy + Default>(window: IntegerAVT<T>, offset: isize) -> IntegerArray<T> {
1745 let (arr, win_offset, win_len) = window;
1746 if offset == 0 {
1747 return IntegerArray {
1748 data: Vec64::from_slice(&arr.data[win_offset..win_offset + win_len]).into(),
1749 null_mask: if win_len != arr.data.len() {
1750 arr.null_mask
1751 .as_ref()
1752 .map(|m| m.slice_clone(win_offset, win_len))
1753 } else {
1754 arr.null_mask.clone()
1755 },
1756 };
1757 } else if offset > 0 {
1758 lead_int((arr, win_offset, win_len), offset as usize)
1759 } else {
1760 lag_int((arr, win_offset, win_len), offset.unsigned_abs())
1761 }
1762}
1763
1764#[inline(always)]
1790pub fn shift_float<T: Copy + num_traits::Zero>(
1791 window: FloatAVT<T>,
1792 offset: isize,
1793) -> FloatArray<T> {
1794 let (arr, win_offset, win_len) = window;
1795 if offset == 0 {
1796 return FloatArray {
1797 data: Vec64::from_slice(&arr.data[win_offset..win_offset + win_len]).into(),
1798 null_mask: if win_len != arr.data.len() {
1799 arr.null_mask
1800 .as_ref()
1801 .map(|m| m.slice_clone(win_offset, win_len))
1802 } else {
1803 arr.null_mask.clone()
1804 },
1805 };
1806 } else if offset > 0 {
1807 lead_float((arr, win_offset, win_len), offset as usize)
1808 } else {
1809 lag_float((arr, win_offset, win_len), offset.unsigned_abs())
1810 }
1811}
1812
1813#[inline(always)]
1848pub fn shift_str<T: Integer>(
1849 arr: StringAVT<T>,
1850 shift_offset: isize,
1851) -> Result<StringArray<T>, KernelError> {
1852 if shift_offset == 0 {
1853 let (array, off, len) = arr;
1855 Ok(array.slice_clone(off, len))
1856 } else if shift_offset > 0 {
1857 lead_str(arr, shift_offset as usize)
1858 } else {
1859 lag_str(arr, shift_offset.unsigned_abs())
1860 }
1861}
1862
1863#[cfg(test)]
1864mod tests {
1865 use minarrow::structs::variants::float::FloatArray;
1866 use minarrow::structs::variants::integer::IntegerArray;
1867 use minarrow::structs::variants::string::StringArray;
1868 use minarrow::{Bitmask, BooleanArray};
1869
1870 use super::*;
1871
1872 fn bm(bits: &[bool]) -> Bitmask {
1876 let mut m = Bitmask::new_set_all(bits.len(), false);
1877 for (i, &b) in bits.iter().enumerate() {
1878 m.set(i, b);
1879 }
1880 m
1881 }
1882
1883 fn expect_int<T: PartialEq + std::fmt::Debug + Clone>(
1885 arr: &IntegerArray<T>,
1886 values: &[T],
1887 valid: &[bool],
1888 ) {
1889 assert_eq!(arr.data.as_slice(), values);
1890 let mask = arr.null_mask.as_ref().expect("mask missing");
1891 for (i, &v) in valid.iter().enumerate() {
1892 assert_eq!(mask.get(i), v, "mask bit {}", i);
1893 }
1894 }
1895
1896 fn expect_float<T: num_traits::Float + std::fmt::Debug>(
1898 arr: &FloatArray<T>,
1899 values: &[T],
1900 valid: &[bool],
1901 eps: T,
1902 ) {
1903 let data = arr.data.as_slice();
1904 assert_eq!(data.len(), values.len());
1905 for (a, b) in data.iter().zip(values.iter()) {
1906 assert!((*a - *b).abs() <= eps, "value mismatch {:?} vs {:?}", a, b);
1907 }
1908 let mask = arr.null_mask.as_ref().expect("mask missing");
1909 for (i, &v) in valid.iter().enumerate() {
1910 assert_eq!(mask.get(i), v);
1911 }
1912 }
1913
1914 #[test]
1917 fn test_rolling_sum_int_basic() {
1918 let arr = IntegerArray::<i32>::from_slice(&[1, 2, 3, 4, 5]);
1919 let res = rolling_sum_int((&arr, 0, arr.len()), 3);
1920 expect_int(&res, &[0, 0, 6, 9, 12], &[false, false, true, true, true]);
1921 }
1922
1923 #[test]
1924 fn test_rolling_sum_int_masked() {
1925 let mut arr = IntegerArray::<i32>::from_slice(&[1, 2, 3, 4]);
1926 arr.null_mask = Some(bm(&[true, false, true, true]));
1927 let res = rolling_sum_int((&arr, 0, arr.len()), 2);
1928 expect_int(
1929 &res,
1930 &[0, 0, 3, 7],
1931 &[false, false, false, true], );
1933 }
1934
1935 #[test]
1936 fn test_rolling_sum_float() {
1937 let arr = FloatArray::<f32>::from_slice(&[1.0, 2.0, 3.0]);
1938 let res = rolling_sum_float((&arr, 0, arr.len()), 2);
1939 expect_float(&res, &[0.0, 0.0, 5.0], &[false, false, true], 1e-6f32);
1940 }
1941
1942 #[test]
1943 fn test_rolling_sum_bool() {
1944 let bools = BooleanArray::from_slice(&[true, true, false, true]);
1945 let res = rolling_sum_bool((&bools, 0, bools.len()), 2); expect_int(&res, &[0, 0, 1, 1], &[false, false, true, true]);
1947 }
1948
1949 #[test]
1950 fn test_rolling_min_max_mean_count() {
1951 let arr = IntegerArray::<i32>::from_slice(&[3, 1, 4, 1, 5]);
1952 let rmin = rolling_min_int((&arr, 0, arr.len()), 2);
1954 expect_int(&rmin, &[0, 0, 1, 1, 1], &[false, false, true, true, true]);
1955
1956 let rmax = rolling_max_int((&arr, 0, arr.len()), 3);
1958 expect_int(&rmax, &[0, 0, 4, 4, 5], &[false, false, true, true, true]);
1959
1960 let rmean = rolling_mean_int((&arr, 0, arr.len()), 2);
1962 expect_float(
1963 &rmean,
1964 &[0.0, 0.0, 2.5, 2.5, 3.0],
1965 &[false, false, true, true, true],
1966 1e-12,
1967 );
1968
1969 let cnt = rolling_count((0, 5), 3);
1971 expect_int(&cnt, &[0, 0, 3, 3, 3], &[false, false, true, true, true]);
1972 }
1973
1974 #[test]
1977 fn test_rank_int_basic() {
1978 let arr = IntegerArray::<i32>::from_slice(&[30, 10, 20]);
1979 let res = rank_int((&arr, 0, arr.len()));
1980 expect_int(&res, &[3, 1, 2], &[true, true, true]);
1981 }
1982
1983 #[test]
1984 fn test_rank_float_with_nulls() {
1985 let mut arr = FloatArray::<f64>::from_slice(&[2.0, 1.0, 3.0]);
1986 arr.null_mask = Some(bm(&[true, false, true]));
1987 let res = rank_float((&arr, 0, arr.len()));
1988 expect_int(&res, &[2, 0, 3], &[true, false, true]);
1989 }
1990
1991 #[test]
1992 fn test_dense_rank_str_duplicates() {
1993 let arr = StringArray::<u32>::from_slice(&["b", "a", "b", "c"]);
1994 let res = dense_rank_str((&arr, 0, arr.len())).unwrap();
1995 expect_int(&res, &[2, 1, 2, 3], &[true, true, true, true]);
1996 }
1997
1998 #[test]
1999 fn test_dense_rank_str_duplicates_chunk() {
2000 let arr = StringArray::<u32>::from_slice(&["x", "b", "a", "b", "c", "y"]);
2002 let res = dense_rank_str((&arr, 1, 4)).unwrap(); expect_int(&res, &[2, 1, 2, 3], &[true, true, true, true]);
2004 }
2005
2006 #[test]
2009 fn test_lag_lead_int() {
2010 let arr = IntegerArray::<i32>::from_slice(&[10, 20, 30, 40]);
2011 let lag1 = lag_int((&arr, 0, arr.len()), 1);
2012 expect_int(&lag1, &[0, 10, 20, 30], &[false, true, true, true]);
2013
2014 let lead2 = lead_int((&arr, 0, arr.len()), 2);
2015 expect_int(&lead2, &[30, 40, 0, 0], &[true, true, false, false]);
2016 }
2017
2018 #[test]
2019 fn test_shift_float_positive_negative_zero() {
2020 let arr = FloatArray::<f32>::from_slice(&[1.0, 2.0, 3.0]);
2021 let s0 = shift_float((&arr, 0, arr.len()), 0);
2022 assert_eq!(s0.data, arr.data); let s1 = shift_float((&arr, 0, arr.len()), 1);
2025 expect_float(&s1, &[2.0, 3.0, 0.0], &[true, true, false], 1e-6f32);
2026
2027 let s_neg = shift_float((&arr, 0, arr.len()), -1);
2028 expect_float(&s_neg, &[0.0, 1.0, 2.0], &[false, true, true], 1e-6f32);
2029 }
2030
2031 #[test]
2032 fn test_lag_lead_str() {
2033 let arr = StringArray::<u32>::from_slice(&["a", "b", "c"]);
2034 let l1 = lag_str((&arr, 0, arr.len()), 1).unwrap();
2035 assert_eq!(l1.get(0), None);
2036 assert_eq!(l1.get(1), Some("a"));
2037 assert_eq!(l1.get(2), Some("b"));
2038
2039 let d1 = lead_str((&arr, 0, arr.len()), 1).unwrap();
2040 assert_eq!(d1.get(0), Some("b"));
2041 assert_eq!(d1.get(1), Some("c"));
2042 assert_eq!(d1.get(2), None);
2043 }
2044
2045 #[test]
2046 fn test_lag_lead_str_chunk() {
2047 let arr = StringArray::<u32>::from_slice(&["x", "a", "b", "c", "y"]);
2049 let l1 = lag_str((&arr, 1, 3), 1).unwrap();
2050 assert_eq!(l1.get(0), None);
2051 assert_eq!(l1.get(1), Some("a"));
2052 assert_eq!(l1.get(2), Some("b"));
2053
2054 let d1 = lead_str((&arr, 1, 3), 1).unwrap();
2055 assert_eq!(d1.get(0), Some("b"));
2056 assert_eq!(d1.get(1), Some("c"));
2057 assert_eq!(d1.get(2), None);
2058 }
2059
2060 #[test]
2061 fn test_rolling_sum_int_edge_windows() {
2062 let arr = IntegerArray::<i32>::from_slice(&[1, 2, 3, 4, 5]);
2063
2064 let r0 = rolling_sum_int((&arr, 0, arr.len()), 0);
2066 assert_eq!(r0.data.as_slice(), &[0, 0, 0, 0, 0]);
2067 assert_eq!(r0.null_mask.as_ref().unwrap().all_unset(), true);
2068
2069 let r1 = rolling_sum_int((&arr, 0, arr.len()), 1);
2071 assert_eq!(r1.data.as_slice(), &[1, 2, 3, 4, 5]);
2072 assert!(r1.null_mask.as_ref().unwrap().all_set());
2073
2074 let r_large = rolling_sum_int((&arr, 0, arr.len()), 10);
2076 assert_eq!(r_large.data.as_slice(), &[0, 0, 0, 0, 0]);
2077 assert_eq!(r_large.null_mask.as_ref().unwrap().all_unset(), true);
2078 }
2079
2080 #[test]
2081 fn test_rolling_sum_float_masked_nulls_propagate() {
2082 let mut arr = FloatArray::<f32>::from_slice(&[1.0, 2.0, 3.0, 4.0]);
2083 arr.null_mask = Some(bm(&[true, true, false, true]));
2085 let r = rolling_sum_float((&arr, 0, arr.len()), 2);
2086 expect_float(
2091 &r,
2092 &[0.0, 0.0, 2.0, 4.0],
2093 &[false, false, false, false],
2094 1e-6,
2095 );
2096 }
2097
2098 #[test]
2099 fn test_rolling_sum_bool_with_nulls() {
2100 let mut b = BooleanArray::from_slice(&[true, false, true, true]);
2101 b.null_mask = Some(bm(&[true, false, true, true]));
2102 let r = rolling_sum_bool((&b, 0, b.len()), 2);
2103 expect_int(&r, &[0, 0, 1, 2], &[false, false, false, true]);
2105 }
2106
2107 #[test]
2108 fn test_lag_str_null_propagation() {
2109 let mut arr = StringArray::<u32>::from_slice(&["x", "y", "z"]);
2110 arr.null_mask = Some(bm(&[true, false, true])); let lag1 = lag_str((&arr, 0, arr.len()), 1).unwrap();
2112 assert_eq!(lag1.get(0), None); assert_eq!(lag1.get(1), Some("x"));
2114 assert_eq!(lag1.get(2), None); let m = lag1.null_mask.unwrap();
2116 assert_eq!(m.get(0), false);
2117 assert_eq!(m.get(1), true);
2118 assert_eq!(m.get(2), false);
2119 }
2120
2121 #[test]
2122 fn test_lag_str_null_propagation_chunk() {
2123 let mut arr = StringArray::<u32>::from_slice(&["w", "x", "y", "z", "q"]);
2125 arr.null_mask = Some(bm(&[true, true, false, true, true]));
2126 let lag1 = lag_str((&arr, 1, 3), 1).unwrap();
2127 assert_eq!(lag1.get(0), None); assert_eq!(lag1.get(1), Some("x")); assert_eq!(lag1.get(2), None); let m = lag1.null_mask.unwrap();
2131 assert_eq!(m.get(0), false);
2132 assert_eq!(m.get(1), true);
2133 assert_eq!(m.get(2), false);
2134 }
2135
2136 #[test]
2137 fn test_shift_str_large_offset() {
2138 let arr = StringArray::<u32>::from_slice(&["a", "b", "c"]);
2139 let shifted = shift_str((&arr, 0, arr.len()), 10).unwrap(); assert_eq!(shifted.len(), 3);
2141 for i in 0..3 {
2142 assert_eq!(shifted.get(i), None);
2143 assert_eq!(shifted.null_mask.as_ref().unwrap().get(i), false);
2144 }
2145 }
2146
2147 #[test]
2148 fn test_shift_str_large_offset_chunk() {
2149 let arr = StringArray::<u32>::from_slice(&["w", "a", "b", "c", "x"]);
2151 let shifted = shift_str((&arr, 1, 3), 10).unwrap(); assert_eq!(shifted.len(), 3);
2153 for i in 0..3 {
2154 assert_eq!(shifted.get(i), None);
2155 assert_eq!(shifted.null_mask.as_ref().unwrap().get(i), false);
2156 }
2157 }
2158
2159 #[test]
2160 fn test_rank_str_ties_and_nulls() {
2161 let mut arr = StringArray::<u32>::from_slice(&["a", "b", "a", "c"]);
2162 arr.null_mask = Some(bm(&[true, true, false, true]));
2163 let r = rank_str((&arr, 0, arr.len())).unwrap();
2164 expect_int(&r, &[1, 2, 0, 3], &[true, true, false, true]);
2166 }
2167
2168 #[test]
2169 fn test_rank_str_ties_and_nulls_chunk() {
2170 let mut arr = StringArray::<u32>::from_slice(&["w", "a", "b", "a", "c"]);
2172 arr.null_mask = Some(bm(&[true, true, true, false, true]));
2173 let r = rank_str((&arr, 1, 4)).unwrap(); expect_int(&r, &[1, 2, 0, 3], &[true, true, false, true]);
2175 }
2176}