arrow_arith/
boolean.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//! Defines boolean kernels on Arrow `BooleanArray`'s, e.g. `AND`, `OR` and `NOT`.
19//!
20//! These kernels can leverage SIMD if available on your system.  Currently no runtime
21//! detection is provided, you should enable the specific SIMD intrinsics using
22//! `RUSTFLAGS="-C target-feature=+avx2"` for example.  See the documentation
23//! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.
24
25use arrow_array::*;
26use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_quaternary_op_helper};
27use arrow_buffer::{BooleanBuffer, NullBuffer, buffer_bin_and_not};
28use arrow_schema::ArrowError;
29
30/// Logical 'and' boolean values with Kleene logic
31///
32/// # Behavior
33///
34/// This function behaves as follows with nulls:
35///
36/// * `true` and `null` = `null`
37/// * `null` and `true` = `null`
38/// * `false` and `null` = `false`
39/// * `null` and `false` = `false`
40/// * `null` and `null` = `null`
41///
42/// In other words, in this context a null value really means \"unknown\",
43/// and an unknown value 'and' false is always false.
44/// For a different null behavior, see function \"and\".
45///
46/// # Example
47///
48/// ```rust
49/// # use arrow_array::BooleanArray;
50/// # use arrow_arith::boolean::and_kleene;
51/// let a = BooleanArray::from(vec![Some(true), Some(false), None]);
52/// let b = BooleanArray::from(vec![None, None, None]);
53/// let and_ab = and_kleene(&a, &b).unwrap();
54/// assert_eq!(and_ab, BooleanArray::from(vec![None, Some(false), None]));
55/// ```
56///
57/// # Fails
58///
59/// If the operands have different lengths
60pub fn and_kleene(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
61    if left.len() != right.len() {
62        return Err(ArrowError::ComputeError(
63            "Cannot perform bitwise operation on arrays of different length".to_string(),
64        ));
65    }
66
67    let left_values = left.values();
68    let right_values = right.values();
69
70    let buffer = match (left.nulls(), right.nulls()) {
71        (None, None) => None,
72        (Some(left_null_buffer), None) => {
73            // The right side has no null values.
74            // The final null bit is set only if:
75            // 1. left null bit is set, or
76            // 2. right data bit is false (because null AND false = false).
77            Some(bitwise_bin_op_helper(
78                left_null_buffer.buffer(),
79                left_null_buffer.offset(),
80                right_values.inner(),
81                right_values.offset(),
82                left.len(),
83                |a, b| a | !b,
84            ))
85        }
86        (None, Some(right_null_buffer)) => {
87            // Same as above
88            Some(bitwise_bin_op_helper(
89                right_null_buffer.buffer(),
90                right_null_buffer.offset(),
91                left_values.inner(),
92                left_values.offset(),
93                left.len(),
94                |a, b| a | !b,
95            ))
96        }
97        (Some(left_null_buffer), Some(right_null_buffer)) => {
98            // Follow the same logic above. Both sides have null values.
99            // Assume a is left null bits, b is left data bits, c is right null bits,
100            // d is right data bits.
101            // The final null bits are:
102            // (a | (c & !d)) & (c | (a & !b))
103            Some(bitwise_quaternary_op_helper(
104                [
105                    left_null_buffer.buffer(),
106                    left_values.inner(),
107                    right_null_buffer.buffer(),
108                    right_values.inner(),
109                ],
110                [
111                    left_null_buffer.offset(),
112                    left_values.offset(),
113                    right_null_buffer.offset(),
114                    right_values.offset(),
115                ],
116                left.len(),
117                |a, b, c, d| (a | (c & !d)) & (c | (a & !b)),
118            ))
119        }
120    };
121    let nulls = buffer.map(|b| NullBuffer::new(BooleanBuffer::new(b, 0, left.len())));
122    Ok(BooleanArray::new(left_values & right_values, nulls))
123}
124
125/// Logical 'or' boolean values with Kleene logic
126///
127/// # Behavior
128///
129/// This function behaves as follows with nulls:
130///
131/// * `true` or `null` = `true`
132/// * `null` or `true` = `true`
133/// * `false` or `null` = `null`
134/// * `null` or `false` = `null`
135/// * `null` or `null` = `null`
136///
137/// In other words, in this context a null value really means \"unknown\",
138/// and an unknown value 'or' true is always true.
139/// For a different null behavior, see function \"or\".
140///
141/// # Example
142///
143/// ```rust
144/// # use arrow_array::BooleanArray;
145/// # use arrow_arith::boolean::or_kleene;
146/// let a = BooleanArray::from(vec![Some(true), Some(false), None]);
147/// let b = BooleanArray::from(vec![None, None, None]);
148/// let or_ab = or_kleene(&a, &b).unwrap();
149/// assert_eq!(or_ab, BooleanArray::from(vec![Some(true), None, None]));
150/// ```
151///
152/// # Fails
153///
154/// If the operands have different lengths
155pub fn or_kleene(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
156    if left.len() != right.len() {
157        return Err(ArrowError::ComputeError(
158            "Cannot perform bitwise operation on arrays of different length".to_string(),
159        ));
160    }
161
162    let left_values = left.values();
163    let right_values = right.values();
164
165    let buffer = match (left.nulls(), right.nulls()) {
166        (None, None) => None,
167        (Some(left_nulls), None) => {
168            // The right side has no null values.
169            // The final null bit is set only if:
170            // 1. left null bit is set, or
171            // 2. right data bit is true (because null OR true = true).
172            Some(bitwise_bin_op_helper(
173                left_nulls.buffer(),
174                left_nulls.offset(),
175                right_values.inner(),
176                right_values.offset(),
177                left.len(),
178                |a, b| a | b,
179            ))
180        }
181        (None, Some(right_nulls)) => {
182            // Same as above
183            Some(bitwise_bin_op_helper(
184                right_nulls.buffer(),
185                right_nulls.offset(),
186                left_values.inner(),
187                left_values.offset(),
188                left.len(),
189                |a, b| a | b,
190            ))
191        }
192        (Some(left_nulls), Some(right_nulls)) => {
193            // Follow the same logic above. Both sides have null values.
194            // Assume a is left null bits, b is left data bits, c is right null bits,
195            // d is right data bits.
196            // The final null bits are:
197            // (a | (c & d)) & (c | (a & b))
198            Some(bitwise_quaternary_op_helper(
199                [
200                    left_nulls.buffer(),
201                    left_values.inner(),
202                    right_nulls.buffer(),
203                    right_values.inner(),
204                ],
205                [
206                    left_nulls.offset(),
207                    left_values.offset(),
208                    right_nulls.offset(),
209                    right_values.offset(),
210                ],
211                left.len(),
212                |a, b, c, d| (a | (c & d)) & (c | (a & b)),
213            ))
214        }
215    };
216
217    let nulls = buffer.map(|b| NullBuffer::new(BooleanBuffer::new(b, 0, left.len())));
218    Ok(BooleanArray::new(left_values | right_values, nulls))
219}
220
221/// Helper function to implement binary kernels
222pub(crate) fn binary_boolean_kernel<F>(
223    left: &BooleanArray,
224    right: &BooleanArray,
225    op: F,
226) -> Result<BooleanArray, ArrowError>
227where
228    F: Fn(&BooleanBuffer, &BooleanBuffer) -> BooleanBuffer,
229{
230    if left.len() != right.len() {
231        return Err(ArrowError::ComputeError(
232            "Cannot perform bitwise operation on arrays of different length".to_string(),
233        ));
234    }
235
236    let nulls = NullBuffer::union(left.nulls(), right.nulls());
237    let values = op(left.values(), right.values());
238    Ok(BooleanArray::new(values, nulls))
239}
240
241/// Performs `AND` operation on two arrays. If either left or right value is null then the
242/// result is also null.
243/// # Error
244/// This function errors when the arrays have different lengths.
245/// # Example
246/// ```rust
247/// # use arrow_array::BooleanArray;
248/// # use arrow_arith::boolean::and;
249/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
250/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
251/// let and_ab = and(&a, &b).unwrap();
252/// assert_eq!(and_ab, BooleanArray::from(vec![Some(false), Some(true), None]));
253/// ```
254pub fn and(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
255    binary_boolean_kernel(left, right, |a, b| a & b)
256}
257
258/// Performs `OR` operation on two arrays. If either left or right value is null then the
259/// result is also null.
260/// # Error
261/// This function errors when the arrays have different lengths.
262/// # Example
263/// ```rust
264/// # use arrow_array::BooleanArray;
265/// # use arrow_arith::boolean::or;
266/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
267/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
268/// let or_ab = or(&a, &b).unwrap();
269/// assert_eq!(or_ab, BooleanArray::from(vec![Some(true), Some(true), None]));
270/// ```
271pub fn or(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
272    binary_boolean_kernel(left, right, |a, b| a | b)
273}
274
275/// Performs `AND_NOT` operation on two arrays. If either left or right value is null then the
276/// result is also null.
277/// # Error
278/// This function errors when the arrays have different lengths.
279/// # Example
280/// ```rust
281/// # use arrow_array::BooleanArray;
282/// # use arrow_arith::boolean::{and, not, and_not};
283/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
284/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
285/// let andn_ab = and_not(&a, &b).unwrap();
286/// assert_eq!(andn_ab, BooleanArray::from(vec![Some(false), Some(false), None]));
287/// // It's equal to and(left, not(right))
288/// assert_eq!(andn_ab, and(&a, &not(&b).unwrap()).unwrap());
289pub fn and_not(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
290    binary_boolean_kernel(left, right, |a, b| {
291        let buffer = buffer_bin_and_not(a.inner(), b.offset(), b.inner(), a.offset(), a.len());
292        BooleanBuffer::new(buffer, left.offset(), left.len())
293    })
294}
295
296/// Performs unary `NOT` operation on an arrays. If value is null then the result is also
297/// null.
298/// # Error
299/// This function never errors. It returns an error for consistency.
300/// # Example
301/// ```rust
302/// # use arrow_array::BooleanArray;
303/// # use arrow_arith::boolean::not;
304/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
305/// let not_a = not(&a).unwrap();
306/// assert_eq!(not_a, BooleanArray::from(vec![Some(true), Some(false), None]));
307/// ```
308pub fn not(left: &BooleanArray) -> Result<BooleanArray, ArrowError> {
309    let nulls = left.nulls().cloned();
310    let values = !left.values();
311    Ok(BooleanArray::new(values, nulls))
312}
313
314/// Returns a non-null [BooleanArray] with whether each value of the array is null.
315/// # Error
316/// This function never errors.
317/// # Example
318/// ```rust
319/// # use arrow_array::BooleanArray;
320/// # use arrow_arith::boolean::is_null;
321/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
322/// let a_is_null = is_null(&a).unwrap();
323/// assert_eq!(a_is_null, BooleanArray::from(vec![false, false, true]));
324/// ```
325pub fn is_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
326    let values = match input.logical_nulls() {
327        None => BooleanBuffer::new_unset(input.len()),
328        Some(nulls) => !nulls.inner(),
329    };
330
331    Ok(BooleanArray::new(values, None))
332}
333
334/// Returns a non-null [BooleanArray] with whether each value of the array is not null.
335/// # Error
336/// This function never errors.
337/// # Example
338/// ```rust
339/// # use arrow_array::BooleanArray;
340/// # use arrow_arith::boolean::is_not_null;
341/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
342/// let a_is_not_null = is_not_null(&a).unwrap();
343/// assert_eq!(a_is_not_null, BooleanArray::from(vec![true, true, false]));
344/// ```
345pub fn is_not_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
346    let values = match input.logical_nulls() {
347        None => BooleanBuffer::new_set(input.len()),
348        Some(n) => n.inner().clone(),
349    };
350    Ok(BooleanArray::new(values, None))
351}
352
353#[cfg(test)]
354mod tests {
355    use arrow_buffer::ScalarBuffer;
356    use arrow_schema::{DataType, Field, UnionFields};
357
358    use super::*;
359    use std::sync::Arc;
360
361    #[test]
362    fn test_bool_array_and() {
363        let a = BooleanArray::from(vec![false, false, true, true]);
364        let b = BooleanArray::from(vec![false, true, false, true]);
365        let c = and(&a, &b).unwrap();
366
367        let expected = BooleanArray::from(vec![false, false, false, true]);
368
369        assert_eq!(c, expected);
370    }
371
372    #[test]
373    fn test_bool_array_or() {
374        let a = BooleanArray::from(vec![false, false, true, true]);
375        let b = BooleanArray::from(vec![false, true, false, true]);
376        let c = or(&a, &b).unwrap();
377
378        let expected = BooleanArray::from(vec![false, true, true, true]);
379
380        assert_eq!(c, expected);
381    }
382
383    #[test]
384    fn test_bool_array_and_not() {
385        let a = BooleanArray::from(vec![false, false, true, true]);
386        let b = BooleanArray::from(vec![false, true, false, true]);
387        let c = and_not(&a, &b).unwrap();
388
389        let expected = BooleanArray::from(vec![false, false, true, false]);
390
391        assert_eq!(c, expected);
392        assert_eq!(c, and(&a, &not(&b).unwrap()).unwrap());
393    }
394
395    #[test]
396    fn test_bool_array_or_nulls() {
397        let a = BooleanArray::from(vec![
398            None,
399            None,
400            None,
401            Some(false),
402            Some(false),
403            Some(false),
404            Some(true),
405            Some(true),
406            Some(true),
407        ]);
408        let b = BooleanArray::from(vec![
409            None,
410            Some(false),
411            Some(true),
412            None,
413            Some(false),
414            Some(true),
415            None,
416            Some(false),
417            Some(true),
418        ]);
419        let c = or(&a, &b).unwrap();
420
421        let expected = BooleanArray::from(vec![
422            None,
423            None,
424            None,
425            None,
426            Some(false),
427            Some(true),
428            None,
429            Some(true),
430            Some(true),
431        ]);
432
433        assert_eq!(c, expected);
434    }
435
436    #[test]
437    fn test_boolean_array_kleene_no_remainder() {
438        let n = 1024;
439        let a = BooleanArray::from(vec![true; n]);
440        let b = BooleanArray::from(vec![None; n]);
441        let result = or_kleene(&a, &b).unwrap();
442
443        assert_eq!(result, a);
444    }
445
446    #[test]
447    fn test_bool_array_and_kleene_nulls() {
448        let a = BooleanArray::from(vec![
449            None,
450            None,
451            None,
452            Some(false),
453            Some(false),
454            Some(false),
455            Some(true),
456            Some(true),
457            Some(true),
458        ]);
459        let b = BooleanArray::from(vec![
460            None,
461            Some(false),
462            Some(true),
463            None,
464            Some(false),
465            Some(true),
466            None,
467            Some(false),
468            Some(true),
469        ]);
470        let c = and_kleene(&a, &b).unwrap();
471
472        let expected = BooleanArray::from(vec![
473            None,
474            Some(false),
475            None,
476            Some(false),
477            Some(false),
478            Some(false),
479            None,
480            Some(false),
481            Some(true),
482        ]);
483
484        assert_eq!(c, expected);
485    }
486
487    #[test]
488    fn test_bool_array_or_kleene_nulls() {
489        let a = BooleanArray::from(vec![
490            None,
491            None,
492            None,
493            Some(false),
494            Some(false),
495            Some(false),
496            Some(true),
497            Some(true),
498            Some(true),
499        ]);
500        let b = BooleanArray::from(vec![
501            None,
502            Some(false),
503            Some(true),
504            None,
505            Some(false),
506            Some(true),
507            None,
508            Some(false),
509            Some(true),
510        ]);
511        let c = or_kleene(&a, &b).unwrap();
512
513        let expected = BooleanArray::from(vec![
514            None,
515            None,
516            Some(true),
517            None,
518            Some(false),
519            Some(true),
520            Some(true),
521            Some(true),
522            Some(true),
523        ]);
524
525        assert_eq!(c, expected);
526    }
527
528    #[test]
529    fn test_bool_array_or_kleene_right_sided_nulls() {
530        let a = BooleanArray::from(vec![false, false, false, true, true, true]);
531
532        // ensure null bitmap of a is absent
533        assert!(a.nulls().is_none());
534
535        let b = BooleanArray::from(vec![
536            Some(true),
537            Some(false),
538            None,
539            Some(true),
540            Some(false),
541            None,
542        ]);
543
544        // ensure null bitmap of b is present
545        assert!(b.nulls().is_some());
546
547        let c = or_kleene(&a, &b).unwrap();
548
549        let expected = BooleanArray::from(vec![
550            Some(true),
551            Some(false),
552            None,
553            Some(true),
554            Some(true),
555            Some(true),
556        ]);
557
558        assert_eq!(c, expected);
559    }
560
561    #[test]
562    fn test_bool_array_or_kleene_left_sided_nulls() {
563        let a = BooleanArray::from(vec![
564            Some(true),
565            Some(false),
566            None,
567            Some(true),
568            Some(false),
569            None,
570        ]);
571
572        // ensure null bitmap of b is absent
573        assert!(a.nulls().is_some());
574
575        let b = BooleanArray::from(vec![false, false, false, true, true, true]);
576
577        // ensure null bitmap of a is present
578        assert!(b.nulls().is_none());
579
580        let c = or_kleene(&a, &b).unwrap();
581
582        let expected = BooleanArray::from(vec![
583            Some(true),
584            Some(false),
585            None,
586            Some(true),
587            Some(true),
588            Some(true),
589        ]);
590
591        assert_eq!(c, expected);
592    }
593
594    #[test]
595    fn test_bool_array_not() {
596        let a = BooleanArray::from(vec![false, true]);
597        let c = not(&a).unwrap();
598
599        let expected = BooleanArray::from(vec![true, false]);
600
601        assert_eq!(c, expected);
602    }
603
604    #[test]
605    fn test_bool_array_not_sliced() {
606        let a = BooleanArray::from(vec![None, Some(true), Some(false), None, Some(true)]);
607        let a = a.slice(1, 4);
608        let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
609        let c = not(a).unwrap();
610
611        let expected = BooleanArray::from(vec![Some(false), Some(true), None, Some(false)]);
612
613        assert_eq!(c, expected);
614    }
615
616    #[test]
617    fn test_bool_array_and_nulls() {
618        let a = BooleanArray::from(vec![
619            None,
620            None,
621            None,
622            Some(false),
623            Some(false),
624            Some(false),
625            Some(true),
626            Some(true),
627            Some(true),
628        ]);
629        let b = BooleanArray::from(vec![
630            None,
631            Some(false),
632            Some(true),
633            None,
634            Some(false),
635            Some(true),
636            None,
637            Some(false),
638            Some(true),
639        ]);
640        let c = and(&a, &b).unwrap();
641
642        let expected = BooleanArray::from(vec![
643            None,
644            None,
645            None,
646            None,
647            Some(false),
648            Some(false),
649            None,
650            Some(false),
651            Some(true),
652        ]);
653
654        assert_eq!(c, expected);
655    }
656
657    #[test]
658    fn test_bool_array_and_sliced_same_offset() {
659        let a = BooleanArray::from(vec![
660            false, false, false, false, false, false, false, false, false, false, true, true,
661        ]);
662        let b = BooleanArray::from(vec![
663            false, false, false, false, false, false, false, false, false, true, false, true,
664        ]);
665
666        let a = a.slice(8, 4);
667        let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
668        let b = b.slice(8, 4);
669        let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
670
671        let c = and(a, b).unwrap();
672
673        let expected = BooleanArray::from(vec![false, false, false, true]);
674
675        assert_eq!(expected, c);
676    }
677
678    #[test]
679    fn test_bool_array_and_sliced_same_offset_mod8() {
680        let a = BooleanArray::from(vec![
681            false, false, true, true, false, false, false, false, false, false, false, false,
682        ]);
683        let b = BooleanArray::from(vec![
684            false, false, false, false, false, false, false, false, false, true, false, true,
685        ]);
686
687        let a = a.slice(0, 4);
688        let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
689        let b = b.slice(8, 4);
690        let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
691
692        let c = and(a, b).unwrap();
693
694        let expected = BooleanArray::from(vec![false, false, false, true]);
695
696        assert_eq!(expected, c);
697    }
698
699    #[test]
700    fn test_bool_array_and_sliced_offset1() {
701        let a = BooleanArray::from(vec![
702            false, false, false, false, false, false, false, false, false, false, true, true,
703        ]);
704        let b = BooleanArray::from(vec![false, true, false, true]);
705
706        let a = a.slice(8, 4);
707        let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
708
709        let c = and(a, &b).unwrap();
710
711        let expected = BooleanArray::from(vec![false, false, false, true]);
712
713        assert_eq!(expected, c);
714    }
715
716    #[test]
717    fn test_bool_array_and_sliced_offset2() {
718        let a = BooleanArray::from(vec![false, false, true, true]);
719        let b = BooleanArray::from(vec![
720            false, false, false, false, false, false, false, false, false, true, false, true,
721        ]);
722
723        let b = b.slice(8, 4);
724        let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
725
726        let c = and(&a, b).unwrap();
727
728        let expected = BooleanArray::from(vec![false, false, false, true]);
729
730        assert_eq!(expected, c);
731    }
732
733    #[test]
734    fn test_bool_array_and_nulls_offset() {
735        let a = BooleanArray::from(vec![None, Some(false), Some(true), None, Some(true)]);
736        let a = a.slice(1, 4);
737        let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
738
739        let b = BooleanArray::from(vec![
740            None,
741            None,
742            Some(true),
743            Some(false),
744            Some(true),
745            Some(true),
746        ]);
747
748        let b = b.slice(2, 4);
749        let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
750
751        let c = and(a, b).unwrap();
752
753        let expected = BooleanArray::from(vec![Some(false), Some(false), None, Some(true)]);
754
755        assert_eq!(expected, c);
756    }
757
758    #[test]
759    fn test_nonnull_array_is_null() {
760        let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
761
762        let res = is_null(a.as_ref()).unwrap();
763
764        let expected = BooleanArray::from(vec![false, false, false, false]);
765
766        assert_eq!(expected, res);
767        assert!(res.nulls().is_none());
768    }
769
770    #[test]
771    fn test_nonnull_array_with_offset_is_null() {
772        let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
773        let a = a.slice(8, 4);
774
775        let res = is_null(&a).unwrap();
776
777        let expected = BooleanArray::from(vec![false, false, false, false]);
778
779        assert_eq!(expected, res);
780        assert!(res.nulls().is_none());
781    }
782
783    #[test]
784    fn test_nonnull_array_is_not_null() {
785        let a = Int32Array::from(vec![1, 2, 3, 4]);
786
787        let res = is_not_null(&a).unwrap();
788
789        let expected = BooleanArray::from(vec![true, true, true, true]);
790
791        assert_eq!(expected, res);
792        assert!(res.nulls().is_none());
793    }
794
795    #[test]
796    fn test_nonnull_array_with_offset_is_not_null() {
797        let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
798        let a = a.slice(8, 4);
799
800        let res = is_not_null(&a).unwrap();
801
802        let expected = BooleanArray::from(vec![true, true, true, true]);
803
804        assert_eq!(expected, res);
805        assert!(res.nulls().is_none());
806    }
807
808    #[test]
809    fn test_nullable_array_is_null() {
810        let a = Int32Array::from(vec![Some(1), None, Some(3), None]);
811
812        let res = is_null(&a).unwrap();
813
814        let expected = BooleanArray::from(vec![false, true, false, true]);
815
816        assert_eq!(expected, res);
817        assert!(res.nulls().is_none());
818    }
819
820    #[test]
821    fn test_nullable_array_with_offset_is_null() {
822        let a = Int32Array::from(vec![
823            None,
824            None,
825            None,
826            None,
827            None,
828            None,
829            None,
830            None,
831            // offset 8, previous None values are skipped by the slice
832            Some(1),
833            None,
834            Some(2),
835            None,
836            Some(3),
837            Some(4),
838            None,
839            None,
840        ]);
841        let a = a.slice(8, 4);
842
843        let res = is_null(&a).unwrap();
844
845        let expected = BooleanArray::from(vec![false, true, false, true]);
846
847        assert_eq!(expected, res);
848        assert!(res.nulls().is_none());
849    }
850
851    #[test]
852    fn test_nullable_array_is_not_null() {
853        let a = Int32Array::from(vec![Some(1), None, Some(3), None]);
854
855        let res = is_not_null(&a).unwrap();
856
857        let expected = BooleanArray::from(vec![true, false, true, false]);
858
859        assert_eq!(expected, res);
860        assert!(res.nulls().is_none());
861    }
862
863    #[test]
864    fn test_nullable_array_with_offset_is_not_null() {
865        let a = Int32Array::from(vec![
866            None,
867            None,
868            None,
869            None,
870            None,
871            None,
872            None,
873            None,
874            // offset 8, previous None values are skipped by the slice
875            Some(1),
876            None,
877            Some(2),
878            None,
879            Some(3),
880            Some(4),
881            None,
882            None,
883        ]);
884        let a = a.slice(8, 4);
885
886        let res = is_not_null(&a).unwrap();
887
888        let expected = BooleanArray::from(vec![true, false, true, false]);
889
890        assert_eq!(expected, res);
891        assert!(res.nulls().is_none());
892    }
893
894    #[test]
895    fn test_null_array_is_null() {
896        let a = NullArray::new(3);
897
898        let res = is_null(&a).unwrap();
899
900        let expected = BooleanArray::from(vec![true, true, true]);
901
902        assert_eq!(expected, res);
903        assert!(res.nulls().is_none());
904    }
905
906    #[test]
907    fn test_null_array_is_not_null() {
908        let a = NullArray::new(3);
909
910        let res = is_not_null(&a).unwrap();
911
912        let expected = BooleanArray::from(vec![false, false, false]);
913
914        assert_eq!(expected, res);
915        assert!(res.nulls().is_none());
916    }
917
918    #[test]
919    fn test_dense_union_is_null() {
920        // union of [{A=1}, {A=}, {B=3.2}, {B=}, {C="a"}, {C=}]
921        let int_array = Int32Array::from(vec![Some(1), None]);
922        let float_array = Float64Array::from(vec![Some(3.2), None]);
923        let str_array = StringArray::from(vec![Some("a"), None]);
924        let type_ids = [0, 0, 1, 1, 2, 2].into_iter().collect::<ScalarBuffer<i8>>();
925        let offsets = [0, 1, 0, 1, 0, 1]
926            .into_iter()
927            .collect::<ScalarBuffer<i32>>();
928
929        let children = vec![
930            Arc::new(int_array) as Arc<dyn Array>,
931            Arc::new(float_array),
932            Arc::new(str_array),
933        ];
934
935        let array = UnionArray::try_new(union_fields(), type_ids, Some(offsets), children).unwrap();
936
937        let result = is_null(&array).unwrap();
938
939        let expected = &BooleanArray::from(vec![false, true, false, true, false, true]);
940        assert_eq!(expected, &result);
941    }
942
943    #[test]
944    fn test_sparse_union_is_null() {
945        // union of [{A=1}, {A=}, {B=3.2}, {B=}, {C="a"}, {C=}]
946        let int_array = Int32Array::from(vec![Some(1), None, None, None, None, None]);
947        let float_array = Float64Array::from(vec![None, None, Some(3.2), None, None, None]);
948        let str_array = StringArray::from(vec![None, None, None, None, Some("a"), None]);
949        let type_ids = [0, 0, 1, 1, 2, 2].into_iter().collect::<ScalarBuffer<i8>>();
950
951        let children = vec![
952            Arc::new(int_array) as Arc<dyn Array>,
953            Arc::new(float_array),
954            Arc::new(str_array),
955        ];
956
957        let array = UnionArray::try_new(union_fields(), type_ids, None, children).unwrap();
958
959        let result = is_null(&array).unwrap();
960
961        let expected = &BooleanArray::from(vec![false, true, false, true, false, true]);
962        assert_eq!(expected, &result);
963    }
964
965    fn union_fields() -> UnionFields {
966        [
967            (0, Arc::new(Field::new("A", DataType::Int32, true))),
968            (1, Arc::new(Field::new("B", DataType::Float64, true))),
969            (2, Arc::new(Field::new("C", DataType::Utf8, true))),
970        ]
971        .into_iter()
972        .collect()
973    }
974}
arrow_arith/boolean.rs

arrow_arith/
boolean.rs
