polars_arrow/legacy/
bit_util.rs

1/// Forked from Arrow until their API stabilizes.
2///
3/// Note that the bound checks are optimized away.
4///
5use crate::bitmap::utils::{BitChunkIterExact, BitChunks};
6
7pub fn find_first_true_false_null(
8    mut bit_chunks: BitChunks<u64>,
9    mut validity_chunks: BitChunks<u64>,
10) -> (Option<usize>, Option<usize>, Option<usize>) {
11    let (mut true_index, mut false_index, mut null_index) = (None, None, None);
12    let (mut true_not_found_mask, mut false_not_found_mask, mut null_not_found_mask) =
13        (!0u64, !0u64, !0u64); // All ones while not found.
14    let mut offset: usize = 0;
15    let mut all_found = false;
16    for (truth_mask, null_mask) in (&mut bit_chunks).zip(&mut validity_chunks) {
17        let mask = null_mask & truth_mask & true_not_found_mask;
18        if mask > 0 {
19            true_index = Some(offset + mask.trailing_zeros() as usize);
20            true_not_found_mask = 0;
21        }
22        let mask = null_mask & !truth_mask & false_not_found_mask;
23        if mask > 0 {
24            false_index = Some(offset + mask.trailing_zeros() as usize);
25            false_not_found_mask = 0;
26        }
27        if !null_mask & null_not_found_mask > 0 {
28            null_index = Some(offset + null_mask.trailing_ones() as usize);
29            null_not_found_mask = 0;
30        }
31        if null_not_found_mask | true_not_found_mask | false_not_found_mask == 0 {
32            all_found = true;
33            break;
34        }
35        offset += 64;
36    }
37    if !all_found {
38        for (val, not_null) in bit_chunks
39            .remainder_iter()
40            .zip(validity_chunks.remainder_iter())
41        {
42            if true_index.is_none() && not_null && val {
43                true_index = Some(offset);
44            } else if false_index.is_none() && not_null && !val {
45                false_index = Some(offset);
46            } else if null_index.is_none() && !not_null {
47                null_index = Some(offset);
48            }
49            offset += 1;
50        }
51    }
52    (true_index, false_index, null_index)
53}
54
55pub fn find_first_true_false_no_null(
56    mut bit_chunks: BitChunks<u64>,
57) -> (Option<usize>, Option<usize>) {
58    let (mut true_index, mut false_index) = (None, None);
59    let (mut true_not_found_mask, mut false_not_found_mask) = (!0u64, !0u64); // All ones while not found.
60    let mut offset: usize = 0;
61    let mut all_found = false;
62    for truth_mask in &mut bit_chunks {
63        let mask = truth_mask & true_not_found_mask;
64        if mask > 0 {
65            true_index = Some(offset + mask.trailing_zeros() as usize);
66            true_not_found_mask = 0;
67        }
68        let mask = !truth_mask & false_not_found_mask;
69        if mask > 0 {
70            false_index = Some(offset + mask.trailing_zeros() as usize);
71            false_not_found_mask = 0;
72        }
73        if true_not_found_mask | false_not_found_mask == 0 {
74            all_found = true;
75            break;
76        }
77        offset += 64;
78    }
79    if !all_found {
80        for val in bit_chunks.remainder_iter() {
81            if true_index.is_none() && val {
82                true_index = Some(offset);
83            } else if false_index.is_none() && !val {
84                false_index = Some(offset);
85            }
86            offset += 1;
87        }
88    }
89    (true_index, false_index)
90}
polars_arrow/legacy/bit_util.rs

polars_arrow/legacy/
bit_util.rs
