polars_arrow/compute/
arity_assign.rs

1//! Defines generics suitable to perform operations to [`PrimitiveArray`] in-place.
2
3use either::Either;
4
5use super::utils::check_same_len;
6use crate::array::PrimitiveArray;
7use crate::types::NativeType;
8
9/// Applies an unary function to a [`PrimitiveArray`], optionally in-place.
10///
11/// # Implementation
12/// This function tries to apply the function directly to the values of the array.
13/// If that region is shared, this function creates a new region and writes to it.
14///
15/// # Panics
16/// This function panics iff
17/// * the arrays have a different length.
18/// * the function itself panics.
19#[inline]
20pub fn unary<I, F>(array: &mut PrimitiveArray<I>, op: F)
21where
22    I: NativeType,
23    F: Fn(I) -> I,
24{
25    if let Some(values) = array.get_mut_values() {
26        // mutate in place
27        values.iter_mut().for_each(|l| *l = op(*l));
28    } else {
29        // alloc and write to new region
30        let values = array.values().iter().map(|l| op(*l)).collect::<Vec<_>>();
31        array.set_values(values.into());
32    }
33}
34
35/// Applies a binary function to two [`PrimitiveArray`]s, optionally in-place, returning
36/// a new [`PrimitiveArray`].
37///
38/// # Implementation
39/// This function tries to apply the function directly to the values of the array.
40/// If that region is shared, this function creates a new region and writes to it.
41/// # Panics
42/// This function panics iff
43/// * the arrays have a different length.
44/// * the function itself panics.
45#[inline]
46pub fn binary<T, D, F>(lhs: &mut PrimitiveArray<T>, rhs: &PrimitiveArray<D>, op: F)
47where
48    T: NativeType,
49    D: NativeType,
50    F: Fn(T, D) -> T,
51{
52    check_same_len(lhs, rhs).unwrap();
53
54    // both for the validity and for the values
55    // we branch to check if we can mutate in place
56    // if we can, great that is fastest.
57    // if we cannot, we allocate a new buffer and assign values to that
58    // new buffer, that is benchmarked to be ~2x faster than first memcpy and assign in place
59    // for the validity bits it can be much faster as we might need to iterate all bits if the
60    // bitmap has an offset.
61    if let Some(rhs) = rhs.validity() {
62        if lhs.validity().is_none() {
63            lhs.set_validity(Some(rhs.clone()));
64        } else {
65            lhs.apply_validity(|bitmap| {
66                match bitmap.into_mut() {
67                    Either::Left(immutable) => {
68                        // alloc new region
69                        &immutable & rhs
70                    },
71                    Either::Right(mutable) => {
72                        // mutate in place
73                        (mutable & rhs).into()
74                    },
75                }
76            });
77        }
78    };
79
80    if let Some(values) = lhs.get_mut_values() {
81        // mutate values in place
82        values
83            .iter_mut()
84            .zip(rhs.values().iter())
85            .for_each(|(l, r)| *l = op(*l, *r));
86    } else {
87        // alloc new region
88        let values = lhs
89            .values()
90            .iter()
91            .zip(rhs.values().iter())
92            .map(|(l, r)| op(*l, *r))
93            .collect::<Vec<_>>();
94        lhs.set_values(values.into());
95    }
96}