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}