arrow_select/
nullif.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//! Implements the `nullif` function for Arrow arrays.
19
20use arrow_array::{Array, ArrayRef, BooleanArray, UInt64Array, make_array};
21use arrow_buffer::buffer::bitwise_bin_op_helper;
22use arrow_buffer::{BooleanBuffer, NullBuffer, bitwise_unary_op_helper};
23use arrow_schema::{ArrowError, DataType};
24
25/// Returns a new array with the same values and the validity bit to false where
26/// the corresponding element of `right` is true.
27///
28/// This can be used to implement SQL `NULLIF`
29///
30/// # Example
31/// ```
32/// # use arrow_array::{Int32Array, BooleanArray};
33/// # use arrow_array::cast::AsArray;
34/// # use arrow_array::types::Int32Type;
35/// # use arrow_select::nullif::nullif;
36/// // input is [null, 8, 1, 9]
37/// let a = Int32Array::from(vec![None, Some(8), Some(1), Some(9)]);
38/// // use nullif to set index 1 to null
39/// let bool_array = BooleanArray::from(vec![Some(false), Some(true), Some(false), None]);
40/// let nulled = nullif(&a, &bool_array).unwrap();
41/// // The resulting array is [null, null, 1, 9]
42/// assert_eq!(nulled.as_primitive(), &Int32Array::from(vec![None, None, Some(1), Some(9)]));
43/// ```
44pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result<ArrayRef, ArrowError> {
45    if left.len() != right.len() {
46        return Err(ArrowError::ComputeError(
47            "Cannot perform comparison operation on arrays of different length".to_string(),
48        ));
49    }
50    let len = left.len();
51
52    if len == 0 || left.data_type() == &DataType::Null {
53        return Ok(make_array(left.to_data()));
54    }
55
56    // Compute right_values & right_bitmap. True bits are positions that should
57    // become null; nulls in `right` do not introduce nulls.
58    let should_null = match right.nulls() {
59        Some(nulls) => right.values() & nulls.inner(),
60        None => right.values().clone(),
61    };
62
63    // REE & Union have cases where they technically shouldn't be allowed to hold
64    // nulls based on nullability of their children
65    match left.data_type() {
66        DataType::RunEndEncoded(_, values) => {
67            if !values.is_nullable() && should_null.has_true() {
68                return Err(ArrowError::ComputeError(
69                    "Cannot introduce nulls into a RunEndEncoded array with a non-nullable values field".into(),
70                ));
71            }
72            return nullif_take(left, &should_null);
73        }
74        DataType::Union(_, _) => return nullif_take(left, &should_null),
75        _ => {}
76    }
77
78    let left_data = left.to_data();
79
80    // left=0 (null)   right=null       output bitmap=null
81    // left=0          right=1          output bitmap=null
82    // left=1 (set)    right=null       output bitmap=set   (passthrough)
83    // left=1          right=1 & comp=true    output bitmap=null
84    // left=1          right=1 & comp=false   output bitmap=set
85    //
86    // Thus: result = left null bitmap & (!right_values | !right_bitmap)
87    //              OR left null bitmap & !(right_values & right_bitmap)
88
89    // Compute left null bitmap & !should_null
90
91    let (combined, null_count) = match left_data.nulls() {
92        Some(left) => {
93            let mut valid_count = 0;
94            let b = bitwise_bin_op_helper(
95                left.buffer(),
96                left.offset(),
97                should_null.inner(),
98                should_null.offset(),
99                len,
100                |l, r| {
101                    let t = l & !r;
102                    valid_count += t.count_ones() as usize;
103                    t
104                },
105            );
106            (b, len - valid_count)
107        }
108        None => {
109            let mut null_count = 0;
110            let buffer =
111                bitwise_unary_op_helper(should_null.inner(), should_null.offset(), len, |b| {
112                    let t = !b;
113                    null_count += t.count_zeros() as usize;
114                    t
115                });
116            (buffer, null_count)
117        }
118    };
119
120    let combined = BooleanBuffer::new(combined, 0, len);
121    // Safety:
122    // Counted nulls whilst computing
123    let nulls = unsafe { NullBuffer::new_unchecked(combined, null_count) };
124    let data = left_data.into_builder().nulls(Some(nulls));
125
126    // SAFETY:
127    // Only altered null mask
128    Ok(make_array(unsafe { data.build_unchecked() }))
129}
130
131/// Applies `nullif` to arrays that represent logical nulls in their children.
132fn nullif_take(left: &dyn Array, should_null: &BooleanBuffer) -> Result<ArrayRef, ArrowError> {
133    let indices = UInt64Array::from_iter((0..left.len()).map(|index| {
134        if should_null.value(index) {
135            None
136        } else {
137            Some(index as u64)
138        }
139    }));
140    crate::take::take(left, &indices, None)
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use arrow_array::builder::{BooleanBuilder, Int32Builder, StructBuilder};
147    use arrow_array::cast::AsArray;
148    use arrow_array::types::Int16Type;
149    use arrow_array::types::Int32Type;
150    use arrow_array::{
151        Int16Array, Int32Array, NullArray, RunArray, StringArray, StructArray, UnionArray,
152    };
153    use arrow_buffer::{RunEndBuffer, ScalarBuffer};
154    use arrow_data::ArrayData;
155    use arrow_schema::{Field, Fields, UnionFields, UnionMode};
156    use rand::prelude::StdRng;
157    use rand::{RngExt, SeedableRng};
158    use std::sync::Arc;
159
160    #[test]
161    fn test_nullif_int_array() {
162        let a = Int32Array::from(vec![Some(15), None, Some(8), Some(1), Some(9)]);
163        let comp = BooleanArray::from(vec![Some(false), None, Some(true), Some(false), None]);
164        let res = nullif(&a, &comp).unwrap();
165
166        let expected = Int32Array::from(vec![
167            Some(15),
168            None,
169            None, // comp true, slot 2 turned into null
170            Some(1),
171            // Even though comp array / right is null, should still pass through original value
172            // comp true, slot 2 turned into null
173            Some(9),
174        ]);
175
176        let res = res.as_primitive::<Int32Type>();
177        assert_eq!(&expected, res);
178    }
179
180    #[test]
181    fn test_nullif_null_array() {
182        assert_eq!(
183            nullif(&NullArray::new(0), &BooleanArray::new_null(0))
184                .unwrap()
185                .as_ref(),
186            &NullArray::new(0)
187        );
188
189        assert_eq!(
190            nullif(
191                &NullArray::new(3),
192                &BooleanArray::from(vec![Some(false), Some(true), None]),
193            )
194            .unwrap()
195            .as_ref(),
196            &NullArray::new(3)
197        );
198    }
199
200    #[test]
201    fn test_nullif_run_end_encoded() {
202        let values = Int32Array::from(vec![10, 20]);
203        let ree = RunArray::<Int16Type>::try_new(&Int16Array::from(vec![1, 2]), &values).unwrap();
204        let mask = BooleanArray::from(vec![Some(false), Some(true)]);
205
206        let result = nullif(&ree, &mask).unwrap();
207        let result = result.as_run::<Int16Type>();
208        assert_eq!(result.logical_null_count(), 1);
209        assert_eq!(
210            result
211                .values()
212                .as_primitive::<Int32Type>()
213                .iter()
214                .collect::<Vec<_>>(),
215            vec![Some(10), None]
216        );
217    }
218
219    #[test]
220    fn test_nullif_run_end_encoded_mask_nulls() {
221        let values = Int32Array::from(vec![10, 20, 30]);
222        let ree =
223            RunArray::<Int16Type>::try_new(&Int16Array::from(vec![1, 2, 3]), &values).unwrap();
224        let mask = BooleanArray::from(vec![Some(true), Some(false), None, Some(true)]);
225        let mask = mask.slice(1, 3); // Some(false), None, Some(true)
226
227        let result = nullif(&ree, &mask).unwrap();
228        let result = result.as_run::<Int16Type>();
229        assert_eq!(result.logical_null_count(), 1);
230        assert_eq!(
231            result
232                .values()
233                .as_primitive::<Int32Type>()
234                .iter()
235                .collect::<Vec<_>>(),
236            vec![Some(10), Some(20), None]
237        );
238    }
239
240    #[test]
241    fn test_nullif_run_end_encoded_non_nullable_values() {
242        let ree = unsafe {
243            RunArray::<Int16Type>::new_unchecked(
244                DataType::RunEndEncoded(
245                    Arc::new(Field::new("run_ends", DataType::Int16, false)),
246                    Arc::new(Field::new("values", DataType::Int32, false)),
247                ),
248                RunEndBuffer::new(vec![1_i16, 2].into(), 0, 2),
249                Arc::new(Int32Array::from(vec![10, 20])),
250            )
251        };
252        let mask = BooleanArray::from(vec![Some(false), Some(true)]);
253
254        let error = nullif(&ree, &mask).unwrap_err();
255        assert_eq!(
256            error.to_string(),
257            "Compute error: Cannot introduce nulls into a RunEndEncoded array with a non-nullable values field"
258        );
259    }
260
261    #[test]
262    fn test_nullif_union() {
263        let fields =
264            UnionFields::try_new(vec![0], vec![Field::new("i", DataType::Int32, true)]).unwrap();
265        let union = UnionArray::try_new(
266            fields,
267            ScalarBuffer::from(vec![0_i8, 0]),
268            Some(ScalarBuffer::from(vec![0_i32, 1])),
269            vec![Arc::new(Int32Array::from(vec![10, 20]))],
270        )
271        .unwrap();
272        let mask = BooleanArray::from(vec![Some(false), Some(true)]);
273
274        let result = nullif(&union, &mask).unwrap();
275        let result = result.as_union();
276        assert_eq!(result.logical_null_count(), 1);
277        assert_eq!(result.type_ids(), &ScalarBuffer::from(vec![0_i8, 0]));
278        assert_eq!(
279            result
280                .child(0)
281                .as_primitive::<Int32Type>()
282                .iter()
283                .collect::<Vec<_>>(),
284            vec![Some(10), None]
285        );
286        assert!(matches!(
287            result.data_type(),
288            DataType::Union(_, UnionMode::Dense)
289        ));
290    }
291
292    #[test]
293    fn test_nullif_union_without_nullable_child() {
294        let fields =
295            UnionFields::try_new(vec![0], vec![Field::new("i", DataType::Int32, false)]).unwrap();
296        let union = UnionArray::try_new(
297            fields,
298            ScalarBuffer::from(vec![0_i8, 0]),
299            None,
300            vec![Arc::new(Int32Array::from(vec![10, 20]))],
301        )
302        .unwrap();
303        let mask = BooleanArray::from(vec![Some(false), Some(true)]);
304
305        let error = nullif(&union, &mask).unwrap_err();
306        assert_eq!(
307            error.to_string(),
308            "Compute error: Cannot take null indices from a union with no field that can represent nulls"
309        );
310    }
311
312    #[test]
313    fn test_nullif_int_array_offset() {
314        let a = Int32Array::from(vec![None, Some(15), Some(8), Some(1), Some(9)]);
315        let a = a.slice(1, 3); // Some(15), Some(8), Some(1)
316        let comp = BooleanArray::from(vec![
317            Some(false),
318            Some(false),
319            Some(false),
320            None,
321            Some(true),
322            Some(false),
323            None,
324        ]);
325        let comp = comp.slice(2, 3); // Some(false), None, Some(true)
326        let res = nullif(&a, &comp).unwrap();
327
328        let expected = Int32Array::from(vec![
329            Some(15), // False => keep it
330            Some(8),  // None => keep it
331            None,     // true => None
332        ]);
333        let res = res.as_primitive::<Int32Type>();
334        assert_eq!(&expected, res)
335    }
336
337    #[test]
338    fn test_nullif_string() {
339        let s = StringArray::from_iter([
340            Some("hello"),
341            None,
342            Some("world"),
343            Some("a"),
344            Some("b"),
345            None,
346            None,
347        ]);
348        let select = BooleanArray::from_iter([
349            Some(true),
350            Some(true),
351            Some(false),
352            Some(true),
353            Some(false),
354            Some(false),
355            None,
356        ]);
357
358        let a = nullif(&s, &select).unwrap();
359        let r: Vec<_> = a.as_string::<i32>().iter().collect();
360        assert_eq!(
361            r,
362            vec![None, None, Some("world"), None, Some("b"), None, None]
363        );
364
365        let s = s.slice(2, 3);
366        let select = select.slice(1, 3);
367        let a = nullif(&s, &select).unwrap();
368        let r: Vec<_> = a.as_string::<i32>().iter().collect();
369        assert_eq!(r, vec![None, Some("a"), None]);
370    }
371
372    #[test]
373    fn test_nullif_int_large_left_offset() {
374        let a = Int32Array::from(vec![
375            Some(-1), // 0
376            Some(-1),
377            Some(-1),
378            Some(-1),
379            Some(-1),
380            Some(-1),
381            Some(-1),
382            Some(-1),
383            Some(-1), // 8
384            Some(-1),
385            Some(-1),
386            Some(-1),
387            Some(-1),
388            Some(-1),
389            Some(-1),
390            Some(-1),
391            None,     // 16
392            Some(15), // 17
393            Some(8),
394            Some(1),
395            Some(9),
396        ]);
397        let a = a.slice(17, 3); // Some(15), Some(8), Some(1)
398
399        let comp = BooleanArray::from(vec![
400            Some(false),
401            Some(false),
402            Some(false),
403            None,
404            Some(true),
405            Some(false),
406            None,
407        ]);
408        let comp = comp.slice(2, 3); // Some(false), None, Some(true)
409        let res = nullif(&a, &comp).unwrap();
410        let res = res.as_primitive::<Int32Type>();
411
412        let expected = Int32Array::from(vec![
413            Some(15), // False => keep it
414            Some(8),  // None => keep it
415            None,     // true => None
416        ]);
417        assert_eq!(&expected, res)
418    }
419
420    #[test]
421    fn test_nullif_int_large_right_offset() {
422        let a = Int32Array::from(vec![
423            None,     // 0
424            Some(15), // 1
425            Some(8),
426            Some(1),
427            Some(9),
428        ]);
429        let a = a.slice(1, 3); // Some(15), Some(8), Some(1)
430
431        let comp = BooleanArray::from(vec![
432            Some(false), // 0
433            Some(false),
434            Some(false),
435            Some(false),
436            Some(false),
437            Some(false),
438            Some(false),
439            Some(false),
440            Some(false), // 8
441            Some(false),
442            Some(false),
443            Some(false),
444            Some(false),
445            Some(false),
446            Some(false),
447            Some(false),
448            Some(false), // 16
449            Some(false), // 17
450            Some(false), // 18
451            None,
452            Some(true),
453            Some(false),
454            None,
455        ]);
456        let comp = comp.slice(18, 3); // Some(false), None, Some(true)
457        let res = nullif(&a, &comp).unwrap();
458        let res = res.as_primitive::<Int32Type>();
459
460        let expected = Int32Array::from(vec![
461            Some(15), // False => keep it
462            Some(8),  // None => keep it
463            None,     // true => None
464        ]);
465        assert_eq!(&expected, res)
466    }
467
468    #[test]
469    fn test_nullif_boolean_offset() {
470        let a = BooleanArray::from(vec![
471            None,       // 0
472            Some(true), // 1
473            Some(false),
474            Some(true),
475            Some(true),
476        ]);
477        let a = a.slice(1, 3); // Some(true), Some(false), Some(true)
478
479        let comp = BooleanArray::from(vec![
480            Some(false), // 0
481            Some(false), // 1
482            Some(false), // 2
483            None,
484            Some(true),
485            Some(false),
486            None,
487        ]);
488        let comp = comp.slice(2, 3); // Some(false), None, Some(true)
489        let res = nullif(&a, &comp).unwrap();
490        let res = res.as_boolean();
491
492        let expected = BooleanArray::from(vec![
493            Some(true),  // False => keep it
494            Some(false), // None => keep it
495            None,        // true => None
496        ]);
497        assert_eq!(&expected, res)
498    }
499
500    struct Foo {
501        a: Option<i32>,
502        b: Option<bool>,
503        /// Whether the entry should be valid.
504        is_valid: bool,
505    }
506
507    impl Foo {
508        fn new_valid(a: i32, b: bool) -> Foo {
509            Self {
510                a: Some(a),
511                b: Some(b),
512                is_valid: true,
513            }
514        }
515
516        fn new_null() -> Foo {
517            Self {
518                a: None,
519                b: None,
520                is_valid: false,
521            }
522        }
523    }
524
525    /// Struct Array equality is a bit weird -- we need to have the *child values*
526    /// correct even if the enclosing struct indicates it is null. But we
527    /// also need the top level is_valid bits to be correct.
528    fn create_foo_struct(values: Vec<Foo>) -> StructArray {
529        let mut struct_array = StructBuilder::new(
530            Fields::from(vec![
531                Field::new("a", DataType::Int32, true),
532                Field::new("b", DataType::Boolean, true),
533            ]),
534            vec![
535                Box::new(Int32Builder::with_capacity(values.len())),
536                Box::new(BooleanBuilder::with_capacity(values.len())),
537            ],
538        );
539
540        for value in values {
541            struct_array
542                .field_builder::<Int32Builder>(0)
543                .unwrap()
544                .append_option(value.a);
545            struct_array
546                .field_builder::<BooleanBuilder>(1)
547                .unwrap()
548                .append_option(value.b);
549            struct_array.append(value.is_valid);
550        }
551
552        struct_array.finish()
553    }
554
555    #[test]
556    fn test_nullif_struct_slices() {
557        let struct_array = create_foo_struct(vec![
558            Foo::new_valid(7, true),
559            Foo::new_valid(15, false),
560            Foo::new_valid(8, true),
561            Foo::new_valid(12, false),
562            Foo::new_null(),
563            Foo::new_null(),
564            Foo::new_valid(42, true),
565        ]);
566
567        // Some({a: 15, b: false}), Some({a: 8, b: true}), Some({a: 12, b: false}),
568        // None, None
569        let struct_array = struct_array.slice(1, 5);
570        let comp = BooleanArray::from(vec![
571            Some(false), // 0
572            Some(false), // 1
573            Some(false), // 2
574            None,
575            Some(true),
576            Some(false),
577            None,
578        ]);
579        let comp = comp.slice(2, 5); // Some(false), None, Some(true), Some(false), None
580        let res = nullif(&struct_array, &comp).unwrap();
581        let res = res.as_struct();
582
583        let expected = create_foo_struct(vec![
584            // Some(false) -> keep
585            Foo::new_valid(15, false),
586            // None -> keep
587            Foo::new_valid(8, true),
588            // Some(true) -> null out. But child values are still there.
589            Foo {
590                a: Some(12),
591                b: Some(false),
592                is_valid: false,
593            },
594            // Some(false) -> keep, but was null
595            Foo::new_null(),
596            // None -> keep, but was null
597            Foo::new_null(),
598        ]);
599
600        assert_eq!(&expected, res);
601    }
602
603    #[test]
604    fn test_nullif_no_nulls() {
605        let a = Int32Array::from(vec![Some(15), Some(7), Some(8), Some(1), Some(9)]);
606        let comp = BooleanArray::from(vec![Some(false), None, Some(true), Some(false), None]);
607        let res = nullif(&a, &comp).unwrap();
608        let res = res.as_primitive::<Int32Type>();
609
610        let expected = Int32Array::from(vec![Some(15), Some(7), None, Some(1), Some(9)]);
611        assert_eq!(res, &expected);
612    }
613
614    #[test]
615    fn nullif_empty() {
616        let a = Int32Array::from(ArrayData::new_empty(&DataType::Int32));
617        let mask = BooleanArray::from(ArrayData::new_empty(&DataType::Boolean));
618        let res = nullif(&a, &mask).unwrap();
619        assert_eq!(res.as_ref(), &a);
620    }
621
622    fn test_nullif(values: &Int32Array, filter: &BooleanArray) {
623        let expected: Int32Array = values
624            .iter()
625            .zip(filter.iter())
626            .map(|(a, b)| match b {
627                Some(true) => None,
628                Some(false) | None => a,
629            })
630            .collect();
631
632        let r = nullif(values, filter).unwrap();
633        let r_data = r.to_data();
634        r_data.validate().unwrap();
635
636        assert_eq!(
637            r.as_ref(),
638            &expected,
639            "expected nulls: {:#?}\n\n\
640        result nulls:   {:#?}\n\n\\
641        expected values: {:#?}\n\n\
642        result values:   {:#?}",
643            expected.nulls(),
644            r.nulls(),
645            expected.values(),
646            r.as_primitive::<Int32Type>().values()
647        );
648        validate_nulls(expected.nulls());
649        validate_nulls(r.nulls());
650    }
651
652    /// Ensures that the null count matches the actual number of nulls.
653    fn validate_nulls(nulls: Option<&NullBuffer>) {
654        let Some(nulls) = nulls else {
655            return;
656        };
657        let mut actual_null_count = 0;
658        for i in 0..nulls.len() {
659            if nulls.is_null(i) {
660                actual_null_count += 1;
661            }
662        }
663        assert_eq!(actual_null_count, nulls.null_count());
664    }
665
666    #[test]
667    #[cfg_attr(miri, ignore)] // Takes too long
668    fn nullif_fuzz() {
669        let mut rng = StdRng::seed_from_u64(7337);
670
671        let arrays = [
672            Int32Array::from(vec![0; 1024]), // no nulls
673            (0..1024) // 50% nulls
674                .map(|_| rng.random_bool(0.5).then_some(1))
675                .collect(),
676        ];
677
678        for a in arrays {
679            let a_slices = [
680                (0, 128),
681                (0, 129),
682                (64, 64),
683                (0, 64),
684                (32, 32),
685                (0, 0),
686                (32, 0),
687                (5, 800),
688                (33, 53),
689                (77, 101),
690            ];
691            for (a_offset, a_length) in a_slices {
692                let a = a.slice(a_offset, a_length);
693
694                for i in 1..65 {
695                    let b_start_offset = rng.random_range(0..i);
696                    let b_end_offset = rng.random_range(0..i);
697
698                    // b with 50% nulls
699                    let b: BooleanArray = (0..a_length + b_start_offset + b_end_offset)
700                        .map(|_| rng.random_bool(0.5).then(|| rng.random_bool(0.5)))
701                        .collect();
702                    let b_sliced = b.slice(b_start_offset, a_length);
703                    test_nullif(&a, &b_sliced);
704
705                    // b with no nulls (and no null buffer)
706                    let b = remove_null_buffer(&b);
707                    let b_sliced = b.slice(b_start_offset, a_length);
708                    test_nullif(&a, &b_sliced);
709
710                    // b with no nulls (but with a null buffer)
711                    let b = remove_null_values(&b);
712                    let b_sliced = b.slice(b_start_offset, a_length);
713                    test_nullif(&a, &b_sliced);
714                }
715            }
716        }
717    }
718
719    /// Returns a new BooleanArray with no null buffer
720    fn remove_null_buffer(array: &BooleanArray) -> BooleanArray {
721        BooleanArray::new(array.values().clone(), None)
722    }
723
724    /// Returns a new BooleanArray with a null buffer where all values are valid
725    fn remove_null_values(array: &BooleanArray) -> BooleanArray {
726        let len = array.len();
727        let new_nulls = NullBuffer::from_iter(std::iter::repeat_n(true, len));
728        BooleanArray::new(array.values().clone(), Some(new_nulls))
729    }
730}
arrow_select/nullif.rs

arrow_select/
nullif.rs
