Repository: pola-rs/polars
File: crates/polars-expr/src/expressions/gather.rs (Rust)

Code before:
            ac.with_update_groups(UpdateGroups::No);
            return Ok(ac);
        }

        // Cast the indices to
        // - IdxSize, if the idx only contains positive integers.
        // - Int64,   if the idx contains negative numbers.
        // This may give false positives if there are masked out elements.
        let idx = idx.aggregated_as_list();
        let idx = idx.apply_to_inner(&|s| match s.dtype() {
            dtype if dtype == &IDX_DTYPE => Ok(s),
            dtype if dtype.is_unsigned_integer() => {
                s.cast_with_options(&IDX_DTYPE, CastOptions::Strict)
            },

            dtype if dtype.is_signed_integer() => {
                let has_negative_integers = s.lt(0)?.any();
                if has_negative_integers && dtype == &DataType::Int64 {
                    Ok(s)
                } else if has_negative_integers {
                    s.cast_with_options(&DataType::Int64, CastOptions::Strict)
                } else {
                    s.cast_with_options(&IDX_DTYPE, CastOptions::Overflowing)
                }
            },
            _ => polars_bail!(opq = get, got = s.dtype(), expected = "integer type"),
        })?;

        let taken = if idx.inner_dtype() == &IDX_DTYPE {
            // Fast path: all indices are positive.

            ac_list
                .amortized_iter()
                .zip(idx.amortized_iter())
                .map(|(s, idx)| Some(s?.as_ref().take(idx?.as_ref().idx().unwrap())))
                .map(|opt_res| opt_res.transpose())
                .collect::<PolarsResult<ListChunked>>()?
                .with_name(ac.get_values().name().clone())
        } else {
            // Slower path: some indices may be negative.
            assert!(idx.inner_dtype() == &DataType::Int64);

            ac_list
                .amortized_iter()
                .zip(idx.amortized_iter())
                .map(|(s, idx)| {
                    let s = s?;
                    let idx = idx?;
                    let idx = idx.as_ref().i64().unwrap();
                    let target_len = s.as_ref().len() as u64;
                    let idx = unary_elementwise_values(idx, |v| v.to_idx(target_len));

Review comment (suggestion):
I would really rather have this refer to `gather` and `get`. It is confusing to get an error message about `get` when you are using gather.

Code after:
            ac.with_update_groups(UpdateGroups::No);
            return Ok(ac);
        }

        // Cast the indices to
        // - IdxSize, if the idx only contains positive integers.
        // - Int64,   if the idx contains negative numbers.
        // This may give false positives if there are masked out elements.
        let idx = idx.aggregated_as_list();
        let idx = idx.apply_to_inner(&|s| match s.dtype() {
            dtype if dtype == &IDX_DTYPE => Ok(s),
            dtype if dtype.is_unsigned_integer() => {
                s.cast_with_options(&IDX_DTYPE, CastOptions::Strict)
            },

            dtype if dtype.is_signed_integer() => {
                let has_negative_integers = s.lt(0)?.any();
                if has_negative_integers && dtype == &DataType::Int64 {
                    Ok(s)
                } else if has_negative_integers {
                    s.cast_with_options(&DataType::Int64, CastOptions::Strict)
                } else {
                    s.cast_with_options(&IDX_DTYPE, CastOptions::Overflowing)
                }
            },
            _ => polars_bail!(
                op = "gather/get",
                got = s.dtype(),
                expected = "integer type"
            ),
        })?;

        let taken = if idx.inner_dtype() == &IDX_DTYPE {
            // Fast path: all indices are positive.

            ac_list
                .amortized_iter()
                .zip(idx.amortized_iter())
                .map(|(s, idx)| Some(s?.as_ref().take(idx?.as_ref().idx().unwrap())))
                .map(|opt_res| opt_res.transpose())
                .collect::<PolarsResult<ListChunked>>()?
                .with_name(ac.get_values().name().clone())
        } else {
            // Slower path: some indices may be negative.
            assert!(idx.inner_dtype() == &DataType::Int64);

            ac_list
                .amortized_iter()
                .zip(idx.amortized_iter())
                .map(|(s, idx)| {
                    let s = s?;