polars_compute/gather/sublist/
fixed_size_list.rs

1use arrow::array::{ArrayRef, FixedSizeListArray, PrimitiveArray};
2use arrow::legacy::prelude::*;
3use arrow::legacy::utils::CustomIterTools;
4use polars_error::{polars_bail, PolarsResult};
5use polars_utils::index::NullCount;
6use polars_utils::IdxSize;
7
8use crate::gather::take_unchecked;
9
10fn sub_fixed_size_list_get_indexes_literal(width: usize, len: usize, index: i64) -> IdxArr {
11    (0..len)
12        .map(|i| {
13            if index >= width as i64 {
14                return None;
15            }
16
17            index
18                .negative_to_usize(width)
19                .map(|idx| (idx + i * width) as IdxSize)
20        })
21        .collect_trusted()
22}
23
24fn sub_fixed_size_list_get_indexes(width: usize, index: &PrimitiveArray<i64>) -> IdxArr {
25    index
26        .iter()
27        .enumerate()
28        .map(|(i, idx)| {
29            if let Some(idx) = idx {
30                if *idx >= width as i64 {
31                    return None;
32                }
33
34                idx.negative_to_usize(width)
35                    .map(|idx| (idx + i * width) as IdxSize)
36            } else {
37                None
38            }
39        })
40        .collect_trusted()
41}
42
43pub fn sub_fixed_size_list_get_literal(
44    arr: &FixedSizeListArray,
45    index: i64,
46    null_on_oob: bool,
47) -> PolarsResult<ArrayRef> {
48    let take_by = sub_fixed_size_list_get_indexes_literal(arr.size(), arr.len(), index);
49    if !null_on_oob && take_by.null_count() > 0 {
50        polars_bail!(ComputeError: "get index is out of bounds");
51    }
52
53    let values = arr.values();
54    // SAFETY:
55    // the indices we generate are in bounds
56    unsafe { Ok(take_unchecked(&**values, &take_by)) }
57}
58
59pub fn sub_fixed_size_list_get(
60    arr: &FixedSizeListArray,
61    index: &PrimitiveArray<i64>,
62    null_on_oob: bool,
63) -> PolarsResult<ArrayRef> {
64    let take_by = sub_fixed_size_list_get_indexes(arr.size(), index);
65    if !null_on_oob && take_by.null_count() > 0 {
66        polars_bail!(ComputeError: "get index is out of bounds");
67    }
68
69    let values = arr.values();
70    // SAFETY:
71    // the indices we generate are in bounds
72    unsafe { Ok(take_unchecked(&**values, &take_by)) }
73}
polars_compute/gather/sublist/fixed_size_list.rs

polars_compute/gather/sublist/
fixed_size_list.rs
