polars_arrow/compute/
boolean_kleene.rs

1//! Boolean operators of [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics).
2use crate::array::{Array, BooleanArray};
3use crate::bitmap::{binary, quaternary, ternary, unary, Bitmap};
4use crate::datatypes::ArrowDataType;
5use crate::scalar::BooleanScalar;
6
7/// Logical 'or' operation on two arrays with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics)
8/// # Panics
9/// This function panics iff the arrays have a different length
10/// # Example
11///
12/// ```rust
13/// use polars_arrow::array::BooleanArray;
14/// use polars_arrow::compute::boolean_kleene::or;
15///
16/// let a = BooleanArray::from(&[Some(true), Some(false), None]);
17/// let b = BooleanArray::from(&[None, None, None]);
18/// let or_ab = or(&a, &b);
19/// assert_eq!(or_ab, BooleanArray::from(&[Some(true), None, None]));
20/// ```
21pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
22    assert_eq!(
23        lhs.len(),
24        rhs.len(),
25        "lhs and rhs must have the same length"
26    );
27
28    let lhs_values = lhs.values();
29    let rhs_values = rhs.values();
30
31    let lhs_validity = lhs.validity();
32    let rhs_validity = rhs.validity();
33
34    let validity = match (lhs_validity, rhs_validity) {
35        (Some(lhs_validity), Some(rhs_validity)) => {
36            Some(quaternary(
37                lhs_values,
38                rhs_values,
39                lhs_validity,
40                rhs_validity,
41                // see https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics
42                |lhs, rhs, lhs_v, rhs_v| {
43                    // A = T
44                    (lhs & lhs_v) |
45                    // B = T
46                    (rhs & rhs_v) |
47                    // A = F & B = F
48                    (!lhs & lhs_v) & (!rhs & rhs_v)
49                },
50            ))
51        },
52        (Some(lhs_validity), None) => {
53            // B != U
54            Some(ternary(
55                lhs_values,
56                rhs_values,
57                lhs_validity,
58                // see https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics
59                |lhs, rhs, lhs_v| {
60                    // A = T
61                    (lhs & lhs_v) |
62                    // B = T
63                    rhs |
64                    // A = F & B = F
65                    (!lhs & lhs_v) & !rhs
66                },
67            ))
68        },
69        (None, Some(rhs_validity)) => {
70            Some(ternary(
71                lhs_values,
72                rhs_values,
73                rhs_validity,
74                // see https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics
75                |lhs, rhs, rhs_v| {
76                    // A = T
77                    lhs |
78                    // B = T
79                    (rhs & rhs_v) |
80                    // A = F & B = F
81                    !lhs & (!rhs & rhs_v)
82                },
83            ))
84        },
85        (None, None) => None,
86    };
87    BooleanArray::new(ArrowDataType::Boolean, lhs_values | rhs_values, validity)
88}
89
90/// Logical 'and' operation on two arrays with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics)
91/// # Panics
92/// This function panics iff the arrays have a different length
93/// # Example
94///
95/// ```rust
96/// use polars_arrow::array::BooleanArray;
97/// use polars_arrow::compute::boolean_kleene::and;
98///
99/// let a = BooleanArray::from(&[Some(true), Some(false), None]);
100/// let b = BooleanArray::from(&[None, None, None]);
101/// let and_ab = and(&a, &b);
102/// assert_eq!(and_ab, BooleanArray::from(&[None, Some(false), None]));
103/// ```
104pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
105    assert_eq!(
106        lhs.len(),
107        rhs.len(),
108        "lhs and rhs must have the same length"
109    );
110
111    let lhs_values = lhs.values();
112    let rhs_values = rhs.values();
113
114    let lhs_validity = lhs.validity();
115    let rhs_validity = rhs.validity();
116
117    let validity = match (lhs_validity, rhs_validity) {
118        (Some(lhs_validity), Some(rhs_validity)) => {
119            Some(quaternary(
120                lhs_values,
121                rhs_values,
122                lhs_validity,
123                rhs_validity,
124                // see https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics
125                |lhs, rhs, lhs_v, rhs_v| {
126                    // B = F
127                    (!rhs & rhs_v) |
128                    // A = F
129                    (!lhs & lhs_v) |
130                    // A = T & B = T
131                    (lhs & lhs_v) & (rhs & rhs_v)
132                },
133            ))
134        },
135        (Some(lhs_validity), None) => {
136            Some(ternary(
137                lhs_values,
138                rhs_values,
139                lhs_validity,
140                // see https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics
141                |lhs, rhs, lhs_v| {
142                    // B = F
143                    !rhs |
144                    // A = F
145                    (!lhs & lhs_v) |
146                    // A = T & B = T
147                    (lhs & lhs_v) & rhs
148                },
149            ))
150        },
151        (None, Some(rhs_validity)) => {
152            Some(ternary(
153                lhs_values,
154                rhs_values,
155                rhs_validity,
156                // see https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics
157                |lhs, rhs, rhs_v| {
158                    // B = F
159                    (!rhs & rhs_v) |
160                    // A = F
161                    !lhs |
162                    // A = T & B = T
163                    lhs & (rhs & rhs_v)
164                },
165            ))
166        },
167        (None, None) => None,
168    };
169    BooleanArray::new(ArrowDataType::Boolean, lhs_values & rhs_values, validity)
170}
171
172/// Logical 'or' operation on an array and a scalar value with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics)
173/// # Example
174///
175/// ```rust
176/// use polars_arrow::array::BooleanArray;
177/// use polars_arrow::scalar::BooleanScalar;
178/// use polars_arrow::compute::boolean_kleene::or_scalar;
179///
180/// let array = BooleanArray::from(&[Some(true), Some(false), None]);
181/// let scalar = BooleanScalar::new(Some(false));
182/// let result = or_scalar(&array, &scalar);
183/// assert_eq!(result, BooleanArray::from(&[Some(true), Some(false), None]));
184/// ```
185pub fn or_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray {
186    match scalar.value() {
187        Some(true) => BooleanArray::new(
188            ArrowDataType::Boolean,
189            Bitmap::new_with_value(true, array.len()),
190            None,
191        ),
192        Some(false) => array.clone(),
193        None => {
194            let values = array.values();
195            let validity = match array.validity() {
196                Some(validity) => binary(values, validity, |value, validity| validity & value),
197                None => unary(values, |value| value),
198            };
199            BooleanArray::new(ArrowDataType::Boolean, values.clone(), Some(validity))
200        },
201    }
202}
203
204/// Logical 'and' operation on an array and a scalar value with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics)
205/// # Example
206///
207/// ```rust
208/// use polars_arrow::array::BooleanArray;
209/// use polars_arrow::scalar::BooleanScalar;
210/// use polars_arrow::compute::boolean_kleene::and_scalar;
211///
212/// let array = BooleanArray::from(&[Some(true), Some(false), None]);
213/// let scalar = BooleanScalar::new(None);
214/// let result = and_scalar(&array, &scalar);
215/// assert_eq!(result, BooleanArray::from(&[None, Some(false), None]));
216/// ```
217pub fn and_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray {
218    match scalar.value() {
219        Some(true) => array.clone(),
220        Some(false) => {
221            let values = Bitmap::new_zeroed(array.len());
222            BooleanArray::new(ArrowDataType::Boolean, values, None)
223        },
224        None => {
225            let values = array.values();
226            let validity = match array.validity() {
227                Some(validity) => binary(values, validity, |value, validity| validity & !value),
228                None => unary(values, |value| !value),
229            };
230            BooleanArray::new(
231                ArrowDataType::Boolean,
232                array.values().clone(),
233                Some(validity),
234            )
235        },
236    }
237}
238
239/// Returns whether any of the values in the array are `true`.
240///
241/// The output is unknown (`None`) if the array contains any null values and
242/// no `true` values.
243///
244/// # Example
245///
246/// ```
247/// use polars_arrow::array::BooleanArray;
248/// use polars_arrow::compute::boolean_kleene::any;
249///
250/// let a = BooleanArray::from(&[Some(true), Some(false)]);
251/// let b = BooleanArray::from(&[Some(false), Some(false)]);
252/// let c = BooleanArray::from(&[None, Some(false)]);
253///
254/// assert_eq!(any(&a), Some(true));
255/// assert_eq!(any(&b), Some(false));
256/// assert_eq!(any(&c), None);
257/// ```
258pub fn any(array: &BooleanArray) -> Option<bool> {
259    if array.is_empty() {
260        Some(false)
261    } else if array.null_count() > 0 {
262        if array.into_iter().any(|v| v == Some(true)) {
263            Some(true)
264        } else {
265            None
266        }
267    } else {
268        let vals = array.values();
269        Some(vals.unset_bits() != vals.len())
270    }
271}
272
273/// Returns whether all values in the array are `true`.
274///
275/// The output is unknown (`None`) if the array contains any null values and
276/// no `false` values.
277///
278/// # Example
279///
280/// ```
281/// use polars_arrow::array::BooleanArray;
282/// use polars_arrow::compute::boolean_kleene::all;
283///
284/// let a = BooleanArray::from(&[Some(true), Some(true)]);
285/// let b = BooleanArray::from(&[Some(false), Some(true)]);
286/// let c = BooleanArray::from(&[None, Some(true)]);
287///
288/// assert_eq!(all(&a), Some(true));
289/// assert_eq!(all(&b), Some(false));
290/// assert_eq!(all(&c), None);
291/// ```
292pub fn all(array: &BooleanArray) -> Option<bool> {
293    if array.is_empty() {
294        Some(true)
295    } else if array.null_count() > 0 {
296        if array.into_iter().any(|v| v == Some(false)) {
297            Some(false)
298        } else {
299            None
300        }
301    } else {
302        let vals = array.values();
303        Some(vals.unset_bits() == 0)
304    }
305}