diff --git a/crates/polars-compute/src/filter/boolean.rs b/crates/polars-compute/src/filter/boolean.rs new file mode 100644 index 000000000000..0050477ac5d7 --- /dev/null +++ b/crates/polars-compute/src/filter/boolean.rs @@ -0,0 +1,160 @@ +use super::*; + +pub(super) fn filter_bitmap_and_validity( + values: &Bitmap, + validity: Option<&Bitmap>, + mask: &Bitmap, +) -> (MutableBitmap, Option) { + if let Some(validity) = validity { + let (values, validity) = null_filter(values, validity, mask); + (values, Some(validity)) + } else { + (nonnull_filter(values, mask), None) + } +} + +/// # Safety +/// This assumes that the `mask_chunks` contains a number of set/true items equal +/// to `filter_count` +unsafe fn nonnull_filter_impl( + values: &Bitmap, + mut mask_chunks: I, + filter_count: usize, +) -> MutableBitmap +where + I: BitChunkIterExact, +{ + // TODO! we might use ChunksExact here if offset = 0. + let mut chunks = values.chunks::(); + let mut new = MutableBitmap::with_capacity(filter_count); + + chunks + .by_ref() + .zip(mask_chunks.by_ref()) + .for_each(|(chunk, mask_chunk)| { + let ones = mask_chunk.count_ones(); + let leading_ones = get_leading_ones(mask_chunk); + + if ones == leading_ones { + let size = leading_ones as usize; + unsafe { new.extend_from_slice_unchecked(chunk.to_ne_bytes().as_ref(), 0, size) }; + return; + } + + let ones_iter = BitChunkOnes::from_known_count(mask_chunk, ones as usize); + for pos in ones_iter { + new.push_unchecked(chunk & (1 << pos) > 0); + } + }); + + chunks + .remainder_iter() + .zip(mask_chunks.remainder_iter()) + .for_each(|(value, is_selected)| { + if is_selected { + unsafe { + new.push_unchecked(value); + }; + } + }); + + new +} + +/// # Safety +/// This assumes that the `mask_chunks` contains a number of set/true items equal +/// to `filter_count` +unsafe fn null_filter_impl( + values: &Bitmap, + validity: &Bitmap, + mut mask_chunks: I, + filter_count: usize, +) -> (MutableBitmap, MutableBitmap) +where + I: BitChunkIterExact, +{ + let mut chunks = values.chunks::(); + let mut validity_chunks = validity.chunks::(); + + let mut new = MutableBitmap::with_capacity(filter_count); + let mut new_validity = MutableBitmap::with_capacity(filter_count); + + chunks + .by_ref() + .zip(validity_chunks.by_ref()) + .zip(mask_chunks.by_ref()) + .for_each(|((chunk, validity_chunk), mask_chunk)| { + let ones = mask_chunk.count_ones(); + let leading_ones = get_leading_ones(mask_chunk); + + if ones == leading_ones { + let size = leading_ones as usize; + + unsafe { + new.extend_from_slice_unchecked(chunk.to_ne_bytes().as_ref(), 0, size); + + // safety: invariant offset + length <= slice.len() + new_validity.extend_from_slice_unchecked( + validity_chunk.to_ne_bytes().as_ref(), + 0, + size, + ); + } + return; + } + + // this triggers a bitcount + let ones_iter = BitChunkOnes::from_known_count(mask_chunk, ones as usize); + for pos in ones_iter { + new.push_unchecked(chunk & (1 << pos) > 0); + new_validity.push_unchecked(validity_chunk & (1 << pos) > 0); + } + }); + + chunks + .remainder_iter() + .zip(validity_chunks.remainder_iter()) + .zip(mask_chunks.remainder_iter()) + .for_each(|((value, is_valid), is_selected)| { + if is_selected { + unsafe { + new.push_unchecked(value); + new_validity.push_unchecked(is_valid); + }; + } + }); + + (new, new_validity) +} + +fn null_filter( + values: &Bitmap, + validity: &Bitmap, + mask: &Bitmap, +) -> (MutableBitmap, MutableBitmap) { + assert_eq!(values.len(), mask.len()); + let filter_count = mask.len() - mask.unset_bits(); + + let (slice, offset, length) = mask.as_slice(); + if offset == 0 { + let mask_chunks = BitChunksExact::::new(slice, length); + unsafe { null_filter_impl(values, validity, mask_chunks, filter_count) } + } else { + let mask_chunks = mask.chunks::(); + unsafe { null_filter_impl(values, validity, mask_chunks, filter_count) } + } +} + +fn nonnull_filter(values: &Bitmap, mask: &Bitmap) -> MutableBitmap { + assert_eq!(values.len(), mask.len()); + let filter_count = mask.len() - mask.unset_bits(); + + let (slice, offset, length) = mask.as_slice(); + if offset == 0 { + let mask_chunks = BitChunksExact::::new(slice, length); + unsafe { nonnull_filter_impl(values, mask_chunks, filter_count) } + } else { + let mask_chunks = mask.chunks::(); + unsafe { nonnull_filter_impl(values, mask_chunks, filter_count) } + } +} diff --git a/crates/polars-compute/src/filter/mod.rs b/crates/polars-compute/src/filter/mod.rs index 22da1d35770d..6277c7ed5b20 100644 --- a/crates/polars-compute/src/filter/mod.rs +++ b/crates/polars-compute/src/filter/mod.rs @@ -1,4 +1,7 @@ //! Contains operators to filter arrays such as [`filter`]. +mod boolean; +mod primitive; + use arrow::array::growable::make_growable; use arrow::array::*; use arrow::bitmap::utils::{BitChunkIterExact, BitChunksExact, SlicesIterator}; @@ -7,7 +10,9 @@ use arrow::datatypes::ArrowDataType; use arrow::types::simd::Simd; use arrow::types::{BitChunkOnes, NativeType}; use arrow::with_match_primitive_type_full; +use boolean::*; use polars_error::*; +use primitive::*; /// Function that can filter arbitrary arrays pub type Filter<'a> = Box Box + 'a + Send + Sync>; @@ -21,188 +26,6 @@ fn get_leading_ones(chunk: u64) -> u32 { } } -/// # Safety -/// This assumes that the `mask_chunks` contains a number of set/true items equal -/// to `filter_count` -unsafe fn nonnull_filter_impl(values: &[T], mut mask_chunks: I, filter_count: usize) -> Vec -where - T: NativeType, - I: BitChunkIterExact, -{ - let mut chunks = values.chunks_exact(64); - let mut new = Vec::::with_capacity(filter_count); - let mut dst = new.as_mut_ptr(); - - chunks - .by_ref() - .zip(mask_chunks.by_ref()) - .for_each(|(chunk, mask_chunk)| { - let ones = mask_chunk.count_ones(); - let leading_ones = get_leading_ones(mask_chunk); - - if ones == leading_ones { - let size = leading_ones as usize; - unsafe { - std::ptr::copy(chunk.as_ptr(), dst, size); - dst = dst.add(size); - } - return; - } - - let ones_iter = BitChunkOnes::from_known_count(mask_chunk, ones as usize); - for pos in ones_iter { - dst.write(*chunk.get_unchecked(pos)); - dst = dst.add(1); - } - }); - - chunks - .remainder() - .iter() - .zip(mask_chunks.remainder_iter()) - .for_each(|(value, b)| { - if b { - unsafe { - dst.write(*value); - dst = dst.add(1); - }; - } - }); - - unsafe { new.set_len(filter_count) }; - new -} - -/// # Safety -/// This assumes that the `mask_chunks` contains a number of set/true items equal -/// to `filter_count` -unsafe fn null_filter_impl( - values: &[T], - validity: &Bitmap, - mut mask_chunks: I, - filter_count: usize, -) -> (Vec, MutableBitmap) -where - T: NativeType, - I: BitChunkIterExact, -{ - let mut chunks = values.chunks_exact(64); - - let mut validity_chunks = validity.chunks::(); - - let mut new = Vec::::with_capacity(filter_count); - let mut dst = new.as_mut_ptr(); - let mut new_validity = MutableBitmap::with_capacity(filter_count); - - chunks - .by_ref() - .zip(validity_chunks.by_ref()) - .zip(mask_chunks.by_ref()) - .for_each(|((chunk, validity_chunk), mask_chunk)| { - let ones = mask_chunk.count_ones(); - let leading_ones = get_leading_ones(mask_chunk); - - if ones == leading_ones { - let size = leading_ones as usize; - unsafe { - std::ptr::copy(chunk.as_ptr(), dst, size); - dst = dst.add(size); - - // safety: invariant offset + length <= slice.len() - new_validity.extend_from_slice_unchecked( - validity_chunk.to_ne_bytes().as_ref(), - 0, - size, - ); - } - return; - } - - // this triggers a bitcount - let ones_iter = BitChunkOnes::from_known_count(mask_chunk, ones as usize); - for pos in ones_iter { - dst.write(*chunk.get_unchecked(pos)); - dst = dst.add(1); - new_validity.push_unchecked(validity_chunk & (1 << pos) > 0); - } - }); - - chunks - .remainder() - .iter() - .zip(validity_chunks.remainder_iter()) - .zip(mask_chunks.remainder_iter()) - .for_each(|((value, is_valid), is_selected)| { - if is_selected { - unsafe { - dst.write(*value); - dst = dst.add(1); - new_validity.push_unchecked(is_valid); - }; - } - }); - - unsafe { new.set_len(filter_count) }; - (new, new_validity) -} - -fn null_filter( - values: &[T], - validity: &Bitmap, - mask: &Bitmap, -) -> (Vec, MutableBitmap) { - assert_eq!(values.len(), mask.len()); - let filter_count = mask.len() - mask.unset_bits(); - - let (slice, offset, length) = mask.as_slice(); - if offset == 0 { - let mask_chunks = BitChunksExact::::new(slice, length); - unsafe { null_filter_impl(values, validity, mask_chunks, filter_count) } - } else { - let mask_chunks = mask.chunks::(); - unsafe { null_filter_impl(values, validity, mask_chunks, filter_count) } - } -} - -fn nonnull_filter(values: &[T], mask: &Bitmap) -> Vec { - assert_eq!(values.len(), mask.len()); - let filter_count = mask.len() - mask.unset_bits(); - - let (slice, offset, length) = mask.as_slice(); - if offset == 0 { - let mask_chunks = BitChunksExact::::new(slice, length); - unsafe { nonnull_filter_impl(values, mask_chunks, filter_count) } - } else { - let mask_chunks = mask.chunks::(); - unsafe { nonnull_filter_impl(values, mask_chunks, filter_count) } - } -} - -fn filter_values_and_validity( - values: &[T], - validity: Option<&Bitmap>, - mask: &Bitmap, -) -> (Vec, Option) { - if let Some(validity) = validity { - let (values, validity) = null_filter(values, validity, mask); - (values, Some(validity)) - } else { - (nonnull_filter(values, mask), None) - } -} - -fn filter_primitive( - array: &PrimitiveArray, - mask: &Bitmap, -) -> PrimitiveArray { - assert_eq!(array.len(), mask.len()); - let (values, validity) = filter_values_and_validity(array.values(), array.validity(), mask); - let validity = validity.map(|validity| validity.freeze()); - unsafe { - PrimitiveArray::::new_unchecked(array.data_type().clone(), values.into(), validity) - } -} - pub fn filter(array: &dyn Array, mask: &BooleanArray) -> PolarsResult> { // The validities may be masking out `true` bits, making the filter operation // based on the values incorrect @@ -229,6 +52,17 @@ pub fn filter(array: &dyn Array, mask: &BooleanArray) -> PolarsResult(array, mask.values()))) }), + Boolean => { + let array = array.as_any().downcast_ref::().unwrap(); + let (values, validity) = + filter_bitmap_and_validity(array.values(), array.validity(), mask.values()); + Ok(BooleanArray::new( + array.data_type().clone(), + values.freeze(), + validity.map(|v| v.freeze()), + ) + .boxed()) + }, BinaryView => { let array = array.as_any().downcast_ref::().unwrap(); let views = array.views(); diff --git a/crates/polars-compute/src/filter/primitive.rs b/crates/polars-compute/src/filter/primitive.rs new file mode 100644 index 000000000000..336009b8a233 --- /dev/null +++ b/crates/polars-compute/src/filter/primitive.rs @@ -0,0 +1,183 @@ +use super::*; + +pub(super) fn filter_values_and_validity( + values: &[T], + validity: Option<&Bitmap>, + mask: &Bitmap, +) -> (Vec, Option) { + if let Some(validity) = validity { + let (values, validity) = null_filter(values, validity, mask); + (values, Some(validity)) + } else { + (nonnull_filter(values, mask), None) + } +} + +pub(super) fn filter_primitive( + array: &PrimitiveArray, + mask: &Bitmap, +) -> PrimitiveArray { + assert_eq!(array.len(), mask.len()); + let (values, validity) = filter_values_and_validity(array.values(), array.validity(), mask); + let validity = validity.map(|validity| validity.freeze()); + unsafe { + PrimitiveArray::::new_unchecked(array.data_type().clone(), values.into(), validity) + } +} + +/// # Safety +/// This assumes that the `mask_chunks` contains a number of set/true items equal +/// to `filter_count` +unsafe fn nonnull_filter_impl(values: &[T], mut mask_chunks: I, filter_count: usize) -> Vec +where + T: NativeType, + I: BitChunkIterExact, +{ + let mut chunks = values.chunks_exact(64); + let mut new = Vec::::with_capacity(filter_count); + let mut dst = new.as_mut_ptr(); + + chunks + .by_ref() + .zip(mask_chunks.by_ref()) + .for_each(|(chunk, mask_chunk)| { + let ones = mask_chunk.count_ones(); + let leading_ones = get_leading_ones(mask_chunk); + + if ones == leading_ones { + let size = leading_ones as usize; + unsafe { + std::ptr::copy(chunk.as_ptr(), dst, size); + dst = dst.add(size); + } + return; + } + + let ones_iter = BitChunkOnes::from_known_count(mask_chunk, ones as usize); + for pos in ones_iter { + dst.write(*chunk.get_unchecked(pos)); + dst = dst.add(1); + } + }); + + chunks + .remainder() + .iter() + .zip(mask_chunks.remainder_iter()) + .for_each(|(value, b)| { + if b { + unsafe { + dst.write(*value); + dst = dst.add(1); + }; + } + }); + + unsafe { new.set_len(filter_count) }; + new +} + +/// # Safety +/// This assumes that the `mask_chunks` contains a number of set/true items equal +/// to `filter_count` +unsafe fn null_filter_impl( + values: &[T], + validity: &Bitmap, + mut mask_chunks: I, + filter_count: usize, +) -> (Vec, MutableBitmap) +where + T: NativeType, + I: BitChunkIterExact, +{ + let mut chunks = values.chunks_exact(64); + + let mut validity_chunks = validity.chunks::(); + + let mut new = Vec::::with_capacity(filter_count); + let mut dst = new.as_mut_ptr(); + let mut new_validity = MutableBitmap::with_capacity(filter_count); + + chunks + .by_ref() + .zip(validity_chunks.by_ref()) + .zip(mask_chunks.by_ref()) + .for_each(|((chunk, validity_chunk), mask_chunk)| { + let ones = mask_chunk.count_ones(); + let leading_ones = get_leading_ones(mask_chunk); + + if ones == leading_ones { + let size = leading_ones as usize; + unsafe { + std::ptr::copy(chunk.as_ptr(), dst, size); + dst = dst.add(size); + + // safety: invariant offset + length <= slice.len() + new_validity.extend_from_slice_unchecked( + validity_chunk.to_ne_bytes().as_ref(), + 0, + size, + ); + } + return; + } + + // this triggers a bitcount + let ones_iter = BitChunkOnes::from_known_count(mask_chunk, ones as usize); + for pos in ones_iter { + dst.write(*chunk.get_unchecked(pos)); + dst = dst.add(1); + new_validity.push_unchecked(validity_chunk & (1 << pos) > 0); + } + }); + + chunks + .remainder() + .iter() + .zip(validity_chunks.remainder_iter()) + .zip(mask_chunks.remainder_iter()) + .for_each(|((value, is_valid), is_selected)| { + if is_selected { + unsafe { + dst.write(*value); + dst = dst.add(1); + new_validity.push_unchecked(is_valid); + }; + } + }); + + unsafe { new.set_len(filter_count) }; + (new, new_validity) +} + +fn null_filter( + values: &[T], + validity: &Bitmap, + mask: &Bitmap, +) -> (Vec, MutableBitmap) { + assert_eq!(values.len(), mask.len()); + let filter_count = mask.len() - mask.unset_bits(); + + let (slice, offset, length) = mask.as_slice(); + if offset == 0 { + let mask_chunks = BitChunksExact::::new(slice, length); + unsafe { null_filter_impl(values, validity, mask_chunks, filter_count) } + } else { + let mask_chunks = mask.chunks::(); + unsafe { null_filter_impl(values, validity, mask_chunks, filter_count) } + } +} + +fn nonnull_filter(values: &[T], mask: &Bitmap) -> Vec { + assert_eq!(values.len(), mask.len()); + let filter_count = mask.len() - mask.unset_bits(); + + let (slice, offset, length) = mask.as_slice(); + if offset == 0 { + let mask_chunks = BitChunksExact::::new(slice, length); + unsafe { nonnull_filter_impl(values, mask_chunks, filter_count) } + } else { + let mask_chunks = mask.chunks::(); + unsafe { nonnull_filter_impl(values, mask_chunks, filter_count) } + } +}