Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/nulls.rs
Line
Count
Source (jump to first uncovered line)
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`], and [`filtered_null_mask`], utilities for working with nulls
19
20
use arrow::array::{Array, ArrowNumericType, BooleanArray, PrimitiveArray};
21
use arrow::buffer::NullBuffer;
22
23
/// Sets the validity mask for a `PrimitiveArray` to `nulls`
24
/// replacing any existing null mask
25
0
pub fn set_nulls<T: ArrowNumericType + Send>(
26
0
    array: PrimitiveArray<T>,
27
0
    nulls: Option<NullBuffer>,
28
0
) -> PrimitiveArray<T> {
29
0
    let (dt, values, _old_nulls) = array.into_parts();
30
0
    PrimitiveArray::<T>::new(values, nulls).with_data_type(dt)
31
0
}
32
33
/// Converts a `BooleanBuffer` representing a filter to a `NullBuffer.
34
///
35
/// The `NullBuffer` is
36
/// * `true` (representing valid) for values that were `true` in filter
37
/// * `false` (representing null) for values that were `false` or `null` in filter
38
0
fn filter_to_nulls(filter: &BooleanArray) -> Option<NullBuffer> {
39
0
    let (filter_bools, filter_nulls) = filter.clone().into_parts();
40
0
    let filter_bools = NullBuffer::from(filter_bools);
41
0
    NullBuffer::union(Some(&filter_bools), filter_nulls.as_ref())
42
0
}
43
44
/// Compute an output validity mask for an array that has been filtered
45
///
46
/// This can be used to compute nulls for the output of
47
/// [`GroupsAccumulator::convert_to_state`], which quickly applies an optional
48
/// filter to the input rows by setting any filtered rows to NULL in the output.
49
/// Subsequent applications of  aggregate functions that ignore NULLs (most of
50
/// them) will thus ignore the filtered rows as well.
51
///
52
/// # Output element is `true` (and thus output is non-null)
53
///
54
/// A `true` in the output represents non null output for all values that were *both*:
55
///
56
/// * `true` in any `opt_filter` (aka values that passed the filter)
57
///
58
/// * `non null` in `input`
59
///
60
/// # Output element is `false` (and thus output is null)
61
///
62
/// A `false` in the output represents an input that was *either*:
63
///
64
/// * `null`
65
///
66
/// * filtered (aka the value was `false` or `null` in the filter)
67
///
68
/// # Example
69
///
70
/// ```text
71
/// ┌─────┐           ┌─────┐            ┌─────┐
72
/// │true │           │NULL │            │false│
73
/// │true │    │      │true │            │true │
74
/// │true │ ───┼───   │false│  ────────▶ │false│       filtered_nulls
75
/// │false│    │      │NULL │            │false│
76
/// │false│           │true │            │false│
77
/// └─────┘           └─────┘            └─────┘
78
/// array           opt_filter           output
79
///  .nulls()
80
///
81
/// false = NULL    true  = pass          false = NULL       Meanings
82
/// true  = valid   false = filter        true  = valid
83
///                 NULL  = filter
84
/// ```
85
///
86
/// [`GroupsAccumulator::convert_to_state`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator
87
0
pub fn filtered_null_mask(
88
0
    opt_filter: Option<&BooleanArray>,
89
0
    input: &dyn Array,
90
0
) -> Option<NullBuffer> {
91
0
    let opt_filter = opt_filter.and_then(filter_to_nulls);
92
0
    NullBuffer::union(opt_filter.as_ref(), input.nulls())
93
0
}