Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-parquet/src/parquet/encoding/bitpacked/encode.rs
7887 views
1
use super::{Unpackable, Unpacked};
2
3
/// Encodes (packs) a slice of [`Unpackable`] into bitpacked bytes `packed`, using `num_bits` per value.
4
///
5
/// This function assumes that the maximum value in `unpacked` fits in `num_bits` bits
6
/// and saturates higher values.
7
///
8
/// Only the first `ceil8(unpacked.len() * num_bits)` of `packed` are populated.
9
pub fn encode<T: Unpackable>(unpacked: &[T], num_bits: usize, packed: &mut [u8]) {
10
let chunks = unpacked.chunks_exact(T::Unpacked::LENGTH);
11
12
let remainder = chunks.remainder();
13
14
let packed_size = (T::Unpacked::LENGTH * num_bits).div_ceil(8);
15
if !remainder.is_empty() {
16
let packed_chunks = packed.chunks_mut(packed_size);
17
let mut last_chunk = T::Unpacked::zero();
18
for i in 0..remainder.len() {
19
last_chunk[i] = remainder[i]
20
}
21
22
chunks
23
.chain(std::iter::once(last_chunk.as_ref()))
24
.zip(packed_chunks)
25
.for_each(|(unpacked, packed)| {
26
T::pack(&unpacked.try_into().unwrap(), num_bits, packed);
27
});
28
} else {
29
let packed_chunks = packed.chunks_exact_mut(packed_size);
30
chunks.zip(packed_chunks).for_each(|(unpacked, packed)| {
31
T::pack(&unpacked.try_into().unwrap(), num_bits, packed);
32
});
33
}
34
}
35
36
/// Encodes (packs) a potentially incomplete pack of [`Unpackable`] into bitpacked
37
/// bytes `packed`, using `num_bits` per value.
38
///
39
/// This function assumes that the maximum value in `unpacked` fits in `num_bits` bits
40
/// and saturates higher values.
41
///
42
/// Only the first `ceil8(unpacked.len() * num_bits)` of `packed` are populated.
43
#[inline]
44
pub fn encode_pack<T: Unpackable>(unpacked: &[T], num_bits: usize, packed: &mut [u8]) {
45
if unpacked.len() < T::Unpacked::LENGTH {
46
let mut complete_unpacked = T::Unpacked::zero();
47
complete_unpacked.as_mut()[..unpacked.len()].copy_from_slice(unpacked);
48
T::pack(&complete_unpacked, num_bits, packed)
49
} else {
50
T::pack(&unpacked.try_into().unwrap(), num_bits, packed)
51
}
52
}
53
54