parquet/arrow/buffer/
bit_util.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use arrow_buffer::bit_chunk_iterator::UnalignedBitChunk;
6use std::ops::Range;
7
8/// Counts the number of set bits in the provided range
9pub fn count_set_bits(bytes: &[u8], range: Range<usize>) -> usize {
10    let unaligned = UnalignedBitChunk::new(bytes, range.start, range.end - range.start);
11    unaligned.count_ones()
12}
13
14/// Iterates through the set bit positions in `bytes` in reverse order
15pub fn iter_set_bits_rev(bytes: &[u8]) -> impl Iterator<Item = usize> + '_ {
16    let bit_length = bytes.len() * 8;
17    let unaligned = UnalignedBitChunk::new(bytes, 0, bit_length);
18    let mut chunk_end_idx = bit_length + unaligned.lead_padding() + unaligned.trailing_padding();
19
20    let iter = unaligned
21        .prefix()
22        .into_iter()
23        .chain(unaligned.chunks().iter().copied())
24        .chain(unaligned.suffix());
25
26    iter.rev().flat_map(move |chunk| {
27        let chunk_idx = chunk_end_idx - 64;
28        chunk_end_idx = chunk_idx;
29        let mut rev_chunk = chunk.reverse_bits();
30        std::iter::from_fn(move || {
31            if rev_chunk != 0 {
32                let bit_pos = rev_chunk.trailing_zeros();
33                rev_chunk &= rev_chunk - 1;
34                return Some(chunk_idx + (63 - bit_pos as usize));
35            }
36            None
37        })
38    })
39}
40
41/// Performs big endian sign extension
42pub fn sign_extend_be<const N: usize>(b: &[u8]) -> [u8; N] {
43    assert!(b.len() <= N, "Array too large, expected less than {N}");
44    let is_negative = (b[0] & 128u8) == 128u8;
45    let mut result = if is_negative { [255u8; N] } else { [0u8; N] };
46    for (d, s) in result.iter_mut().skip(N - b.len()).zip(b) {
47        *d = *s;
48    }
49    result
50}
51
52/// Sign extends `b`, a big-endian two's complement value of `W` bytes, to `N` bytes.
53///
54/// Equivalent to [`sign_extend_be`] for a width known at compile time, which keeps the
55/// copy a constant-size move.
56#[inline]
57pub fn sign_extend_be_fixed<const W: usize, const N: usize>(b: &[u8; W]) -> [u8; N] {
58    const { assert!(W > 0 && W <= N, "invalid width for sign extension") };
59    let mut result = if b[0] & 0x80 != 0 {
60        [u8::MAX; N]
61    } else {
62        [0; N]
63    };
64    result[N - W..].copy_from_slice(b);
65    result
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use arrow_array::builder::BooleanBufferBuilder;
72    use rand::{prelude::*, rng};
73
74    #[test]
75    fn test_bit_fns() {
76        let mut rng = rng();
77        let mask_length = rng.random_range(1..1024);
78        let bools: Vec<_> = std::iter::from_fn(|| Some(rng.next_u32() & 1 == 0))
79            .take(mask_length)
80            .collect();
81
82        let mut nulls = BooleanBufferBuilder::new(mask_length);
83        bools.iter().for_each(|b| nulls.append(*b));
84
85        let actual: Vec<_> = iter_set_bits_rev(nulls.as_slice()).collect();
86        let expected: Vec<_> = bools
87            .iter()
88            .enumerate()
89            .rev()
90            .filter_map(|(x, y)| y.then_some(x))
91            .collect();
92        assert_eq!(actual, expected);
93
94        assert_eq!(iter_set_bits_rev(&[]).count(), 0);
95        assert_eq!(count_set_bits(&[], 0..0), 0);
96        assert_eq!(count_set_bits(&[0xFF], 1..1), 0);
97
98        for _ in 0..20 {
99            let start = rng.random_range(0..bools.len());
100            let end = rng.random_range(start..bools.len());
101
102            let actual = count_set_bits(nulls.as_slice(), start..end);
103            let expected = bools[start..end].iter().filter(|x| **x).count();
104
105            assert_eq!(actual, expected);
106        }
107    }
108}
parquet/arrow/buffer/bit_util.rs

parquet/arrow/buffer/
bit_util.rs
