polars_arrow/bitmap/
aligned.rs

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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
use std::iter::Copied;
use std::slice::Iter;

use crate::bitmap::utils::BitChunk;

fn load_chunk_le<T: BitChunk>(src: &[u8]) -> T {
    if let Ok(chunk) = src.try_into() {
        return T::from_le_bytes(chunk);
    }

    let mut chunk = T::Bytes::default();
    let len = src.len().min(chunk.as_ref().len());
    chunk.as_mut()[..len].copy_from_slice(&src[..len]);
    T::from_le_bytes(chunk)
}

/// Represents a bitmap split in three portions, a prefix, a suffix and an
/// aligned bulk section in the middle.
#[derive(Default, Clone, Debug)]
pub struct AlignedBitmapSlice<'a, T: BitChunk> {
    prefix: T,
    prefix_len: u32,
    bulk: &'a [T],
    suffix: T,
    suffix_len: u32,
}

impl<'a, T: BitChunk> AlignedBitmapSlice<'a, T> {
    #[inline(always)]
    pub fn prefix(&self) -> T {
        self.prefix
    }

    #[inline(always)]
    pub fn bulk_iter(&self) -> Copied<Iter<'a, T>> {
        self.bulk.iter().copied()
    }

    #[inline(always)]
    pub fn bulk(&self) -> &'a [T] {
        self.bulk
    }

    #[inline(always)]
    pub fn suffix(&self) -> T {
        self.suffix
    }

    /// The length (in bits) of the portion of the bitmap found in prefix.
    #[inline(always)]
    pub fn prefix_bitlen(&self) -> usize {
        self.prefix_len as usize
    }

    /// The length (in bits) of the portion of the bitmap found in bulk.
    #[inline(always)]
    pub fn bulk_bitlen(&self) -> usize {
        8 * size_of::<T>() * self.bulk.len()
    }

    /// The length (in bits) of the portion of the bitmap found in suffix.
    #[inline(always)]
    pub fn suffix_bitlen(&self) -> usize {
        self.suffix_len as usize
    }

    pub fn new(mut bytes: &'a [u8], mut offset: usize, len: usize) -> Self {
        if len == 0 {
            return Self::default();
        }

        assert!(bytes.len() * 8 >= offset + len);

        // Strip off useless bytes from start.
        let start_byte_idx = offset / 8;
        bytes = &bytes[start_byte_idx..];
        offset %= 8;

        // Fast-path: fits entirely in one chunk.
        let chunk_len = size_of::<T>();
        let chunk_len_bits = 8 * chunk_len;
        if offset + len <= chunk_len_bits {
            let mut prefix = load_chunk_le::<T>(bytes) >> offset;
            if len < chunk_len_bits {
                prefix &= (T::one() << len) - T::one();
            }
            return Self {
                prefix,
                prefix_len: len as u32,
                ..Self::default()
            };
        }

        // Find how many bytes from the start our aligned section would start.
        let mut align_offset = bytes.as_ptr().align_offset(chunk_len);
        let mut align_offset_bits = 8 * align_offset;

        // Oops, the original pointer was already aligned, but our offset means
        // we can't start there, start one chunk later.
        if offset > align_offset_bits {
            align_offset_bits += chunk_len_bits;
            align_offset += chunk_len;
        }

        // Calculate based on this the lengths of our sections (in bits).
        let prefix_len = (align_offset_bits - offset).min(len);
        let rest_len = len - prefix_len;
        let suffix_len = rest_len % chunk_len_bits;
        let bulk_len = rest_len - suffix_len;
        debug_assert!(prefix_len < chunk_len_bits);
        debug_assert!(bulk_len % chunk_len_bits == 0);
        debug_assert!(suffix_len < chunk_len_bits);

        // Now we just have to load.
        let (prefix_bytes, rest_bytes) = bytes.split_at(align_offset);
        let (bulk_bytes, suffix_bytes) = rest_bytes.split_at(bulk_len / 8);
        let mut prefix = load_chunk_le::<T>(prefix_bytes) >> offset;
        let mut suffix = load_chunk_le::<T>(suffix_bytes);
        prefix &= (T::one() << prefix_len) - T::one();
        suffix &= (T::one() << suffix_len) - T::one();
        Self {
            prefix,
            bulk: bytemuck::cast_slice(bulk_bytes),
            suffix,
            prefix_len: prefix_len as u32,
            suffix_len: suffix_len as u32,
        }
    }
}