polars_arrow/compute/
utils.rs

1use std::borrow::Borrow;
2use std::ops::{BitAnd, BitOr};
3
4use polars_error::{PolarsResult, polars_ensure};
5
6use crate::array::Array;
7use crate::bitmap::{Bitmap, and_not, push_bitchunk, ternary};
8
9pub fn combine_validities_and3(
10    opt1: Option<&Bitmap>,
11    opt2: Option<&Bitmap>,
12    opt3: Option<&Bitmap>,
13) -> Option<Bitmap> {
14    match (opt1, opt2, opt3) {
15        (Some(a), Some(b), Some(c)) => Some(ternary(a, b, c, |x, y, z| x & y & z)),
16        (Some(a), Some(b), None) => Some(a.bitand(b)),
17        (Some(a), None, Some(c)) => Some(a.bitand(c)),
18        (None, Some(b), Some(c)) => Some(b.bitand(c)),
19        (Some(a), None, None) => Some(a.clone()),
20        (None, Some(b), None) => Some(b.clone()),
21        (None, None, Some(c)) => Some(c.clone()),
22        (None, None, None) => None,
23    }
24}
25
26pub fn combine_validities_and(opt_l: Option<&Bitmap>, opt_r: Option<&Bitmap>) -> Option<Bitmap> {
27    match (opt_l, opt_r) {
28        (Some(l), Some(r)) => Some(l.bitand(r)),
29        (None, Some(r)) => Some(r.clone()),
30        (Some(l), None) => Some(l.clone()),
31        (None, None) => None,
32    }
33}
34
35pub fn combine_validities_or(opt_l: Option<&Bitmap>, opt_r: Option<&Bitmap>) -> Option<Bitmap> {
36    match (opt_l, opt_r) {
37        (Some(l), Some(r)) => Some(l.bitor(r)),
38        _ => None,
39    }
40}
41
42pub fn combine_validities_and_not(
43    opt_l: Option<&Bitmap>,
44    opt_r: Option<&Bitmap>,
45) -> Option<Bitmap> {
46    match (opt_l, opt_r) {
47        (Some(l), Some(r)) => Some(and_not(l, r)),
48        (None, Some(r)) => Some(!r),
49        (Some(l), None) => Some(l.clone()),
50        (None, None) => None,
51    }
52}
53
54pub fn combine_validities_and_many<B: Borrow<Bitmap>>(bitmaps: &[Option<B>]) -> Option<Bitmap> {
55    let mut bitmaps = bitmaps
56        .iter()
57        .flatten()
58        .map(|b| b.borrow())
59        .collect::<Vec<_>>();
60
61    match bitmaps.len() {
62        0 => None,
63        1 => bitmaps.pop().cloned(),
64        2 => combine_validities_and(bitmaps.pop(), bitmaps.pop()),
65        3 => combine_validities_and3(bitmaps.pop(), bitmaps.pop(), bitmaps.pop()),
66        _ => {
67            let mut iterators = bitmaps
68                .iter()
69                .map(|v| v.fast_iter_u64())
70                .collect::<Vec<_>>();
71            let mut buffer = Vec::with_capacity(iterators.first().unwrap().size_hint().0 + 2);
72
73            'rows: loop {
74                // All ones so as identity for & operation
75                let mut out = u64::MAX;
76                for iter in iterators.iter_mut() {
77                    if let Some(v) = iter.next() {
78                        out &= v
79                    } else {
80                        break 'rows;
81                    }
82                }
83                push_bitchunk(&mut buffer, out);
84            }
85
86            // All ones so as identity for & operation
87            let mut out = [u64::MAX, u64::MAX];
88            let mut len = 0;
89            for iter in iterators.into_iter() {
90                let (rem, rem_len) = iter.remainder();
91                len = rem_len;
92
93                for (out, rem) in out.iter_mut().zip(rem) {
94                    *out &= rem;
95                }
96            }
97            push_bitchunk(&mut buffer, out[0]);
98            if len > 64 {
99                push_bitchunk(&mut buffer, out[1]);
100            }
101            let bitmap = Bitmap::from_u8_vec(buffer, bitmaps[0].len());
102            if bitmap.unset_bits() == bitmap.len() {
103                None
104            } else {
105                Some(bitmap)
106            }
107        },
108    }
109}
110
111// Errors iff the two arrays have a different length.
112#[inline]
113pub fn check_same_len(lhs: &dyn Array, rhs: &dyn Array) -> PolarsResult<()> {
114    polars_ensure!(lhs.len() == rhs.len(), ComputeError:
115            "arrays must have the same length"
116    );
117    Ok(())
118}
polars_arrow/compute/utils.rs

polars_arrow/compute/
utils.rs
