1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

use crate::array::PrimitiveArray;
use crate::compute::utils::combine_validities_and;
use crate::datatypes::ArrowDataType;
use crate::types::NativeType;

#[cfg(feature = "dtype-decimal")]
pub mod decimal;
// pub mod take;
pub mod tile;

#[inline]
pub fn binary_mut<T, D, F>(
    lhs: &PrimitiveArray<T>,
    rhs: &PrimitiveArray<D>,
    data_type: ArrowDataType,
    mut op: F,
) -> PrimitiveArray<T>
where
    T: NativeType,
    D: NativeType,
    F: FnMut(T, D) -> T,
{
    assert_eq!(lhs.len(), rhs.len());
    let validity = combine_validities_and(lhs.validity(), rhs.validity());

    let values = lhs
        .values()
        .iter()
        .zip(rhs.values().iter())
        .map(|(l, r)| op(*l, *r))
        .collect::<Vec<_>>()
        .into();

    PrimitiveArray::<T>::new(data_type, values, validity)
}

#[inline]
pub fn unary_mut<I, F, O>(
    array: &PrimitiveArray<I>,
    mut op: F,
    data_type: ArrowDataType,
) -> PrimitiveArray<O>
where
    I: NativeType,
    O: NativeType,
    F: FnMut(I) -> O,
{
    let values = array.values().iter().map(|v| op(*v)).collect::<Vec<_>>();

    PrimitiveArray::<O>::new(data_type, values.into(), array.validity().cloned())
}