diff --git a/parquet/src/column/reader.rs b/parquet/src/column/reader.rs index 29cb50185a58..d1a328237f26 100644 --- a/parquet/src/column/reader.rs +++ b/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(); diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 4c3dbabc2132..cdf489f3b619 100644 --- a/parquet/src/column/writer/mod.rs +++ b/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_stats: Vec, - // Reused buffers - def_levels_sink: Vec, - rep_levels_sink: Vec, + // Streaming level encoders for definition/repetition levels. + def_levels_encoder: LevelEncoder, + rep_levels_encoder: LevelEncoder, data_pages: VecDeque, // column index and offset index column_index_builder: ColumnIndexBuilder, @@ -411,6 +411,8 @@ 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, @@ -418,8 +420,6 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { 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 { - 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 { - 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<()> { diff --git a/parquet/src/encodings/levels.rs b/parquet/src/encodings/levels.rs index 8c95e5ca51aa..b761a5ac5daf 100644 --- a/parquet/src/encodings/levels.rs +++ b/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 { 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(&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::()) 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::()); + } + LevelEncoder::RleV2(encoder) => encoder.clear(), + LevelEncoder::BitPacked(_, encoder) => encoder.clear(), + } + result + } } diff --git a/parquet/src/encodings/rle.rs b/parquet/src/encodings/rle.rs index c95a46c634d2..2815c20dab56 100644 --- a/parquet/src/encodings/rle.rs +++ b/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] diff --git a/parquet/src/util/bit_util.rs b/parquet/src/util/bit_util.rs index 262564825864..7d7907f6f543 100644 --- a/parquet/src/util/bit_util.rs +++ b/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) { diff --git a/parquet/src/util/test_common/page_util.rs b/parquet/src/util/test_common/page_util.rs index 5b64eb54133c..6a99beaea151 100644 --- a/parquet/src/util/test_common/page_util.rs +++ b/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)