[Parquet] Vectorise Pack Bits within BitWriter

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

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

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

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;

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]

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

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

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 {