Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Fixed filter of predicate with validity #653

Merged
merged 2 commits into from
Dec 2, 2021
Merged
Show file tree
Hide file tree
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
16 changes: 14 additions & 2 deletions src/compute/filter.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
//! Contains operators to filter arrays such as [`filter`].
use crate::array::growable::{make_growable, Growable};
use crate::bitmap::{utils::SlicesIterator, Bitmap, MutableBitmap};
use crate::datatypes::DataType;
use crate::record_batch::RecordBatch;
use crate::{array::*, types::NativeType};
use crate::{buffer::MutableBuffer, error::Result};
Expand Down Expand Up @@ -110,8 +111,10 @@ pub fn build_filter(filter: &BooleanArray) -> Result<Filter> {
}

/// Filters an [Array], returning elements matching the filter (i.e. where the values are true).
/// 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.
///
/// Note that the nulls of `filter` are interpreted as `false` will lead to these elements being
/// masked out.
///
/// # Example
/// ```rust
/// # use arrow2::array::{Int32Array, PrimitiveArray, BooleanArray};
Expand All @@ -127,6 +130,15 @@ pub fn build_filter(filter: &BooleanArray) -> Result<Filter> {
/// # }
/// ```
pub fn filter(array: &dyn Array, filter: &BooleanArray) -> Result<Box<dyn Array>> {
// The validities may be masking out `true` bits, making the filter operation
// based on the values incorrect
if let Some(validities) = filter.validity() {
let values = filter.values();
let new_values = values & validities;
let filter = BooleanArray::from_data(DataType::Boolean, new_values, None);
return crate::compute::filter::filter(array, &filter);
}

use crate::datatypes::PhysicalType::*;
match array.data_type().to_physical_type() {
Primitive(primitive) => with_match_primitive_type!(primitive, |$T| {
Expand Down
16 changes: 16 additions & 0 deletions tests/it/compute/filter.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use arrow2::array::*;
use arrow2::bitmap::Bitmap;
use arrow2::compute::filter::*;

#[test]
Expand Down Expand Up @@ -112,6 +113,21 @@ fn binary_array_with_null() {
assert!(d.is_null(1));
}

#[test]
fn masked_true_values() {
let a = Int32Array::from_slice(&[1, 2, 3]);
let b = BooleanArray::from_slice(&[true, false, true]);
let validity = Bitmap::from(&[true, false, false]);

let b = b.with_validity(Some(validity));

let c = filter(&a, &b).unwrap();

let expected = Int32Array::from_slice(&[1]);

assert_eq!(expected, c.as_ref());
}

/*
#[test]
fn dictionary_array() {
Expand Down