1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

use arrow::bitmap::Bitmap;

use super::*;
use crate::chunked_array::metadata::MetadataProperties;

impl<T: PolarsDataType> ChunkedArray<T> {
    /// Get a mask of the null values.
    pub fn is_null(&self) -> BooleanChunked {
        if !self.has_nulls() {
            return BooleanChunked::full(self.name().clone(), false, self.len());
        }
        // dispatch to non-generic function
        is_null(self.name().clone(), &self.chunks)
    }

    /// Get a mask of the valid values.
    pub fn is_not_null(&self) -> BooleanChunked {
        if self.null_count() == 0 {
            return BooleanChunked::full(self.name().clone(), true, self.len());
        }
        // dispatch to non-generic function
        is_not_null(self.name().clone(), &self.chunks)
    }

    pub(crate) fn coalesce_nulls(&self, other: &[ArrayRef]) -> Self {
        let chunks = coalesce_nulls(&self.chunks, other);
        let mut ca = unsafe { self.copy_with_chunks(chunks) };
        ca.copy_metadata(self, MetadataProperties::SORTED);
        ca
    }
}

pub fn is_not_null(name: PlSmallStr, chunks: &[ArrayRef]) -> BooleanChunked {
    let chunks = chunks.iter().map(|arr| {
        let bitmap = arr
            .validity()
            .cloned()
            .unwrap_or_else(|| !(&Bitmap::new_zeroed(arr.len())));
        BooleanArray::from_data_default(bitmap, None)
    });
    BooleanChunked::from_chunk_iter(name, chunks)
}

pub fn is_null(name: PlSmallStr, chunks: &[ArrayRef]) -> BooleanChunked {
    let chunks = chunks.iter().map(|arr| {
        let bitmap = arr
            .validity()
            .map(|bitmap| !bitmap)
            .unwrap_or_else(|| Bitmap::new_zeroed(arr.len()));
        BooleanArray::from_data_default(bitmap, None)
    });
    BooleanChunked::from_chunk_iter(name, chunks)
}

pub fn replace_non_null(name: PlSmallStr, chunks: &[ArrayRef], default: bool) -> BooleanChunked {
    BooleanChunked::from_chunk_iter(
        name,
        chunks.iter().map(|el| {
            BooleanArray::from_data_default(
                Bitmap::new_with_value(default, el.len()),
                el.validity().cloned(),
            )
        }),
    )
}

pub(crate) fn coalesce_nulls(chunks: &[ArrayRef], other: &[ArrayRef]) -> Vec<ArrayRef> {
    assert_eq!(chunks.len(), other.len());
    chunks
        .iter()
        .zip(other)
        .map(|(a, b)| {
            assert_eq!(a.len(), b.len());
            let validity = match (a.validity(), b.validity()) {
                (None, Some(b)) => Some(b.clone()),
                (Some(a), Some(b)) => Some(a & b),
                (Some(a), None) => Some(a.clone()),
                (None, None) => None,
            };

            a.with_validity(validity)
        })
        .collect()
}