|
| 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 | +pub fn set_nulls<T: ArrowNumericType + Send>( |
| 26 | + array: PrimitiveArray<T>, |
| 27 | + nulls: Option<NullBuffer>, |
| 28 | +) -> PrimitiveArray<T> { |
| 29 | + let (dt, values, _old_nulls) = array.into_parts(); |
| 30 | + PrimitiveArray::<T>::new(values, nulls).with_data_type(dt) |
| 31 | +} |
| 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 | +fn filter_to_nulls(filter: &BooleanArray) -> Option<NullBuffer> { |
| 39 | + let (filter_bools, filter_nulls) = filter.clone().into_parts(); |
| 40 | + let filter_bools = NullBuffer::from(filter_bools); |
| 41 | + NullBuffer::union(Some(&filter_bools), filter_nulls.as_ref()) |
| 42 | +} |
| 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 | +pub fn filtered_null_mask( |
| 88 | + opt_filter: Option<&BooleanArray>, |
| 89 | + input: &dyn Array, |
| 90 | +) -> Option<NullBuffer> { |
| 91 | + let opt_filter = opt_filter.and_then(filter_to_nulls); |
| 92 | + NullBuffer::union(opt_filter.as_ref(), input.nulls()) |
| 93 | +} |
0 commit comments