polars_compute/if_then_else/
list.rs

1use arrow::array::builder::{ShareStrategy, StaticArrayBuilder, make_builder};
2use arrow::array::{Array, ArrayCollectIterExt, ListArray, ListArrayBuilder};
3use arrow::bitmap::Bitmap;
4
5use super::{IfThenElseKernel, if_then_else_extend};
6
7impl IfThenElseKernel for ListArray<i64> {
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 = ListArrayBuilder::new(if_true.dtype().clone(), make_builder(inner_dt));
13        builder.reserve(mask.len());
14        if_then_else_extend(
15            &mut builder,
16            mask,
17            |b, off, len| b.subslice_extend(if_true, off, len, ShareStrategy::Always),
18            |b, off, len| b.subslice_extend(if_false, off, len, ShareStrategy::Always),
19        );
20        builder.freeze()
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: ListArray<i64> =
29            std::iter::once(if_true).collect_arr_trusted_with_dtype(if_false.dtype().clone());
30        let inner_dt = if_false.dtype().inner_dtype().unwrap();
31        let mut builder = ListArrayBuilder::new(if_false.dtype().clone(), make_builder(inner_dt));
32        builder.reserve(mask.len());
33        if_then_else_extend(
34            &mut builder,
35            mask,
36            |b, _, len| b.subslice_extend_repeated(&if_true_list, 0, 1, len, ShareStrategy::Always),
37            |b, off, len| b.subslice_extend(if_false, off, len, ShareStrategy::Always),
38        );
39        builder.freeze()
40    }
41
42    fn if_then_else_broadcast_false(
43        mask: &Bitmap,
44        if_true: &Self,
45        if_false: Self::Scalar<'_>,
46    ) -> Self {
47        let if_false_list: ListArray<i64> =
48            std::iter::once(if_false).collect_arr_trusted_with_dtype(if_true.dtype().clone());
49        let inner_dt = if_true.dtype().inner_dtype().unwrap();
50        let mut builder = ListArrayBuilder::new(if_true.dtype().clone(), make_builder(inner_dt));
51        builder.reserve(mask.len());
52        if_then_else_extend(
53            &mut builder,
54            mask,
55            |b, off, len| b.subslice_extend(if_true, off, len, ShareStrategy::Always),
56            |b, _, len| {
57                b.subslice_extend_repeated(&if_false_list, 0, 1, len, ShareStrategy::Always)
58            },
59        );
60        builder.freeze()
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: ListArray<i64> =
70            std::iter::once(if_true).collect_arr_trusted_with_dtype(dtype.clone());
71        let if_false_list: ListArray<i64> =
72            std::iter::once(if_false).collect_arr_trusted_with_dtype(dtype.clone());
73        let inner_dt = dtype.inner_dtype().unwrap();
74        let mut builder = ListArrayBuilder::new(dtype.clone(), make_builder(inner_dt));
75        builder.reserve(mask.len());
76        if_then_else_extend(
77            &mut builder,
78            mask,
79            |b, _, len| b.subslice_extend_repeated(&if_true_list, 0, 1, len, ShareStrategy::Always),
80            |b, _, len| {
81                b.subslice_extend_repeated(&if_false_list, 0, 1, len, ShareStrategy::Always)
82            },
83        );
84        builder.freeze()
85    }
86}
polars_compute/if_then_else/list.rs

polars_compute/if_then_else/
list.rs
