Alamb/test encoding speed

parquet/src/column/reader.rs

@@ -1401,11 +1401,11 @@ mod tests {
         // Helper: build a DataPage v2 for this list column.
         let make_v2_page =
             |rep_levels: &[i16], def_levels: &[i16], values: &[i32], num_rows: u32| -> Page {
-                let mut rep_enc = LevelEncoder::v2(max_rep_level, rep_levels.len());
+                let mut rep_enc = LevelEncoder::v2_streaming(max_rep_level);
                 rep_enc.put(rep_levels);
                 let rep_bytes = rep_enc.consume();
 
-                let mut def_enc = LevelEncoder::v2(max_def_level, def_levels.len());
+                let mut def_enc = LevelEncoder::v2_streaming(max_def_level);
                 def_enc.put(def_levels);
                 let def_bytes = def_enc.consume();
 

parquet/src/column/writer/mod.rs

@@ -351,9 +351,9 @@ pub struct GenericColumnWriter<'a, E: ColumnValueEncoder> {
     /// but we use a BTreeSet so that the output is deterministic
     encodings: BTreeSet<Encoding>,
     encoding_stats: Vec<PageEncodingStats>,
-    // Reused buffers
-    def_levels_sink: Vec<i16>,
-    rep_levels_sink: Vec<i16>,
+    // Streaming level encoders for definition/repetition levels.
+    def_levels_encoder: LevelEncoder,
+    rep_levels_encoder: LevelEncoder,
     data_pages: VecDeque<CompressedPage>,
     // column index and offset index
     column_index_builder: ColumnIndexBuilder,
@@ -411,15 +411,15 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
         };
 
         Self {
+            def_levels_encoder: Self::create_level_encoder(descr.max_def_level(), &props),
+            rep_levels_encoder: Self::create_level_encoder(descr.max_rep_level(), &props),
             descr,
             props,
             statistics_enabled,
             page_writer,
             codec,
             compressor,
             encoder,
-            def_levels_sink: vec![],
-            rep_levels_sink: vec![],
             data_pages: VecDeque::new(),
             page_metrics,
             column_metrics,
@@ -647,6 +647,14 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
         })
     }
 
+    /// Creates a new streaming level encoder appropriate for the writer version.
+    fn create_level_encoder(max_level: i16, props: &WriterProperties) -> LevelEncoder {
+        match props.writer_version() {
+            WriterVersion::PARQUET_1_0 => LevelEncoder::v1_streaming(Encoding::RLE, max_level),
+            WriterVersion::PARQUET_2_0 => LevelEncoder::v2_streaming(max_level),
+        }
+    }
+
     /// Writes mini batch of values, definition and repetition levels.
     /// This allows fine-grained processing of values and maintaining a reasonable
     /// page size.
@@ -677,7 +685,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
             // Update histogram
             self.page_metrics.update_definition_level_histogram(levels);
 
