529 lines
17 KiB
Rust
529 lines
17 KiB
Rust
// Licensed to the Apache Software Foundation (ASF) under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing,
|
|
// software distributed under the License is distributed on an
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
// KIND, either express or implied. See the License for the
|
|
// specific language governing permissions and limitations
|
|
// under the License.
|
|
|
|
//! Kernels for operating on [`PrimitiveArray`]s
|
|
|
|
use arrow_array::builder::BufferBuilder;
|
|
use arrow_array::*;
|
|
use arrow_buffer::ArrowNativeType;
|
|
use arrow_buffer::MutableBuffer;
|
|
use arrow_buffer::buffer::NullBuffer;
|
|
use arrow_data::ArrayData;
|
|
use arrow_schema::ArrowError;
|
|
|
|
/// See [`PrimitiveArray::unary`]
|
|
pub fn unary<I, F, O>(array: &PrimitiveArray<I>, op: F) -> PrimitiveArray<O>
|
|
where
|
|
I: ArrowPrimitiveType,
|
|
O: ArrowPrimitiveType,
|
|
F: Fn(I::Native) -> O::Native,
|
|
{
|
|
array.unary(op)
|
|
}
|
|
|
|
/// See [`PrimitiveArray::unary_mut`]
|
|
pub fn unary_mut<I, F>(
|
|
array: PrimitiveArray<I>,
|
|
op: F,
|
|
) -> Result<PrimitiveArray<I>, PrimitiveArray<I>>
|
|
where
|
|
I: ArrowPrimitiveType,
|
|
F: Fn(I::Native) -> I::Native,
|
|
{
|
|
array.unary_mut(op)
|
|
}
|
|
|
|
/// See [`PrimitiveArray::try_unary`]
|
|
pub fn try_unary<I, F, O>(array: &PrimitiveArray<I>, op: F) -> Result<PrimitiveArray<O>, ArrowError>
|
|
where
|
|
I: ArrowPrimitiveType,
|
|
O: ArrowPrimitiveType,
|
|
F: Fn(I::Native) -> Result<O::Native, ArrowError>,
|
|
{
|
|
array.try_unary(op)
|
|
}
|
|
|
|
/// See [`PrimitiveArray::try_unary_mut`]
|
|
pub fn try_unary_mut<I, F>(
|
|
array: PrimitiveArray<I>,
|
|
op: F,
|
|
) -> Result<Result<PrimitiveArray<I>, ArrowError>, PrimitiveArray<I>>
|
|
where
|
|
I: ArrowPrimitiveType,
|
|
F: Fn(I::Native) -> Result<I::Native, ArrowError>,
|
|
{
|
|
array.try_unary_mut(op)
|
|
}
|
|
|
|
/// Allies a binary infallable function to two [`PrimitiveArray`]s,
|
|
/// producing a new [`PrimitiveArray`]
|
|
///
|
|
/// # Details
|
|
///
|
|
/// Given two arrays of length `len`, calls `op(a[i], b[i])` for `i` in `0..len`, collecting
|
|
/// the results in a [`PrimitiveArray`].
|
|
///
|
|
/// If any index is null in either `a` or `b`, the
|
|
/// corresponding index in the result will also be null
|
|
///
|
|
/// Like [`unary`], the `op` is evaluated for every element in the two arrays,
|
|
/// including those elements which are NULL. This is beneficial as the cost of
|
|
/// the operation is low compared to the cost of branching, and especially when
|
|
/// the operation can be vectorised, however, requires `op` to be infallible for
|
|
/// all possible values of its inputs
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// * if the arrays have different lengths.
|
|
///
|
|
/// # Example
|
|
/// ```
|
|
/// # use arrow_arith::arity::binary;
|
|
/// # use arrow_array::{Float32Array, Int32Array};
|
|
/// # use arrow_array::types::Int32Type;
|
|
/// let a = Float32Array::from(vec![Some(5.1f32), None, Some(6.8), Some(7.2)]);
|
|
/// let b = Int32Array::from(vec![1, 2, 4, 9]);
|
|
/// // compute int(a) + b for each element
|
|
/// let c = binary(&a, &b, |a, b| a as i32 + b).unwrap();
|
|
/// assert_eq!(c, Int32Array::from(vec![Some(6), None, Some(10), Some(16)]));
|
|
/// ```
|
|
pub fn binary<A, B, F, O>(
|
|
a: &PrimitiveArray<A>,
|
|
b: &PrimitiveArray<B>,
|
|
op: F,
|
|
) -> Result<PrimitiveArray<O>, ArrowError>
|
|
where
|
|
A: ArrowPrimitiveType,
|
|
B: ArrowPrimitiveType,
|
|
O: ArrowPrimitiveType,
|
|
F: Fn(A::Native, B::Native) -> O::Native,
|
|
{
|
|
if a.len() != b.len() {
|
|
return Err(ArrowError::ComputeError(
|
|
"Cannot perform binary operation on arrays of different length".to_string(),
|
|
));
|
|
}
|
|
|
|
if a.is_empty() {
|
|
return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)));
|
|
}
|
|
|
|
let nulls = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref());
|
|
|
|
let values = a
|
|
.values()
|
|
.into_iter()
|
|
.zip(b.values())
|
|
.map(|(l, r)| op(*l, *r));
|
|
|
|
let buffer: Vec<_> = values.collect();
|
|
Ok(PrimitiveArray::new(buffer.into(), nulls))
|
|
}
|
|
|
|
/// Applies a binary and infallible function to values in two arrays, replacing
|
|
/// the values in the first array in place.
|
|
///
|
|
/// # Details
|
|
///
|
|
/// Given two arrays of length `len`, calls `op(a[i], b[i])` for `i` in
|
|
/// `0..len`, modifying the [`PrimitiveArray`] `a` in place, if possible.
|
|
///
|
|
/// If any index is null in either `a` or `b`, the corresponding index in the
|
|
/// result will also be null.
|
|
///
|
|
/// # Buffer Reuse
|
|
///
|
|
/// If the underlying buffers in `a` are not shared with other arrays, mutates
|
|
/// the underlying buffer in place, without allocating.
|
|
///
|
|
/// If the underlying buffer in `a` are shared, returns Err(self)
|
|
///
|
|
/// Like [`unary`] the provided function is evaluated for every index, ignoring validity. This
|
|
/// is beneficial when the cost of the operation is low compared to the cost of branching, and
|
|
/// especially when the operation can be vectorised, however, requires `op` to be infallible
|
|
/// for all possible values of its inputs
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// * If the arrays have different lengths
|
|
/// * If the array is not mutable (see "Buffer Reuse")
|
|
///
|
|
/// # See Also
|
|
///
|
|
/// * Documentation on [`PrimitiveArray::unary_mut`] for operating on [`ArrayRef`].
|
|
///
|
|
/// # Example
|
|
/// ```
|
|
/// # use arrow_arith::arity::binary_mut;
|
|
/// # use arrow_array::{Float32Array, Int32Array};
|
|
/// # use arrow_array::types::Int32Type;
|
|
/// // compute a + b for each element
|
|
/// let a = Float32Array::from(vec![Some(5.1f32), None, Some(6.8)]);
|
|
/// let b = Int32Array::from(vec![Some(1), None, Some(2)]);
|
|
/// // compute a + b, updating the value in a in place if possible
|
|
/// let a = binary_mut(a, &b, |a, b| a + b as f32).unwrap().unwrap();
|
|
/// // a is updated in place
|
|
/// assert_eq!(a, Float32Array::from(vec![Some(6.1), None, Some(8.8)]));
|
|
/// ```
|
|
///
|
|
/// # Example with shared buffers
|
|
/// ```
|
|
/// # use arrow_arith::arity::binary_mut;
|
|
/// # use arrow_array::Float32Array;
|
|
/// # use arrow_array::types::Int32Type;
|
|
/// let a = Float32Array::from(vec![Some(5.1f32), None, Some(6.8)]);
|
|
/// let b = Float32Array::from(vec![Some(1.0f32), None, Some(2.0)]);
|
|
/// // a_clone shares the buffer with a
|
|
/// let a_cloned = a.clone();
|
|
/// // try to update a in place, but it is shared. Returns Err(a)
|
|
/// let a = binary_mut(a, &b, |a, b| a + b).unwrap_err();
|
|
/// assert_eq!(a_cloned, a);
|
|
/// // drop shared reference
|
|
/// drop(a_cloned);
|
|
/// // now a is not shared, so we can update it in place
|
|
/// let a = binary_mut(a, &b, |a, b| a + b).unwrap().unwrap();
|
|
/// assert_eq!(a, Float32Array::from(vec![Some(6.1), None, Some(8.8)]));
|
|
/// ```
|
|
pub fn binary_mut<T, U, F>(
|
|
a: PrimitiveArray<T>,
|
|
b: &PrimitiveArray<U>,
|
|
op: F,
|
|
) -> Result<Result<PrimitiveArray<T>, ArrowError>, PrimitiveArray<T>>
|
|
where
|
|
T: ArrowPrimitiveType,
|
|
U: ArrowPrimitiveType,
|
|
F: Fn(T::Native, U::Native) -> T::Native,
|
|
{
|
|
if a.len() != b.len() {
|
|
return Ok(Err(ArrowError::ComputeError(
|
|
"Cannot perform binary operation on arrays of different length".to_string(),
|
|
)));
|
|
}
|
|
|
|
if a.is_empty() {
|
|
return Ok(Ok(PrimitiveArray::from(ArrayData::new_empty(
|
|
&T::DATA_TYPE,
|
|
))));
|
|
}
|
|
|
|
let mut builder = a.into_builder()?;
|
|
|
|
builder
|
|
.values_slice_mut()
|
|
.iter_mut()
|
|
.zip(b.values())
|
|
.for_each(|(l, r)| *l = op(*l, *r));
|
|
|
|
let array = builder.finish();
|
|
|
|
// The builder has the null buffer from `a`, it is not changed.
|
|
let nulls = NullBuffer::union(array.logical_nulls().as_ref(), b.logical_nulls().as_ref());
|
|
|
|
let array_builder = array.into_data().into_builder().nulls(nulls);
|
|
|
|
let array_data = unsafe { array_builder.build_unchecked() };
|
|
Ok(Ok(PrimitiveArray::<T>::from(array_data)))
|
|
}
|
|
|
|
/// Applies the provided fallible binary operation across `a` and `b`.
|
|
///
|
|
/// This will return any error encountered, or collect the results into
|
|
/// a [`PrimitiveArray`]. If any index is null in either `a`
|
|
/// or `b`, the corresponding index in the result will also be null
|
|
///
|
|
/// Like [`try_unary`] the function is only evaluated for non-null indices
|
|
///
|
|
/// # Error
|
|
///
|
|
/// Return an error if the arrays have different lengths or
|
|
/// the operation is under erroneous
|
|
pub fn try_binary<A: ArrayAccessor, B: ArrayAccessor, F, O>(
|
|
a: A,
|
|
b: B,
|
|
op: F,
|
|
) -> Result<PrimitiveArray<O>, ArrowError>
|
|
where
|
|
O: ArrowPrimitiveType,
|
|
F: Fn(A::Item, B::Item) -> Result<O::Native, ArrowError>,
|
|
{
|
|
if a.len() != b.len() {
|
|
return Err(ArrowError::ComputeError(
|
|
"Cannot perform a binary operation on arrays of different length".to_string(),
|
|
));
|
|
}
|
|
if a.is_empty() {
|
|
return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)));
|
|
}
|
|
let len = a.len();
|
|
|
|
if a.null_count() == 0 && b.null_count() == 0 {
|
|
try_binary_no_nulls(len, a, b, op)
|
|
} else {
|
|
let nulls =
|
|
NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()).unwrap();
|
|
|
|
let mut buffer = BufferBuilder::<O::Native>::new(len);
|
|
buffer.append_n_zeroed(len);
|
|
let slice = buffer.as_slice_mut();
|
|
|
|
nulls.try_for_each_valid_idx(|idx| {
|
|
unsafe {
|
|
*slice.get_unchecked_mut(idx) = op(a.value_unchecked(idx), b.value_unchecked(idx))?
|
|
};
|
|
Ok::<_, ArrowError>(())
|
|
})?;
|
|
|
|
let values = buffer.finish().into();
|
|
Ok(PrimitiveArray::new(values, Some(nulls)))
|
|
}
|
|
}
|
|
|
|
/// Applies the provided fallible binary operation across `a` and `b` by mutating the mutable
|
|
/// [`PrimitiveArray`] `a` with the results.
|
|
///
|
|
/// Returns any error encountered, or collects the results into a [`PrimitiveArray`] as return
|
|
/// value. If any index is null in either `a` or `b`, the corresponding index in the result will
|
|
/// also be null.
|
|
///
|
|
/// Like [`try_unary`] the function is only evaluated for non-null indices.
|
|
///
|
|
/// See [`binary_mut`] for errors and buffer reuse information.
|
|
pub fn try_binary_mut<T, F>(
|
|
a: PrimitiveArray<T>,
|
|
b: &PrimitiveArray<T>,
|
|
op: F,
|
|
) -> Result<Result<PrimitiveArray<T>, ArrowError>, PrimitiveArray<T>>
|
|
where
|
|
T: ArrowPrimitiveType,
|
|
F: Fn(T::Native, T::Native) -> Result<T::Native, ArrowError>,
|
|
{
|
|
if a.len() != b.len() {
|
|
return Ok(Err(ArrowError::ComputeError(
|
|
"Cannot perform binary operation on arrays of different length".to_string(),
|
|
)));
|
|
}
|
|
let len = a.len();
|
|
|
|
if a.is_empty() {
|
|
return Ok(Ok(PrimitiveArray::from(ArrayData::new_empty(
|
|
&T::DATA_TYPE,
|
|
))));
|
|
}
|
|
|
|
if a.null_count() == 0 && b.null_count() == 0 {
|
|
try_binary_no_nulls_mut(len, a, b, op)
|
|
} else {
|
|
let nulls =
|
|
create_union_null_buffer(a.logical_nulls().as_ref(), b.logical_nulls().as_ref())
|
|
.unwrap();
|
|
|
|
let mut builder = a.into_builder()?;
|
|
|
|
let slice = builder.values_slice_mut();
|
|
|
|
let r = nulls.try_for_each_valid_idx(|idx| {
|
|
unsafe {
|
|
*slice.get_unchecked_mut(idx) =
|
|
op(*slice.get_unchecked(idx), b.value_unchecked(idx))?
|
|
};
|
|
Ok::<_, ArrowError>(())
|
|
});
|
|
if let Err(err) = r {
|
|
return Ok(Err(err));
|
|
}
|
|
let array_builder = builder.finish().into_data().into_builder();
|
|
let array_data = unsafe { array_builder.nulls(Some(nulls)).build_unchecked() };
|
|
Ok(Ok(PrimitiveArray::<T>::from(array_data)))
|
|
}
|
|
}
|
|
|
|
/// Computes the union of the nulls in two optional [`NullBuffer`] which
|
|
/// is not shared with the input buffers.
|
|
///
|
|
/// The union of the nulls is the same as `NullBuffer::union(lhs, rhs)` but
|
|
/// it does not increase the reference count of the null buffer.
|
|
fn create_union_null_buffer(
|
|
lhs: Option<&NullBuffer>,
|
|
rhs: Option<&NullBuffer>,
|
|
) -> Option<NullBuffer> {
|
|
match (lhs, rhs) {
|
|
(Some(lhs), Some(rhs)) => Some(NullBuffer::new(lhs.inner() & rhs.inner())),
|
|
(Some(n), None) | (None, Some(n)) => Some(NullBuffer::new(n.inner() & n.inner())),
|
|
(None, None) => None,
|
|
}
|
|
}
|
|
|
|
/// This intentional inline(never) attribute helps LLVM optimize the loop.
|
|
#[inline(never)]
|
|
fn try_binary_no_nulls<A: ArrayAccessor, B: ArrayAccessor, F, O>(
|
|
len: usize,
|
|
a: A,
|
|
b: B,
|
|
op: F,
|
|
) -> Result<PrimitiveArray<O>, ArrowError>
|
|
where
|
|
O: ArrowPrimitiveType,
|
|
F: Fn(A::Item, B::Item) -> Result<O::Native, ArrowError>,
|
|
{
|
|
let mut buffer = MutableBuffer::new(len * O::Native::get_byte_width());
|
|
for idx in 0..len {
|
|
unsafe {
|
|
buffer.push_unchecked(op(a.value_unchecked(idx), b.value_unchecked(idx))?);
|
|
};
|
|
}
|
|
Ok(PrimitiveArray::new(buffer.into(), None))
|
|
}
|
|
|
|
/// This intentional inline(never) attribute helps LLVM optimize the loop.
|
|
#[inline(never)]
|
|
fn try_binary_no_nulls_mut<T, F>(
|
|
len: usize,
|
|
a: PrimitiveArray<T>,
|
|
b: &PrimitiveArray<T>,
|
|
op: F,
|
|
) -> Result<Result<PrimitiveArray<T>, ArrowError>, PrimitiveArray<T>>
|
|
where
|
|
T: ArrowPrimitiveType,
|
|
F: Fn(T::Native, T::Native) -> Result<T::Native, ArrowError>,
|
|
{
|
|
let mut builder = a.into_builder()?;
|
|
let slice = builder.values_slice_mut();
|
|
|
|
for idx in 0..len {
|
|
unsafe {
|
|
match op(*slice.get_unchecked(idx), b.value_unchecked(idx)) {
|
|
Ok(value) => *slice.get_unchecked_mut(idx) = value,
|
|
Err(err) => return Ok(Err(err)),
|
|
};
|
|
};
|
|
}
|
|
Ok(Ok(builder.finish()))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use arrow_array::types::*;
|
|
use std::sync::Arc;
|
|
|
|
#[test]
|
|
#[allow(deprecated)]
|
|
fn test_unary_f64_slice() {
|
|
let input = Float64Array::from(vec![Some(5.1f64), None, Some(6.8), None, Some(7.2)]);
|
|
let input_slice = input.slice(1, 4);
|
|
let result = unary(&input_slice, |n| n.round());
|
|
assert_eq!(
|
|
result,
|
|
Float64Array::from(vec![None, Some(7.0), None, Some(7.0)])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_binary_mut() {
|
|
let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
|
|
let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
|
|
let c = binary_mut(a, &b, |l, r| l + r).unwrap().unwrap();
|
|
|
|
let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]);
|
|
assert_eq!(c, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_binary_mut_null_buffer() {
|
|
let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
|
|
|
|
let b = Int32Array::from(vec![Some(10), Some(11), Some(12), Some(13), Some(14)]);
|
|
|
|
let r1 = binary_mut(a, &b, |a, b| a + b).unwrap();
|
|
|
|
let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
|
|
let b = Int32Array::new(
|
|
vec![10, 11, 12, 13, 14].into(),
|
|
Some(vec![true, true, true, true, true].into()),
|
|
);
|
|
|
|
// unwrap here means that no copying occured
|
|
let r2 = binary_mut(a, &b, |a, b| a + b).unwrap();
|
|
assert_eq!(r1.unwrap(), r2.unwrap());
|
|
}
|
|
|
|
#[test]
|
|
fn test_try_binary_mut() {
|
|
let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
|
|
let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
|
|
let c = try_binary_mut(a, &b, |l, r| Ok(l + r)).unwrap().unwrap();
|
|
|
|
let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]);
|
|
assert_eq!(c, expected);
|
|
|
|
let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
|
|
let b = Int32Array::from(vec![1, 2, 3, 4, 5]);
|
|
let c = try_binary_mut(a, &b, |l, r| Ok(l + r)).unwrap().unwrap();
|
|
let expected = Int32Array::from(vec![16, 16, 12, 12, 6]);
|
|
assert_eq!(c, expected);
|
|
|
|
let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
|
|
let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
|
|
let _ = try_binary_mut(a, &b, |l, r| {
|
|
if l == 1 {
|
|
Err(ArrowError::InvalidArgumentError(
|
|
"got error".parse().unwrap(),
|
|
))
|
|
} else {
|
|
Ok(l + r)
|
|
}
|
|
})
|
|
.unwrap()
|
|
.expect_err("should got error");
|
|
}
|
|
|
|
#[test]
|
|
fn test_try_binary_mut_null_buffer() {
|
|
let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
|
|
|
|
let b = Int32Array::from(vec![Some(10), Some(11), Some(12), Some(13), Some(14)]);
|
|
|
|
let r1 = try_binary_mut(a, &b, |a, b| Ok(a + b)).unwrap();
|
|
|
|
let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
|
|
let b = Int32Array::new(
|
|
vec![10, 11, 12, 13, 14].into(),
|
|
Some(vec![true, true, true, true, true].into()),
|
|
);
|
|
|
|
// unwrap here means that no copying occured
|
|
let r2 = try_binary_mut(a, &b, |a, b| Ok(a + b)).unwrap();
|
|
assert_eq!(r1.unwrap(), r2.unwrap());
|
|
}
|
|
|
|
#[test]
|
|
fn test_unary_dict_mut() {
|
|
let values = Int32Array::from(vec![Some(10), Some(20), None]);
|
|
let keys = Int8Array::from_iter_values([0, 0, 1, 2]);
|
|
let dictionary = DictionaryArray::new(keys, Arc::new(values));
|
|
|
|
let updated = dictionary.unary_mut::<_, Int32Type>(|x| x + 1).unwrap();
|
|
let typed = updated.downcast_dict::<Int32Array>().unwrap();
|
|
assert_eq!(typed.value(0), 11);
|
|
assert_eq!(typed.value(1), 11);
|
|
assert_eq!(typed.value(2), 21);
|
|
|
|
let values = updated.values();
|
|
assert!(values.is_null(2));
|
|
}
|
|
}
|