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