arrow_array/
scalar.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::Array;
6
7/// A possibly [`Scalar`] [`Array`]
8///
9/// This allows optimised binary kernels where one or more arguments are constant
10///
11/// ```
12/// # use arrow_array::*;
13/// # use arrow_buffer::{BooleanBuffer, MutableBuffer, NullBuffer};
14/// # use arrow_schema::ArrowError;
15/// #
16/// fn eq_impl<T: ArrowPrimitiveType>(
17///     a: &PrimitiveArray<T>,
18///     a_scalar: bool,
19///     b: &PrimitiveArray<T>,
20///     b_scalar: bool,
21/// ) -> BooleanArray {
22///     let (array, scalar) = match (a_scalar, b_scalar) {
23///         (true, true) | (false, false) => {
24///             let len = a.len().min(b.len());
25///             let nulls = NullBuffer::union(a.nulls(), b.nulls());
26///             let buffer = BooleanBuffer::collect_bool(len, |idx| a.value(idx) == b.value(idx));
27///             return BooleanArray::new(buffer, nulls);
28///         }
29///         (true, false) => (b, (a.null_count() == 0).then(|| a.value(0))),
30///         (false, true) => (a, (b.null_count() == 0).then(|| b.value(0))),
31///     };
32///     match scalar {
33///         Some(v) => {
34///             let len = array.len();
35///             let nulls = array.nulls().cloned();
36///             let buffer = BooleanBuffer::collect_bool(len, |idx| array.value(idx) == v);
37///             BooleanArray::new(buffer, nulls)
38///         }
39///         None => BooleanArray::new_null(array.len()),
40///     }
41/// }
42///
43/// pub fn eq(l: &dyn Datum, r: &dyn Datum) -> Result<BooleanArray, ArrowError> {
44///     let (l_array, l_scalar) = l.get();
45///     let (r_array, r_scalar) = r.get();
46///     downcast_primitive_array!(
47///         (l_array, r_array) => Ok(eq_impl(l_array, l_scalar, r_array, r_scalar)),
48///         (a, b) => Err(ArrowError::NotYetImplemented(format!("{a} == {b}"))),
49///     )
50/// }
51///
52/// // Comparison of two arrays
53/// let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
54/// let b = Int32Array::from(vec![1, 2, 4, 7, 3]);
55/// let r = eq(&a, &b).unwrap();
56/// let values: Vec<_> = r.values().iter().collect();
57/// assert_eq!(values, &[true, true, false, false, false]);
58///
59/// // Comparison of an array and a scalar
60/// let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
61/// let b = Int32Array::new_scalar(1);
62/// let r = eq(&a, &b).unwrap();
63/// let values: Vec<_> = r.values().iter().collect();
64/// assert_eq!(values, &[true, false, false, false, false]);
65/// ```
66pub trait Datum {
67    /// Returns the value for this [`Datum`] and a boolean indicating if the value is scalar
68    fn get(&self) -> (&dyn Array, bool);
69}
70
71impl<T: Array> Datum for T {
72    fn get(&self) -> (&dyn Array, bool) {
73        (self, false)
74    }
75}
76
77impl Datum for dyn Array {
78    fn get(&self) -> (&dyn Array, bool) {
79        (self, false)
80    }
81}
82
83impl Datum for &dyn Array {
84    fn get(&self) -> (&dyn Array, bool) {
85        (*self, false)
86    }
87}
88
89/// A wrapper around a single value [`Array`] that implements
90/// [`Datum`] and indicates [compute] kernels should treat this array
91/// as a scalar value (a single value).
92///
93/// Using a [`Scalar`] is often much more efficient than creating an
94/// [`Array`] with the same (repeated) value.
95///
96/// See [`Datum`] for more information.
97///
98/// # Example
99///
100/// ```rust
101/// # use arrow_array::{Scalar, Int32Array, ArrayRef};
102/// # fn get_array() -> ArrayRef { std::sync::Arc::new(Int32Array::from(vec![42])) }
103/// // Create a (typed) scalar for Int32Array for the value 42
104/// let scalar = Scalar::new(Int32Array::from(vec![42]));
105///
106/// // Create a scalar using PrimitiveArray::scalar
107/// let scalar = Int32Array::new_scalar(42);
108///
109/// // create a scalar from an ArrayRef (for dynamic typed Arrays)
110/// let array: ArrayRef = get_array();
111/// let scalar = Scalar::new(array);
112/// ```
113///
114/// [compute]: https://docs.rs/arrow/latest/arrow/compute/index.html
115#[derive(Debug, Copy, Clone)]
116pub struct Scalar<T: Array>(T);
117
118impl<T: Array> Scalar<T> {
119    /// Create a new [`Scalar`] from an [`Array`]
120    ///
121    /// # Panics
122    ///
123    /// Panics if `array.len() != 1`
124    pub fn new(array: T) -> Self {
125        assert_eq!(array.len(), 1);
126        Self(array)
127    }
128
129    /// Returns the inner array
130    #[inline]
131    pub fn into_inner(self) -> T {
132        self.0
133    }
134}
135
136impl<T: Array> Datum for Scalar<T> {
137    fn get(&self) -> (&dyn Array, bool) {
138        (&self.0, true)
139    }
140}
arrow_array/scalar.rs

arrow_array/
scalar.rs
