polars_compute/if_then_else/
array.rs

1use arrow::array::growable::{Growable, GrowableFixedSizeList};
2use arrow::array::{Array, ArrayCollectIterExt, FixedSizeListArray};
3use arrow::bitmap::Bitmap;
4
5use super::{if_then_else_extend, IfThenElseKernel};
6
7impl IfThenElseKernel for FixedSizeListArray {
8    type Scalar<'a> = Box<dyn Array>;
9
10    fn if_then_else(mask: &Bitmap, if_true: &Self, if_false: &Self) -> Self {
11        let mut growable = GrowableFixedSizeList::new(vec![if_true, if_false], false, mask.len());
12        unsafe {
13            if_then_else_extend(
14                &mut growable,
15                mask,
16                |g, off, len| g.extend(0, off, len),
17                |g, off, len| g.extend(1, off, len),
18            )
19        };
20        growable.to()
21    }
22
23    fn if_then_else_broadcast_true(
24        mask: &Bitmap,
25        if_true: Self::Scalar<'_>,
26        if_false: &Self,
27    ) -> Self {
28        let if_true_list: FixedSizeListArray =
29            std::iter::once(if_true).collect_arr_trusted_with_dtype(if_false.dtype().clone());
30        let mut growable =
31            GrowableFixedSizeList::new(vec![&if_true_list, if_false], false, mask.len());
32        unsafe {
33            if_then_else_extend(
34                &mut growable,
35                mask,
36                |g, _, len| g.extend_copies(0, 0, 1, len),
37                |g, off, len| g.extend(1, off, len),
38            )
39        };
40        growable.to()
41    }
42
43    fn if_then_else_broadcast_false(
44        mask: &Bitmap,
45        if_true: &Self,
46        if_false: Self::Scalar<'_>,
47    ) -> Self {
48        let if_false_list: FixedSizeListArray =
49            std::iter::once(if_false).collect_arr_trusted_with_dtype(if_true.dtype().clone());
50        let mut growable =
51            GrowableFixedSizeList::new(vec![if_true, &if_false_list], false, mask.len());
52        unsafe {
53            if_then_else_extend(
54                &mut growable,
55                mask,
56                |g, off, len| g.extend(0, off, len),
57                |g, _, len| g.extend_copies(1, 0, 1, len),
58            )
59        };
60        growable.to()
61    }
62
63    fn if_then_else_broadcast_both(
64        dtype: arrow::datatypes::ArrowDataType,
65        mask: &Bitmap,
66        if_true: Self::Scalar<'_>,
67        if_false: Self::Scalar<'_>,
68    ) -> Self {
69        let if_true_list: FixedSizeListArray =
70            std::iter::once(if_true).collect_arr_trusted_with_dtype(dtype.clone());
71        let if_false_list: FixedSizeListArray =
72            std::iter::once(if_false).collect_arr_trusted_with_dtype(dtype.clone());
73        let mut growable =
74            GrowableFixedSizeList::new(vec![&if_true_list, &if_false_list], false, mask.len());
75        unsafe {
76            if_then_else_extend(
77                &mut growable,
78                mask,
79                |g, _, len| g.extend_copies(0, 0, 1, len),
80                |g, _, len| g.extend_copies(1, 0, 1, len),
81            )
82        };
83        growable.to()
84    }
85}
polars_compute/if_then_else/array.rs

polars_compute/if_then_else/
array.rs
