perf(parquet): splice buffered pages with write_all instead of io::copy (adapts #10052)

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<dyn PageStore>,
     /// Page handles in final file order: the dictionary page first (if any),
     /// then the data pages.
     keys: IntoIter<PageKey>,
-    /// 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<usize> {
-        // 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<Bytes>;
 
-        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<Self::Item> {
+        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)
     }
 }
 

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<R: 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<I>(
+        &mut self,
+        pages: I,
+        close: ColumnCloseResult,
+    ) -> Result<()>
+    where
+        I: IntoIterator<Item = Result<Bytes>>,
+    {
+        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())