use super::utils::{check_same_len, combine_validities_and};
use crate::array::PrimitiveArray;
use crate::datatypes::ArrowDataType;
use crate::types::NativeType;
#[inline]
pub fn unary<I, F, O>(array: &PrimitiveArray<I>, op: F, dtype: ArrowDataType) -> PrimitiveArray<O>
where
I: NativeType,
O: NativeType,
F: Fn(I) -> O,
{
let values = array.values().iter().map(|v| op(*v)).collect::<Vec<_>>();
PrimitiveArray::<O>::new(dtype, values.into(), array.validity().cloned())
}
#[inline]
pub fn binary<T, D, F>(
lhs: &PrimitiveArray<T>,
rhs: &PrimitiveArray<D>,
dtype: ArrowDataType,
op: F,
) -> PrimitiveArray<T>
where
T: NativeType,
D: NativeType,
F: Fn(T, D) -> T,
{
check_same_len(lhs, rhs).unwrap();
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(dtype, values, validity)
}