polars_ops/series/ops/
index.rs

1use num_traits::{Bounded, ToPrimitive, Zero};
2use polars_arrow::array::Array;
3use polars_arrow::bitmap::BitmapBuilder;
4use polars_arrow::compute::utils::combine_validities_and;
5use polars_arrow::datatypes::IdxArr;
6use polars_core::error::{PolarsResult, polars_bail, polars_ensure};
7use polars_core::prelude::{ChunkedArray, IdxCa, IdxSize, PolarsIntegerType, Series};
8use polars_core::with_match_physical_integer_polars_type;
9use polars_utils::select::select_unpredictable;
10use polars_utils::vec::PushUnchecked;
11
12/// UNSIGNED conversion:
13/// - `0 <= v < target_len`  → `Some(v)`
14/// - `v >= target_len`      → `None`
15///
16/// SIGNED conversion with Python-style negative semantics:
17/// - `v < -target_len`              → `None`
18/// - `-target_len <= v < 0`         → `Some(target_len + v)`
19/// - `0 <= v < target_len`          → `Some(v)`
20/// - `v >= target_len`              → `None`
21pub fn convert_and_bound_idx_ca<T>(
22    ca: &ChunkedArray<T>,
23    target_len: usize,
24    null_on_oob: bool,
25) -> PolarsResult<IdxCa>
26where
27    T: PolarsIntegerType,
28    T::Native: ToPrimitive,
29{
30    let mut out = Vec::with_capacity(ca.len());
31    let mut in_bounds = BitmapBuilder::with_capacity(ca.len());
32    assert!(target_len < IdxSize::MAX as usize);
33
34    let unsigned = T::Native::min_value() == T::Native::zero(); // Optimized to constant by compiler.
35    if unsigned {
36        let len_u64 = target_len as u64;
37        for arr in ca.downcast_iter() {
38            for v in arr.values().iter() {
39                // SAFETY: we reserved.
40                unsafe {
41                    if let Some(v_u64) = v.to_u64() {
42                        // Usually infallible.
43                        out.push_unchecked(v_u64 as IdxSize);
44                        in_bounds.push_unchecked(v_u64 < len_u64);
45                    } else {
46                        out.push_unchecked(0);
47                        in_bounds.push_unchecked(false);
48                    }
49                }
50            }
51        }
52    } else {
53        let len_i64 = target_len as i64;
54        for arr in ca.downcast_iter() {
55            for v in arr.values().iter() {
56                // SAFETY: we reserved.
57                unsafe {
58                    if let Some(v_i64) = v.to_i64() {
59                        // Usually infallible.
60                        let mut shifted = v_i64;
61                        shifted += select_unpredictable(v_i64 < 0, len_i64, 0);
62                        out.push_unchecked(shifted as IdxSize);
63                        in_bounds.push_unchecked((v_i64 >= -len_i64) & (v_i64 < len_i64));
64                    } else {
65                        out.push_unchecked(0);
66                        in_bounds.push_unchecked(false);
67                    }
68                }
69            }
70        }
71    }
72
73    let idx_arr = IdxArr::from_vec(out);
74    let in_bounds_valid = in_bounds.into_opt_validity();
75    let ca_valid = ca.rechunk_validity();
76    let valid = combine_validities_and(in_bounds_valid.as_ref(), ca_valid.as_ref());
77    let out = idx_arr.with_validity(valid);
78
79    if !null_on_oob && out.null_count() != ca.null_count() {
80        polars_bail!(
81            OutOfBounds: "gather indices are out of bounds"
82        );
83    }
84
85    Ok(out.into())
86}
87
88/// Convert arbitrary integer Series into IdxCa, using `target_len` as logical length.
89///
90/// - All OOB indices are mapped to null in `convert_*`.
91/// - We track null counts before and after:
92///   - if `null_on_oob == true`, extra nulls are expected and we just return.
93///   - if `null_on_oob == false` and new nulls appear, we raise OutOfBounds.
94pub fn convert_and_bound_index(
95    s: &Series,
96    target_len: usize,
97    null_on_oob: bool,
98) -> PolarsResult<IdxCa> {
99    let dtype = s.dtype();
100    polars_ensure!(
101        dtype.is_integer(),
102        InvalidOperation: "expected integers as index, got `{}`", dtype
103    );
104
105    with_match_physical_integer_polars_type!(dtype, |$T| {
106        let ca: &ChunkedArray<$T> = s.as_ref().as_ref();
107        convert_and_bound_idx_ca(ca, target_len, null_on_oob)
108    })
109}
polars_ops/series/ops/index.rs

polars_ops/series/ops/
index.rs
