polars_compute/if_then_else/
array.rs

1use arrow::array::builder::{ShareStrategy, StaticArrayBuilder, make_builder};
2use arrow::array::{Array, ArrayCollectIterExt, FixedSizeListArray, FixedSizeListArrayBuilder};
3use arrow::bitmap::Bitmap;
4
5use super::{IfThenElseKernel, if_then_else_extend};
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 inner_dt = if_true.dtype().inner_dtype().unwrap();
12        let mut builder =
13            FixedSizeListArrayBuilder::new(if_true.dtype().clone(), make_builder(inner_dt));
14        builder.reserve(mask.len());
15        if_then_else_extend(
16            &mut builder,
17            mask,
18            |b, off, len| b.subslice_extend(if_true, off, len, ShareStrategy::Always),
19            |b, off, len| b.subslice_extend(if_false, off, len, ShareStrategy::Always),
20        );
21        builder.freeze()
22    }
23
24    fn if_then_else_broadcast_true(
25        mask: &Bitmap,
26        if_true: Self::Scalar<'_>,
27        if_false: &Self,
28    ) -> Self {
29        let if_true_list: FixedSizeListArray =
30            std::iter::once(if_true).collect_arr_trusted_with_dtype(if_false.dtype().clone());
31        let inner_dt = if_false.dtype().inner_dtype().unwrap();
32        let mut builder =
33            FixedSizeListArrayBuilder::new(if_false.dtype().clone(), make_builder(inner_dt));
34        builder.reserve(mask.len());
35        if_then_else_extend(
36            &mut builder,
37            mask,
38            |b, _, len| b.subslice_extend_repeated(&if_true_list, 0, 1, len, ShareStrategy::Always),
39            |b, off, len| b.subslice_extend(if_false, off, len, ShareStrategy::Always),
40        );
41        builder.freeze()
42    }
43
44    fn if_then_else_broadcast_false(
45        mask: &Bitmap,
46        if_true: &Self,
47        if_false: Self::Scalar<'_>,
48    ) -> Self {
49        let if_false_list: FixedSizeListArray =
50            std::iter::once(if_false).collect_arr_trusted_with_dtype(if_true.dtype().clone());
51        let inner_dt = if_true.dtype().inner_dtype().unwrap();
52        let mut builder =
53            FixedSizeListArrayBuilder::new(if_true.dtype().clone(), make_builder(inner_dt));
54        builder.reserve(mask.len());
55        if_then_else_extend(
56            &mut builder,
57            mask,
58            |b, off, len| b.subslice_extend(if_true, off, len, ShareStrategy::Always),
59            |b, _, len| {
60                b.subslice_extend_repeated(&if_false_list, 0, 1, len, ShareStrategy::Always)
61            },
62        );
63        builder.freeze()
64    }
65
66    fn if_then_else_broadcast_both(
67        dtype: arrow::datatypes::ArrowDataType,
68        mask: &Bitmap,
69        if_true: Self::Scalar<'_>,
70        if_false: Self::Scalar<'_>,
71    ) -> Self {
72        let if_true_list: FixedSizeListArray =
73            std::iter::once(if_true).collect_arr_trusted_with_dtype(dtype.clone());
74        let if_false_list: FixedSizeListArray =
75            std::iter::once(if_false).collect_arr_trusted_with_dtype(dtype.clone());
76        let inner_dt = dtype.inner_dtype().unwrap();
77        let mut builder = FixedSizeListArrayBuilder::new(dtype.clone(), make_builder(inner_dt));
78        builder.reserve(mask.len());
79        if_then_else_extend(
80            &mut builder,
81            mask,
82            |b, _, len| b.subslice_extend_repeated(&if_true_list, 0, 1, len, ShareStrategy::Always),
83            |b, _, len| {
84                b.subslice_extend_repeated(&if_false_list, 0, 1, len, ShareStrategy::Always)
85            },
86        );
87        builder.freeze()
88    }
89}
polars_compute/if_then_else/array.rs

polars_compute/if_then_else/
array.rs
