/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 | } |