polars-expr 0.55.1

use std::borrow::Cow;

use arrow::array::PrimitiveArray;
use num_traits::Zero;
use polars_core::error::constants::LENGTH_LIMIT_MSG;
use polars_core::prelude::sum_output_dtype;
use polars_core::with_match_physical_numeric_polars_type;
use polars_utils::float::IsFloat;
use polars_utils::index::{idxsize_to_u64, idxsize_try_from};

use super::*;

pub fn new_sum_reduction(dtype: DataType) -> PolarsResult<Box<dyn GroupedReduction>> {
    // TODO: Move the error checks up and make this function infallible
    use DataType::*;
    use VecGroupedReduction as VGR;
    Ok(match dtype {
        Boolean => Box::new(VGR::new(dtype, BoolSumReducer)),
        _ if dtype.is_primitive_numeric() => {
            with_match_physical_numeric_polars_type!(dtype.to_physical(), |$T| {
                Box::new(VGR::new(dtype, NumSumReducer::<$T>(PhantomData)))
            })
        },
        #[cfg(feature = "dtype-decimal")]
        Decimal(_, _) => Box::new(VGR::new(dtype, NumSumReducer::<Int128Type>(PhantomData))),
        Duration(_) => Box::new(VGR::new(dtype, NumSumReducer::<Int64Type>(PhantomData))),
        Null => Box::new(super::NullGroupedReduction::new(Scalar::null(
            DataType::Null,
        ))),
        String => {
            polars_bail!(
                op = "`sum`",
                DataType::String,
                hint = "you may mean to call `str.join` or `list.join`"
            );
        },
        _ => polars_bail!(op = "`sum`", dtype),
    })
}

struct NumSumReducer<T>(PhantomData<T>);
impl<T> Clone for NumSumReducer<T> {
    fn clone(&self) -> Self {
        Self(PhantomData)
    }
}

impl<T> Reducer for NumSumReducer<T>
where
    T: PolarsNumericType,
    <T as PolarsNumericType>::Native: SumCast,
    ChunkedArray<T>: ChunkAgg<T::Native>,
{
    type Dtype = T;
    type Value = <T::Native as SumCast>::Sum;

    #[inline(always)]
    fn init(&self) -> Self::Value {
        Zero::zero()
    }

    fn cast_series<'a>(&self, s: &'a Series) -> Cow<'a, Series> {
        s.to_physical_repr()
    }

    #[inline(always)]
    fn combine(&self, a: &mut Self::Value, b: &Self::Value) {
        *a += *b;
    }

    #[inline(always)]
    fn reduce_one(&self, a: &mut Self::Value, b: Option<T::Native>, _seq_id: u64) {
        *a += b.map(Into::into).unwrap_or(Zero::zero());
    }

    fn reduce_ca(&self, v: &mut Self::Value, ca: &ChunkedArray<Self::Dtype>, _seq_id: u64) {
        if T::Native::is_float() {
            *v += ChunkAgg::sum(ca).map(Into::into).unwrap_or(Zero::zero());
        } else {
            for arr in ca.downcast_iter() {
                if arr.has_nulls() {
                    for x in arr.iter() {
                        *v += x.copied().map(Into::into).unwrap_or(Zero::zero());
                    }
                } else {
                    for x in arr.values_iter().copied() {
                        *v += x.into();
                    }
                }
            }
        }
    }

    fn finish(
        &self,
        v: Vec<Self::Value>,
        m: Option<Bitmap>,
        dtype: &DataType,
    ) -> PolarsResult<Series> {
        assert!(m.is_none());
        let arr = Box::new(PrimitiveArray::from_vec(v));
        Ok(unsafe {
            Series::from_chunks_and_dtype_unchecked(
                PlSmallStr::EMPTY,
                vec![arr],
                &sum_output_dtype(dtype),
            )
        })
    }
}

/// Reduces as u64. Converts to IdxSize on `finish()`, raising bigidx error if the
/// result doesn't fit into the configured `IdxSize`.
#[derive(Default)]
pub struct IdxTypeCheckedSumReducer(PhantomData<(IdxType, UInt64Type)>);

impl IdxTypeCheckedSumReducer {
    pub fn new_grouped_reduction() -> VecGroupedReduction<Self> {
        VecGroupedReduction::new(DataType::IDX_DTYPE, Self::default())
    }
}

impl Clone for IdxTypeCheckedSumReducer {
    fn clone(&self) -> Self {
        Self(PhantomData)
    }
}

impl Reducer for IdxTypeCheckedSumReducer {
    type Dtype = IdxType;
    type Value = <<UInt64Type as PolarsNumericType>::Native as SumCast>::Sum;

    #[inline(always)]
    fn init(&self) -> Self::Value {
        Zero::zero()
    }

    fn cast_series<'a>(&self, s: &'a Series) -> Cow<'a, Series> {
        s.to_physical_repr()
    }

    #[inline(always)]
    fn combine(&self, a: &mut Self::Value, b: &Self::Value) {
        *a += *b;
    }

    #[inline(always)]
    fn reduce_one(
        &self,
        a: &mut Self::Value,
        b: Option<<Self::Dtype as PolarsNumericType>::Native>,
        _seq_id: u64,
    ) {
        *a += b.map(idxsize_to_u64).unwrap_or(0);
    }

    fn reduce_ca(&self, v: &mut Self::Value, ca: &ChunkedArray<Self::Dtype>, _seq_id: u64) {
        for arr in ca.downcast_iter() {
            if arr.has_nulls() {
                for x in arr.iter() {
                    *v += x.copied().map(idxsize_to_u64).unwrap_or(0);
                }
            } else {
                for x in arr.values_iter().copied() {
                    *v += idxsize_to_u64(x);
                }
            }
        }
    }

    fn finish(
        &self,
        v: Vec<Self::Value>,
        m: Option<Bitmap>,
        dtype: &DataType,
    ) -> PolarsResult<Series> {
        assert!(m.is_none());

        let len = v.len();
        let v: Vec<IdxSize> = v
            .into_iter()
            .filter_map(|x| idxsize_try_from(x).ok())
            .collect();

        polars_ensure!(
            v.len() == len,
            ComputeError:
            LENGTH_LIMIT_MSG
        );

        let arr = PrimitiveArray::from_vec(v);

        Ok(unsafe {
            Series::from_chunks_and_dtype_unchecked(
                PlSmallStr::EMPTY,
                vec![Box::new(arr)],
                &sum_output_dtype(dtype),
            )
        })
    }
}

#[derive(Clone)]
struct BoolSumReducer;

impl Reducer for BoolSumReducer {
    type Dtype = BooleanType;
    type Value = IdxSize;

    #[inline(always)]
    fn init(&self) -> Self::Value {
        0
    }

    #[inline(always)]
    fn combine(&self, a: &mut Self::Value, b: &Self::Value) {
        *a += *b;
    }

    #[inline(always)]
    fn reduce_one(&self, a: &mut Self::Value, b: Option<bool>, _seq_id: u64) {
        *a += b.unwrap_or(false) as IdxSize;
    }

    fn reduce_ca(&self, v: &mut Self::Value, ca: &ChunkedArray<Self::Dtype>, _seq_id: u64) {
        *v += ca.sum().unwrap_or(0) as IdxSize;
    }

    fn finish(
        &self,
        v: Vec<Self::Value>,
        m: Option<Bitmap>,
        dtype: &DataType,
    ) -> PolarsResult<Series> {
        assert!(m.is_none());
        assert!(dtype == &DataType::Boolean);
        Ok(IdxCa::from_vec(PlSmallStr::EMPTY, v).into_series())
    }
}