polars_compute/
sum.rs

1use std::ops::Add;
2#[cfg(feature = "simd")]
3use std::simd::Select;
4#[cfg(feature = "simd")]
5use std::simd::prelude::*;
6
7use arrow::array::{Array, PrimitiveArray};
8use arrow::bitmap::bitmask::BitMask;
9use arrow::types::NativeType;
10use num_traits::Zero;
11use polars_utils::float16::pf16;
12
13macro_rules! wrapping_impl {
14    ($trait_name:ident, $method:ident, $t:ty) => {
15        impl $trait_name for $t {
16            #[inline(always)]
17            fn wrapping_add(&self, v: &Self) -> Self {
18                <$t>::$method(*self, *v)
19            }
20        }
21    };
22}
23
24/// Performs addition that wraps around on overflow.
25///
26/// Differs from num::WrappingAdd in that this is also implemented for floats.
27pub trait WrappingAdd: Sized {
28    /// Wrapping (modular) addition. Computes `self + other`, wrapping around at
29    /// the boundary of the type.
30    fn wrapping_add(&self, v: &Self) -> Self;
31}
32
33wrapping_impl!(WrappingAdd, wrapping_add, u8);
34wrapping_impl!(WrappingAdd, wrapping_add, u16);
35wrapping_impl!(WrappingAdd, wrapping_add, u32);
36wrapping_impl!(WrappingAdd, wrapping_add, u64);
37wrapping_impl!(WrappingAdd, wrapping_add, usize);
38wrapping_impl!(WrappingAdd, wrapping_add, u128);
39
40wrapping_impl!(WrappingAdd, wrapping_add, i8);
41wrapping_impl!(WrappingAdd, wrapping_add, i16);
42wrapping_impl!(WrappingAdd, wrapping_add, i32);
43wrapping_impl!(WrappingAdd, wrapping_add, i64);
44wrapping_impl!(WrappingAdd, wrapping_add, isize);
45wrapping_impl!(WrappingAdd, wrapping_add, i128);
46
47wrapping_impl!(WrappingAdd, add, pf16);
48wrapping_impl!(WrappingAdd, add, f32);
49wrapping_impl!(WrappingAdd, add, f64);
50
51#[cfg(feature = "simd")]
52const STRIPE: usize = 16;
53
54fn wrapping_sum_with_mask_scalar<T: Zero + WrappingAdd + Copy>(vals: &[T], mask: &BitMask) -> T {
55    assert!(vals.len() == mask.len());
56    vals.iter()
57        .enumerate()
58        .map(|(i, x)| {
59            // No filter but rather select of 0 for cmov opt.
60            if mask.get(i) { *x } else { T::zero() }
61        })
62        .fold(T::zero(), |a, b| a.wrapping_add(&b))
63}
64
65fn wrapping_sum_with_mask_scalar_upcast<T, S>(vals: &[T], mask: &BitMask) -> S
66where
67    T: NativeType + Zero + Into<S>,
68    S: Zero + WrappingAdd + Copy,
69{
70    assert!(vals.len() == mask.len());
71    vals.iter()
72        .enumerate()
73        .map(|(i, x)| {
74            // No filter but rather select of 0 for cmov opt.
75            if mask.get(i) { *x } else { T::zero() }
76        })
77        .fold(S::zero(), |a, b| a.wrapping_add(&b.into()))
78}
79
80#[cfg(not(feature = "simd"))]
81impl<T> WrappingSum for T
82where
83    T: NativeType + WrappingAdd + Zero,
84{
85    fn wrapping_sum(vals: &[Self]) -> Self {
86        vals.iter()
87            .copied()
88            .fold(T::zero(), |a, b| a.wrapping_add(&b))
89    }
90
91    fn wrapping_sum_with_validity(vals: &[Self], mask: &BitMask) -> Self {
92        wrapping_sum_with_mask_scalar(vals, mask)
93    }
94}
95
96#[cfg(feature = "simd")]
97impl<T> WrappingSum for T
98where
99    T: NativeType + WrappingAdd + Zero + crate::SimdPrimitive,
100{
101    fn wrapping_sum(vals: &[Self]) -> Self {
102        vals.iter()
103            .copied()
104            .fold(T::zero(), |a, b| a.wrapping_add(&b))
105    }
106
107    fn wrapping_sum_with_validity(vals: &[Self], mask: &BitMask) -> Self {
108        assert!(vals.len() == mask.len());
109        let remainder = vals.len() % STRIPE;
110        let (rest, main) = vals.split_at(remainder);
111        let (rest_mask, main_mask) = mask.split_at(remainder);
112        let zero: Simd<T, STRIPE> = Simd::default();
113
114        let vsum = main
115            .chunks_exact(STRIPE)
116            .enumerate()
117            .map(|(i, a)| {
118                let m: Mask<T::Mask, STRIPE> = main_mask.get_simd(i * STRIPE);
119                m.select(Simd::from_slice(a), zero)
120            })
121            .fold(zero, |a, b| {
122                let a = a.to_array();
123                let b = b.to_array();
124                Simd::from_array(std::array::from_fn(|i| a[i].wrapping_add(&b[i])))
125            });
126
127        let mainsum = vsum
128            .to_array()
129            .into_iter()
130            .fold(T::zero(), |a, b| a.wrapping_add(&b));
131
132        // TODO: faster remainder.
133        let restsum = wrapping_sum_with_mask_scalar(rest, &rest_mask);
134        mainsum.wrapping_add(&restsum)
135    }
136}
137
138#[cfg(feature = "simd")]
139impl WrappingSum for u128 {
140    fn wrapping_sum(vals: &[Self]) -> Self {
141        vals.iter().copied().fold(0, |a, b| a.wrapping_add(b))
142    }
143
144    fn wrapping_sum_with_validity(vals: &[Self], mask: &BitMask) -> Self {
145        wrapping_sum_with_mask_scalar(vals, mask)
146    }
147}
148
149#[cfg(feature = "simd")]
150impl WrappingSum for i128 {
151    fn wrapping_sum(vals: &[Self]) -> Self {
152        vals.iter().copied().fold(0, |a, b| a.wrapping_add(b))
153    }
154
155    fn wrapping_sum_with_validity(vals: &[Self], mask: &BitMask) -> Self {
156        wrapping_sum_with_mask_scalar(vals, mask)
157    }
158}
159
160#[cfg(feature = "simd")]
161impl WrappingSum for pf16 {
162    fn wrapping_sum(_vals: &[Self]) -> Self {
163        unimplemented!("should have been dispatched to other sum kernel")
164    }
165
166    fn wrapping_sum_with_validity(_vals: &[Self], _mask: &BitMask) -> Self {
167        unimplemented!("should have been dispatched to other sum kernel")
168    }
169}
170
171pub trait WrappingSum: Sized {
172    fn wrapping_sum(vals: &[Self]) -> Self;
173    fn wrapping_sum_with_validity(vals: &[Self], mask: &BitMask) -> Self;
174}
175
176pub fn wrapping_sum_arr<T>(arr: &PrimitiveArray<T>) -> T
177where
178    T: NativeType + WrappingSum,
179{
180    let validity = arr.validity().filter(|_| arr.null_count() > 0);
181    if let Some(mask) = validity {
182        WrappingSum::wrapping_sum_with_validity(arr.values(), &BitMask::from_bitmap(mask))
183    } else {
184        WrappingSum::wrapping_sum(arr.values())
185    }
186}
187
188pub fn wrapping_sum_arr_upcast<T, S>(arr: &PrimitiveArray<T>) -> S
189where
190    T: NativeType + Zero + Into<S>,
191    S: Zero + WrappingAdd + Copy,
192{
193    let validity = arr.validity().filter(|_| arr.null_count() > 0);
194    if let Some(mask) = validity {
195        wrapping_sum_with_mask_scalar_upcast(arr.values(), &BitMask::from_bitmap(mask))
196    } else {
197        arr.values()
198            .iter()
199            .fold(S::zero(), |a, b| a.wrapping_add(&(*b).into()))
200    }
201}
polars_compute/sum.rs

polars_compute/
sum.rs
