Function try_binary_mut

arrow-arith/src/arity.rs:305–352

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 indice

(
    a: PrimitiveArray<T>,
    b: &PrimitiveArray<T>,
    op: F,
)

///
/// 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.