polars_ops/series/ops/
index.rs

1use num_traits::ToPrimitive;
2use polars_core::error::{PolarsResult, polars_bail, polars_ensure};
3use polars_core::prelude::{ChunkedArray, DataType, IdxCa, IdxSize, PolarsIntegerType, Series};
4
5/// UNSIGNED conversion:
6///
7/// - `0 <= v < target_len`  → `Some(v as IdxSize)`
8/// - `v >= target_len`      → `None`
9/// - values that cannot be represented as `u64` → `None`
10fn convert_unsigned<T>(ca: &ChunkedArray<T>, target_len: usize) -> PolarsResult<IdxCa>
11where
12    T: PolarsIntegerType,
13    T::Native: ToPrimitive,
14{
15    let len_u64 = target_len as u64;
16
17    let out: IdxCa = ca
18        .into_iter()
19        .map(|opt_v| {
20            opt_v.and_then(|v| {
21                // cannot represent as u64 → None (treated as OOB)
22                let v_u64 = v.to_u64()?;
23
24                if v_u64 < len_u64 {
25                    Some(v_u64 as IdxSize)
26                } else {
27                    None
28                }
29            })
30        })
31        .collect();
32
33    Ok(out)
34}
35
36/// SIGNED conversion with Python-style negative semantics:
37///
38/// - `0 <= v < target_len`          → `Some(v as IdxSize)`
39/// - `v >= target_len`              → `None`
40/// - `-target_len <= v < 0`         → `Some(target_len + v)`
41/// - `v < -target_len`              → `None`
42/// - values that cannot be represented as `i64` → `None`
43fn convert_signed<T>(ca: &ChunkedArray<T>, target_len: usize) -> PolarsResult<IdxCa>
44where
45    T: PolarsIntegerType,
46    T::Native: ToPrimitive,
47{
48    let len_i64 = target_len as i64;
49
50    let out: IdxCa = ca
51        .into_iter()
52        .map(|opt_v| {
53            opt_v.and_then(|v| {
54                // cannot represent as i64 → None (treated as OOB)
55                let v_i64 = v.to_i64()?;
56
57                if v_i64 >= 0 {
58                    // 0 <= v < len
59                    if v_i64 < len_i64 {
60                        Some(v_i64 as IdxSize)
61                    } else {
62                        None
63                    }
64                } else {
65                    // negative index: valid iff -len <= v < 0
66                    if v_i64 >= -len_i64 {
67                        let pos = len_i64 + v_i64; // in [0, len)
68                        debug_assert!(pos >= 0 && pos < len_i64);
69                        Some(pos as IdxSize)
70                    } else {
71                        None
72                    }
73                }
74            })
75        })
76        .collect();
77
78    Ok(out)
79}
80
81/// Convert arbitrary integer Series into IdxCa, using `target_len` as logical length.
82///
83/// - All OOB indices are mapped to null in `convert_*`.
84/// - We track null counts before and after:
85///   - if `null_on_oob == true`, extra nulls are expected and we just return.
86///   - if `null_on_oob == false` and new nulls appear, we raise OutOfBounds.
87pub fn convert_to_unsigned_index(
88    s: &Series,
89    target_len: usize,
90    null_on_oob: bool,
91) -> PolarsResult<IdxCa> {
92    let dtype = s.dtype();
93    polars_ensure!(
94        dtype.is_integer(),
95        InvalidOperation: "expected integers as index"
96    );
97
98    assert!(
99        (target_len as u128) <= (IdxSize::MAX as u128),
100        "internal error: target_len does not fit in IdxSize"
101    );
102
103    let in_nulls = s.null_count();
104
105    // Normalize to IdxCa with all OOB indices already mapped to nulls.
106    let idx: IdxCa = match dtype {
107        // ----- SIGNED -----
108        DataType::Int64 => {
109            let ca = s.i64().unwrap();
110            convert_signed(ca, target_len)?
111        },
112        DataType::Int32 => {
113            let ca = s.i32().unwrap();
114            convert_signed(ca, target_len)?
115        },
116        #[cfg(feature = "dtype-i16")]
117        DataType::Int16 => {
118            let ca = s.i16().unwrap();
119            convert_signed(ca, target_len)?
120        },
121        #[cfg(feature = "dtype-i8")]
122        DataType::Int8 => {
123            let ca = s.i8().unwrap();
124            convert_signed(ca, target_len)?
125        },
126        #[cfg(feature = "dtype-i128")]
127        DataType::Int128 => {
128            let ca = s.i128().unwrap();
129            convert_signed(ca, target_len)?
130        },
131
132        // ----- UNSIGNED -----
133        DataType::UInt64 => {
134            let ca = s.u64().unwrap();
135            convert_unsigned(ca, target_len)?
136        },
137        DataType::UInt32 => {
138            let ca = s.u32().unwrap();
139            convert_unsigned(ca, target_len)?
140        },
141        #[cfg(feature = "dtype-u16")]
142        DataType::UInt16 => {
143            let ca = s.u16().unwrap();
144            convert_unsigned(ca, target_len)?
145        },
146        #[cfg(feature = "dtype-u8")]
147        DataType::UInt8 => {
148            let ca = s.u8().unwrap();
149            convert_unsigned(ca, target_len)?
150        },
151        #[cfg(feature = "dtype-u128")]
152        DataType::UInt128 => {
153            let ca = s.u128().unwrap();
154            convert_unsigned(ca, target_len)?
155        },
156
157        _ => unreachable!(),
158    };
159
160    let out_nulls = idx.null_count();
161
162    if !null_on_oob && out_nulls > in_nulls {
163        polars_bail!(
164            OutOfBounds: "gather indices are out of bounds"
165        );
166    }
167
168    Ok(idx)
169}
polars_ops/series/ops/index.rs

polars_ops/series/ops/
index.rs
