## Coverage Report

#### Created: 2025-11-22 06:00

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

    /Users/andrewlamb/Software/datafusion/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/nulls.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
    //! [`set_nulls`], other utilities for working with nulls
    19
    20
    use arrow::array::{
    21
        Array, ArrayRef, ArrowNumericType, AsArray, BinaryArray, BinaryViewArray,
    22
        BooleanArray, LargeBinaryArray, LargeStringArray, PrimitiveArray, StringArray,
    23
        StringViewArray, StructArray,
    24
    };
    25
    use arrow::buffer::NullBuffer;
    26
    use arrow::datatypes::DataType;
    27
    use datafusion_common::{not_impl_err, Result};
    28
    use std::sync::Arc;
    29
    30
    /// Sets the validity mask for a `PrimitiveArray` to `nulls`
    31
    /// replacing any existing null mask
    32
    ///
    33
    /// See [`set_nulls_dyn`] for a version that works with `Array`
    34
    0
    pub fn set_nulls<T: ArrowNumericType + Send>(
    35
    0
        array: PrimitiveArray<T>,
    36
    0
        nulls: Option<NullBuffer>,
    37
    0
    ) -> PrimitiveArray<T> {
    38
    0
        let (dt, values, _old_nulls) = array.into_parts();
    39
    0
        PrimitiveArray::<T>::new(values, nulls).with_data_type(dt)
    40
    0
    }
    41
    42
    /// Converts a `BooleanBuffer` representing a filter to a `NullBuffer.
    43
    ///
    44
    /// The `NullBuffer` is
    45
    /// * `true` (representing valid) for values that were `true` in filter
    46
    /// * `false` (representing null) for values that were `false` or `null` in filter
    47
    0
    fn filter_to_nulls(filter: &BooleanArray) -> Option<NullBuffer> {
    48
    0
        let (filter_bools, filter_nulls) = filter.clone().into_parts();
    49
    0
        let filter_bools = NullBuffer::from(filter_bools);
    50
    0
        NullBuffer::union(Some(&filter_bools), filter_nulls.as_ref())
    51
    0
    }
    52
    53
    /// Compute an output validity mask for an array that has been filtered
    54
    ///
    55
    /// This can be used to compute nulls for the output of
    56
    /// [`GroupsAccumulator::convert_to_state`], which quickly applies an optional
    57
    /// filter to the input rows by setting any filtered rows to NULL in the output.
    58
    /// Subsequent applications of  aggregate functions that ignore NULLs (most of
    59
    /// them) will thus ignore the filtered rows as well.
    60
    ///
    61
    /// # Output element is `true` (and thus output is non-null)
    62
    ///
    63
    /// A `true` in the output represents non null output for all values that were *both*:
    64
    ///
    65
    /// * `true` in any `opt_filter` (aka values that passed the filter)
    66
    ///
    67
    /// * `non null` in `input`
    68
    ///
    69
    /// # Output element is `false` (and thus output is null)
    70
    ///
    71
    /// A `false` in the output represents an input that was *either*:
    72
    ///
    73
    /// * `null`
    74
    ///
    75
    /// * filtered (aka the value was `false` or `null` in the filter)
    76
    ///
    77
    /// # Example
    78
    ///
    79
    /// ```text
    80
    /// ┌─────┐           ┌─────┐            ┌─────┐
    81
    /// │true │           │NULL │            │false│
    82
    /// │true │    │      │true │            │true │
    83
    /// │true │ ───┼───   │false│  ────────▶ │false│       filtered_nulls
    84
    /// │false│    │      │NULL │            │false│
    85
    /// │false│           │true │            │false│
    86
    /// └─────┘           └─────┘            └─────┘
    87
    /// array           opt_filter           output
    88
    ///  .nulls()
    89
    ///
    90
    /// false = NULL    true  = pass          false = NULL       Meanings
    91
    /// true  = valid   false = filter        true  = valid
    92
    ///                 NULL  = filter
    93
    /// ```
    94
    ///
    95
    /// [`GroupsAccumulator::convert_to_state`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator
    96
    0
    pub fn filtered_null_mask(
    97
    0
        opt_filter: Option<&BooleanArray>,
    98
    0
        input: &dyn Array,
    99
    0
    ) -> Option<NullBuffer> {
    100
    0
        let opt_filter = opt_filter.and_then(filter_to_nulls);
    101
    0
        NullBuffer::union(opt_filter.as_ref(), input.nulls())
    102
    0
    }
    103
    104
    /// Applies optional filter to input, returning a new array of the same type
    105
    /// with the same data, but with any values that were filtered out set to null
    106
    0
    pub fn apply_filter_as_nulls(
    107
    0
        input: &dyn Array,
    108
    0
        opt_filter: Option<&BooleanArray>,
    109
    0
    ) -> Result<ArrayRef> {
    110
    0
        let nulls = filtered_null_mask(opt_filter, input);
    111
    0
        set_nulls_dyn(input, nulls)
    112
    0
    }
    113
    114
    /// Replaces the nulls in the input array with the given `NullBuffer`
    115
    ///
    116
    /// TODO: replace when upstreamed in arrow-rs: <https://github.com/apache/arrow-rs/issues/6528>
    117
    0
    pub fn set_nulls_dyn(input: &dyn Array, nulls: Option<NullBuffer>) -> Result<ArrayRef> {
    118
    0
        if let Some(nulls) = nulls.as_ref() {
    119
    0
            assert_eq!(nulls.len(), input.len());
    120
    0
        }
    121
    122
    0
        let output: ArrayRef = match input.data_type() {
    123
            DataType::Utf8 => {
    124
    0
                let input = input.as_string::<i32>();
    125
                // safety: values / offsets came from a valid string array, so are valid utf8
    126
                // and we checked nulls has the same length as values
    127
    0
                unsafe {
    128
    0
                    Arc::new(StringArray::new_unchecked(
    129
    0
                        input.offsets().clone(),
    130
    0
                        input.values().clone(),
    131
    0
                        nulls,
    132
    0
                    ))
    133
    0
                }
    134
            }
    135
            DataType::LargeUtf8 => {
    136
    0
                let input = input.as_string::<i64>();
    137
                // safety: values / offsets came from a valid string array, so are valid utf8
    138
                // and we checked nulls has the same length as values
    139
                unsafe {
    140
    0
                    Arc::new(LargeStringArray::new_unchecked(
    141
    0
                        input.offsets().clone(),
    142
    0
                        input.values().clone(),
    143
    0
                        nulls,
    144
    0
                    ))
    145
                }
    146
            }
    147
            DataType::Utf8View => {
    148
    0
                let input = input.as_string_view();
    149
                // safety: values / views came from a valid string view array, so are valid utf8
    150
                // and we checked nulls has the same length as values
    151
                unsafe {
    152
    0
                    Arc::new(StringViewArray::new_unchecked(
    153
    0
                        input.views().clone(),
    154
    0
                        input.data_buffers().to_vec(),
    155
    0
                        nulls,
    156
    0
                    ))
    157
                }
    158
            }
    159
    160
            DataType::Binary => {
    161
    0
                let input = input.as_binary::<i32>();
    162
                // safety: values / offsets came from a valid binary array
    163
                // and we checked nulls has the same length as values
    164
                unsafe {
    165
    0
                    Arc::new(BinaryArray::new_unchecked(
    166
    0
                        input.offsets().clone(),
    167
    0
                        input.values().clone(),
    168
    0
                        nulls,
    169
    0
                    ))
    170
                }
    171
            }
    172
            DataType::LargeBinary => {
    173
    0
                let input = input.as_binary::<i64>();
    174
                // safety: values / offsets came from a valid large binary array
    175
                // and we checked nulls has the same length as values
    176
                unsafe {
    177
    0
                    Arc::new(LargeBinaryArray::new_unchecked(
    178
    0
                        input.offsets().clone(),
    179
    0
                        input.values().clone(),
    180
    0
                        nulls,
    181
    0
                    ))
    182
                }
    183
            }
    184
            DataType::BinaryView => {
    185
    0
                let input = input.as_binary_view();
    186
                // safety: values / views came from a valid binary view array
    187
                // and we checked nulls has the same length as values
    188
                unsafe {
    189
    0
                    Arc::new(BinaryViewArray::new_unchecked(
    190
    0
                        input.views().clone(),
    191
    0
                        input.data_buffers().to_vec(),
    192
    0
                        nulls,
    193
    0
                    ))
    194
                }
    195
            }
    196
            DataType::Struct(_) => {
    197
    0
                let input = input.as_struct();
    198
                // safety: values / offsets came from a valid struct array
    199
                // and we checked nulls has the same length as values
    200
                unsafe {
    201
    0
                    Arc::new(StructArray::new_unchecked(
    202
    0
                        input.fields().clone(),
    203
    0
                        input.columns().to_vec(),
    204
    0
                        nulls,
    205
    0
                    ))
    206
                }
    207
            }
    208
            _ => {
    209
    0
                return not_impl_err!("Applying nulls {:?}", input.data_type());
    210
            }
    211
        };
    212
    0
        assert_eq!(input.len(), output.len());
    213
    0
        assert_eq!(input.data_type(), output.data_type());
    214
    215
    0
        Ok(output)
    216
    0
    }