Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 40 additions & 3 deletions arrow-select/src/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ use arrow_buffer::{bit_util, ArrowNativeType, BooleanBuffer, NullBuffer, RunEndB
use arrow_buffer::{Buffer, MutableBuffer};
use arrow_data::bit_iterator::{BitIndexIterator, BitSliceIterator};
use arrow_data::transform::MutableArrayData;
use arrow_data::ArrayDataBuilder;
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::*;

/// If the filter selects more than this fraction of rows, use
Expand Down Expand Up @@ -112,6 +112,43 @@ fn filter_count(filter: &BooleanArray) -> usize {
filter.values().count_set_bits()
}

/// Function that can filter arbitrary arrays
///
/// Deprecated: Use [`FilterPredicate`] instead
#[deprecated]
pub type Filter<'a> = Box<dyn Fn(&ArrayData) -> ArrayData + 'a>;

/// Returns a prepared function optimized to filter multiple arrays.
///
/// Creating this function requires time, but using it is faster than [filter] when the
/// same filter needs to be applied to multiple arrays (e.g. a multi-column `RecordBatch`).
/// WARNING: the nulls of `filter` are ignored and the value on its slot is considered.
/// Therefore, it is considered undefined behavior to pass `filter` with null values.
///
/// Deprecated: Use [`FilterBuilder`] instead
#[deprecated]
#[allow(deprecated)]
pub fn build_filter(filter: &BooleanArray) -> Result<Filter, ArrowError> {
let iter = SlicesIterator::new(filter);
let filter_count = filter_count(filter);
let chunks = iter.collect::<Vec<_>>();

Ok(Box::new(move |array: &ArrayData| {
match filter_count {
// return all
len if len == array.len() => array.clone(),
0 => ArrayData::new_empty(array.data_type()),
_ => {
let mut mutable = MutableArrayData::new(vec![array], false, filter_count);
chunks
.iter()
.for_each(|(start, end)| mutable.extend(0, *start, *end));
mutable.freeze()
}
}
}))
}

/// Remove null values by do a bitmask AND operation with null bits and the boolean bits.
pub fn prep_null_mask_filter(filter: &BooleanArray) -> BooleanArray {
let nulls = filter.nulls().unwrap();
Expand Down Expand Up @@ -853,16 +890,16 @@ fn filter_sparse_union(

#[cfg(test)]
mod tests {
use super::*;
use arrow_array::builder::*;
use arrow_array::cast::as_run_array;
use arrow_array::types::*;
use arrow_data::ArrayData;
use rand::distr::uniform::{UniformSampler, UniformUsize};
use rand::distr::{Alphanumeric, StandardUniform};
use rand::prelude::*;
use rand::rng;

use super::*;

macro_rules! def_temporal_test {
($test:ident, $array_type: ident, $data: expr) => {
#[test]
Expand Down
Loading