-            self.def_levels_sink.extend_from_slice(levels);
+            self.def_levels_encoder.put(levels);
             values_to_write
         } else {
             num_levels
@@ -708,7 +716,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
             // Update histogram
             self.page_metrics.update_repetition_level_histogram(levels);
 
-            self.rep_levels_sink.extend_from_slice(levels);
+            self.rep_levels_encoder.put(levels);
         } else {
             // Each value is exactly one row.
             // Equals to the number of values, we count nulls as well.
@@ -1060,23 +1068,13 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
                 let mut buffer = vec![];
 
                 if max_rep_level > 0 {
-                    buffer.extend_from_slice(
-                        &self.encode_levels_v1(
-                            Encoding::RLE,
-                            &self.rep_levels_sink[..],
-                            max_rep_level,
-                        )[..],
-                    );
+                    self.rep_levels_encoder
+                        .flush_to(|data| buffer.extend_from_slice(data));
                 }
 
                 if max_def_level > 0 {
-                    buffer.extend_from_slice(
-                        &self.encode_levels_v1(
-                            Encoding::RLE,
-                            &self.def_levels_sink[..],
-                            max_def_level,
-                        )[..],
-                    );
+                    self.def_levels_encoder
+                        .flush_to(|data| buffer.extend_from_slice(data));
                 }
 
                 buffer.extend_from_slice(&values_data.buf);
@@ -1106,15 +1104,15 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
                 let mut buffer = vec![];
 
                 if max_rep_level > 0 {
-                    let levels = self.encode_levels_v2(&self.rep_levels_sink[..], max_rep_level);
-                    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 levels = self.encode_levels_v2(&self.def_levels_sink[..], max_def_level);
-                    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 =
@@ -1164,10 +1162,6 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
 
         // Update total number of rows.
         self.column_metrics.total_rows_written += self.page_metrics.num_buffered_rows as u64;
-
-        // Reset state.
-        self.rep_levels_sink.clear();
-        self.def_levels_sink.clear();
         self.page_metrics.new_page();
 
         Ok(())
@@ -1244,23 +1238,6 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
         Ok(metadata)
     }
 
-    /// Encodes definition or repetition levels for Data Page v1.
-    #[inline]
-    fn encode_levels_v1(&self, encoding: Encoding, levels: &[i16], max_level: i16) -> Vec<u8> {
-        let mut encoder = LevelEncoder::v1(encoding, max_level, levels.len());
-        encoder.put(levels);
-        encoder.consume()
-    }
-
-    /// Encodes definition or repetition levels for Data Page v2.
-    /// Encoding is always RLE.
-    #[inline]
-    fn encode_levels_v2(&self, levels: &[i16], max_level: i16) -> Vec<u8> {
-        let mut encoder = LevelEncoder::v2(max_level, levels.len());
-        encoder.put(levels);
-        encoder.consume()
-    }
-
     /// Writes compressed data page into underlying sink and updates global metrics.
     #[inline]
     fn write_data_page(&mut self, page: CompressedPage) -> Result<()> {

parquet/src/encodings/levels.rs

@@ -21,21 +21,7 @@ use super::rle::RleEncoder;
 
 use crate::basic::Encoding;
 use crate::data_type::AsBytes;
-use crate::util::bit_util::{BitWriter, ceil, num_required_bits};
-
-/// Computes max buffer size for level encoder/decoder based on encoding, max
-/// repetition/definition level and number of total buffered values (includes null
-/// values).
-#[inline]
-pub fn max_buffer_size(encoding: Encoding, max_level: i16, num_buffered_values: usize) -> usize {
-    let bit_width = num_required_bits(max_level as u64);
-    match encoding {
-        Encoding::RLE => RleEncoder::max_buffer_size(bit_width, num_buffered_values),
-        #[allow(deprecated)]
-        Encoding::BIT_PACKED => ceil(num_buffered_values * bit_width as usize, 8),
-        _ => panic!("Unsupported encoding type {encoding}"),
-    }
-}
+use crate::util::bit_util::{BitWriter, num_required_bits};
 
 /// Encoder for definition/repetition levels.
 /// Currently only supports Rle and BitPacked (dev/null) encoding, including v2.
@@ -46,46 +32,44 @@ pub enum LevelEncoder {
 }
 
 impl LevelEncoder {
-    /// Creates new level encoder based on encoding, max level and underlying byte buffer.
-    /// For bit packed encoding it is assumed that buffer is already allocated with
-    /// `levels::max_buffer_size` method.
-    ///
-    /// Used to encode levels for Data Page v1.
+    /// Creates a new streaming level encoder for Data Page v1.
     ///
-    /// Panics, if encoding is not supported.
-    pub fn v1(encoding: Encoding, max_level: i16, capacity: usize) -> Self {
-        let capacity_bytes = max_buffer_size(encoding, max_level, capacity);
-        let mut buffer = Vec::with_capacity(capacity_bytes);
+    /// Unlike [`v1`](Self::v1), this does not require knowing the number of values
+    /// upfront, making it suitable for incremental encoding where levels are fed in
+    /// as they arrive via [`put`](Self::put).
+    pub fn v1_streaming(encoding: Encoding, max_level: i16) -> Self {
         let bit_width = num_required_bits(max_level as u64);
         match encoding {
             Encoding::RLE => {
                 // Reserve space for length header
-                buffer.extend_from_slice(&[0; 4]);
+                let buffer = vec![0u8; 4];
                 LevelEncoder::Rle(RleEncoder::new_from_buf(bit_width, buffer))
             }
             #[allow(deprecated)]
             Encoding::BIT_PACKED => {
-                // Here we set full byte buffer without adjusting for num_buffered_values,
-                // because byte buffer will already be allocated with size from
-                // `max_buffer_size()` method.
-                LevelEncoder::BitPacked(bit_width, BitWriter::new_from_buf(buffer))
+                LevelEncoder::BitPacked(bit_width, BitWriter::new_from_buf(Vec::new()))
             }
             _ => panic!("Unsupported encoding type {encoding}"),
         }
     }
 
-    /// Creates new level encoder based on RLE encoding. Used to encode Data Page v2
-    /// repetition and definition levels.
-    pub fn v2(max_level: i16, capacity: usize) -> Self {
-        let capacity_bytes = max_buffer_size(Encoding::RLE, max_level, capacity);
-        let buffer = Vec::with_capacity(capacity_bytes);
+    /// Creates a new streaming RLE level encoder for Data Page v2.
+    ///
+    /// Unlike [`v2`](Self::v2), this does not require knowing the number of values
+    /// upfront, making it suitable for incremental encoding where levels are fed in
+    /// as they arrive via [`put`](Self::put).
+    pub fn v2_streaming(max_level: i16) -> Self {
         let bit_width = num_required_bits(max_level as u64);
-        LevelEncoder::RleV2(RleEncoder::new_from_buf(bit_width, buffer))
+        LevelEncoder::RleV2(RleEncoder::new_from_buf(bit_width, Vec::new()))
     }
 
     /// Put/encode levels vector into this level encoder.
     /// Returns number of encoded values that are less than or equal to length of the
     /// input buffer.
+    ///
+    /// This method does **not** flush the underlying encoder, so it can be called
+    /// incrementally across multiple batches without forcing run boundaries.
+    /// The encoder is flushed automatically when [`consume`](Self::consume) is called.
     #[inline]
     pub fn put(&mut self, buffer: &[i16]) -> usize {
         let mut num_encoded = 0;
@@ -95,14 +79,12 @@ impl LevelEncoder {
                     encoder.put(*value as u64);
                     num_encoded += 1;
                 }
-                encoder.flush();
             }
             LevelEncoder::BitPacked(bit_width, ref mut encoder) => {
                 for value in buffer {
                     encoder.put_value(*value as u64, bit_width as usize);
                     num_encoded += 1;
                 }
-                encoder.flush();
             }
         }
         num_encoded
@@ -111,6 +93,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) => {
@@ -126,4 +109,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) {

parquet/src/util/test_common/page_util.rs

@@ -75,7 +75,7 @@ impl DataPageBuilderImpl {
         if max_level <= 0 {
             return 0;
         }
-        let mut level_encoder = LevelEncoder::v1(Encoding::RLE, max_level, levels.len());
+        let mut level_encoder = LevelEncoder::v1_streaming(Encoding::RLE, max_level);
         level_encoder.put(levels);
         let encoded_levels = level_encoder.consume();
         // Actual encoded bytes (without length offset)