polars_core/chunked_array/ops/
nulls.rs

1use arrow::bitmap::Bitmap;
2
3use super::*;
4use crate::chunked_array::flags::StatisticsFlags;
5
6impl<T: PolarsDataType> ChunkedArray<T> {
7    /// Get a mask of the null values.
8    pub fn is_null(&self) -> BooleanChunked {
9        if !self.has_nulls() {
10            return BooleanChunked::full(self.name().clone(), false, self.len());
11        }
12        // dispatch to non-generic function
13        is_null(self.name().clone(), &self.chunks)
14    }
15
16    /// Get a mask of the valid values.
17    pub fn is_not_null(&self) -> BooleanChunked {
18        if self.null_count() == 0 {
19            return BooleanChunked::full(self.name().clone(), true, self.len());
20        }
21        // dispatch to non-generic function
22        is_not_null(self.name().clone(), &self.chunks)
23    }
24
25    pub(crate) fn coalesce_nulls(&self, other: &[ArrayRef]) -> Self {
26        let chunks = coalesce_nulls(&self.chunks, other);
27        let mut ca = unsafe { self.copy_with_chunks(chunks) };
28        use StatisticsFlags as F;
29        ca.retain_flags_from(self, F::IS_SORTED_ANY);
30        ca
31    }
32}
33
34pub fn is_not_null(name: PlSmallStr, chunks: &[ArrayRef]) -> BooleanChunked {
35    let chunks = chunks.iter().map(|arr| {
36        let bitmap = arr
37            .validity()
38            .cloned()
39            .unwrap_or_else(|| !(&Bitmap::new_zeroed(arr.len())));
40        BooleanArray::from_data_default(bitmap, None)
41    });
42    BooleanChunked::from_chunk_iter(name, chunks)
43}
44
45pub fn is_null(name: PlSmallStr, chunks: &[ArrayRef]) -> BooleanChunked {
46    let chunks = chunks.iter().map(|arr| {
47        let bitmap = arr
48            .validity()
49            .map(|bitmap| !bitmap)
50            .unwrap_or_else(|| Bitmap::new_zeroed(arr.len()));
51        BooleanArray::from_data_default(bitmap, None)
52    });
53    BooleanChunked::from_chunk_iter(name, chunks)
54}
55
56pub fn replace_non_null(name: PlSmallStr, chunks: &[ArrayRef], default: bool) -> BooleanChunked {
57    BooleanChunked::from_chunk_iter(
58        name,
59        chunks.iter().map(|el| {
60            BooleanArray::from_data_default(
61                Bitmap::new_with_value(default, el.len()),
62                el.validity().cloned(),
63            )
64        }),
65    )
66}
67
68pub(crate) fn coalesce_nulls(chunks: &[ArrayRef], other: &[ArrayRef]) -> Vec<ArrayRef> {
69    assert_eq!(chunks.len(), other.len());
70    chunks
71        .iter()
72        .zip(other)
73        .map(|(a, b)| {
74            assert_eq!(a.len(), b.len());
75            let validity = match (a.validity(), b.validity()) {
76                (None, Some(b)) => Some(b.clone()),
77                (Some(a), Some(b)) => Some(a & b),
78                (Some(a), None) => Some(a.clone()),
79                (None, None) => None,
80            };
81
82            a.with_validity(validity)
83        })
84        .collect()
85}
polars_core/chunked_array/ops/nulls.rs

polars_core/chunked_array/ops/
nulls.rs
