polars-compute 0.53.0

Private compute kernels for the Polars DataFrame library
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

use arrow::array::{Array, FixedSizeListArray, ListArray, StructArray};
use arrow::offset::OffsetsBuffer;
use arrow::types::Offset;

/// Trim all lists of unused start and end elements recursively.
pub fn trim_lists_to_normalized_offsets(arr: &dyn Array) -> Option<Box<dyn Array>> {
    let arr = arr.as_any();
    if let Some(arr) = arr.downcast_ref::<ListArray<i32>>() {
        return trim_lists_to_normalized_offsets_list(arr).map(|arr| Box::new(arr) as _);
    }
    if let Some(arr) = arr.downcast_ref::<ListArray<i64>>() {
        return trim_lists_to_normalized_offsets_list(arr).map(|arr| Box::new(arr) as _);
    }
    if let Some(arr) = arr.downcast_ref::<FixedSizeListArray>() {
        return trim_lists_to_normalized_offsets_fsl(arr).map(|arr| Box::new(arr) as _);
    }
    if let Some(arr) = arr.downcast_ref::<StructArray>() {
        return trim_lists_to_normalized_offsets_struct(arr).map(|arr| Box::new(arr) as _);
    }

    None
}

pub fn trim_lists_to_normalized_offsets_list<O: Offset>(
    arr: &ListArray<O>,
) -> Option<ListArray<O>> {
    let offsets = arr.offsets();
    let values = arr.values();

    let len = offsets.range().to_usize();

    let (values, offsets) = if values.len() == len {
        let values = trim_lists_to_normalized_offsets(values.as_ref())?;
        (values, offsets.clone())
    } else {
        let first_idx: O = *offsets.first();

        let offsets = if first_idx.to_usize() == 0 {
            offsets.clone()
        } else {
            let v: Vec<O> = offsets.iter().map(|x| *x - first_idx).collect();
            unsafe { OffsetsBuffer::<O>::new_unchecked(v.into()) }
        };

        let values = values.sliced(first_idx.to_usize(), len);
        let values = trim_lists_to_normalized_offsets(values.as_ref()).unwrap_or(values);
        (values, offsets)
    };

    assert_eq!(offsets.first().to_usize(), 0);
    assert_eq!(values.len(), offsets.range().to_usize());

    Some(ListArray::new(
        arr.dtype().clone(),
        offsets,
        values,
        arr.validity().cloned(),
    ))
}

pub fn trim_lists_to_normalized_offsets_fsl(
    arr: &FixedSizeListArray,
) -> Option<FixedSizeListArray> {
    let values = trim_lists_to_normalized_offsets(arr.values().as_ref())?;

    Some(FixedSizeListArray::new(
        arr.dtype().clone(),
        arr.len(),
        values,
        arr.validity().cloned(),
    ))
}

pub fn trim_lists_to_normalized_offsets_struct(arr: &StructArray) -> Option<StructArray> {
    let mut new_values = Vec::new();
    for (i, field_array) in arr.values().iter().enumerate() {
        let Some(field_array) = trim_lists_to_normalized_offsets(field_array.as_ref()) else {
            // Nothing was changed. Return the original array.
            continue;
        };

        new_values.reserve(arr.values().len());
        new_values.extend(arr.values()[..i].iter().cloned());
        new_values.push(field_array);
        break;
    }

    if new_values.is_empty() {
        return None;
    }

    new_values.extend(arr.values()[new_values.len()..].iter().map(|field_array| {
        trim_lists_to_normalized_offsets(field_array.as_ref())
            .unwrap_or_else(|| field_array.clone())
    }));

    Some(StructArray::new(
        arr.dtype().clone(),
        arr.len(),
        new_values,
        arr.validity().cloned(),
    ))
}