## Coverage Report

#### Created: 2025-11-17 14:14

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/arrow-rs/arrow-select/src/nullif.rs

    Line
    Count
    Source
    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
    20
    use arrow_array::{Array, ArrayRef, BooleanArray, make_array};
    21
    use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_unary_op_helper};
    22
    use arrow_buffer::{BooleanBuffer, NullBuffer};
    23
    use 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
    /// ```
    44
    780
    pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result<ArrayRef, ArrowError> {
    45
    780
        let left_data = left.to_data();
    46
    47
    780
        if left_data.len() != right.len() {
    48
    0
            return Err(ArrowError::ComputeError(
    49
    0
                "Cannot perform comparison operation on arrays of different length".to_string(),
    50
    0
            ));
    51
    780
        }
    52
    780
        let len = left_data.len();
    53
    54
    780
        if len == 0 || left_data.data_type() == &DataType::Null522 {
    55
    259
            return Ok(make_array(left_data));
    56
    521
        }
    57
    58
        // left=0 (null)   right=null       output bitmap=null
    59
        // left=0          right=1          output bitmap=null
    60
        // left=1 (set)    right=null       output bitmap=set   (passthrough)
    61
        // left=1          right=1 & comp=true    output bitmap=null
    62
        // left=1          right=1 & comp=false   output bitmap=set
    63
        //
    64
        // Thus: result = left null bitmap & (!right_values | !right_bitmap)
    65
        //              OR left null bitmap & !(right_values & right_bitmap)
    66
    67
        // Compute right_values & right_bitmap
    68
    521
        let right = match right.nulls() {
    69
    521
            Some(nulls) => right.values() & nulls.inner(),
    70
    0
            None => right.values().clone(),
    71
        };
    72
    73
        // Compute left null bitmap & !right
    74
    75
    521
        let (combined, null_count) = match left_data.nulls() {
    76
    259
            Some(left) => {
    77
    259
                let mut valid_count = 0;
    78
    259
                let b = bitwise_bin_op_helper(
    79
    259
                    left.buffer(),
    80
    259
                    left.offset(),
    81
    259
                    right.inner(),
    82
    259
                    right.offset(),
    83
    259
                    len,
    84
    515
                    |l, r| {
    85
    515
                        let t = l & !r;
    86
    515
                        valid_count += t.count_ones() as usize;
    87
    515
                        t
    88
    515
                    },
    89
                );
    90
    259
                (b, len - valid_count)
    91
            }
    92
            None => {
    93
    262
                let mut null_count = 0;
    94
    518
                let buffer262 = bitwise_unary_op_helper262(right.inner()262, right262.offset262(), len262, |b| {
    95
    518
                    let t = !b;
    96
    518
                    null_count += t.count_zeros() as usize;
    97
    518
                    t
    98
    518
                });
    99
    262
                (buffer, null_count)
    100
            }
    101
        };
    102
    103
    521
        let combined = BooleanBuffer::new(combined, 0, len);
    104
        // Safety:
    105
        // Counted nulls whilst computing
    106
    521
        let nulls = unsafe { NullBuffer::new_unchecked(combined, null_count) };
    107
    521
        let data = left_data.into_builder().nulls(Some(nulls));
    108
    109
        // SAFETY:
    110
        // Only altered null mask
    111
    521
        Ok(make_array(unsafe { data.build_unchecked() }))
    112
    780
    }
    113
    114
    #[cfg(test)]
    115
    mod tests {
    116
        use super::*;
    117
        use arrow_array::builder::{BooleanBuilder, Int32Builder, StructBuilder};
    118
        use arrow_array::cast::AsArray;
    119
        use arrow_array::types::Int32Type;
    120
        use arrow_array::{Int32Array, NullArray, StringArray, StructArray};
    121
        use arrow_data::ArrayData;
    122
        use arrow_schema::{Field, Fields};
    123
        use rand::{Rng, rng};
    124
    125
        #[test]
    126
    1
        fn test_nullif_int_array() {
    127
    1
            let a = Int32Array::from(vec![Some(15), None, Some(8), Some(1), Some(9)]);
    128
    1
            let comp = BooleanArray::from(vec![Some(false), None, Some(true), Some(false), None]);
    129
    1
            let res = nullif(&a, &comp).unwrap();
    130
    131
    1
            let expected = Int32Array::from(vec![
    132
    1
                Some(15),
    133
    1
                None,
    134
    1
                None, // comp true, slot 2 turned into null
    135
    1
                Some(1),
    136
                // Even though comp array / right is null, should still pass through original value
    137
                // comp true, slot 2 turned into null
    138
    1
                Some(9),
    139
            ]);
    140
    141
    1
            let res = res.as_primitive::<Int32Type>();
    142
    1
            assert_eq!(&expected, res);
    143
    1
        }
    144
    145
        #[test]
    146
    1
        fn test_nullif_null_array() {
    147
    1
            assert_eq!(
    148
    1
                nullif(&NullArray::new(0), &BooleanArray::new_null(0))
    149
    1
                    .unwrap()
    150
    1
                    .as_ref(),
    151
    1
                &NullArray::new(0)
    152
            );
    153
    154
    1
            assert_eq!(
    155
    1
                nullif(
    156
    1
                    &NullArray::new(3),
    157
    1
                    &BooleanArray::from(vec![Some(false), Some(true), None]),
    158
    1
                )
    159
    1
                .unwrap()
    160
    1
                .as_ref(),
    161
    1
                &NullArray::new(3)
    162
            );
    163
    1
        }
    164
    165
        #[test]
    166
    1
        fn test_nullif_int_array_offset() {
    167
    1
            let a = Int32Array::from(vec![None, Some(15), Some(8), Some(1), Some(9)]);
    168
    1
            let a = a.slice(1, 3); // Some(15), Some(8), Some(1)
    169
    1
            let a = a.as_any().downcast_ref::<Int32Array>().unwrap();
    170
    1
            let comp = BooleanArray::from(vec![
    171
    1
                Some(false),
    172
    1
                Some(false),
    173
    1
                Some(false),
    174
    1
                None,
    175
    1
                Some(true),
    176
    1
                Some(false),
    177
    1
                None,
    178
            ]);
    179
    1
            let comp = comp.slice(2, 3); // Some(false), None, Some(true)
    180
    1
            let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
    181
    1
            let res = nullif(a, comp).unwrap();
    182
    183
    1
            let expected = Int32Array::from(vec![
    184
    1
                Some(15), // False => keep it
    185
    1
                Some(8),  // None => keep it
    186
    1
                None,     // true => None
    187
            ]);
    188
    1
            let res = res.as_primitive::<Int32Type>();
    189
    1
            assert_eq!(&expected, res)
    190
    1
        }
    191
    192
        #[test]
    193
    1
        fn test_nullif_string() {
    194
    1
            let s = StringArray::from_iter([
    195
    1
                Some("hello"),
    196
    1
                None,
    197
    1
                Some("world"),
    198
    1
                Some("a"),
    199
    1
                Some("b"),
    200
    1
                None,
    201
    1
                None,
    202
    1
            ]);
    203
    1
            let select = BooleanArray::from_iter([
    204
    1
                Some(true),
    205
    1
                Some(true),
    206
    1
                Some(false),
    207
    1
                Some(true),
    208
    1
                Some(false),
    209
    1
                Some(false),
    210
    1
                None,
    211
    1
            ]);
    212
    213
    1
            let a = nullif(&s, &select).unwrap();
    214
    1
            let r: Vec<_> = a.as_string::<i32>().iter().collect();
    215
    1
            assert_eq!(
    216
                r,
    217
    1
                vec![None, None, Some("world"), None, Some("b"), None, None]
    218
            );
    219
    220
    1
            let s = s.slice(2, 3);
    221
    1
            let select = select.slice(1, 3);
    222
    1
            let a = nullif(&s, &select).unwrap();
    223
    1
            let r: Vec<_> = a.as_string::<i32>().iter().collect();
    224
    1
            assert_eq!(r, vec![None, Some("a"), None]);
    225
    1
        }
    226
    227
        #[test]
    228
    1
        fn test_nullif_int_large_left_offset() {
    229
    1
            let a = Int32Array::from(vec![
    230
    1
                Some(-1), // 0
    231
    1
                Some(-1),
    232
    1
                Some(-1),
    233
    1
                Some(-1),
    234
    1
                Some(-1),
    235
    1
                Some(-1),
    236
    1
                Some(-1),
    237
    1
                Some(-1),
    238
    1
                Some(-1), // 8
    239
    1
                Some(-1),
    240
    1
                Some(-1),
    241
    1
                Some(-1),
    242
    1
                Some(-1),
    243
    1
                Some(-1),
    244
    1
                Some(-1),
    245
    1
                Some(-1),
    246
    1
                None,     // 16
    247
    1
                Some(15), // 17
    248
    1
                Some(8),
    249
    1
                Some(1),
    250
    1
                Some(9),
    251
            ]);
    252
    1
            let a = a.slice(17, 3); // Some(15), Some(8), Some(1)
    253
    254
    1
            let comp = BooleanArray::from(vec![
    255
    1
                Some(false),
    256
    1
                Some(false),
    257
    1
                Some(false),
    258
    1
                None,
    259
    1
                Some(true),
    260
    1
                Some(false),
    261
    1
                None,
    262
            ]);
    263
    1
            let comp = comp.slice(2, 3); // Some(false), None, Some(true)
    264
    1
            let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
    265
    1
            let res = nullif(&a, comp).unwrap();
    266
    1
            let res = res.as_any().downcast_ref::<Int32Array>().unwrap();
    267
    268
    1
            let expected = Int32Array::from(vec![
    269
    1
                Some(15), // False => keep it
    270
    1
                Some(8),  // None => keep it
    271
    1
                None,     // true => None
    272
            ]);
    273
    1
            assert_eq!(&expected, res)
    274
    1
        }
    275
    276
        #[test]
    277
    1
        fn test_nullif_int_large_right_offset() {
    278
    1
            let a = Int32Array::from(vec![
    279
    1
                None,     // 0
    280
    1
                Some(15), // 1
    281
    1
                Some(8),
    282
    1
                Some(1),
    283
    1
                Some(9),
    284
            ]);
    285
    1
            let a = a.slice(1, 3); // Some(15), Some(8), Some(1)
    286
    287
    1
            let comp = BooleanArray::from(vec![
    288
    1
                Some(false), // 0
    289
    1
                Some(false),
    290
    1
                Some(false),
    291
    1
                Some(false),
    292
    1
                Some(false),
    293
    1
                Some(false),
    294
    1
                Some(false),
    295
    1
                Some(false),
    296
    1
                Some(false), // 8
    297
    1
                Some(false),
    298
    1
                Some(false),
    299
    1
                Some(false),
    300
    1
                Some(false),
    301
    1
                Some(false),
    302
    1
                Some(false),
    303
    1
                Some(false),
    304
    1
                Some(false), // 16
    305
    1
                Some(false), // 17
    306
    1
                Some(false), // 18
    307
    1
                None,
    308
    1
                Some(true),
    309
    1
                Some(false),
    310
    1
                None,
    311
            ]);
    312
    1
            let comp = comp.slice(18, 3); // Some(false), None, Some(true)
    313
    1
            let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
    314
    1
            let res = nullif(&a, comp).unwrap();
    315
    1
            let res = res.as_any().downcast_ref::<Int32Array>().unwrap();
    316
    317
    1
            let expected = Int32Array::from(vec![
    318
    1
                Some(15), // False => keep it
    319
    1
                Some(8),  // None => keep it
    320
    1
                None,     // true => None
    321
            ]);
    322
    1
            assert_eq!(&expected, res)
    323
    1
        }
    324
    325
        #[test]
    326
    1
        fn test_nullif_boolean_offset() {
    327
    1
            let a = BooleanArray::from(vec![
    328
    1
                None,       // 0
    329
    1
                Some(true), // 1
    330
    1
                Some(false),
    331
    1
                Some(true),
    332
    1
                Some(true),
    333
            ]);
    334
    1
            let a = a.slice(1, 3); // Some(true), Some(false), Some(true)
    335
    336
    1
            let comp = BooleanArray::from(vec![
    337
    1
                Some(false), // 0
    338
    1
                Some(false), // 1
    339
    1
                Some(false), // 2
    340
    1
                None,
    341
    1
                Some(true),
    342
    1
                Some(false),
    343
    1
                None,
    344
            ]);
    345
    1
            let comp = comp.slice(2, 3); // Some(false), None, Some(true)
    346
    1
            let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
    347
    1
            let res = nullif(&a, comp).unwrap();
    348
    1
            let res = res.as_any().downcast_ref::<BooleanArray>().unwrap();
    349
    350
    1
            let expected = BooleanArray::from(vec![
    351
    1
                Some(true),  // False => keep it
    352
    1
                Some(false), // None => keep it
    353
    1
                None,        // true => None
    354
            ]);
    355
    1
            assert_eq!(&expected, res)
    356
    1
        }
    357
    358
        struct Foo {
    359
            a: Option<i32>,
    360
            b: Option<bool>,
    361
            /// Whether the entry should be valid.
    362
            is_valid: bool,
    363
        }
    364
    365
        impl Foo {
    366
    7
            fn new_valid(a: i32, b: bool) -> Foo {
    367
    7
                Self {
    368
    7
                    a: Some(a),
    369
    7
                    b: Some(b),
    370
    7
                    is_valid: true,
    371
    7
                }
    372
    7
            }
    373
    374
    4
            fn new_null() -> Foo {
    375
    4
                Self {
    376
    4
                    a: None,
    377
    4
                    b: None,
    378
    4
                    is_valid: false,
    379
    4
                }
    380
    4
            }
    381
        }
    382
    383
        /// Struct Array equality is a bit weird -- we need to have the *child values*
    384
        /// correct even if the enclosing struct indicates it is null. But we
    385
        /// also need the top level is_valid bits to be correct.
    386
    2
        fn create_foo_struct(values: Vec<Foo>) -> StructArray {
    387
    2
            let mut struct_array = StructBuilder::new(
    388
    2
                Fields::from(vec![
    389
    2
                    Field::new("a", DataType::Int32, true),
    390
    2
                    Field::new("b", DataType::Boolean, true),
    391
                ]),
    392
    2
                vec![
    393
    2
                    Box::new(Int32Builder::with_capacity(values.len())),
    394
    2
                    Box::new(BooleanBuilder::with_capacity(values.len())),
    395
                ],
    396
            );
    397
    398
    14
            for value12 in values {
    399
    12
                struct_array
    400
    12
                    .field_builder::<Int32Builder>(0)
    401
    12
                    .unwrap()
    402
    12
                    .append_option(value.a);
    403
    12
                struct_array
    404
    12
                    .field_builder::<BooleanBuilder>(1)
    405
    12
                    .unwrap()
    406
    12
                    .append_option(value.b);
    407
    12
                struct_array.append(value.is_valid);
    408
    12
            }
    409
    410
    2
            struct_array.finish()
    411
    2
        }
    412
    413
        #[test]
    414
    1
        fn test_nullif_struct_slices() {
    415
    1
            let struct_array = create_foo_struct(vec![
    416
    1
                Foo::new_valid(7, true),
    417
    1
                Foo::new_valid(15, false),
    418
    1
                Foo::new_valid(8, true),
    419
    1
                Foo::new_valid(12, false),
    420
    1
                Foo::new_null(),
    421
    1
                Foo::new_null(),
    422
    1
                Foo::new_valid(42, true),
    423
            ]);
    424
    425
            // Some({a: 15, b: false}), Some({a: 8, b: true}), Some({a: 12, b: false}),
    426
            // None, None
    427
    1
            let struct_array = struct_array.slice(1, 5);
    428
    1
            let comp = BooleanArray::from(vec![
    429
    1
                Some(false), // 0
    430
    1
                Some(false), // 1
    431
    1
                Some(false), // 2
    432
    1
                None,
    433
    1
                Some(true),
    434
    1
                Some(false),
    435
    1
                None,
    436
            ]);
    437
    1
            let comp = comp.slice(2, 5); // Some(false), None, Some(true), Some(false), None
    438
    1
            let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
    439
    1
            let res = nullif(&struct_array, comp).unwrap();
    440
    1
            let res = res.as_any().downcast_ref::<StructArray>().unwrap();
    441
    442
    1
            let expected = create_foo_struct(vec![
    443
                // Some(false) -> keep
    444
    1
                Foo::new_valid(15, false),
    445
                // None -> keep
    446
    1
                Foo::new_valid(8, true),
    447
                // Some(true) -> null out. But child values are still there.
    448
    1
                Foo {
    449
    1
                    a: Some(12),
    450
    1
                    b: Some(false),
    451
    1
                    is_valid: false,
    452
    1
                },
    453
                // Some(false) -> keep, but was null
    454
    1
                Foo::new_null(),
    455
                // None -> keep, but was null
    456
    1
                Foo::new_null(),
    457
            ]);
    458
    459
    1
            assert_eq!(&expected, res);
    460
    1
        }
    461
    462
        #[test]
    463
    1
        fn test_nullif_no_nulls() {
    464
    1
            let a = Int32Array::from(vec![Some(15), Some(7), Some(8), Some(1), Some(9)]);
    465
    1
            let comp = BooleanArray::from(vec![Some(false), None, Some(true), Some(false), None]);
    466
    1
            let res = nullif(&a, &comp).unwrap();
    467
    1
            let res = res.as_primitive::<Int32Type>();
    468
    469
    1
            let expected = Int32Array::from(vec![Some(15), Some(7), None, Some(1), Some(9)]);
    470
    1
            assert_eq!(res, &expected);
    471
    1
        }
    472
    473
        #[test]
    474
    1
        fn nullif_empty() {
    475
    1
            let a = Int32Array::from(ArrayData::new_empty(&DataType::Int32));
    476
    1
            let mask = BooleanArray::from(ArrayData::new_empty(&DataType::Boolean));
    477
    1
            let res = nullif(&a, &mask).unwrap();
    478
    1
            assert_eq!(res.as_ref(), &a);
    479
    1
        }
    480
    481
    768
        fn test_nullif(values: &Int32Array, filter: &BooleanArray) {
    482
    768
            let expected: Int32Array = values
    483
    768
                .iter()
    484
    768
                .zip(filter.iter())
    485
    36.8k
                .map768(|(a, b)| match b {
    486
    9.28k
                    Some(true) => None,
    487
    27.5k
                    Some(false) | None => a,
    488
    36.8k
                })
    489
    768
                .collect();
    490
    491
    768
            let r = nullif(values, filter).unwrap();
    492
    768
            let r_data = r.to_data();
    493
    768
            r_data.validate().unwrap();
    494
    495
    768
            assert_eq!(r.as_ref(), &expected);
    496
    768
        }
    497
    498
        #[test]
    499
    1
        fn nullif_fuzz() {
    500
    1
            let mut rng = rng();
    501
    502
    1
            let arrays = [
    503
    1
                Int32Array::from(vec![0; 128]),
    504
    1
                (0..128)
    505
    128
                    .map1(|_| rng.random_bool(0.5).then_some(0))
    506
    1
                    .collect(),
    507
            ];
    508
    509
    3
            for a2 in arrays {
    510
    2
                let a_slices = [(0, 128), (64, 64), (0, 64), (32, 32), (0, 0), (32, 0)];
    511
    512
    14
                for (a_offset12, a_length12) in a_slices {
    513
    12
                    let a = a.slice(a_offset, a_length);
    514
    515
    780
                    for i768 in 1..65 {
    516
    768
                        let b_start_offset = rng.random_range(0..i);
    517
    768
                        let b_end_offset = rng.random_range(0..i);
    518
    519
    768
                        let b: BooleanArray = (0..a_length + b_start_offset + b_end_offset)
    520
    61.0k
                            .map768(|_| rng.random_bool(0.5).then(|| rng30.3k.random_bool30.3k(0.5)))
    521
    768
                            .collect();
    522
    768
                        let b = b.slice(b_start_offset, a_length);
    523
    524
    768
                        test_nullif(&a, &b);
    525
                    }
    526
                }
    527
            }
    528
    1
        }
    529
    }