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

arrow_arith/
boolean.rs
