datafusion_functions_aggregate_common/aggregate/groups_accumulator/nulls.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//! [`set_nulls`], other utilities for working with nulls
19
20use arrow::array::{
21 Array, ArrayRef, ArrowNumericType, AsArray, BinaryArray, BinaryViewArray,
22 BooleanArray, LargeBinaryArray, LargeStringArray, PrimitiveArray, StringArray,
23 StringViewArray, StructArray,
24};
25use arrow::buffer::NullBuffer;
26use arrow::datatypes::DataType;
27use datafusion_common::{not_impl_err, Result};
28use 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`
34pub fn set_nulls<T: ArrowNumericType + Send>(
35 array: PrimitiveArray<T>,
36 nulls: Option<NullBuffer>,
37) -> PrimitiveArray<T> {
38 let (dt, values, _old_nulls) = array.into_parts();
39 PrimitiveArray::<T>::new(values, nulls).with_data_type(dt)
40}
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
47fn filter_to_nulls(filter: &BooleanArray) -> Option<NullBuffer> {
48 let (filter_bools, filter_nulls) = filter.clone().into_parts();
49 let filter_bools = NullBuffer::from(filter_bools);
50 NullBuffer::union(Some(&filter_bools), filter_nulls.as_ref())
51}
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
96pub fn filtered_null_mask(
97 opt_filter: Option<&BooleanArray>,
98 input: &dyn Array,
99) -> Option<NullBuffer> {
100 let opt_filter = opt_filter.and_then(filter_to_nulls);
101 NullBuffer::union(opt_filter.as_ref(), input.nulls())
102}
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
106pub fn apply_filter_as_nulls(
107 input: &dyn Array,
108 opt_filter: Option<&BooleanArray>,
109) -> Result<ArrayRef> {
110 let nulls = filtered_null_mask(opt_filter, input);
111 set_nulls_dyn(input, nulls)
112}
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>
117pub fn set_nulls_dyn(input: &dyn Array, nulls: Option<NullBuffer>) -> Result<ArrayRef> {
118 if let Some(nulls) = nulls.as_ref() {
119 assert_eq!(nulls.len(), input.len());
120 }
121
122 let output: ArrayRef = match input.data_type() {
123 DataType::Utf8 => {
124 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 unsafe {
128 Arc::new(StringArray::new_unchecked(
129 input.offsets().clone(),
130 input.values().clone(),
131 nulls,
132 ))
133 }
134 }
135 DataType::LargeUtf8 => {
136 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 Arc::new(LargeStringArray::new_unchecked(
141 input.offsets().clone(),
142 input.values().clone(),
143 nulls,
144 ))
145 }
146 }
147 DataType::Utf8View => {
148 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 Arc::new(StringViewArray::new_unchecked(
153 input.views().clone(),
154 input.data_buffers().to_vec(),
155 nulls,
156 ))
157 }
158 }
159
160 DataType::Binary => {
161 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 Arc::new(BinaryArray::new_unchecked(
166 input.offsets().clone(),
167 input.values().clone(),
168 nulls,
169 ))
170 }
171 }
172 DataType::LargeBinary => {
173 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 Arc::new(LargeBinaryArray::new_unchecked(
178 input.offsets().clone(),
179 input.values().clone(),
180 nulls,
181 ))
182 }
183 }
184 DataType::BinaryView => {
185 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 Arc::new(BinaryViewArray::new_unchecked(
190 input.views().clone(),
191 input.data_buffers().to_vec(),
192 nulls,
193 ))
194 }
195 }
196 DataType::Struct(_) => {
197 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 Arc::new(StructArray::new_unchecked(
202 input.fields().clone(),
203 input.columns().to_vec(),
204 nulls,
205 ))
206 }
207 }
208 _ => {
209 return not_impl_err!("Applying nulls {:?}", input.data_type());
210 }
211 };
212 assert_eq!(input.len(), output.len());
213 assert_eq!(input.data_type(), output.data_type());
214
215 Ok(output)
216}