parquet: reuse level encoder allocations across page flushes

Previously, flushing a data page would `mem::replace` each level encoder
with a freshly allocated one, consuming the old encoder to get its
buffer. This allocated new internal `Vec`s on every page boundary.

We now preserve the internal state of the encoder and reuse memory
across pages.

Signed-off-by: Hippolyte Barraud <hippolyte.barraud@datadoghq.com>

Author: HippoBaro

parquet/src/column/writer/mod.rs

@@ -24,7 +24,6 @@ use crate::bloom_filter::Sbbf;
 use crate::file::page_index::column_index::ColumnIndexMetaData;
 use crate::file::page_index::offset_index::OffsetIndexMetaData;
 use std::collections::{BTreeSet, VecDeque};
-use std::mem;
 use std::str;
 
 use crate::basic::{
@@ -1060,19 +1059,13 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
                 let mut buffer = vec![];
 
                 if max_rep_level > 0 {
-                    let encoder = mem::replace(
-                        &mut self.rep_levels_encoder,
-                        Self::create_level_encoder(max_rep_level, &self.props),
-                    );
-                    buffer.extend_from_slice(&encoder.consume());
+                    self.rep_levels_encoder
+                        .flush_to(|data| buffer.extend_from_slice(data));
                 }
 
                 if max_def_level > 0 {
-                    let encoder = mem::replace(
-                        &mut self.def_levels_encoder,
-                        Self::create_level_encoder(max_def_level, &self.props),
-                    );
-                    buffer.extend_from_slice(&encoder.consume());
+                    self.def_levels_encoder
+                        .flush_to(|data| buffer.extend_from_slice(data));
                 }
 
                 buffer.extend_from_slice(&values_data.buf);
@@ -1102,23 +1095,15 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
                 let mut buffer = vec![];
 
                 if max_rep_level > 0 {
-                    let encoder = mem::replace(
-                        &mut self.rep_levels_encoder,
-                        Self::create_level_encoder(max_rep_level, &self.props),
-                    );
-                    let levels = encoder.consume();
-                    rep_levels_byte_len = levels.len();
-                    buffer.extend_from_slice(&levels);
+                    self.rep_levels_encoder
+                        .flush_to(|data| buffer.extend_from_slice(data));
+                    rep_levels_byte_len = buffer.len();
                 }
 
                 if max_def_level > 0 {
-                    let encoder = mem::replace(
-                        &mut self.def_levels_encoder,
-                        Self::create_level_encoder(max_def_level, &self.props),
-                    );
-                    let levels = encoder.consume();
-                    def_levels_byte_len = levels.len();
-                    buffer.extend_from_slice(&levels);
+                    self.def_levels_encoder
+                        .flush_to(|data| buffer.extend_from_slice(data));
+                    def_levels_byte_len = buffer.len() - rep_levels_byte_len;
                 }
 
                 let uncompressed_size =

parquet/src/encodings/levels.rs

@@ -147,6 +147,7 @@ impl LevelEncoder {
     /// Finalizes level encoder, flush all intermediate buffers and return resulting
     /// encoded buffer. Returned buffer is already truncated to encoded bytes only.
     #[inline]
+    #[allow(unused)]
     pub fn consume(self) -> Vec<u8> {
         match self {
             LevelEncoder::Rle(encoder) => {
@@ -162,4 +163,34 @@ impl LevelEncoder {
             LevelEncoder::BitPacked(_, encoder) => encoder.consume(),
         }
     }
+
+    /// Flushes all intermediate buffers, passes the encoded data to `f`, then
+    /// resets the encoder for reuse while retaining the buffer allocation.
+    #[inline]
+    pub fn flush_to<F, R>(&mut self, f: F) -> R
+    where
+        F: FnOnce(&[u8]) -> R,
+    {
+        let result = match self {
+            LevelEncoder::Rle(encoder) => {
+                let data = encoder.flush_buffer_mut();
+                // Patch the 4-byte length header reserved at the start of the buffer
+                let encoded_len = (data.len() - mem::size_of::<i32>()) as i32;
+                data[..4].copy_from_slice(&encoded_len.to_le_bytes());
+                f(data)
+            }
+            LevelEncoder::RleV2(encoder) => f(encoder.flush_buffer()),
+            LevelEncoder::BitPacked(_, encoder) => f(encoder.flush_buffer()),
+        };
+        match self {
+            LevelEncoder::Rle(encoder) => {
+                encoder.clear();
+                // Re-reserve the 4-byte length header for the next page
+                encoder.skip(mem::size_of::<i32>());
+            }
+            LevelEncoder::RleV2(encoder) => encoder.clear(),
+            LevelEncoder::BitPacked(_, encoder) => encoder.clear(),
+        }
+        result
+    }
 }

parquet/src/encodings/rle.rs

@@ -177,16 +177,22 @@ impl RleEncoder {
     /// Borrow equivalent of the `consume` method.
     /// Call `clear()` after invoking this method.
     #[inline]
-    #[allow(unused)]
     pub fn flush_buffer(&mut self) -> &[u8] {
         self.flush();
         self.bit_writer.flush_buffer()
     }
 
+    /// Like `flush_buffer`, but returns mutable access to the internal buffer.
+    /// Call `clear()` after invoking this method.
+    #[inline]
+    pub fn flush_buffer_mut(&mut self) -> &mut [u8] {
+        self.flush();
+        self.bit_writer.flush_buffer_mut()
+    }
+
     /// Clears the internal state so this encoder can be reused (e.g., after becoming
     /// full).
     #[inline]
-    #[allow(unused)]
     pub fn clear(&mut self) {
         self.bit_writer.clear();
         self.num_buffered_values = 0;
@@ -196,6 +202,13 @@ impl RleEncoder {
         self.indicator_byte_pos = -1;
     }
 
+    /// Advances the buffer by `num_bytes` zero bytes, delegating to the
+    /// underlying [`BitWriter::skip`].
+    #[inline]
+    pub fn skip(&mut self, num_bytes: usize) {
+        self.bit_writer.skip(num_bytes);
+    }
+
     /// Flushes all remaining values and return the final byte buffer maintained by the
     /// internal writer.
     #[inline]

parquet/src/util/bit_util.rs

@@ -219,6 +219,13 @@ impl BitWriter {
         self.buffer()
     }
 
+    /// Like `flush_buffer`, but returns mutable access to the buffer.
+    #[inline]
+    pub fn flush_buffer_mut(&mut self) -> &mut [u8] {
+        self.flush();
+        &mut self.buffer
+    }
+
     /// Clears the internal state so the buffer can be reused.
     #[inline]
     pub fn clear(&mut self) {