use arrow::array::{Array, PrimitiveArray};
use arrow::bitmap::Bitmap;
use arrow::legacy::utils::CustomIterTools;
use arrow::types::NativeType;
use num_traits::{NumCast, ToPrimitive};
use polars_core::prelude::*;
use polars_utils::float16::pf16;
use crate::chunked_array::sum::sum_slice;
fn dispatch_sum<T, S>(arr: &dyn Array, width: usize, validity: Option<&Bitmap>) -> ArrayRef
where
T: NativeType + ToPrimitive,
S: NativeType + NumCast + std::iter::Sum,
{
let values = arr.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
let values = values.values().as_slice();
let summed: Vec<_> = (0..values.len())
.step_by(width)
.map(|start| {
let slice = unsafe { values.get_unchecked(start..start + width) };
sum_slice::<T, S>(slice)
})
.collect_trusted();
Box::new(PrimitiveArray::from_data_default(
summed.into(),
validity.cloned(),
)) as ArrayRef
}
pub(super) fn sum_array_numerical(ca: &ArrayChunked, inner_type: &DataType) -> Series {
let width = ca.width();
use DataType::*;
let chunks = ca
.downcast_iter()
.map(|arr| {
let values = arr.values().as_ref();
match inner_type {
Int8 => dispatch_sum::<i8, i64>(values, width, arr.validity()),
Int16 => dispatch_sum::<i16, i64>(values, width, arr.validity()),
Int32 => dispatch_sum::<i32, i32>(values, width, arr.validity()),
Int64 => dispatch_sum::<i64, i64>(values, width, arr.validity()),
Int128 => dispatch_sum::<i128, i128>(values, width, arr.validity()),
UInt8 => dispatch_sum::<u8, i64>(values, width, arr.validity()),
UInt16 => dispatch_sum::<u16, i64>(values, width, arr.validity()),
UInt32 => dispatch_sum::<u32, u32>(values, width, arr.validity()),
UInt64 => dispatch_sum::<u64, u64>(values, width, arr.validity()),
UInt128 => dispatch_sum::<u128, u128>(values, width, arr.validity()),
Float16 => dispatch_sum::<pf16, pf16>(values, width, arr.validity()),
Float32 => dispatch_sum::<f32, f32>(values, width, arr.validity()),
Float64 => dispatch_sum::<f64, f64>(values, width, arr.validity()),
_ => unimplemented!(),
}
})
.collect::<Vec<_>>();
Series::try_from((ca.name().clone(), chunks)).unwrap()
}
pub(super) fn sum_with_nulls(ca: &ArrayChunked, inner_dtype: &DataType) -> PolarsResult<Series> {
use DataType::*;
let mut out = {
match inner_dtype {
Boolean => {
let out: IdxCa = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
UInt8 => {
let out: Int64Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
UInt16 => {
let out: Int64Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
UInt32 => {
let out: UInt32Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
UInt64 => {
let out: UInt64Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
#[cfg(feature = "dtype-u128")]
UInt128 => {
let out: UInt128Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
Int8 => {
let out: Int64Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
Int16 => {
let out: Int64Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
Int32 => {
let out: Int32Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
Int64 => {
let out: Int64Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
#[cfg(feature = "dtype-i128")]
Int128 => {
let out: Int128Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
#[cfg(feature = "dtype-f16")]
Float16 => {
let out: Float16Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
Float32 => {
let out: Float32Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
Float64 => {
let out: Float64Chunked = ca
.amortized_iter()
.map(|s| s.and_then(|s| s.as_ref().sum().ok()))
.collect();
out.into_series()
},
_ => {
polars_bail!(ComputeError: "summing array with dtype: {} not yet supported", ca.dtype())
},
}
};
out.rename(ca.name().clone());
Ok(out)
}