polars_ops/series/ops/
index.rs

1use polars_core::error::{PolarsResult, polars_ensure};
2use polars_core::prelude::arity::unary_elementwise_values;
3use polars_core::prelude::{ChunkedArray, DataType, IDX_DTYPE, IdxCa, PolarsIntegerType, Series};
4use polars_utils::index::ToIdx;
5
6fn convert<T>(ca: &ChunkedArray<T>, target_len: usize) -> PolarsResult<IdxCa>
7where
8    T: PolarsIntegerType,
9    T::Native: ToIdx,
10{
11    let target_len = target_len as u64;
12    Ok(unary_elementwise_values(ca, |v| v.to_idx(target_len)))
13}
14
15pub fn convert_to_unsigned_index(s: &Series, target_len: usize) -> PolarsResult<IdxCa> {
16    let dtype = s.dtype();
17    polars_ensure!(dtype.is_integer(), InvalidOperation: "expected integers as index");
18    if dtype.is_unsigned_integer() {
19        let nulls_before_cast = s.null_count();
20        let out = s.cast(&IDX_DTYPE).unwrap();
21        polars_ensure!(out.null_count() == nulls_before_cast, OutOfBounds: "some integers did not fit polars' index size");
22        return Ok(out.idx().unwrap().clone());
23    }
24    match dtype {
25        DataType::Int64 => {
26            let ca = s.i64().unwrap();
27            convert(ca, target_len)
28        },
29        DataType::Int32 => {
30            let ca = s.i32().unwrap();
31            convert(ca, target_len)
32        },
33        #[cfg(feature = "dtype-i16")]
34        DataType::Int16 => {
35            let ca = s.i16().unwrap();
36            convert(ca, target_len)
37        },
38        #[cfg(feature = "dtype-i8")]
39        DataType::Int8 => {
40            let ca = s.i8().unwrap();
41            convert(ca, target_len)
42        },
43        _ => unreachable!(),
44    }
45}
polars_ops/series/ops/index.rs

polars_ops/series/ops/
index.rs
