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)
                }
            },
            _ => Err(PolarsError::InvalidOperation(
                format!(
                    "Invalid dtype {:?} for gather indices; expected integer type",
                    s.dtype()
                ).into()
            )),
        })?;

        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)| {

Review comment (suggestion):
This should use `polars_bail!(opq = "get/gather", s.dtype())`

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?;