[Parquet] Vectorise Pack Bits within BitWriter

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::<f32>()).to_le_bytes().to_vec(),
@@ -93,12 +97,20 @@ fn criterion_benchmark(c: &mut Criterion) {
         d128s.push(FixedLenByteArray::from(
             rng.random::<i128>().to_be_bytes().to_vec(),
         ));
+        bools.push(rng.random::<bool>());
+        i32s.push(rng.random::<i32>());
+        // Keep deltas below 32 bits, mimicking timestamp-like data
+        i64s.push(rng.random_range(0..1i64 << 28));
     }
 
     bench_typed::<FloatType>(c, &f32s, Encoding::BYTE_STREAM_SPLIT, 0);
     bench_typed::<DoubleType>(c, &f64s, Encoding::BYTE_STREAM_SPLIT, 0);
     bench_typed::<FixedLenByteArrayType>(c, &f16s, Encoding::BYTE_STREAM_SPLIT, 2);
     bench_typed::<FixedLenByteArrayType>(c, &d128s, Encoding::BYTE_STREAM_SPLIT, 16);
+    bench_typed::<BoolType>(c, &bools, Encoding::PLAIN, 0);
+    bench_typed::<BoolType>(c, &bools, Encoding::RLE, 0);
+    bench_typed::<Int32Type>(c, &i32s, Encoding::DELTA_BINARY_PACKED, 0);
+    bench_typed::<Int64Type>(c, &i64s, Encoding::DELTA_BINARY_PACKED, 0);
 }
 
 criterion_group!(benches, criterion_benchmark);

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<K: ArrowNativeType, V: OffsetSizeTrait> {
 
 impl<K, V> ByteArrayDictionaryReader<K, V>
 where
-    K: FromBitpacked + Ord + ArrowNativeType,
+    K: BitPacking + Ord + ArrowNativeType,
     V: OffsetSizeTrait,
 {
     fn new(
@@ -152,7 +152,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 {
@@ -234,7 +234,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>;

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();
 

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(())
         }
 

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<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<()> {
@@ -661,7 +661,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]

parquet/src/encodings/encoding/dict_encoder.rs

@@ -138,9 +138,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())
     }

parquet/src/encodings/encoding/mod.rs

@@ -218,9 +218,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(())
     }
@@ -411,12 +415,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 {