This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 98de410920 [Parquet] Vectorise Pack Bits within BitWriter (#10432)
98de410920 is described below
commit 98de410920d1bdd1f98a2527109c49dd5ffc528f
Author: Kosta Tarasov <[email protected]>
AuthorDate: Fri Sep 18 12:07:06 2026 -0400
[Parquet] Vectorise Pack Bits within BitWriter (#10432)
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- Closes #2257.
# Rationale for this change
Perf :rocket:
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
- Renamed the `FromBitpacked` to `BitPacking` to store the Packing APIs
- Added vectorised `pack_batch` matching `unpack_batch` kernels
- Added unit tests
- Added bench arms
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
Yes!
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
If this PR claims a performance improvement, please include evidence
such as benchmark results.
-->
# Are there any user-facing changes?
No
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
.../arrow/array_reader/byte_array_dictionary.rs | 8 +-
parquet/src/arrow/arrow_writer/byte_array.rs | 4 +-
parquet/src/data_type.rs | 4 +-
parquet/src/encodings/alp.rs | 6 +-
parquet/src/encodings/decoding.rs | 6 +-
parquet/src/encodings/encoding/alp_encoder.rs | 12 +-
parquet/src/encodings/encoding/dict_encoder.rs | 4 +-
parquet/src/encodings/encoding/mod.rs | 23 +-
parquet/src/encodings/rle.rs | 269 +++++++++++-----
parquet/src/util/bit_pack.rs | 198 ++++++++++++
parquet/src/util/bit_util.rs | 353 +++++++++++++++++++--
11 files changed, 743 insertions(+), 144 deletions(-)
diff --git a/parquet/src/arrow/array_reader/byte_array_dictionary.rs
b/parquet/src/arrow/array_reader/byte_array_dictionary.rs
index ae00eac6e3..4b4e1ad852 100644
--- a/parquet/src/arrow/array_reader/byte_array_dictionary.rs
+++ b/parquet/src/arrow/array_reader/byte_array_dictionary.rs
@@ -38,7 +38,7 @@ use crate::column::reader::decoder::ColumnValueDecoder;
use crate::encodings::rle::{MAX_RLE_DICTIONARY_BIT_WIDTH, RleDecoder};
use crate::errors::{ParquetError, Result};
use crate::schema::types::ColumnDescPtr;
-use crate::util::bit_util::FromBitpacked;
+use crate::util::bit_util::BitPacking;
/// A macro to reduce verbosity of [`make_byte_array_dictionary_reader`]
macro_rules! make_reader {
@@ -168,7 +168,7 @@ struct ByteArrayDictionaryReader<K: ArrowNativeType, V:
OffsetSizeTrait> {
impl<K, V> ByteArrayDictionaryReader<K, V>
where
- K: FromBitpacked + Ord + ArrowNativeType,
+ K: BitPacking + Ord + ArrowNativeType,
V: OffsetSizeTrait,
{
fn try_new(
@@ -209,7 +209,7 @@ where
impl<K, V> ArrayReader for ByteArrayDictionaryReader<K, V>
where
- K: FromBitpacked + Ord + ArrowNativeType,
+ K: BitPacking + Ord + ArrowNativeType,
V: OffsetSizeTrait,
{
fn as_any(&self) -> &dyn Any {
@@ -300,7 +300,7 @@ struct DictionaryDecoder<K, V> {
impl<K, V> ColumnValueDecoder for DictionaryDecoder<K, V>
where
- K: FromBitpacked + Ord + ArrowNativeType,
+ K: BitPacking + Ord + ArrowNativeType,
V: OffsetSizeTrait,
{
type Buffer = DictionaryBuffer<K, V>;
diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs
b/parquet/src/arrow/arrow_writer/byte_array.rs
index bd8ccbb36a..3a5a561218 100644
--- a/parquet/src/arrow/arrow_writer/byte_array.rs
+++ b/parquet/src/arrow/arrow_writer/byte_array.rs
@@ -401,9 +401,7 @@ impl DictEncoder {
buffer.push(self.bit_width());
let mut encoder = RleEncoder::new_from_buf(self.bit_width(), buffer);
- for index in &self.indices {
- encoder.put(*index)
- }
+ encoder.put_batch(&self.indices);
self.indices.clear();
diff --git a/parquet/src/data_type.rs b/parquet/src/data_type.rs
index 1f53505d52..3af6d426f3 100644
--- a/parquet/src/data_type.rs
+++ b/parquet/src/data_type.rs
@@ -791,9 +791,7 @@ pub(crate) mod private {
_: &mut W,
bit_writer: &mut BitWriter,
) -> Result<()> {
- for value in values {
- bit_writer.put_value(*value as u64, 1)
- }
+ bit_writer.put_batch(values, 1);
Ok(())
}
diff --git a/parquet/src/encodings/alp.rs b/parquet/src/encodings/alp.rs
index ba7d56f31f..8c51039a2a 100644
--- a/parquet/src/encodings/alp.rs
+++ b/parquet/src/encodings/alp.rs
@@ -38,7 +38,7 @@
//! `[AlpInfo][ForInfo][PackedValues][ExceptionPositions][ExceptionValues]`.
use crate::errors::{ParquetError, Result};
-use crate::util::bit_util::{FromBitpacked, FromBytes};
+use crate::util::bit_util::{BitPacking, FromBytes};
pub(crate) const ALP_HEADER_SIZE: usize = 7;
pub(crate) const ALP_COMPRESSION_MODE: u8 = 0;
@@ -246,10 +246,10 @@ impl<Exact: AlpExact> ForInfo<Exact> {
/// range exceeds the signed maximum, and unpacking needs no sign extension.
/// Signed interpretation is applied later during decimal reconstruction.
pub(crate) trait AlpExact:
- Copy + std::fmt::Debug + PartialEq + FromBitpacked + Default
+ Copy + std::fmt::Debug + PartialEq + BitPacking + Default
{
const WIDTH: usize;
- type Signed: Copy + Ord + std::fmt::Debug + Send;
+ type Signed: BitPacking + Copy + Ord + std::fmt::Debug + Send;
fn from_le_slice(slice: &[u8]) -> Self;
fn wrapping_add(self, rhs: Self) -> Self;
fn wrapping_sub(self, rhs: Self) -> Self;
diff --git a/parquet/src/encodings/decoding.rs
b/parquet/src/encodings/decoding.rs
index 9830cd9d5e..379819d55a 100644
--- a/parquet/src/encodings/decoding.rs
+++ b/parquet/src/encodings/decoding.rs
@@ -32,7 +32,7 @@ use crate::encodings::decoding::byte_stream_split_decoder::{
};
use crate::errors::{ParquetError, Result};
use crate::schema::types::ColumnDescPtr;
-use crate::util::bit_util::{self, BitReader, FromBitpacked};
+use crate::util::bit_util::{self, BitPacking, BitReader};
pub(crate) mod alp_decoder;
mod byte_stream_split_decoder;
@@ -461,7 +461,7 @@ impl<T: DataType> RleValueDecoder<T> {
impl<T: DataType> Decoder<T> for RleValueDecoder<T>
where
- T::T: FromBitpacked,
+ T::T: BitPacking,
{
#[inline]
fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> {
@@ -670,7 +670,7 @@ where
impl<T: DataType> Decoder<T> for DeltaBitPackDecoder<T>
where
- T::T: Default + FromPrimitive + FromBitpacked + WrappingAdd + Copy,
+ T::T: Default + FromPrimitive + BitPacking + WrappingAdd + Copy,
{
// # of total values is derived from encoding
#[inline]
diff --git a/parquet/src/encodings/encoding/alp_encoder.rs
b/parquet/src/encodings/encoding/alp_encoder.rs
index 95c3a55c84..9b2629d410 100644
--- a/parquet/src/encodings/encoding/alp_encoder.rs
+++ b/parquet/src/encodings/encoding/alp_encoder.rs
@@ -383,12 +383,16 @@ fn encode_vector<F: AlpFloat>(
// width means every value equals the frame of reference, so nothing is
// stored.
if bit_width > 0 {
- let mut writer = BitWriter::new_from_buf(std::mem::take(out));
- for &encoded_value in encoded.iter() {
+ // The encoded values are dead after choosing the frame and bit width.
+ // Reuse their allocation for deltas; signed BitPacking preserves the
+ // underlying two's-complement bit patterns.
+ for encoded_value in encoded.iter_mut() {
let delta =
-
F::Exact::reinterpret_from_signed(encoded_value).wrapping_sub(frame_of_reference);
- writer.put_value(delta.to_u64(), bit_width as usize);
+
F::Exact::reinterpret_from_signed(*encoded_value).wrapping_sub(frame_of_reference);
+ *encoded_value = delta.reinterpret_as_signed();
}
+ let mut writer = BitWriter::new_from_buf(std::mem::take(out));
+ writer.put_batch(encoded, bit_width as usize);
// Pads to a byte boundary, giving exactly the ceil(n * bit_width / 8)
// bytes the decoder derives from the metadata.
*out = writer.consume();
diff --git a/parquet/src/encodings/encoding/dict_encoder.rs
b/parquet/src/encodings/encoding/dict_encoder.rs
index 995ade1f51..2747b3c7d6 100644
--- a/parquet/src/encodings/encoding/dict_encoder.rs
+++ b/parquet/src/encodings/encoding/dict_encoder.rs
@@ -143,9 +143,7 @@ impl<T: DataType> DictEncoder<T> {
// Write bit width in the first byte
let mut encoder = RleEncoder::new_from_buf(self.bit_width(), buffer);
- for index in &self.indices {
- encoder.put(*index)
- }
+ encoder.put_batch(&self.indices);
self.indices.clear();
Ok(encoder.consume().into())
}
diff --git a/parquet/src/encodings/encoding/mod.rs
b/parquet/src/encodings/encoding/mod.rs
index 09ad15fd01..05117a625a 100644
--- a/parquet/src/encodings/encoding/mod.rs
+++ b/parquet/src/encodings/encoding/mod.rs
@@ -287,9 +287,13 @@ impl<T: DataType> Encoder<T> for RleValueEncoder<T> {
RleEncoder::new_from_buf(1, buffer)
});
- for value in values {
- let value = value.as_u64()?;
- rle_encoder.put(value)
+ let mut buf = [0_u64; 64];
+ for chunk in values.chunks(buf.len()) {
+ let buf = &mut buf[..chunk.len()];
+ for (b, value) in buf.iter_mut().zip(chunk) {
+ *b = value.as_u64()?;
+ }
+ rle_encoder.put_batch(buf);
}
Ok(())
}
@@ -480,12 +484,15 @@ impl<T: DataType> DeltaBitPackEncoder<T> {
let bit_width = num_required_bits(self.subtract_u64(max_delta,
min_delta)) as usize;
self.bit_writer.write_at(offset + i, bit_width as u8);
- // Encode values in current mini block using min_delta and
bit_width
- for j in 0..n {
- let packed_value =
- self.subtract_u64(self.deltas[i * self.mini_block_size +
j], min_delta);
- self.bit_writer.put_value(packed_value, bit_width);
+ // Encode values in current mini block using min_delta and
bit_width. This
+ // mini block's deltas are not read again, so they can be
rewritten in place
+ // with the values to pack
+ let start = i * self.mini_block_size;
+ for j in start..start + n {
+ self.deltas[j] = self.subtract_u64(self.deltas[j], min_delta)
as i64;
}
+ self.bit_writer
+ .put_batch(&self.deltas[start..start + n], bit_width);
// Pad the last block (n < mini_block_size)
for _ in n..self.mini_block_size {
diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs
index e4673a3c6c..2d56a30918 100644
--- a/parquet/src/encodings/rle.rs
+++ b/parquet/src/encodings/rle.rs
@@ -39,7 +39,7 @@ use std::{cmp, mem::size_of};
use bytes::Bytes;
use crate::errors::{ParquetError, Result};
-use crate::util::bit_util::{self, BitReader, BitWriter, FromBitpacked};
+use crate::util::bit_util::{self, BitPacking, BitReader, BitWriter};
/// Maximum bit width for dictionary page indices encoded with RLE /
Bit-Packing hybrid encoding.
///
@@ -65,11 +65,11 @@ pub struct RleEncoder {
// Underlying writer which holds an internal buffer.
bit_writer: BitWriter,
- // Buffered values for bit-packed runs.
- buffered_values: [u64; BIT_PACK_GROUP_SIZE],
-
- // Number of current buffered values. Must be less than
BIT_PACK_GROUP_SIZE.
- num_buffered_values: usize,
+ // Values of the current in-progress bit-packed run, including the
partially
+ // filled trailing group. Only materialized into `bit_writer` when the run
+ // closes, so the whole run can be packed in one vectorised `put_batch`
call.
+ // Bounded by `MAX_GROUPS_PER_BIT_PACKED_RUN * BIT_PACK_GROUP_SIZE` values.
+ pending_values: Vec<u64>,
// The current (also last) value that was written and the count of how many
// times in a row that value has been seen.
@@ -78,13 +78,6 @@ pub struct RleEncoder {
// The number of repetitions for `current_value`. If this gets too high
we'd
// switch to use RLE encoding.
repeat_count: usize,
-
- // Number of bit-packed values in the current run. This doesn't include
values
- // in `buffered_values`.
- bit_packed_count: usize,
-
- // The position of the indicator byte in the `bit_writer`.
- indicator_byte_pos: i64,
}
impl RleEncoder {
@@ -101,12 +94,9 @@ impl RleEncoder {
RleEncoder {
bit_width,
bit_writer,
- buffered_values: [0; BIT_PACK_GROUP_SIZE],
- num_buffered_values: 0,
+ pending_values: Vec::new(),
current_value: 0,
repeat_count: 0,
- bit_packed_count: 0,
- indicator_byte_pos: -1,
}
}
@@ -174,19 +164,80 @@ impl RleEncoder {
} else {
if self.repeat_count >= BIT_PACK_GROUP_SIZE {
// The current RLE run has ended and we've gathered enough.
Flush first.
- debug_assert_eq!(self.bit_packed_count, 0);
+ debug_assert!(self.pending_values.is_empty());
self.flush_rle_run();
}
self.repeat_count = 1;
self.current_value = value;
}
- self.buffered_values[self.num_buffered_values] = value;
- self.num_buffered_values += 1;
- if self.num_buffered_values == BIT_PACK_GROUP_SIZE {
- // Buffered values are full. Flush them.
- debug_assert_eq!(self.bit_packed_count % BIT_PACK_GROUP_SIZE, 0);
- self.flush_buffered_values();
+ self.pending_values.push(value);
+ if self
+ .pending_values
+ .len()
+ .is_multiple_of(BIT_PACK_GROUP_SIZE)
+ {
+ // A group of values is complete. Decide its encoding.
+ self.commit_group();
+ }
+ }
+
+ /// Encodes `values`, each of which must be representable with `bit_width`
bits.
+ ///
+ /// Produces output identical to calling [`Self::put`] for each value, but
avoids
+ /// the per-value run tracking wherever a whole group of values can be
examined at
+ /// once.
+ pub fn put_batch(&mut self, values: &[u64]) {
+ let mut i = 0;
+ while i < values.len() {
+ // Extend an active RLE run with its whole continuation at once
+ if self.repeat_count >= BIT_PACK_GROUP_SIZE && values[i] ==
self.current_value {
+ let run_len = values[i..]
+ .iter()
+ .take_while(|&&v| v == self.current_value)
+ .count();
+ self.repeat_count += run_len;
+ i += run_len;
+ continue;
+ }
+
+ // Whole groups bypass the per-value state machine while the
pending run is
+ // group-aligned and no RLE run is being accumulated.
`repeat_count` is
+ // always 0 here: `commit_group` resets it whenever a group
completes, and
+ // that is the only way to reach a group-aligned state
+ if self.repeat_count < BIT_PACK_GROUP_SIZE
+ && self
+ .pending_values
+ .len()
+ .is_multiple_of(BIT_PACK_GROUP_SIZE)
+ && values.len() - i >= BIT_PACK_GROUP_SIZE
+ {
+ debug_assert_eq!(self.repeat_count, 0);
+ let group = &values[i..i + BIT_PACK_GROUP_SIZE];
+ i += BIT_PACK_GROUP_SIZE;
+ if group.iter().all(|&v| v == group[0]) {
+ // The same decision `commit_group` makes for a group
holding a
+ // single repeated value, without appending it first
+ if !self.pending_values.is_empty() {
+ self.close_bit_packed_run();
+ }
+ self.current_value = group[0];
+ self.repeat_count = BIT_PACK_GROUP_SIZE;
+ } else {
+ self.pending_values.extend_from_slice(group);
+ self.current_value = group[BIT_PACK_GROUP_SIZE - 1];
+ let num_groups = self.pending_values.len() /
BIT_PACK_GROUP_SIZE;
+ if num_groups + 1 >= MAX_GROUPS_PER_BIT_PACKED_RUN {
+ self.close_bit_packed_run();
+ }
+ }
+ continue;
+ }
+
+ // Ending an RLE run, filling up a partial group, or a trailing
partial
+ // group: fall back to the per-value path
+ self.put(values[i]);
+ i += 1;
}
}
@@ -198,7 +249,13 @@ impl RleEncoder {
#[inline]
pub fn len(&self) -> usize {
- self.bit_writer.bytes_written()
+ // Include the eventual size of the pending bit-packed run, so size
estimates
+ // do not lag behind by up to a full run
+ let pending_bytes = match self.pending_values.len() {
+ 0 => 0,
+ n => 1 + bit_util::ceil(n * self.bit_width as usize, u8::BITS as
usize),
+ };
+ self.bit_writer.bytes_written() + pending_bytes
}
#[cfg_attr(not(feature = "experimental"), expect(unused))]
@@ -233,11 +290,9 @@ impl RleEncoder {
#[inline]
pub fn clear(&mut self) {
self.bit_writer.clear();
- self.num_buffered_values = 0;
+ self.pending_values.clear();
self.current_value = 0;
self.repeat_count = 0;
- self.bit_packed_count = 0;
- self.indicator_byte_pos = -1;
}
/// Advances the buffer by `num_bytes` zero bytes, delegating to the
@@ -251,21 +306,21 @@ impl RleEncoder {
/// internal writer.
#[inline]
pub fn flush(&mut self) {
- if self.bit_packed_count > 0 || self.repeat_count > 0 ||
self.num_buffered_values > 0 {
- let all_repeat = self.bit_packed_count == 0
- && (self.repeat_count == self.num_buffered_values ||
self.num_buffered_values == 0);
+ if !self.pending_values.is_empty() || self.repeat_count > 0 {
+ let tail = self.pending_values.len() % BIT_PACK_GROUP_SIZE;
+ let all_repeat =
+ self.pending_values.len() == tail && (self.repeat_count ==
tail || tail == 0);
if self.repeat_count > 0 && all_repeat {
self.flush_rle_run();
} else {
- // Buffer the last group of bit-packed values to
BIT_PACK_GROUP_SIZE by padding with 0s.
- if self.num_buffered_values > 0 {
- while self.num_buffered_values < BIT_PACK_GROUP_SIZE {
- self.buffered_values[self.num_buffered_values] = 0;
- self.num_buffered_values += 1;
- }
+ // Pad the last group of bit-packed values to
BIT_PACK_GROUP_SIZE with 0s.
+ if tail > 0 {
+ self.pending_values
+ .resize(self.pending_values.len() +
BIT_PACK_GROUP_SIZE - tail, 0);
+ }
+ if !self.pending_values.is_empty() {
+ self.close_bit_packed_run();
}
- self.bit_packed_count += self.num_buffered_values;
- self.flush_bit_packed_run(true);
self.repeat_count = 0;
}
}
@@ -279,65 +334,54 @@ impl RleEncoder {
self.current_value,
bit_util::ceil(self.bit_width as usize, u8::BITS as usize),
);
- self.num_buffered_values = 0;
+ self.pending_values.clear();
self.repeat_count = 0;
}
- fn flush_bit_packed_run(&mut self, end_current_run: bool) {
- if self.indicator_byte_pos < 0 {
- self.indicator_byte_pos = self.bit_writer.skip(1) as i64;
- }
-
- // Write all buffered values as bit-packed literals
- for v in &self.buffered_values[..self.num_buffered_values] {
- self.bit_writer.put_value(*v, self.bit_width as usize);
- }
- self.num_buffered_values = 0;
- if end_current_run {
- self.finish_bit_packed_run();
- }
- }
-
- // Called when ending a bit-packed run. Writes the indicator byte to the
reserved
- // position in `bit_writer`
- fn finish_bit_packed_run(&mut self) {
- let num_groups = self.bit_packed_count / BIT_PACK_GROUP_SIZE;
- let indicator_byte = ((num_groups << 1) | 1) as u8;
- self.bit_writer
- .put_aligned_offset(indicator_byte, 1, self.indicator_byte_pos as
usize);
- self.indicator_byte_pos = -1;
- self.bit_packed_count = 0;
- }
-
- fn flush_buffered_values(&mut self) {
+ // Called when a group of BIT_PACK_GROUP_SIZE values completes: either
keeps it in
+ // the pending bit-packed run, or drops it and switches to RLE when it
repeats a
+ // single value.
+ fn commit_group(&mut self) {
if self.repeat_count >= BIT_PACK_GROUP_SIZE {
- // Clear buffered values as they are not needed
- self.num_buffered_values = 0;
- if self.bit_packed_count > 0 {
+ // The group repeats a single value, drop it and let
`repeat_count` track
+ // the run in RLE mode instead
+ self.pending_values
+ .truncate(self.pending_values.len() - BIT_PACK_GROUP_SIZE);
+ if !self.pending_values.is_empty() {
// In this case we have chosen to switch to RLE encoding.
Close out the
// previous bit-packed run.
- debug_assert_eq!(self.bit_packed_count % BIT_PACK_GROUP_SIZE,
0);
- self.finish_bit_packed_run();
+ self.close_bit_packed_run();
}
return;
}
- self.bit_packed_count += self.num_buffered_values;
- let num_groups = self.bit_packed_count / BIT_PACK_GROUP_SIZE;
+ let num_groups = self.pending_values.len() / BIT_PACK_GROUP_SIZE;
if num_groups + 1 >= MAX_GROUPS_PER_BIT_PACKED_RUN {
// We've reached the maximum value that can be hold in a single
bit-packed
// run.
- debug_assert!(self.indicator_byte_pos >= 0);
- self.flush_bit_packed_run(true);
- } else {
- self.flush_bit_packed_run(false);
+ self.close_bit_packed_run();
}
self.repeat_count = 0;
}
+ // Called when ending a bit-packed run. Writes the indicator byte followed
by all
+ // pending values in packed form. Deferring the packing until here allows
the whole
+ // run to go through the vectorised `put_batch` in one call.
+ fn close_bit_packed_run(&mut self) {
+ debug_assert_eq!(self.pending_values.len() % BIT_PACK_GROUP_SIZE, 0);
+ let num_groups = self.pending_values.len() / BIT_PACK_GROUP_SIZE;
+ let indicator_byte = ((num_groups << 1) | 1) as u8;
+ self.bit_writer.put_aligned(indicator_byte, 1);
+ self.bit_writer
+ .put_batch(&self.pending_values, self.bit_width as usize);
+ self.pending_values.clear();
+ }
+
/// return the estimated memory size of this encoder.
pub(crate) fn estimated_memory_size(&self) -> usize {
- self.bit_writer.estimated_memory_size() + std::mem::size_of::<Self>()
+ self.bit_writer.estimated_memory_size()
+ + self.pending_values.capacity() * std::mem::size_of::<u64>()
+ + std::mem::size_of::<Self>()
}
}
@@ -401,7 +445,7 @@ impl RleDecoder {
// that damage L1d-cache occupancy. This results in a ~18% performance drop
#[inline(never)]
#[cfg_attr(all(not(feature = "experimental"), not(test)), expect(unused))]
- pub fn get<T: FromBitpacked>(&mut self) -> Result<Option<T>> {
+ pub fn get<T: BitPacking>(&mut self) -> Result<Option<T>> {
assert!(size_of::<T>() <= size_of::<u64>());
while self.rle_left == 0 && self.bit_packed_left == 0 {
@@ -434,7 +478,7 @@ impl RleDecoder {
}
#[inline(never)]
- pub fn get_batch<T: FromBitpacked + Clone>(&mut self, buffer: &mut [T]) ->
Result<usize> {
+ pub fn get_batch<T: BitPacking + Clone>(&mut self, buffer: &mut [T]) ->
Result<usize> {
assert!(size_of::<T>() <= size_of::<u64>());
let mut values_read = 0;
@@ -954,6 +998,69 @@ mod tests {
}
}
+ #[test]
+ fn test_rle_bit_packed_run_cap() {
+ // A bit-packed run holds at most MAX_GROUPS_PER_BIT_PACKED_RUN - 1 =
63 groups
+ // (504 values), which keeps every run header a single byte. 1024
alternating
+ // values must split into runs of 504, 504 and 16 values. A round-trip
cannot
+ // pin this down (longer runs with multi-byte headers also decode
fine), so
+ // check the exact bytes.
+ let values: Vec<i64> = (0..1024).map(|i| i % 2).collect();
+
+ let mut expected_buffer = Vec::new();
+ for num_groups in [63u8, 63, 2] {
+ expected_buffer.push((num_groups << 1) | 1);
+ expected_buffer.resize(expected_buffer.len() + num_groups as
usize, 0b10101010);
+ }
+
+ validate_rle(&values, 1, Some(&expected_buffer), 131);
+ }
+
+ #[test]
+ fn test_put_batch_matches_put() {
+ // `put_batch` must produce byte-identical output to a `put` loop, for
any
+ // input, any split of the input into batches, and mixed use of both
APIs
+ let mut rng = rng();
+ for width in [1, 2, 5, 8, 17, 32] {
+ let mask = (1u64 << width) - 1;
+ for _ in 0..20 {
+ // Mix random stretches with runs long enough to trigger RLE,
sized to
+ // also cross the bit-packed run cap
+ let mut values: Vec<u64> = Vec::new();
+ while values.len() < 1300 {
+ let n = rng.random_range(1..40);
+ if rng.random_bool(0.5) {
+ let v = rng.random::<u64>() & mask;
+ values.extend(std::iter::repeat_n(v, n));
+ } else {
+ values.extend((0..n).map(|_| rng.random::<u64>() &
mask));
+ }
+ }
+
+ let mut expected = RleEncoder::new(width as u8, 1024);
+ for &v in &values {
+ expected.put(v);
+ }
+
+ let mut actual = RleEncoder::new(width as u8, 1024);
+ let mut remaining = &values[..];
+ while !remaining.is_empty() {
+ let n = rng.random_range(1..=remaining.len().min(300));
+ if rng.random_bool(0.25) {
+ for &v in &remaining[..n] {
+ actual.put(v);
+ }
+ } else {
+ actual.put_batch(&remaining[..n]);
+ }
+ remaining = &remaining[n..];
+ }
+
+ assert_eq!(expected.consume(), actual.consume(), "width =
{width}");
+ }
+ }
+ }
+
// `validate_rle` on `num_vals` with width `bit_width`. If `value` is -1,
that value
// is used, otherwise alternating values are used.
fn test_rle_values(bit_width: usize, num_vals: usize, value: i32) {
diff --git a/parquet/src/util/bit_pack.rs b/parquet/src/util/bit_pack.rs
index 94ab9578b9..9ef55ca2b3 100644
--- a/parquet/src/util/bit_pack.rs
+++ b/parquet/src/util/bit_pack.rs
@@ -94,6 +94,106 @@ unpack!(unpack16, u16, 2, 16);
unpack!(unpack32, u32, 4, 32);
unpack!(unpack64, u64, 8, 64);
+/// Macro that generates a pack function taking the number of bits as a const
generic
+macro_rules! pack_impl {
+ ($t:ty, $bytes:literal, $bits:tt) => {
+ #[inline(never)]
+ pub fn pack<const NUM_BITS: usize>(input: &[$t; $bits], output: &mut
[u8]) {
+ if NUM_BITS == 0 {
+ return;
+ }
+
+ assert!(NUM_BITS <= $bytes * 8);
+ assert!(output.len() >= NUM_BITS * $bytes);
+
+ let mask = match NUM_BITS {
+ $bits => <$t>::MAX,
+ _ => ((1 << NUM_BITS) - 1),
+ };
+
+ // Accumulate into locals so the packed words stay in registers.
Only the
+ // first NUM_BITS entries are used, `[$t; NUM_BITS]` needs
generic_const_exprs
+ let mut words = [0; $bits];
+
+ seq_macro::seq!(i in 0..$bits {
+ let value = input[i] & mask;
+
+ let start_bit = i * NUM_BITS;
+ let end_bit = start_bit + NUM_BITS;
+
+ let start_bit_offset = start_bit % $bits;
+ let end_bit_offset = end_bit % $bits;
+ let start_word = start_bit / $bits;
+ let end_word = end_bit / $bits;
+
+ words[start_word] |= value << start_bit_offset;
+ if start_word != end_word && end_bit_offset != 0 {
+ words[end_word] |= value >> (NUM_BITS - end_bit_offset);
+ }
+ });
+
+ seq_macro::seq!(w in 0..$bits {
+ if w < NUM_BITS {
+ output[w * $bytes..(w + 1) *
$bytes].copy_from_slice(&words[w].to_le_bytes());
+ }
+ });
+ }
+
+ pub fn pack_blocks<const NUM_BITS: usize>(input: &[$t], output: &mut
[u8]) {
+ if NUM_BITS == 0 {
+ return;
+ }
+ let block_bytes = NUM_BITS * $bytes;
+ let blocks = input.len() / $bits;
+ assert!(output.len() >= blocks * block_bytes);
+ for (input, output) in input
+ .chunks_exact($bits)
+ .zip(output.chunks_exact_mut(block_bytes))
+ {
+ pack::<NUM_BITS>(input.try_into().unwrap(), output);
+ }
+ }
+ };
+}
+
+/// Macro that generates pack functions that accept num_bits as a parameter
+macro_rules! pack {
+ ($name:ident, $blocks:ident, $t:ty, $bytes:literal, $bits:tt) => {
+ mod $name {
+ pack_impl!($t, $bytes, $bits);
+ }
+
+ /// Pack `input` into `output` with a bit width of `num_bits`
+ ///
+ /// Only the `num_bits` least significant bits of each value are
written,
+ /// and `output` must contain at least `num_bits * size_of::<T>()`
bytes
+ pub fn $name(input: &[$t; $bits], output: &mut [u8], num_bits: usize) {
+ // This will get optimised into a jump table
+ seq_macro::seq!(i in 0..=$bits {
+ if i == num_bits {
+ return $name::pack::<i>(input, output);
+ }
+ });
+ unreachable!("invalid num_bits {}", num_bits);
+ }
+
+ #[inline(never)]
+ pub(crate) fn $blocks(input: &[$t], output: &mut [u8], num_bits:
usize) {
+ seq_macro::seq!(i in 0..=$bits {
+ if i == num_bits {
+ return $name::pack_blocks::<i>(input, output);
+ }
+ });
+ unreachable!("invalid num_bits {}", num_bits);
+ }
+ };
+}
+
+pack!(pack8, pack8_blocks, u8, 1, 8);
+pack!(pack16, pack16_blocks, u16, 2, 16);
+pack!(pack32, pack32_blocks, u32, 4, 32);
+pack!(pack64, pack64_blocks, u64, 8, 64);
+
#[cfg(test)]
mod tests {
use super::*;
@@ -134,4 +234,102 @@ mod tests {
}
}
}
+
+ #[test]
+ fn test_pack_all_ones() {
+ // Packing all-ones values must set every bit of the packed block and
+ // touch nothing beyond it
+ let mut output = [0u8; 4096];
+
+ for i in 0..=8 {
+ output.fill(0);
+ pack8(&[u8::MAX; 8], &mut output, i);
+ assert!(output[..i].iter().all(|&b| b == u8::MAX), "num_bits =
{i}");
+ assert!(output[i..].iter().all(|&b| b == 0), "num_bits = {i}");
+ }
+
+ for i in 0..=16 {
+ output.fill(0);
+ pack16(&[u16::MAX; 16], &mut output, i);
+ assert!(
+ output[..2 * i].iter().all(|&b| b == u8::MAX),
+ "num_bits = {i}"
+ );
+ assert!(output[2 * i..].iter().all(|&b| b == 0), "num_bits = {i}");
+ }
+
+ for i in 0..=32 {
+ output.fill(0);
+ pack32(&[u32::MAX; 32], &mut output, i);
+ assert!(
+ output[..4 * i].iter().all(|&b| b == u8::MAX),
+ "num_bits = {i}"
+ );
+ assert!(output[4 * i..].iter().all(|&b| b == 0), "num_bits = {i}");
+ }
+
+ for i in 0..=64 {
+ output.fill(0);
+ pack64(&[u64::MAX; 64], &mut output, i);
+ assert!(
+ output[..8 * i].iter().all(|&b| b == u8::MAX),
+ "num_bits = {i}"
+ );
+ assert!(output[8 * i..].iter().all(|&b| b == 0), "num_bits = {i}");
+ }
+ }
+
+ #[test]
+ fn test_pack_round_trip() {
+ use crate::util::test_common::rand_gen::random_numbers;
+
+ // Values are deliberately not masked, pack must ignore the high bits
+ for i in 0..=8 {
+ let input: [u8; 8] = random_numbers(8).try_into().unwrap();
+ let mut packed = vec![0u8; i];
+ pack8(&input, &mut packed, i);
+ let mut output = [0; 8];
+ unpack8(&packed, &mut output, i);
+ let mask = ((1u16 << i) - 1) as u8;
+ for (idx, (&v, &out)) in
input.iter().zip(output.iter()).enumerate() {
+ assert_eq!(v & mask, out, "num_bits = {i}, index = {idx}");
+ }
+ }
+
+ for i in 0..=16 {
+ let input: [u16; 16] = random_numbers(16).try_into().unwrap();
+ let mut packed = vec![0u8; 2 * i];
+ pack16(&input, &mut packed, i);
+ let mut output = [0; 16];
+ unpack16(&packed, &mut output, i);
+ let mask = ((1u32 << i) - 1) as u16;
+ for (idx, (&v, &out)) in
input.iter().zip(output.iter()).enumerate() {
+ assert_eq!(v & mask, out, "num_bits = {i}, index = {idx}");
+ }
+ }
+
+ for i in 0..=32 {
+ let input: [u32; 32] = random_numbers(32).try_into().unwrap();
+ let mut packed = vec![0u8; 4 * i];
+ pack32(&input, &mut packed, i);
+ let mut output = [0; 32];
+ unpack32(&packed, &mut output, i);
+ let mask = ((1u64 << i) - 1) as u32;
+ for (idx, (&v, &out)) in
input.iter().zip(output.iter()).enumerate() {
+ assert_eq!(v & mask, out, "num_bits = {i}, index = {idx}");
+ }
+ }
+
+ for i in 0..=64 {
+ let input: [u64; 64] = random_numbers(64).try_into().unwrap();
+ let mut packed = vec![0u8; 8 * i];
+ pack64(&input, &mut packed, i);
+ let mut output = [0; 64];
+ unpack64(&packed, &mut output, i);
+ let mask = ((1u128 << i) - 1) as u64;
+ for (idx, (&v, &out)) in
input.iter().zip(output.iter()).enumerate() {
+ assert_eq!(v & mask, out, "num_bits = {i}, index = {idx}");
+ }
+ }
+ }
}
diff --git a/parquet/src/util/bit_util.rs b/parquet/src/util/bit_util.rs
index ab1524f639..eeb265fa9b 100644
--- a/parquet/src/util/bit_util.rs
+++ b/parquet/src/util/bit_util.rs
@@ -21,7 +21,10 @@ use bytes::Bytes;
use crate::data_type::{AsBytes, ByteArray, FixedLenByteArray, Int96};
use crate::errors::{ParquetError, Result};
-use crate::util::bit_pack::{unpack8, unpack16, unpack32, unpack64};
+use crate::util::bit_pack::{
+ pack8, pack8_blocks, pack16, pack16_blocks, pack32, pack32_blocks, pack64,
pack64_blocks,
+ unpack8, unpack16, unpack32, unpack64,
+};
#[inline]
fn array_from_slice<const N: usize>(bs: &[u8]) -> Result<[u8; N]> {
@@ -44,28 +47,60 @@ pub trait FromBytes: Sized {
fn from_le_bytes(bs: Self::Buffer) -> Self;
}
-/// Types that can be decoded from bitpacked representations.
+/// Types that can be converted to and from bitpacked representations.
///
/// This is implemented for primitive types and bool that can be
-/// directly converted from a u64 value. Types like Int96, ByteArray,
+/// directly converted from and to a u64 value. Types like Int96, ByteArray,
/// and FixedLenByteArray that cannot be represented in 64 bits do not
/// implement this trait.
-pub trait FromBitpacked {
+pub trait BitPacking {
/// The maximum number of bits that are allowed to be converted to this
type.
/// This is at most the size of the type in bits, but could be less, for
example
/// for the boolean type.
const BIT_CAPACITY: usize;
- /// How many values are converted by one call to `unpack_batch`.
+ /// How many values are converted by one call to `unpack_batch` or
`pack_batch`.
const BATCH_SIZE: usize;
+
/// Convert directly from a u64 value by truncation, avoiding byte slice
copies.
fn from_u64(v: u64) -> Self;
+ /// Convert directly to a u64. Values wider than `num_bits` are permitted,
+ /// the packing routines only use the low `num_bits` bits.
+ fn to_u64(&self) -> u64;
+
/// Converts multiple bitpacked values from `input` to `output`.
/// The `output` slice needs to have space for at least `BATCH_SIZE`
elements,
/// otherwise this method will panic.
fn unpack_batch(input: &[u8], output: &mut [Self], num_bits: usize)
where
Self: Sized;
+
+ /// Packs the first `BATCH_SIZE` values of `input` into `output`.
+ /// `input` needs to contain at least `BATCH_SIZE` elements and `output`
needs
+ /// space for at least `num_bits * BATCH_SIZE / 8` bytes, otherwise this
+ /// method will panic.
+ fn pack_batch(input: &[Self], output: &mut [u8], num_bits: usize)
+ where
+ Self: Sized;
+
+ /// Packs all complete `BATCH_SIZE` blocks in `input` into `output`.
+ fn pack_batches(input: &[Self], output: &mut [u8], num_bits: usize)
+ where
+ Self: Sized,
+ {
+ if num_bits == 0 {
+ return;
+ }
+ let block_bytes = num_bits * Self::BATCH_SIZE / 8;
+ let blocks = input.len() / Self::BATCH_SIZE;
+ assert!(output.len() >= blocks * block_bytes);
+ #[expect(clippy::chunks_exact_to_as_chunks)]
+ // Requires using an associated constant in const operations:
generic_const_exprs
+ let input_chunks = input.chunks_exact(Self::BATCH_SIZE);
+ for (input, output) in
input_chunks.zip(output.chunks_exact_mut(block_bytes)) {
+ Self::pack_batch(input, output, num_bits);
+ }
+ }
}
macro_rules! from_le_bytes {
@@ -84,12 +119,12 @@ macro_rules! from_le_bytes {
};
}
-macro_rules! from_bitpacked {
- ($($ty: ty => $unpack: path),*) => {
+macro_rules! bit_packing {
+ ($($ty: ty => ($unpack: path, $pack: path, $pack_blocks: path)),*) => {
$(
- impl FromBitpacked for $ty {
+ impl BitPacking for $ty {
const BIT_CAPACITY: usize = std::mem::size_of::<$ty>() * 8;
- // this has to match the signature of the unpack* functions
+ // this has to match the signature of the unpack*/pack*
functions
const BATCH_SIZE: usize = std::mem::size_of::<$ty>() * 8;
#[inline]
@@ -97,55 +132,92 @@ macro_rules! from_bitpacked {
v as _
}
+ #[inline]
+ fn to_u64(&self) -> u64 {
+ *self as u64
+ }
+
#[inline]
fn unpack_batch(input: &[u8], output: &mut [Self], num_bits:
usize) {
$unpack(input, (&mut
output[..Self::BATCH_SIZE]).try_into().unwrap(), num_bits)
}
+
+ #[inline]
+ fn pack_batch(input: &[Self], output: &mut [u8], num_bits:
usize) {
+ $pack((&input[..Self::BATCH_SIZE]).try_into().unwrap(),
output, num_bits)
+ }
+
+ #[inline]
+ fn pack_batches(input: &[Self], output: &mut [u8], num_bits:
usize) {
+ $pack_blocks(input, output, num_bits)
+ }
}
)*
}
}
-macro_rules! from_bitpacked_delegate {
+macro_rules! bit_packing_delegate {
($($ty: ty => $delegate: ty),*) => {
$(
- impl FromBitpacked for $ty {
- const BIT_CAPACITY: usize = <$delegate as
FromBitpacked>::BIT_CAPACITY;
- const BATCH_SIZE: usize = <$delegate as
FromBitpacked>::BATCH_SIZE;
+ // Guard against misusages of this macro, this fails already at
+ // compile-time if the types are not compatible.
+ const _: () = assert!(
+ std::mem::size_of::<$ty>() == std::mem::size_of::<$delegate>()
+ && std::mem::align_of::<$ty>() ==
std::mem::align_of::<$delegate>(),
+ "types need to have the same size and alignment"
+ );
+
+ impl BitPacking for $ty {
+ const BIT_CAPACITY: usize = <$delegate as
BitPacking>::BIT_CAPACITY;
+ const BATCH_SIZE: usize = <$delegate as
BitPacking>::BATCH_SIZE;
#[inline]
fn from_u64(v: u64) -> Self {
v as _
}
+ #[inline]
+ fn to_u64(&self) -> u64 {
+ *self as u64
+ }
+
#[inline]
fn unpack_batch(input: &[u8], output: &mut [Self], num_bits:
usize) {
- // Guard against misusages of this macro, due to the const
block this will fail
- // already at compile-time if the types are not compatible.
- const {
- assert!(
- std::mem::size_of::<$ty>() ==
std::mem::size_of::<$delegate>()
- && std::mem::align_of::<$ty>() ==
std::mem::align_of::<$delegate>(),
- "types need to have the same size and alignment"
- );
- }
// Safety: ty and delegate have the same size and
alignment, and this macro is only used for types that have transmutable bit
patterns.
let output: &mut [$delegate] = unsafe {
std::slice::from_raw_parts_mut(output.as_mut_ptr().cast::<$delegate>(),
output.len()) };
<$delegate>::unpack_batch(input, output, num_bits);
}
+
+ #[inline]
+ fn pack_batch(input: &[Self], output: &mut [u8], num_bits:
usize) {
+ // Safety: ty and delegate have the same size and
alignment, and this macro is only used for types that have transmutable bit
patterns.
+ let input: &[$delegate] = unsafe {
std::slice::from_raw_parts(input.as_ptr().cast::<$delegate>(), input.len()) };
+ <$delegate>::pack_batch(input, output, num_bits);
+ }
+
+ #[inline]
+ fn pack_batches(input: &[Self], output: &mut [u8], num_bits:
usize) {
+ // Safety: ty and delegate have the same size and
alignment, and this macro is only used for types that have transmutable bit
patterns.
+ let input: &[$delegate] = unsafe {
std::slice::from_raw_parts(input.as_ptr().cast::<$delegate>(), input.len()) };
+ <$delegate>::pack_batches(input, output, num_bits);
+ }
}
)*
}
}
from_le_bytes! { u8, u16, u32, u64, i8, i16, i32, i64 }
-from_bitpacked!(u8 => unpack8, u16 => unpack16, u32 => unpack32);
+bit_packing!(
+ u8 => (unpack8, pack8, pack8_blocks),
+ u16 => (unpack16, pack16, pack16_blocks),
+ u32 => (unpack32, pack32, pack32_blocks)
+);
-// `u64` is written out by hand: the `as` cast the macro uses would be a no-op
here,
+// `u64` is written out by hand: the `as` casts the macro uses would be no-ops
here,
// and it is the only instantiation for which that is true.
-impl FromBitpacked for u64 {
+impl BitPacking for u64 {
const BIT_CAPACITY: usize = std::mem::size_of::<u64>() * 8;
- // this has to match the signature of the unpack* functions
+ // this has to match the signature of the unpack*/pack* functions
const BATCH_SIZE: usize = std::mem::size_of::<u64>() * 8;
#[inline]
@@ -153,6 +225,11 @@ impl FromBitpacked for u64 {
v
}
+ #[inline]
+ fn to_u64(&self) -> u64 {
+ *self
+ }
+
#[inline]
fn unpack_batch(input: &[u8], output: &mut [Self], num_bits: usize) {
unpack64(
@@ -161,18 +238,37 @@ impl FromBitpacked for u64 {
num_bits,
)
}
+
+ #[inline]
+ fn pack_batch(input: &[Self], output: &mut [u8], num_bits: usize) {
+ pack64(
+ (&input[..Self::BATCH_SIZE]).try_into().unwrap(),
+ output,
+ num_bits,
+ )
+ }
+
+ #[inline]
+ fn pack_batches(input: &[Self], output: &mut [u8], num_bits: usize) {
+ pack64_blocks(input, output, num_bits)
+ }
}
-from_bitpacked_delegate!(i8 => u8, i16 => u16, i32 => u32, i64 => u64);
+bit_packing_delegate!(i8 => u8, i16 => u16, i32 => u32, i64 => u64);
-impl FromBitpacked for bool {
+impl BitPacking for bool {
const BIT_CAPACITY: usize = 1;
- const BATCH_SIZE: usize = <u8 as FromBitpacked>::BATCH_SIZE;
+ const BATCH_SIZE: usize = <u8 as BitPacking>::BATCH_SIZE;
#[inline]
fn from_u64(v: u64) -> Self {
v != 0
}
+ #[inline]
+ fn to_u64(&self) -> u64 {
+ *self as u64
+ }
+
#[inline]
fn unpack_batch(input: &[u8], output: &mut [Self], num_bits: usize) {
assert_eq!(num_bits, 1);
@@ -184,6 +280,26 @@ impl FromBitpacked for bool {
};
u8::unpack_batch(input, output, num_bits);
}
+
+ #[inline]
+ fn pack_batch(input: &[Self], output: &mut [u8], num_bits: usize) {
+ assert_eq!(num_bits, 1);
+ // Safety: bool is a single byte that is guaranteed to be 0 or 1, so it
+ // can always be read as a u8.
+ let input: &[u8] =
+ unsafe { std::slice::from_raw_parts(input.as_ptr().cast::<u8>(),
input.len()) };
+ u8::pack_batch(input, output, num_bits);
+ }
+
+ #[inline]
+ fn pack_batches(input: &[Self], output: &mut [u8], num_bits: usize) {
+ assert_eq!(num_bits, 1);
+ // Safety: bool is a single byte that is guaranteed to be 0 or 1, so it
+ // can always be read as a u8.
+ let input: &[u8] =
+ unsafe { std::slice::from_raw_parts(input.as_ptr().cast::<u8>(),
input.len()) };
+ u8::pack_batches(input, output, num_bits);
+ }
}
impl FromBytes for f32 {
@@ -510,6 +626,105 @@ impl BitWriter {
}
}
+ /// Writes all values in `batch` in bit-packed form, `num_bits` bits per
+ /// value.
+ ///
+ /// Equivalent to repeatedly calling [`BitWriter::put_value`] with the same
+ /// `num_bits`, but faster because it dispatches to SIMD-friendly
+ /// fixed-width packing routines whenever possible.
+ ///
+ /// Unlike [`BitWriter::put_value`], values wider than `num_bits` are
+ /// permitted, only the `num_bits` least significant bits of each value are
+ /// written.
+ ///
+ /// # Panics
+ ///
+ /// This function panics if
+ /// - `num_bits` is larger than the bit-capacity of `T`
+ pub fn put_batch<T: BitPacking>(&mut self, batch: &[T], num_bits: usize) {
+ assert_ne!(T::BIT_CAPACITY, 0);
+ assert!(num_bits <= T::BIT_CAPACITY);
+ if batch.is_empty() || num_bits == 0 {
+ return;
+ }
+
+ let mask = match num_bits {
+ 64 => u64::MAX,
+ _ => (1 << num_bits) - 1,
+ };
+
+ let mut i = 0;
+
+ // First fill the accumulator up to a byte boundary
+ while !self.bit_offset.is_multiple_of(8) {
+ match batch.get(i) {
+ Some(v) => self.put_value(v.to_u64() & mask, num_bits),
+ None => return,
+ }
+ i += 1;
+ }
+
+ // Move the accumulator's whole bytes out, so packed blocks can be
+ // appended directly to the buffer. Lossless because the offset is
+ // byte-aligned
+ self.flush();
+
+ // Pack whole blocks directly into the buffer
+ let blocks = (batch.len() - i) / T::BATCH_SIZE;
+ let block_bytes = num_bits * T::BATCH_SIZE / 8;
+ let offset = self.buffer.len();
+ self.buffer.resize(offset + blocks * block_bytes, 0);
+ let block_values = blocks * T::BATCH_SIZE;
+ if blocks != 0 {
+ T::pack_batches(
+ &batch[i..i + block_values],
+ &mut self.buffer[offset..offset + blocks * block_bytes],
+ num_bits,
+ );
+ }
+ i += block_values;
+
+ // Try to write smaller batches if possible
+ if size_of::<T>() > 4 && batch.len() - i >= 32 && num_bits <= 32 {
+ let mut in_buf = [0_u32; 32];
+ for (j, v) in in_buf.iter_mut().enumerate() {
+ *v = (batch[i + j].to_u64() & mask) as u32;
+ }
+ let start = self.buffer.len();
+ self.buffer.resize(start + 4 * num_bits, 0);
+ pack32(&in_buf, &mut self.buffer[start..], num_bits);
+ i += 32;
+ }
+
+ if size_of::<T>() > 2 && batch.len() - i >= 16 && num_bits <= 16 {
+ let mut in_buf = [0_u16; 16];
+ for (j, v) in in_buf.iter_mut().enumerate() {
+ *v = (batch[i + j].to_u64() & mask) as u16;
+ }
+ let start = self.buffer.len();
+ self.buffer.resize(start + 2 * num_bits, 0);
+ pack16(&in_buf, &mut self.buffer[start..], num_bits);
+ i += 16;
+ }
+
+ if size_of::<T>() > 1 && batch.len() - i >= 8 && num_bits <= 8 {
+ let mut in_buf = [0_u8; 8];
+ for (j, v) in in_buf.iter_mut().enumerate() {
+ *v = (batch[i + j].to_u64() & mask) as u8;
+ }
+ let start = self.buffer.len();
+ self.buffer.resize(start + num_bits, 0);
+ pack8(&in_buf, &mut self.buffer[start..], num_bits);
+ i += 8;
+ }
+
+ // Write any trailing values
+ while i < batch.len() {
+ self.put_value(batch[i].to_u64() & mask, num_bits);
+ i += 1;
+ }
+ }
+
/// Writes the first `num_bytes` little-endian bytes of `val` to the
/// writer at the next byte boundary.
///
@@ -665,7 +880,7 @@ impl BitReader {
/// Returns `None` if there are fewer than `num_bits` bits left in the
/// buffer; otherwise `Some(value)`. On `None` the reader's position is
/// left unchanged.
- pub fn get_value<T: FromBitpacked>(&mut self, num_bits: usize) ->
Option<T> {
+ pub fn get_value<T: BitPacking>(&mut self, num_bits: usize) -> Option<T> {
debug_assert!(num_bits <= 64);
debug_assert!(num_bits <= size_of::<T>() * 8);
@@ -715,7 +930,7 @@ impl BitReader {
///
/// This function panics if
/// - `num_bits` is larger than the bit-capacity of `T`
- pub fn get_batch<T: FromBitpacked>(&mut self, batch: &mut [T], num_bits:
usize) -> usize {
+ pub fn get_batch<T: BitPacking>(&mut self, batch: &mut [T], num_bits:
usize) -> usize {
debug_assert!(num_bits <= size_of::<T>() * 8);
let mut values_to_read = batch.len();
@@ -1338,7 +1553,7 @@ mod tests {
fn test_get_batch_helper<T>(total: usize, num_bits: usize)
where
- T: FromBitpacked + Default + Clone + Debug + Eq,
+ T: BitPacking + Default + Clone + Debug + Eq,
{
assert!(num_bits <= 64);
let num_bytes = ceil(num_bits, 8);
@@ -1376,6 +1591,80 @@ mod tests {
}
}
+ #[test]
+ fn test_put_batch() {
+ const SIZE: &[usize] = &[1, 7, 8, 31, 32, 33, 128, 129];
+ for s in SIZE {
+ // Exercise every bit width for each type, the narrow widths on the
+ // wider types cover the smaller-batch packing paths
+ for i in 0..=8 {
+ test_put_batch_helper::<u8>(*s, i);
+ test_put_batch_helper::<i8>(*s, i);
+ }
+ for i in 0..=16 {
+ test_put_batch_helper::<u16>(*s, i);
+ test_put_batch_helper::<i16>(*s, i);
+ }
+ for i in 0..=32 {
+ test_put_batch_helper::<u32>(*s, i);
+ test_put_batch_helper::<i32>(*s, i);
+ }
+ for i in 0..=64 {
+ test_put_batch_helper::<u64>(*s, i);
+ test_put_batch_helper::<i64>(*s, i);
+ }
+ // `bool` only supports a bit width of 1.
+ test_put_batch_helper::<bool>(*s, 1);
+ }
+ }
+
+ fn test_put_batch_helper<T>(total: usize, num_bits: usize)
+ where
+ T: BitPacking + Default + Copy + Debug + Eq,
+ {
+ let mask = match num_bits {
+ 64 => u64::MAX,
+ _ => (1 << num_bits) - 1,
+ };
+
+ let values: Vec<T> = random_numbers::<u64>(total)
+ .iter()
+ .map(|v| T::from_u64(v & mask))
+ .collect();
+
+ // `put_batch` must produce bit-identical output to `put_value`, from
+ // both byte-aligned and unaligned starting positions
+ for misalignment in [0, 1, 3] {
+ let mut expected = BitWriter::new(ceil(num_bits * total, 8));
+ let mut actual = BitWriter::new(ceil(num_bits * total, 8));
+
+ for _ in 0..misalignment {
+ expected.put_value(1, 3);
+ actual.put_value(1, 3);
+ }
+
+ for v in &values {
+ expected.put_value(v.to_u64() & mask, num_bits);
+ }
+ actual.put_batch(&values, num_bits);
+
+ assert_eq!(
+ expected.flush_buffer(),
+ actual.flush_buffer(),
+ "num_bits = {num_bits}, total = {total}, misalignment =
{misalignment}"
+ );
+ }
+
+ // And round-trip through `get_batch`
+ let mut writer = BitWriter::new(ceil(num_bits * total, 8));
+ writer.put_batch(&values, num_bits);
+ let mut reader = BitReader::from(writer.consume());
+ let mut batch = vec![T::default(); total];
+ let values_read = reader.get_batch::<T>(&mut batch, num_bits);
+ assert_eq!(values_read, total);
+ assert_eq!(batch, values, "num_bits = {num_bits}, total = {total}");
+ }
+
#[test]
fn test_put_aligned_roundtrip() {
test_put_aligned_rand_numbers::<u8>(4, 3);
