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

Code before:
use polars_core::chunked_array::cast::CastOptions;
use polars_core::prelude::arity::unary_elementwise_values;
use polars_core::prelude::*;
use polars_ops::prelude::lst_get;
use polars_utils::index::ToIdx;

use super::*;
use crate::expressions::{AggState, AggregationContext, PhysicalExpr, UpdateGroups};

pub struct GatherExpr {
    pub(crate) phys_expr: Arc<dyn PhysicalExpr>,
    pub(crate) idx: Arc<dyn PhysicalExpr>,
    pub(crate) expr: Expr,
    pub(crate) returns_scalar: bool,
    pub(crate) null_on_oob: bool,
}

impl PhysicalExpr for GatherExpr {
    fn as_expression(&self) -> Option<&Expr> {
        Some(&self.expr)
    }

    fn evaluate(&self, df: &DataFrame, state: &ExecutionState) -> PolarsResult<Column> {
        let series = self.phys_expr.evaluate(df, state)?;
        let idx_col = self.idx.evaluate(df, state)?;
        let idx_series = idx_col.as_materialized_series();

        let len = series.len() as i64;

        let casted = idx_series.cast(&DataType::Int64)?;
        let idx_i64 = casted.i64().unwrap();

        // Pre-scan for OOB when null_on_oob == false, like lst_get.
        if !self.null_on_oob {
            if idx_i64.into_iter().any(|opt_v| {
                if let Some(v) = opt_v {
                    let j = if v >= 0 { v } else { len + v };
                    j < 0 || j >= len
                } else {
                    false
                }
            }) {
                polars_bail!(OutOfBounds: "gather/get index is out of bounds");
            }
        }

        let idx_ca: IdxCa = idx_i64
            .into_iter()
            .map(|opt_v| {
                opt_v.and_then(|v| {
                    // Python-style negative indexing: v >= 0 from front, v < 0 from back.

Review comment (suggestion):
Rather than implement it here I think we should add an extra argument to `convert_to_unsigned_index` and improve that implementation in general.

Code after:
use polars_core::chunked_array::cast::CastOptions;
use polars_core::prelude::arity::unary_elementwise_values;
use polars_core::prelude::*;
use polars_ops::prelude::lst_get;
use polars_ops::series::convert_to_unsigned_index;
use polars_utils::index::ToIdx;

use super::*;
use crate::expressions::{AggState, AggregationContext, PhysicalExpr, UpdateGroups};

pub struct GatherExpr {
    pub(crate) phys_expr: Arc<dyn PhysicalExpr>,
    pub(crate) idx: Arc<dyn PhysicalExpr>,
    pub(crate) expr: Expr,
    pub(crate) returns_scalar: bool,
    pub(crate) null_on_oob: bool,
}

impl PhysicalExpr for GatherExpr {
    fn as_expression(&self) -> Option<&Expr> {
        Some(&self.expr)
    }

    fn evaluate(&self, df: &DataFrame, state: &ExecutionState) -> PolarsResult<Column> {
        let series = self.phys_expr.evaluate(df, state)?;
        let idx = self.idx.evaluate(df, state)?;
        let idx = convert_to_unsigned_index(
            idx.as_materialized_series(),
            series.len(),
            self.null_on_oob,
        )?;
        series.take(&idx)
    }

    #[allow(clippy::ptr_arg)]
    fn evaluate_on_groups<'a>(
        &self,
        df: &DataFrame,
        groups: &'a GroupPositions,
        state: &ExecutionState,
    ) -> PolarsResult<AggregationContext<'a>> {
        let mut ac = self.phys_expr.evaluate_on_groups(df, groups, state)?;
        let mut idx = self.idx.evaluate_on_groups(df, groups, state)?;

        let ac_list = ac.aggregated_as_list();

        if self.returns_scalar {
            polars_ensure!(
                !matches!(idx.agg_state(), AggState::AggregatedList(_) | AggState::NotAggregated(_)),
                ComputeError: "expected single index"
            );