arrow_arith/
arity.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Kernels for operating on [`PrimitiveArray`]s
19
20use arrow_array::builder::BufferBuilder;
21use arrow_array::*;
22use arrow_buffer::ArrowNativeType;
23use arrow_buffer::MutableBuffer;
24use arrow_buffer::buffer::NullBuffer;
25use arrow_data::ArrayData;
26use arrow_schema::ArrowError;
27
28/// See [`PrimitiveArray::unary`]
29pub fn unary<I, F, O>(array: &PrimitiveArray<I>, op: F) -> PrimitiveArray<O>
30where
31    I: ArrowPrimitiveType,
32    O: ArrowPrimitiveType,
33    F: Fn(I::Native) -> O::Native,
34{
35    array.unary(op)
36}
37
38/// See [`PrimitiveArray::unary_mut`]
39pub fn unary_mut<I, F>(
40    array: PrimitiveArray<I>,
41    op: F,
42) -> Result<PrimitiveArray<I>, PrimitiveArray<I>>
43where
44    I: ArrowPrimitiveType,
45    F: Fn(I::Native) -> I::Native,
46{
47    array.unary_mut(op)
48}
49
50/// See [`PrimitiveArray::try_unary`]
51pub fn try_unary<I, F, O>(array: &PrimitiveArray<I>, op: F) -> Result<PrimitiveArray<O>, ArrowError>
52where
53    I: ArrowPrimitiveType,
54    O: ArrowPrimitiveType,
55    F: Fn(I::Native) -> Result<O::Native, ArrowError>,
56{
57    array.try_unary(op)
58}
59
60/// See [`PrimitiveArray::try_unary_mut`]
61pub fn try_unary_mut<I, F>(
62    array: PrimitiveArray<I>,
63    op: F,
64) -> Result<Result<PrimitiveArray<I>, ArrowError>, PrimitiveArray<I>>
65where
66    I: ArrowPrimitiveType,
67    F: Fn(I::Native) -> Result<I::Native, ArrowError>,
68{
69    array.try_unary_mut(op)
70}
71
72/// Applies a binary infallible function to two [`PrimitiveArray`]s,
73/// producing a new [`PrimitiveArray`]
74///
75/// # Details
76///
77/// Given two arrays of length `len`, calls `op(a[i], b[i])` for `i` in `0..len`, collecting
78/// the results in a [`PrimitiveArray`].
79///
80/// If any index is null in either `a` or `b`, the
81/// corresponding index in the result will also be null
82///
83/// Like [`unary`], the `op` is evaluated for every element in the two arrays,
84/// including those elements which are NULL. This is beneficial as the cost of
85/// the operation is low compared to the cost of branching, and especially when
86/// the operation can be vectorised, however, requires `op` to be infallible for
87/// all possible values of its inputs
88///
89/// # Errors
90///
91/// * if the arrays have different lengths.
92///
93/// # Example
94/// ```
95/// # use arrow_arith::arity::binary;
96/// # use arrow_array::{Float32Array, Int32Array};
97/// # use arrow_array::types::Int32Type;
98/// let a = Float32Array::from(vec![Some(5.1f32), None, Some(6.8), Some(7.2)]);
99/// let b = Int32Array::from(vec![1, 2, 4, 9]);
100/// // compute int(a) + b for each element
101/// let c = binary(&a, &b, |a, b| a as i32 + b).unwrap();
102/// assert_eq!(c, Int32Array::from(vec![Some(6), None, Some(10), Some(16)]));
103/// ```
104pub fn binary<A, B, F, O>(
105    a: &PrimitiveArray<A>,
106    b: &PrimitiveArray<B>,
107    op: F,
108) -> Result<PrimitiveArray<O>, ArrowError>
109where
110    A: ArrowPrimitiveType,
111    B: ArrowPrimitiveType,
112    O: ArrowPrimitiveType,
113    F: Fn(A::Native, B::Native) -> O::Native,
114{
115    if a.len() != b.len() {
116        return Err(ArrowError::ComputeError(
117            "Cannot perform binary operation on arrays of different length".to_string(),
118        ));
119    }
120
121    if a.is_empty() {
122        return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)));
123    }
124
125    let nulls = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref());
126
127    let values = a
128        .values()
129        .into_iter()
130        .zip(b.values())
131        .map(|(l, r)| op(*l, *r));
132
133    let buffer: Vec<_> = values.collect();
134    Ok(PrimitiveArray::new(buffer.into(), nulls))
135}
136
137/// Applies a binary and infallible function to values in two arrays, replacing
138/// the values in the first array in place.
139///
140/// # Details
141///
142/// Given two arrays of length `len`, calls `op(a[i], b[i])` for `i` in
143/// `0..len`, modifying the [`PrimitiveArray`] `a` in place, if possible.
144///
145/// If any index is null in either `a` or `b`, the corresponding index in the
146/// result will also be null.
147///
148/// # Buffer Reuse
149///
150/// If the underlying buffers in `a` are not shared with other arrays,  mutates
151/// the underlying buffer in place, without allocating.
152///
153/// If the underlying buffer in `a` are shared, returns Err(self)
154///
155/// Like [`unary`] the provided function is evaluated for every index, ignoring validity. This
156/// is beneficial when the cost of the operation is low compared to the cost of branching, and
157/// especially when the operation can be vectorised, however, requires `op` to be infallible
158/// for all possible values of its inputs
159///
160/// # Errors
161///
162/// * If the arrays have different lengths
163/// * If the array is not mutable (see "Buffer Reuse")
164///
165/// # See Also
166///
167/// * Documentation on [`PrimitiveArray::unary_mut`] for operating on [`ArrayRef`].
168///
169/// # Example
170/// ```
171/// # use arrow_arith::arity::binary_mut;
172/// # use arrow_array::{Float32Array, Int32Array};
173/// # use arrow_array::types::Int32Type;
174/// // compute a + b for each element
175/// let a = Float32Array::from(vec![Some(5.1f32), None, Some(6.8)]);
176/// let b = Int32Array::from(vec![Some(1), None, Some(2)]);
177/// // compute a + b, updating the value in a in place if possible
178/// let a = binary_mut(a, &b, |a, b| a + b as f32).unwrap().unwrap();
179/// // a is updated in place
180/// assert_eq!(a, Float32Array::from(vec![Some(6.1), None, Some(8.8)]));
181/// ```
182///
183/// # Example with shared buffers
184/// ```
185/// # use arrow_arith::arity::binary_mut;
186/// # use arrow_array::Float32Array;
187/// # use arrow_array::types::Int32Type;
188/// let a = Float32Array::from(vec![Some(5.1f32), None, Some(6.8)]);
189/// let b = Float32Array::from(vec![Some(1.0f32), None, Some(2.0)]);
190/// // a_clone shares the buffer with a
191/// let a_cloned = a.clone();
192/// // try to update a in place, but it is shared. Returns Err(a)
193/// let a = binary_mut(a, &b, |a, b| a + b).unwrap_err();
194/// assert_eq!(a_cloned, a);
195/// // drop shared reference
196/// drop(a_cloned);
197/// // now a is not shared, so we can update it in place
198/// let a = binary_mut(a, &b, |a, b| a + b).unwrap().unwrap();
199/// assert_eq!(a, Float32Array::from(vec![Some(6.1), None, Some(8.8)]));
200/// ```
201pub fn binary_mut<T, U, F>(
202    a: PrimitiveArray<T>,
203    b: &PrimitiveArray<U>,
204    op: F,
205) -> Result<Result<PrimitiveArray<T>, ArrowError>, PrimitiveArray<T>>
206where
207    T: ArrowPrimitiveType,
208    U: ArrowPrimitiveType,
209    F: Fn(T::Native, U::Native) -> T::Native,
210{
211    if a.len() != b.len() {
212        return Ok(Err(ArrowError::ComputeError(
213            "Cannot perform binary operation on arrays of different length".to_string(),
214        )));
215    }
216
217    if a.is_empty() {
218        return Ok(Ok(PrimitiveArray::from(ArrayData::new_empty(
219            &T::DATA_TYPE,
220        ))));
221    }
222
223    let mut builder = a.into_builder()?;
224
225    builder
226        .values_slice_mut()
227        .iter_mut()
228        .zip(b.values())
229        .for_each(|(l, r)| *l = op(*l, *r));
230
231    let array = builder.finish();
232
233    // The builder has the null buffer from `a`, it is not changed.
234    let nulls = NullBuffer::union(array.logical_nulls().as_ref(), b.logical_nulls().as_ref());
235
236    let array_builder = array.into_data().into_builder().nulls(nulls);
237
238    let array_data = unsafe { array_builder.build_unchecked() };
239    Ok(Ok(PrimitiveArray::<T>::from(array_data)))
240}
241
242/// Applies the provided fallible binary operation across `a` and `b`.
243///
244/// This will return any error encountered, or collect the results into
245/// a [`PrimitiveArray`]. If any index is null in either `a`
246/// or `b`, the corresponding index in the result will also be null
247///
248/// Like [`try_unary`] the function is only evaluated for non-null indices
249///
250/// # Errors
251///
252/// Returns an error if the arrays have different lengths,
253/// or if the operation returns one.
254pub fn try_binary<A: ArrayAccessor, B: ArrayAccessor, F, O>(
255    a: A,
256    b: B,
257    op: F,
258) -> Result<PrimitiveArray<O>, ArrowError>
259where
260    O: ArrowPrimitiveType,
261    F: Fn(A::Item, B::Item) -> Result<O::Native, ArrowError>,
262{
263    if a.len() != b.len() {
264        return Err(ArrowError::ComputeError(
265            "Cannot perform a binary operation on arrays of different length".to_string(),
266        ));
267    }
268    if a.is_empty() {
269        return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)));
270    }
271    let len = a.len();
272
273    // Physical nulls are not the whole story: a `RunArray` or `DictionaryArray` can have
274    // logical nulls in its values while its own null buffer is absent. `is_nullable` covers
275    // those, but is allowed to be conservative, so the union of the logical nulls can still
276    // be empty.
277    if !a.is_nullable() && !b.is_nullable() {
278        try_binary_no_nulls(len, a, b, op)
279    } else {
280        let Some(nulls) = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref())
281        else {
282            return try_binary_no_nulls(len, a, b, op);
283        };
284
285        let mut buffer = BufferBuilder::<O::Native>::new(len);
286        buffer.append_n_zeroed(len);
287        let slice = buffer.as_slice_mut();
288
289        nulls.try_for_each_valid_idx(|idx| {
290            unsafe {
291                *slice.get_unchecked_mut(idx) = op(a.value_unchecked(idx), b.value_unchecked(idx))?
292            };
293            Ok::<_, ArrowError>(())
294        })?;
295
296        let values = buffer.finish().into();
297        Ok(PrimitiveArray::new(values, Some(nulls)))
298    }
299}
300
301/// Applies the provided fallible binary operation across `a` and `b` by mutating the mutable
302/// [`PrimitiveArray`] `a` with the results.
303///
304/// Returns any error encountered, or collects the results into a [`PrimitiveArray`] as return
305/// value. If any index is null in either `a` or `b`, the corresponding index in the result will
306/// also be null.
307///
308/// Like [`try_unary`] the function is only evaluated for non-null indices.
309///
310/// See [`binary_mut`] for errors and buffer reuse information.
311pub fn try_binary_mut<T, F>(
312    a: PrimitiveArray<T>,
313    b: &PrimitiveArray<T>,
314    op: F,
315) -> Result<Result<PrimitiveArray<T>, ArrowError>, PrimitiveArray<T>>
316where
317    T: ArrowPrimitiveType,
318    F: Fn(T::Native, T::Native) -> Result<T::Native, ArrowError>,
319{
320    if a.len() != b.len() {
321        return Ok(Err(ArrowError::ComputeError(
322            "Cannot perform binary operation on arrays of different length".to_string(),
323        )));
324    }
325    let len = a.len();
326
327    if a.is_empty() {
328        return Ok(Ok(PrimitiveArray::from(ArrayData::new_empty(
329            &T::DATA_TYPE,
330        ))));
331    }
332
333    // Physical and logical nulls coincide for a `PrimitiveArray`, but gate on `is_nullable` to
334    // match `try_binary`. That check is allowed to be conservative, so fall back to the no-nulls
335    // path when the union of the logical nulls turns out to be empty, instead of unwrapping it.
336    if !a.is_nullable() && !b.is_nullable() {
337        try_binary_no_nulls_mut(len, a, b, op)
338    } else {
339        let Some(nulls) =
340            create_union_null_buffer(a.logical_nulls().as_ref(), b.logical_nulls().as_ref())
341        else {
342            return try_binary_no_nulls_mut(len, a, b, op);
343        };
344
345        let mut builder = a.into_builder()?;
346
347        let slice = builder.values_slice_mut();
348
349        let r = nulls.try_for_each_valid_idx(|idx| {
350            unsafe {
351                *slice.get_unchecked_mut(idx) =
352                    op(*slice.get_unchecked(idx), b.value_unchecked(idx))?
353            };
354            Ok::<_, ArrowError>(())
355        });
356        if let Err(err) = r {
357            return Ok(Err(err));
358        }
359        let array_builder = builder.finish().into_data().into_builder();
360        let array_data = unsafe { array_builder.nulls(Some(nulls)).build_unchecked() };
361        Ok(Ok(PrimitiveArray::<T>::from(array_data)))
362    }
363}
364
365/// Computes the union of the nulls in two optional [`NullBuffer`] which
366/// is not shared with the input buffers.
367///
368/// The union of the nulls is the same as `NullBuffer::union(lhs, rhs)` but
369/// it does not increase the reference count of the null buffer.
370fn create_union_null_buffer(
371    lhs: Option<&NullBuffer>,
372    rhs: Option<&NullBuffer>,
373) -> Option<NullBuffer> {
374    match (lhs, rhs) {
375        (Some(lhs), Some(rhs)) => Some(NullBuffer::new(lhs.inner() & rhs.inner())),
376        (Some(n), None) | (None, Some(n)) => Some(NullBuffer::new(n.inner() & n.inner())),
377        (None, None) => None,
378    }
379}
380
381/// This intentional inline(never) attribute helps LLVM optimize the loop.
382#[inline(never)]
383fn try_binary_no_nulls<A: ArrayAccessor, B: ArrayAccessor, F, O>(
384    len: usize,
385    a: A,
386    b: B,
387    op: F,
388) -> Result<PrimitiveArray<O>, ArrowError>
389where
390    O: ArrowPrimitiveType,
391    F: Fn(A::Item, B::Item) -> Result<O::Native, ArrowError>,
392{
393    let mut buffer = MutableBuffer::new(len * O::Native::get_byte_width());
394    for idx in 0..len {
395        unsafe {
396            buffer.push_unchecked(op(a.value_unchecked(idx), b.value_unchecked(idx))?);
397        };
398    }
399    Ok(PrimitiveArray::new(buffer.into(), None))
400}
401
402/// This intentional inline(never) attribute helps LLVM optimize the loop.
403#[inline(never)]
404fn try_binary_no_nulls_mut<T, F>(
405    len: usize,
406    a: PrimitiveArray<T>,
407    b: &PrimitiveArray<T>,
408    op: F,
409) -> Result<Result<PrimitiveArray<T>, ArrowError>, PrimitiveArray<T>>
410where
411    T: ArrowPrimitiveType,
412    F: Fn(T::Native, T::Native) -> Result<T::Native, ArrowError>,
413{
414    let mut builder = a.into_builder()?;
415    let slice = builder.values_slice_mut();
416
417    for idx in 0..len {
418        unsafe {
419            match op(*slice.get_unchecked(idx), b.value_unchecked(idx)) {
420                Ok(value) => *slice.get_unchecked_mut(idx) = value,
421                Err(err) => return Ok(Err(err)),
422            }
423        };
424    }
425    Ok(Ok(builder.finish()))
426}
427
428#[cfg(test)]
429mod tests {
430    use super::*;
431    use arrow_array::types::*;
432    use std::sync::Arc;
433
434    #[test]
435    fn test_unary_f64_slice() {
436        let input = Float64Array::from(vec![Some(5.1f64), None, Some(6.8), None, Some(7.2)]);
437        let input_slice = input.slice(1, 4);
438        let result = unary(&input_slice, |n| n.round());
439        assert_eq!(
440            result,
441            Float64Array::from(vec![None, Some(7.0), None, Some(7.0)])
442        );
443    }
444
445    #[test]
446    fn test_try_binary_run_array_logical_nulls() {
447        // A `RunArray` has no null buffer of its own, so its logical nulls live in the values.
448        let run_ends = Int32Array::from(vec![1, 2, 3]);
449        let values = Int32Array::from(vec![Some(10), None, Some(30)]);
450        let run = RunArray::<Int32Type>::try_new(&run_ends, &values).expect("valid run array");
451        assert_eq!(run.null_count(), 0);
452        assert_eq!(run.logical_null_count(), 1);
453
454        let typed = run.downcast::<Int32Array>().expect("Int32 values");
455        let other = Int32Array::from(vec![1, 1, 1]);
456        let result =
457            try_binary::<_, _, _, Int32Type>(typed, &other, |a, b| Ok(a + b)).expect("no overflow");
458        assert_eq!(result, Int32Array::from(vec![Some(11), None, Some(31)]));
459    }
460
461    #[test]
462    fn test_binary_mut() {
463        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
464        let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
465        let c = binary_mut(a, &b, |l, r| l + r).unwrap().unwrap();
466
467        let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]);
468        assert_eq!(c, expected);
469    }
470
471    #[test]
472    fn test_binary_mut_null_buffer() {
473        let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
474
475        let b = Int32Array::from(vec![Some(10), Some(11), Some(12), Some(13), Some(14)]);
476
477        let r1 = binary_mut(a, &b, |a, b| a + b).unwrap();
478
479        let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
480        let b = Int32Array::new(
481            vec![10, 11, 12, 13, 14].into(),
482            Some(vec![true, true, true, true, true].into()),
483        );
484
485        // unwrap here means that no copying occurred
486        let r2 = binary_mut(a, &b, |a, b| a + b).unwrap();
487        assert_eq!(r1.unwrap(), r2.unwrap());
488    }
489
490    #[test]
491    fn test_try_binary_mut() {
492        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
493        let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
494        let c = try_binary_mut(a, &b, |l, r| Ok(l + r)).unwrap().unwrap();
495
496        let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]);
497        assert_eq!(c, expected);
498
499        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
500        let b = Int32Array::from(vec![1, 2, 3, 4, 5]);
501        let c = try_binary_mut(a, &b, |l, r| Ok(l + r)).unwrap().unwrap();
502        let expected = Int32Array::from(vec![16, 16, 12, 12, 6]);
503        assert_eq!(c, expected);
504
505        let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
506        let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
507        let _ = try_binary_mut(a, &b, |l, r| {
508            if l == 1 {
509                Err(ArrowError::InvalidArgumentError(
510                    "got error".parse().unwrap(),
511                ))
512            } else {
513                Ok(l + r)
514            }
515        })
516        .unwrap()
517        .expect_err("should got error");
518    }
519
520    #[test]
521    fn test_try_binary_mut_null_buffer() {
522        let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
523
524        let b = Int32Array::from(vec![Some(10), Some(11), Some(12), Some(13), Some(14)]);
525
526        let r1 = try_binary_mut(a, &b, |a, b| Ok(a + b)).unwrap();
527
528        let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
529        let b = Int32Array::new(
530            vec![10, 11, 12, 13, 14].into(),
531            Some(vec![true, true, true, true, true].into()),
532        );
533
534        // unwrap here means that no copying occurred
535        let r2 = try_binary_mut(a, &b, |a, b| Ok(a + b)).unwrap();
536        assert_eq!(r1.unwrap(), r2.unwrap());
537    }
538
539    #[test]
540    fn test_try_binary_mut_all_valid_null_buffers() {
541        // Both arrays carry a null buffer with no nulls in it: the no-nulls path must still run.
542        let a = Int32Array::new(vec![1, 2].into(), Some(vec![true, true].into()));
543        let b = Int32Array::new(vec![10, 20].into(), Some(vec![true, true].into()));
544        let c = try_binary_mut(a, &b, |a, b| Ok(a + b))
545            .expect("not shared")
546            .expect("no overflow");
547        assert_eq!(c, Int32Array::from(vec![11, 22]));
548        assert_eq!(c.logical_null_count(), 0);
549    }
550
551    #[test]
552    fn test_unary_dict_mut() {
553        let values = Int32Array::from(vec![Some(10), Some(20), None]);
554        let keys = Int8Array::from_iter_values([0, 0, 1, 2]);
555        let dictionary = DictionaryArray::new(keys, Arc::new(values));
556
557        let updated = dictionary.unary_mut::<_, Int32Type>(|x| x + 1).unwrap();
558        let typed = updated.downcast_dict::<Int32Array>().unwrap();
559        assert_eq!(typed.value(0), 11);
560        assert_eq!(typed.value(1), 11);
561        assert_eq!(typed.value(2), 21);
562
563        let values = updated.values();
564        assert!(values.is_null(2));
565    }
566}
arrow_arith/arity.rs

arrow_arith/
arity.rs
