polars-ops 0.46.0

More operations on Polars data structures
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

use arrow::array::{BooleanArray, ListArray};
use arrow::bitmap::Bitmap;

use super::*;

fn list_all_any<F>(arr: &ListArray<i64>, op: F, is_all: bool) -> PolarsResult<BooleanArray>
where
    F: Fn(&BooleanArray) -> bool,
{
    let offsets = arr.offsets().as_slice();
    let values = arr.values();

    polars_ensure!(values.dtype() == &ArrowDataType::Boolean, ComputeError: "expected boolean elements in list");

    let values = values.as_any().downcast_ref::<BooleanArray>().unwrap();
    let validity = arr.validity().cloned();

    // Fast path where all values set (all is free).
    if is_all {
        let all_set = arrow::compute::boolean::all(values);
        if all_set {
            let bits = Bitmap::new_with_value(true, arr.len());
            return Ok(BooleanArray::from_data_default(bits, None).with_validity(validity));
        }
    }

    let mut start = offsets[0] as usize;
    let iter = offsets[1..].iter().map(|&end| {
        let end = end as usize;
        let len = end - start;
        let val = unsafe { values.clone().sliced_unchecked(start, len) };
        start = end;
        op(&val)
    });

    Ok(BooleanArray::from_trusted_len_values_iter(iter).with_validity(validity))
}

pub(super) fn list_all(ca: &ListChunked) -> PolarsResult<Series> {
    let chunks = ca
        .downcast_iter()
        .map(|arr| list_all_any(arr, arrow::compute::boolean::all, true));
    Ok(BooleanChunked::try_from_chunk_iter(ca.name().clone(), chunks)?.into_series())
}

pub(super) fn list_any(ca: &ListChunked) -> PolarsResult<Series> {
    let chunks = ca
        .downcast_iter()
        .map(|arr| list_all_any(arr, arrow::compute::boolean::any, false));
    Ok(BooleanChunked::try_from_chunk_iter(ca.name().clone(), chunks)?.into_series())
}