diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 37d85cf7da78..7dd7fb778909 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -20,7 +20,7 @@ use crate::column::chunker::ContentDefinedChunker; use bytes::Bytes; -use std::io::{Read, Write}; +use std::io::Write; use std::slice::Iter; use std::sync::{Arc, Mutex}; use std::vec::IntoIter; @@ -769,24 +769,22 @@ impl ArrowColumnChunkData { } } -/// A streaming [`Read`] over one column chunk's buffered pages, in final file -/// order: the dictionary page (if any) first, then the data pages. +/// A streaming iterator over one column chunk's buffered page blobs, in final +/// file order: the dictionary page (if any) first, then the data pages. /// /// Each blob is taken back out of the [`PageStore`] *as it is /// consumed* and released immediately afterwards, so splicing a chunk into the /// output file never materializes more than a single page in memory at a time. /// This is what keeps the splice phase within the memory bound for a spilling /// backend (an in-memory store already holds the bytes, so it is unaffected). -struct StreamingColumnChunkReader { +struct StreamingColumnChunkPages { store: Box, /// Page handles in final file order: the dictionary page first (if any), /// then the data pages. keys: IntoIter, - /// The blob currently being drained into the output; emptied as it is read. - current: Bytes, } -impl StreamingColumnChunkReader { +impl StreamingColumnChunkPages { fn new(data: ArrowColumnChunkData) -> Self { // The dictionary page must be emitted first, ahead of the data pages, // even though it was the last page produced. @@ -801,27 +799,16 @@ impl StreamingColumnChunkReader { Self { store: data.store, keys: keys.into_iter(), - current: Bytes::new(), } } } -impl Read for StreamingColumnChunkReader { - fn read(&mut self, out: &mut [u8]) -> std::io::Result { - // Refill from the next blob whenever the current one is drained: the - // dictionary page first, then each data page, all taken from the store. - while self.current.is_empty() { - if let Some(key) = self.keys.next() { - self.current = self.store.take(key).map_err(std::io::Error::other)?; - } else { - return Ok(0); - } - } +impl Iterator for StreamingColumnChunkPages { + type Item = Result; - let len = self.current.len().min(out.len()); - let b = self.current.split_to(len); - out[..len].copy_from_slice(&b); - Ok(len) + fn next(&mut self) -> Option { + let key = self.keys.next()?; + Some(self.store.take(key)) } } @@ -998,8 +985,8 @@ impl ArrowColumnChunk { // it ahead of the data pages in the recorded offsets before the splice. let close = close.update_dictionary_location(data.dictionary_len)?; - let reader = StreamingColumnChunkReader::new(data); - writer.append_column_from_read(reader, close) + let pages = StreamingColumnChunkPages::new(data); + writer.append_column_from_pages(pages, close) } } diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs index 8ec16ba36739..b62d8886cd8d 100644 --- a/parquet/src/file/writer.rs +++ b/parquet/src/file/writer.rs @@ -22,6 +22,8 @@ use crate::file::metadata::thrift::PageHeader; use crate::file::page_index::column_index::ColumnIndexMetaData; use crate::file::page_index::offset_index::OffsetIndexMetaData; use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift}; +#[cfg(feature = "arrow")] +use bytes::Bytes; use std::fmt::Debug; use std::io::{BufWriter, IoSlice, Read}; use std::{io::Write, sync::Arc}; @@ -708,14 +710,75 @@ impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> { pub(crate) fn append_column_from_read( &mut self, read: R, - mut close: ColumnCloseResult, + close: ColumnCloseResult, ) -> Result<()> { + let (src_offset, src_length, write_offset) = self.begin_appended_column(&close)?; + + let mut read = read.take(src_length as _); + let write_length = std::io::copy(&mut read, &mut self.buf)?; + + if src_length as u64 != write_length { + return Err(general_err!( + "Failed to splice column data, expected {src_length} got {write_length}" + )); + } + + self.finish_appended_column(close, src_offset, write_offset) + } + + /// Splice an already-encoded column chunk into the row group from an + /// in-order sequence of byte buffers (typically its serialized pages). + /// + /// This is a lower-overhead alternative to [`Self::append_column`] / + /// [`Self::append_column_from_read`] for callers that already hold the + /// chunk as owned [`Bytes`]: each buffer is written straight to the output + /// with a single `write_all`, skipping the intermediate copy through + /// [`std::io::copy`]'s fixed-size buffer. + /// + /// `pages` must yield the chunk's compressed bytes in final file order + /// (the dictionary page, if any, first) and together total exactly the + /// compressed size recorded in `close`. + #[cfg(feature = "arrow")] + pub(crate) fn append_column_from_pages( + &mut self, + pages: I, + close: ColumnCloseResult, + ) -> Result<()> + where + I: IntoIterator>, + { + let (src_offset, src_length, write_offset) = self.begin_appended_column(&close)?; + + let mut write_length = 0u64; + for page in pages { + let page = page?; + self.buf.write_all(&page)?; + write_length += page.len() as u64; + } + + if src_length as u64 != write_length { + return Err(general_err!( + "Failed to splice column data, expected {src_length} got {write_length}" + )); + } + + self.finish_appended_column(close, src_offset, write_offset) + } + + /// [`Self::append_column_from_read`] / [`Self::append_column_from_pages`] + /// preamble: validates the writer state and that `close` matches the next + /// expected column. + /// + /// Returns `(src_offset, src_length, write_offset)`: the chunk's start + /// offset and length in the source buffer, and the offset at which it will + /// land in the output file. + fn begin_appended_column(&mut self, close: &ColumnCloseResult) -> Result<(i64, i64, usize)> { self.assert_previous_writer_closed()?; let desc = self .next_column_desc() .ok_or_else(|| general_err!("exhausted columns in SerializedRowGroupWriter"))?; - let metadata = close.metadata; + let metadata = &close.metadata; if metadata.column_descr() != desc.as_ref() { return Err(general_err!( @@ -725,20 +788,26 @@ impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> { )); } - let src_dictionary_offset = metadata.dictionary_page_offset(); - let src_data_offset = metadata.data_page_offset(); - let src_offset = src_dictionary_offset.unwrap_or(src_data_offset); + let src_offset = metadata + .dictionary_page_offset() + .unwrap_or_else(|| metadata.data_page_offset()); let src_length = metadata.compressed_size(); - let write_offset = self.buf.bytes_written(); - let mut read = read.take(src_length as _); - let write_length = std::io::copy(&mut read, &mut self.buf)?; + Ok((src_offset, src_length, write_offset)) + } - if src_length as u64 != write_length { - return Err(general_err!( - "Failed to splice column data, expected {read_length} got {write_length}" - )); - } + /// [`Self::append_column_from_read`] / [`Self::append_column_from_pages`] + /// epilogue: rewrites the buffer-relative page offsets recorded in `close` + /// to their final positions in the output file and closes the column. + fn finish_appended_column( + &mut self, + mut close: ColumnCloseResult, + src_offset: i64, + write_offset: usize, + ) -> Result<()> { + let metadata = close.metadata; + let src_dictionary_offset = metadata.dictionary_page_offset(); + let src_data_offset = metadata.data_page_offset(); let map_offset = |x| x - src_offset + write_offset as i64; let mut builder = ColumnChunkMetaData::builder(metadata.column_descr_ptr())