This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 4aee48acaf perf(parquet): splice buffered pages with `write_all`
instead of `io::copy` (adapts #10052) (#10353)
4aee48acaf is described below
commit 4aee48acaf4f37233db478650fd60aab238ae12b
Author: Adrian Garcia Badaracco <[email protected]>
AuthorDate: Sat Jul 18 21:19:50 2026 -0500
perf(parquet): splice buffered pages with `write_all` instead of `io::copy`
(adapts #10052) (#10353)
# Which issue does this PR close?
- Adapts #10052 (by @alamb) to the current `PageStore`-based writer.
That PR predates #10020 / #10142 and no longer applies cleanly; this PR
carries its idea forward. Closes #10052's use case.
- closes https://github.com/apache/arrow-rs/pull/10052
# Rationale for this change
When the `ArrowWriter` flushes a row group, each column chunk's buffered
pages are spliced into the output through `StreamingColumnChunkReader`
(a `Read`) driven by `std::io::copy`. Every byte is copied twice — once
into `io::copy`'s fixed 8 KiB buffer, then again into `TrackedWrite`'s
`BufWriter` — in ~`len / 8 KiB` write calls.
The page blobs are already contiguous owned `Bytes` sitting in the
`PageStore`; we can write each one straight to the sink with a single
`write_all` (large writes bypass the `BufWriter` buffer entirely).
#10052 did exactly this, but its `IntoIterator<Item = Bytes>` signature
can't express the fallible, one-page-at-a-time drain of a (possibly
spilling) `PageStore`, and it assumed the pre-#10020 `Vec<Bytes>` chunk
representation.
# What changes are included in this PR?
- `StreamingColumnChunkReader` (`Read` impl) becomes
`StreamingColumnChunkPages`, an `Iterator<Item = Result<Bytes>>` that
`take()`s each page out of the `PageStore` as it is consumed —
preserving the one-page-in-memory bound for spilling backends.
- New `pub(crate) SerializedRowGroupWriter::append_column_from_pages`,
which writes each page blob with a single `write_all`.
- The shared validation / offset-rewriting code is factored into
`begin_appended_column` / `finish_appended_column` (same refactor as
#10052), also used by the existing `Read`-based path that still backs
the public `append_column`.
# Are these changes tested?
Covered by existing tests. Also verified that the output files are
byte-for-byte identical to main for multi-row-group writes with int /
string / dictionary columns under both default and small-page
properties.
Local benchmark results (`arrow_writer` bench, Apple Silicon),
reproducing the wins from #10052's benchmark run without its
`large_string_non_null` regression:
| benchmark | main | this PR | change |
|---|---|---|---|
| `string_dictionary/default` | 36.9 ms | 27.9 ms | ~24% faster |
| `list_primitive/default` | 179.8 ms | 150.1 ms | ~17% faster |
| `large_string_non_null/default` | 31.4 ms | 30.1 ms | ~4% faster |
# Are there any user-facing changes?
No. `append_column_from_pages` is kept `pub(crate)` for now; it could be
made public in a follow-up if the `Result<Bytes>` iterator contract is
deemed a good public API (#10052 proposed a public variant).
🤖 Generated with [Claude Code](https://claude.com/claude-code)
https://claude.ai/code/session_01Sd8JjF3qQMkn7htofu8bF3
---------
Co-authored-by: Claude Fable 5 <[email protected]>
---
parquet/src/arrow/arrow_writer/mod.rs | 37 +++++---------
parquet/src/file/writer.rs | 95 ++++++++++++++++++++++++++++++-----
2 files changed, 94 insertions(+), 38 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index 37d85cf7da..7dd7fb7789 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<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)
}
}
diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs
index 8ec16ba367..b62d8886cd 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<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())
