polars_core/chunked_array/ops/
filter.rs

1use polars_compute::filter::filter as filter_fn;
2
3#[cfg(feature = "object")]
4use crate::chunked_array::object::builder::ObjectChunkedBuilder;
5use crate::prelude::*;
6
7macro_rules! check_filter_len {
8    ($self:expr, $filter:expr) => {{
9        polars_ensure!(
10            $self.len() == $filter.len(),
11            ShapeMismatch: "filter's length: {} differs from that of the series: {}",
12            $filter.len(), $self.len()
13        )
14    }};
15}
16
17impl<T> ChunkFilter<T> for ChunkedArray<T>
18where
19    T: PolarsDataType<IsObject = FalseT>,
20{
21    fn filter(&self, filter: &BooleanChunked) -> PolarsResult<ChunkedArray<T>> {
22        // Broadcast.
23        if filter.len() == 1 {
24            return match filter.get(0) {
25                Some(true) => Ok(self.clone()),
26                _ => Ok(self.clear()),
27            };
28        }
29        check_filter_len!(self, filter);
30        Ok(unsafe {
31            arity::binary_unchecked_same_type(
32                self,
33                filter,
34                |left, mask| filter_fn(left, mask),
35                true,
36                true,
37            )
38        })
39    }
40}
41
42#[cfg(feature = "object")]
43impl<T> ChunkFilter<ObjectType<T>> for ObjectChunked<T>
44where
45    T: PolarsObject,
46{
47    fn filter(&self, filter: &BooleanChunked) -> PolarsResult<ChunkedArray<ObjectType<T>>>
48    where
49        Self: Sized,
50    {
51        // Broadcast.
52        if filter.len() == 1 {
53            return match filter.get(0) {
54                Some(true) => Ok(self.clone()),
55                _ => Ok(ObjectChunked::new_empty(self.name().clone())),
56            };
57        }
58        check_filter_len!(self, filter);
59        let chunks = self.downcast_iter().collect::<Vec<_>>();
60        let mut builder = ObjectChunkedBuilder::<T>::new(self.name().clone(), self.len());
61        for (idx, mask) in filter.into_iter().enumerate() {
62            if mask.unwrap_or(false) {
63                let (chunk_idx, idx) = self.index_to_chunked_index(idx);
64                unsafe {
65                    let arr = chunks.get_unchecked(chunk_idx);
66                    match arr.is_null(idx) {
67                        true => builder.append_null(),
68                        false => {
69                            let v = arr.value(idx);
70                            builder.append_value(v.clone())
71                        },
72                    }
73                }
74            }
75        }
76        Ok(builder.finish())
77    }
78}
polars_core/chunked_array/ops/filter.rs

polars_core/chunked_array/ops/
filter.rs
