Repository: pola-rs/polars
File: crates/polars-ops/src/series/ops/index.rs (Rust)

Code before:
use num_traits::ToPrimitive;
use polars_core::error::{PolarsResult, polars_bail, polars_ensure};
use polars_core::prelude::{ChunkedArray, DataType, IdxCa, IdxSize, PolarsIntegerType, Series};

/// Generic conversion of an integer ChunkedArray<T> into IdxCa.
///
/// Semantics:
/// - For *all* integer types:
///   - `0 <= v < target_len`          → `v` as `IdxSize`
///   - `v >= target_len`              → `null`
/// - Additionally, for *signed* integer types (Python-style negative indices):
///   - `-target_len <= v < 0`         → `target_len + v` as `IdxSize`
///   - `v < -target_len`              → `null`
fn convert<T>(ca: &ChunkedArray<T>, target_len: usize) -> PolarsResult<IdxCa>
where
    T: PolarsIntegerType,
    T::Native: ToPrimitive,
{
    let dtype = ca.dtype();
    let is_unsigned = dtype.is_unsigned_integer();
    let len_i64 = target_len as i64;
    let len_u64 = target_len as u64;

    let out: IdxCa = ca
        .into_iter()
        .map(|opt_v| {
            opt_v.and_then(|v| {
                if is_unsigned {
                    // ---------- UNSIGNED PATH ----------
                    // Only unsigned semantics: 0 <= v < len.
                    let v_u64 = match v.to_u64() {
                        Some(x) => x,
                        None => return None, // cannot represent → treat as OOB
                    };

                    if v_u64 < len_u64 {
                        Some(v_u64 as IdxSize)
                    } else {
                        None
                    }
                } else {
                    // ---------- SIGNED PATH (Python-style negatives) ----------
                    let v_i64 = match v.to_i64() {
                        Some(x) => x,
                        None => return None, // cannot represent → treat as OOB

Review comment (performance):
Oh oops, I just noticed this. This uses the dynamic `dtype` stored in the `ChunkedArray`, but at this point we know the static dtype `T`. I'd like the compiler to be able to optimize the signed path out for unsigned types entirely. I think we can do that and also simplify the implementation by...

Code after:
use num_traits::ToPrimitive;
use polars_core::error::{PolarsResult, polars_bail, polars_ensure};
use polars_core::prelude::{ChunkedArray, DataType, IdxCa, IdxSize, PolarsIntegerType, Series};

/// UNSIGNED conversion:
///
/// - `0 <= v < target_len`  → `Some(v as IdxSize)`
/// - `v >= target_len`      → `None`
/// - values that cannot be represented as `u64` → `None`
fn convert_unsigned<T>(ca: &ChunkedArray<T>, target_len: usize) -> PolarsResult<IdxCa>
where
    T: PolarsIntegerType,
    T::Native: ToPrimitive,
{
    let len_u64 = target_len as u64;

    let out: IdxCa = ca
        .into_iter()
        .map(|opt_v| {
            opt_v.and_then(|v| {
                // cannot represent as u64 → None (treated as OOB)
                let v_u64 = v.to_u64()?;

                if v_u64 < len_u64 {
                    Some(v_u64 as IdxSize)
                } else {
                    None
                }
            })
        })
        .collect();

    Ok(out)
}

/// SIGNED conversion with Python-style negative semantics:
///
/// - `0 <= v < target_len`          → `Some(v as IdxSize)`
/// - `v >= target_len`              → `None`
/// - `-target_len <= v < 0`         → `Some(target_len + v)`
/// - `v < -target_len`              → `None`
/// - values that cannot be represented as `i64` → `None`
fn convert_signed<T>(ca: &ChunkedArray<T>, target_len: usize) -> PolarsResult<IdxCa>
where
    T: PolarsIntegerType,