polars_arrow/compute/
boolean.rs

1//! null-preserving operators such as [`and`], [`or`] and [`not`].
2use super::utils::combine_validities_and;
3use crate::array::{Array, BooleanArray};
4use crate::bitmap::Bitmap;
5use crate::datatypes::ArrowDataType;
6use crate::scalar::BooleanScalar;
7
8fn assert_lengths(lhs: &BooleanArray, rhs: &BooleanArray) {
9    assert_eq!(
10        lhs.len(),
11        rhs.len(),
12        "lhs and rhs must have the same length"
13    );
14}
15
16/// Helper function to implement binary kernels
17pub(crate) fn binary_boolean_kernel<F>(
18    lhs: &BooleanArray,
19    rhs: &BooleanArray,
20    op: F,
21) -> BooleanArray
22where
23    F: Fn(&Bitmap, &Bitmap) -> Bitmap,
24{
25    assert_lengths(lhs, rhs);
26    let validity = combine_validities_and(lhs.validity(), rhs.validity());
27
28    let left_buffer = lhs.values();
29    let right_buffer = rhs.values();
30
31    let values = op(left_buffer, right_buffer);
32
33    BooleanArray::new(ArrowDataType::Boolean, values, validity)
34}
35
36/// Performs `&&` operation on two [`BooleanArray`], combining the validities.
37/// # Panics
38/// This function panics iff the arrays have different lengths.
39/// # Examples
40/// ```rust
41/// use polars_arrow::array::BooleanArray;
42/// use polars_arrow::compute::boolean::and;
43///
44/// let a = BooleanArray::from(&[Some(false), Some(true), None]);
45/// let b = BooleanArray::from(&[Some(true), Some(true), Some(false)]);
46/// let and_ab = and(&a, &b);
47/// assert_eq!(and_ab, BooleanArray::from(&[Some(false), Some(true), None]));
48/// ```
49pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
50    if lhs.null_count() == 0 && rhs.null_count() == 0 {
51        let left_buffer = lhs.values();
52        let right_buffer = rhs.values();
53
54        match (left_buffer.unset_bits(), right_buffer.unset_bits()) {
55            // all values are `true` on both sides
56            (0, 0) => {
57                assert_lengths(lhs, rhs);
58                return lhs.clone();
59            },
60            // all values are `false` on left side
61            (l, _) if l == lhs.len() => {
62                assert_lengths(lhs, rhs);
63                return lhs.clone();
64            },
65            // all values are `false` on right side
66            (_, r) if r == rhs.len() => {
67                assert_lengths(lhs, rhs);
68                return rhs.clone();
69            },
70            // ignore the rest
71            _ => {},
72        }
73    }
74
75    binary_boolean_kernel(lhs, rhs, |lhs, rhs| lhs & rhs)
76}
77
78/// Performs `||` operation on two [`BooleanArray`], combining the validities.
79/// # Panics
80/// This function panics iff the arrays have different lengths.
81/// # Examples
82/// ```rust
83/// use polars_arrow::array::BooleanArray;
84/// use polars_arrow::compute::boolean::or;
85///
86/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
87/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
88/// let or_ab = or(&a, &b);
89/// assert_eq!(or_ab, BooleanArray::from(vec![Some(true), Some(true), None]));
90/// ```
91pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
92    if lhs.null_count() == 0 && rhs.null_count() == 0 {
93        let left_buffer = lhs.values();
94        let right_buffer = rhs.values();
95
96        match (left_buffer.unset_bits(), right_buffer.unset_bits()) {
97            // all values are `true` on left side
98            (0, _) => {
99                assert_lengths(lhs, rhs);
100                return lhs.clone();
101            },
102            // all values are `true` on right side
103            (_, 0) => {
104                assert_lengths(lhs, rhs);
105                return rhs.clone();
106            },
107            // all values on lhs and rhs are `false`
108            (l, r) if l == lhs.len() && r == rhs.len() => {
109                assert_lengths(lhs, rhs);
110                return rhs.clone();
111            },
112            // ignore the rest
113            _ => {},
114        }
115    }
116
117    binary_boolean_kernel(lhs, rhs, |lhs, rhs| lhs | rhs)
118}
119
120/// Performs unary `NOT` operation on an arrays. If value is null then the result is also
121/// null.
122/// # Example
123/// ```rust
124/// use polars_arrow::array::BooleanArray;
125/// use polars_arrow::compute::boolean::not;
126///
127/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
128/// let not_a = not(&a);
129/// assert_eq!(not_a, BooleanArray::from(vec![Some(true), Some(false), None]));
130/// ```
131pub fn not(array: &BooleanArray) -> BooleanArray {
132    let values = !array.values();
133    let validity = array.validity().cloned();
134    BooleanArray::new(ArrowDataType::Boolean, values, validity)
135}
136
137/// Returns a non-null [`BooleanArray`] with whether each value of the array is null.
138/// # Example
139/// ```rust
140/// use polars_arrow::array::BooleanArray;
141/// use polars_arrow::compute::boolean::is_null;
142/// # fn main() {
143/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
144/// let a_is_null = is_null(&a);
145/// assert_eq!(a_is_null, BooleanArray::from_slice(vec![false, false, true]));
146/// # }
147/// ```
148pub fn is_null(input: &dyn Array) -> BooleanArray {
149    let len = input.len();
150
151    let values = match input.validity() {
152        None => Bitmap::new_zeroed(len),
153        Some(buffer) => !buffer,
154    };
155
156    BooleanArray::new(ArrowDataType::Boolean, values, None)
157}
158
159/// Returns a non-null [`BooleanArray`] with whether each value of the array is not null.
160/// # Example
161/// ```rust
162/// use polars_arrow::array::BooleanArray;
163/// use polars_arrow::compute::boolean::is_not_null;
164///
165/// let a = BooleanArray::from(&vec![Some(false), Some(true), None]);
166/// let a_is_not_null = is_not_null(&a);
167/// assert_eq!(a_is_not_null, BooleanArray::from_slice(&vec![true, true, false]));
168/// ```
169pub fn is_not_null(input: &dyn Array) -> BooleanArray {
170    let values = match input.validity() {
171        None => Bitmap::new_with_value(true, input.len()),
172        Some(buffer) => buffer.clone(),
173    };
174    BooleanArray::new(ArrowDataType::Boolean, values, None)
175}
176
177/// Performs `AND` operation on an array and a scalar value. If either left or right value
178/// is null then the result is also null.
179/// # Example
180/// ```rust
181/// use polars_arrow::array::BooleanArray;
182/// use polars_arrow::compute::boolean::and_scalar;
183/// use polars_arrow::scalar::BooleanScalar;
184///
185/// let array = BooleanArray::from_slice(&[false, false, true, true]);
186/// let scalar = BooleanScalar::new(Some(true));
187/// let result = and_scalar(&array, &scalar);
188/// assert_eq!(result, BooleanArray::from_slice(&[false, false, true, true]));
189///
190/// ```
191pub fn and_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray {
192    match scalar.value() {
193        Some(true) => array.clone(),
194        Some(false) => {
195            let values = Bitmap::new_zeroed(array.len());
196            BooleanArray::new(ArrowDataType::Boolean, values, array.validity().cloned())
197        },
198        None => BooleanArray::new_null(ArrowDataType::Boolean, array.len()),
199    }
200}
201
202/// Performs `OR` operation on an array and a scalar value. If either left or right value
203/// is null then the result is also null.
204/// # Example
205/// ```rust
206/// use polars_arrow::array::BooleanArray;
207/// use polars_arrow::compute::boolean::or_scalar;
208/// use polars_arrow::scalar::BooleanScalar;
209/// # fn main() {
210/// let array = BooleanArray::from_slice(&[false, false, true, true]);
211/// let scalar = BooleanScalar::new(Some(true));
212/// let result = or_scalar(&array, &scalar);
213/// assert_eq!(result, BooleanArray::from_slice(&[true, true, true, true]));
214/// # }
215/// ```
216pub fn or_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray {
217    match scalar.value() {
218        Some(true) => BooleanArray::new(
219            ArrowDataType::Boolean,
220            Bitmap::new_with_value(true, array.len()),
221            array.validity().cloned(),
222        ),
223        Some(false) => array.clone(),
224        None => BooleanArray::new_null(ArrowDataType::Boolean, array.len()),
225    }
226}
227
228/// Returns whether any of the values in the array are `true`.
229///
230/// Null values are ignored.
231///
232/// # Example
233///
234/// ```
235/// use polars_arrow::array::BooleanArray;
236/// use polars_arrow::compute::boolean::any;
237///
238/// let a = BooleanArray::from(&[Some(true), Some(false)]);
239/// let b = BooleanArray::from(&[Some(false), Some(false)]);
240/// let c = BooleanArray::from(&[None, Some(false)]);
241///
242/// assert_eq!(any(&a), true);
243/// assert_eq!(any(&b), false);
244/// assert_eq!(any(&c), false);
245/// ```
246pub fn any(array: &BooleanArray) -> bool {
247    if array.is_empty() {
248        false
249    } else if array.null_count() > 0 {
250        array.into_iter().any(|v| v == Some(true))
251    } else {
252        let vals = array.values();
253        vals.unset_bits() != vals.len()
254    }
255}
256
257/// Returns whether all values in the array are `true`.
258///
259/// Null values are ignored.
260///
261/// # Example
262///
263/// ```
264/// use polars_arrow::array::BooleanArray;
265/// use polars_arrow::compute::boolean::all;
266///
267/// let a = BooleanArray::from(&[Some(true), Some(true)]);
268/// let b = BooleanArray::from(&[Some(false), Some(true)]);
269/// let c = BooleanArray::from(&[None, Some(true)]);
270///
271/// assert_eq!(all(&a), true);
272/// assert_eq!(all(&b), false);
273/// assert_eq!(all(&c), true);
274/// ```
275pub fn all(array: &BooleanArray) -> bool {
276    if array.is_empty() {
277        true
278    } else if array.null_count() > 0 {
279        !array.into_iter().any(|v| v == Some(false))
280    } else {
281        let vals = array.values();
282        vals.unset_bits() == 0
283    }
284}
polars_arrow/compute/boolean.rs

polars_arrow/compute/
boolean.rs
