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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

use arrow::array::Array;
use polars_compute::gather::sublist::fixed_size_list::{
    sub_fixed_size_list_get, sub_fixed_size_list_get_literal,
};
use polars_core::utils::align_chunks_binary;

use super::*;

fn array_get_literal(ca: &ArrayChunked, idx: i64, null_on_oob: bool) -> PolarsResult<Series> {
    let chunks = ca
        .downcast_iter()
        .map(|arr| sub_fixed_size_list_get_literal(arr, idx, null_on_oob))
        .collect::<PolarsResult<Vec<_>>>()?;
    Series::try_from((ca.name().clone(), chunks))
        .unwrap()
        .cast(ca.inner_dtype())
}

/// Get the value by literal index in the array.
/// So index `0` would return the first item of every sub-array
/// and index `-1` would return the last item of every sub-array
/// if an index is out of bounds, it will return a `None`.
pub fn array_get(
    ca: &ArrayChunked,
    index: &Int64Chunked,
    null_on_oob: bool,
) -> PolarsResult<Series> {
    match index.len() {
        1 => {
            let index = index.get(0);
            if let Some(index) = index {
                array_get_literal(ca, index, null_on_oob)
            } else {
                Ok(Series::full_null(
                    ca.name().clone(),
                    ca.len(),
                    ca.inner_dtype(),
                ))
            }
        },
        len if len == ca.len() => {
            let out = binary_to_series_arr_get(ca, index, null_on_oob, |arr, idx, nob| {
                sub_fixed_size_list_get(arr, idx, nob)
            });
            out?.cast(ca.inner_dtype())
        },
        len => polars_bail!(
            ComputeError:
            "`arr.get` expression got an index array of length {} while the array has {} elements",
            len, ca.len()
        ),
    }
}

pub fn binary_to_series_arr_get<T, U, F>(
    lhs: &ChunkedArray<T>,
    rhs: &ChunkedArray<U>,
    null_on_oob: bool,
    mut op: F,
) -> PolarsResult<Series>
where
    T: PolarsDataType,
    U: PolarsDataType,
    F: FnMut(&T::Array, &U::Array, bool) -> PolarsResult<Box<dyn Array>>,
{
    let (lhs, rhs) = align_chunks_binary(lhs, rhs);
    let chunks = lhs
        .downcast_iter()
        .zip(rhs.downcast_iter())
        .map(|(lhs_arr, rhs_arr)| op(lhs_arr, rhs_arr, null_on_oob))
        .collect::<PolarsResult<Vec<_>>>()?;
    Series::try_from((lhs.name().clone(), chunks))
}