diff --git a/parquet/benches/encoding.rs b/parquet/benches/encoding.rs index 65c2ec3d37ee..df605e05e35f 100644 --- a/parquet/benches/encoding.rs +++ b/parquet/benches/encoding.rs @@ -19,7 +19,8 @@ use criterion::*; use half::f16; use parquet::basic::{Encoding, Type as ParquetType}; use parquet::data_type::{ - DataType, DoubleType, FixedLenByteArray, FixedLenByteArrayType, FloatType, + BoolType, DataType, DoubleType, FixedLenByteArray, FixedLenByteArrayType, FloatType, Int32Type, + Int64Type, }; use parquet::decoding::{Decoder, get_decoder}; use parquet::encoding::get_encoder; @@ -84,6 +85,9 @@ fn criterion_benchmark(c: &mut Criterion) { let mut f32s = Vec::new(); let mut f64s = Vec::new(); let mut d128s = Vec::new(); + let mut bools = Vec::new(); + let mut i32s = Vec::new(); + let mut i64s = Vec::new(); for _ in 0..n { f16s.push(FixedLenByteArray::from( f16::from_f32(rng.random::()).to_le_bytes().to_vec(), @@ -93,12 +97,20 @@ fn criterion_benchmark(c: &mut Criterion) { d128s.push(FixedLenByteArray::from( rng.random::().to_be_bytes().to_vec(), )); + bools.push(rng.random::()); + i32s.push(rng.random::()); + // Keep deltas below 32 bits, mimicking timestamp-like data + i64s.push(rng.random_range(0..1i64 << 28)); } bench_typed::(c, &f32s, Encoding::BYTE_STREAM_SPLIT, 0); bench_typed::(c, &f64s, Encoding::BYTE_STREAM_SPLIT, 0); bench_typed::(c, &f16s, Encoding::BYTE_STREAM_SPLIT, 2); bench_typed::(c, &d128s, Encoding::BYTE_STREAM_SPLIT, 16); + bench_typed::(c, &bools, Encoding::PLAIN, 0); + bench_typed::(c, &bools, Encoding::RLE, 0); + bench_typed::(c, &i32s, Encoding::DELTA_BINARY_PACKED, 0); + bench_typed::(c, &i64s, Encoding::DELTA_BINARY_PACKED, 0); } criterion_group!(benches, criterion_benchmark); diff --git a/parquet/src/arrow/array_reader/byte_array_dictionary.rs b/parquet/src/arrow/array_reader/byte_array_dictionary.rs index 65f2d9cb0a7a..04316109378d 100644 --- a/parquet/src/arrow/array_reader/byte_array_dictionary.rs +++ b/parquet/src/arrow/array_reader/byte_array_dictionary.rs @@ -34,7 +34,7 @@ use crate::column::reader::decoder::ColumnValueDecoder; use crate::encodings::rle::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 { @@ -132,7 +132,7 @@ struct ByteArrayDictionaryReader { impl ByteArrayDictionaryReader where - K: FromBitpacked + Ord + ArrowNativeType, + K: BitPacking + Ord + ArrowNativeType, V: OffsetSizeTrait, { fn new( @@ -152,7 +152,7 @@ where impl ArrayReader for ByteArrayDictionaryReader where - K: FromBitpacked + Ord + ArrowNativeType, + K: BitPacking + Ord + ArrowNativeType, V: OffsetSizeTrait, { fn as_any(&self) -> &dyn Any { @@ -234,7 +234,7 @@ struct DictionaryDecoder { impl ColumnValueDecoder for DictionaryDecoder where - K: FromBitpacked + Ord + ArrowNativeType, + K: BitPacking + Ord + ArrowNativeType, V: OffsetSizeTrait, { type Buffer = DictionaryBuffer; diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 145431c26465..47d2877c5f2e 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -399,9 +399,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 d8c7b9201389..56405e6fe331 100644 --- a/parquet/src/data_type.rs +++ b/parquet/src/data_type.rs @@ -774,9 +774,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/decoding.rs b/parquet/src/encodings/decoding.rs index 0fb960412295..75d65140fe39 100644 --- a/parquet/src/encodings/decoding.rs +++ b/parquet/src/encodings/decoding.rs @@ -31,7 +31,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}; mod byte_stream_split_decoder; @@ -457,7 +457,7 @@ impl RleValueDecoder { impl Decoder for RleValueDecoder where - T::T: FromBitpacked, + T::T: BitPacking, { #[inline] fn set_data(&mut self, data: Bytes, num_values: usize) -> Result<()> { @@ -661,7 +661,7 @@ where impl Decoder for DeltaBitPackDecoder 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/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index 89666bbe7313..eac17c0b9009 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -138,9 +138,7 @@ impl DictEncoder { // 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 eeabcf4ba5ce..3db69dc38399 100644 --- a/parquet/src/encodings/encoding/mod.rs +++ b/parquet/src/encodings/encoding/mod.rs @@ -218,9 +218,13 @@ impl Encoder for RleValueEncoder { 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(()) } @@ -411,12 +415,15 @@ impl DeltaBitPackEncoder { 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 5de9dc850a34..bf52310f8358 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}; /// Number of values in one bit-packed group. The Parquet RLE/bit-packing hybrid /// format always bit-packs values in multiples of this count (see the @@ -59,11 +59,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, // The current (also last) value that was written and the count of how many // times in a row that value has been seen. @@ -72,13 +72,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 { @@ -95,12 +88,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, } } @@ -168,19 +158,73 @@ 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() % BIT_PACK_GROUP_SIZE == 0 { + // 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() % BIT_PACK_GROUP_SIZE == 0 + && 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; } } @@ -192,7 +236,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 } #[allow(unused)] @@ -227,11 +277,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 @@ -245,21 +293,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; } } @@ -273,65 +321,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.bit_writer.estimated_memory_size() + + self.pending_values.capacity() * std::mem::size_of::() + + std::mem::size_of::() } } @@ -390,7 +427,7 @@ impl RleDecoder { // that damage L1d-cache occupancy. This results in a ~18% performance drop #[inline(never)] #[allow(unused)] - pub fn get(&mut self) -> Result> { + pub fn get(&mut self) -> Result> { assert!(size_of::() <= size_of::()); while self.rle_left == 0 && self.bit_packed_left == 0 { @@ -423,7 +460,7 @@ impl RleDecoder { } #[inline(never)] - pub fn get_batch(&mut self, buffer: &mut [T]) -> Result { + pub fn get_batch(&mut self, buffer: &mut [T]) -> Result { assert!(size_of::() <= size_of::()); let mut values_read = 0; @@ -947,6 +984,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 = (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 = Vec::new(); + while values.len() < 1300 { + let n = rng.random_range(1..40); + if rng.random_bool(0.5) { + let v = rng.random::() & mask; + values.extend(std::iter::repeat_n(v, n)); + } else { + values.extend((0..n).map(|_| rng.random::() & 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 94ab9578b991..44a62cc59fa3 100644 --- a/parquet/src/util/bit_pack.rs +++ b/parquet/src/util/bit_pack.rs @@ -94,6 +94,80 @@ 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) => { + pub fn pack(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 as $t; $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()); + } + }); + } + }; +} + +/// Macro that generates pack functions that accept num_bits as a parameter +macro_rules! pack { + ($name: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::()` 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::(input, output); + } + }); + unreachable!("invalid num_bits {}", num_bits); + } + }; +} + +pack!(pack8, u8, 1, 8); +pack!(pack16, u16, 2, 16); +pack!(pack32, u32, 4, 32); +pack!(pack64, u64, 8, 64); + #[cfg(test)] mod tests { use super::*; @@ -134,4 +208,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 d5f5945a12a9..777d429eaaf2 100644 --- a/parquet/src/util/bit_util.rs +++ b/parquet/src/util/bit_util.rs @@ -21,7 +21,7 @@ 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, pack16, pack32, pack64, unpack8, unpack16, unpack32, unpack64}; #[inline] fn array_from_slice(bs: &[u8]) -> Result<[u8; N]> { @@ -44,28 +44,41 @@ 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; } macro_rules! from_le_bytes { @@ -84,12 +97,12 @@ macro_rules! from_le_bytes { }; } -macro_rules! from_bitpacked { - ($($ty: ty => $unpack: path),*) => { +macro_rules! bit_packing { + ($($ty: ty => ($unpack: path, $pack: 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,60 +110,91 @@ 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) + } } )* } } -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); + } } )* } } from_le_bytes! { u8, u16, u32, u64, i8, i16, i32, i64 } -from_bitpacked!(u8 => unpack8, u16 => unpack16, u32 => unpack32, u64 => unpack64); -from_bitpacked_delegate!(i8 => u8, i16 => u16, i32 => u32, i64 => u64); - -impl FromBitpacked for bool { +bit_packing!( + u8 => (unpack8, pack8), + u16 => (unpack16, pack16), + u32 => (unpack32, pack32), + u64 => (unpack64, pack64) +); +bit_packing_delegate!(i8 => u8, i16 => u16, i32 => u32, i64 => u64); + +impl BitPacking for bool { const BIT_CAPACITY: usize = 1; - const BATCH_SIZE: usize = ::BATCH_SIZE; + const BATCH_SIZE: usize = ::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!(num_bits == 1); @@ -162,6 +206,16 @@ impl FromBitpacked for bool { }; u8::unpack_batch(input, output, num_bits); } + + #[inline] + fn pack_batch(input: &[Self], output: &mut [u8], num_bits: usize) { + assert!(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::(), input.len()) }; + u8::pack_batch(input, output, num_bits); + } } impl FromBytes for f32 { @@ -488,6 +542,98 @@ 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(&mut self, batch: &[T], num_bits: usize) { + assert_ne!(T::BIT_CAPACITY, 0); + assert!(num_bits <= T::BIT_CAPACITY); + + 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 % 8 != 0 { + 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 mut offset = self.buffer.len(); + self.buffer.resize(offset + blocks * block_bytes, 0); + for _ in 0..blocks { + T::pack_batch(&batch[i..], &mut self.buffer[offset..], num_bits); + offset += block_bytes; + i += T::BATCH_SIZE; + } + + // Try to write smaller batches if possible + if size_of::() > 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::() > 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::() > 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. /// @@ -643,7 +789,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(&mut self, num_bits: usize) -> Option { + pub fn get_value(&mut self, num_bits: usize) -> Option { debug_assert!(num_bits <= 64); debug_assert!(num_bits <= size_of::() * 8); @@ -693,7 +839,7 @@ impl BitReader { /// /// This function panics if /// - `num_bits` is larger than the bit-capacity of `T` - pub fn get_batch(&mut self, batch: &mut [T], num_bits: usize) -> usize { + pub fn get_batch(&mut self, batch: &mut [T], num_bits: usize) -> usize { debug_assert!(num_bits <= size_of::() * 8); let mut values_to_read = batch.len(); @@ -1315,7 +1461,7 @@ mod tests { fn test_get_batch_helper(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); @@ -1353,6 +1499,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::(*s, i); + test_put_batch_helper::(*s, i); + } + for i in 0..=16 { + test_put_batch_helper::(*s, i); + test_put_batch_helper::(*s, i); + } + for i in 0..=32 { + test_put_batch_helper::(*s, i); + test_put_batch_helper::(*s, i); + } + for i in 0..=64 { + test_put_batch_helper::(*s, i); + test_put_batch_helper::(*s, i); + } + // `bool` only supports a bit width of 1. + test_put_batch_helper::(*s, 1); + } + } + + fn test_put_batch_helper(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 = random_numbers::(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::(&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::(4, 3);