This is an automated email from the ASF dual-hosted git repository.
alamb 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 8678fd6ad6 Minor: improve PageStore docs with a temp-file spilling
example (#10074)
8678fd6ad6 is described below
commit 8678fd6ad6e95a6a795982f22edcd4e1bbe380d5
Author: Andrew Lamb <[email protected]>
AuthorDate: Wed Jun 24 14:47:24 2026 -0400
Minor: improve PageStore docs with a temp-file spilling example (#10074)
# Which issue does this PR close?
- Follow-on to #10020, implementing the documentation suggestion from
https://github.com/apache/arrow-rs/pull/10020#discussion_r3350282205.
# Rationale for this change
The original `with_page_store_factory` example used a `HashMap`-backed
store with deliberately sparse keys to demonstrate that the writer
treats `PageKey`s as opaque
I think that was more of a unit test -- but it made the example longer
and harder to follow. A temp-file-backed store is both simpler to read
and much closer to what I think users would actually build, since
spilling pages off the heap is the whole point of the API.
# What changes are included in this PR?
Replaces the doctest with a `TempFilePageStore` (one temp file per
column chunk: `put` appends the page, `take` seeks it back) and reflows
the surrounding prose. Also tidies the `PageStore` trait docs to link
out to the example. This is documentation only — no code or behavior
changes.
# Are these changes tested?
Yes — the example is a runnable doctest that writes a record batch
through the spilling store and asserts the round-tripped data matches.
# Are there any user-facing changes?
Documentation only; no public API or behavior changes.
---
parquet/src/arrow/arrow_writer/mod.rs | 74 +++++++++++++++++++----------------
parquet/src/column/page_store.rs | 9 +++--
2 files changed, 47 insertions(+), 36 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index 0001a83160..063d5abcf1 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -582,22 +582,20 @@ impl ArrowWriterOptions {
/// Sets the [`PageStoreFactory`] used to buffer completed pages while a
row
/// group is being written.
///
- /// By default (an [`InMemoryPageStore`] per column chunk) completed pages
- /// are buffered on the heap until the row group is flushed, so peak memory
- /// grows with the row group size. Supplying a factory that spills to a
temp
- /// file or object storage instead bounds peak write memory, decoupling it
- /// from the row group size while keeping large, read-optimal row groups.
+ /// The default implementation ([`InMemoryPageStore`]) buffers all
completed
+ /// pages on the heap until the row group is flushed, so peak write memory
+ /// grows with the row group size. Using this API, pages can be spilled to
a
+ /// file or object storage instead, reducing peak write memory
substantially
+ /// at the expense of an extra write to and read from secondary storage.
///
- /// # Example: a custom [`PageStore`]
+ /// # Example: spilling pages to a temp file
///
- /// A store only has to map an opaque, store-allocated [`PageKey`] to a
blob
- /// and hand the blob back once. The keys need not be dense or sequential —
- /// here a `HashMap`-backed store mints sparse handles, proving the writer
- /// relies only on the opaque-handle contract. A real spilling backend
would
- /// write the bytes to a temp file in `put` and read them back in `take`.
+ /// A simple spilling backend uses one temp file per column chunk; `put`
+ /// appends the page and `take` reads it back.
///
/// ```
- /// # use std::collections::HashMap;
+ /// # use std::fs::File;
+ /// # use std::io::{Read, Seek, SeekFrom, Write};
/// # use std::sync::Arc;
/// # use bytes::Bytes;
/// # use arrow_array::{ArrayRef, Int64Array, RecordBatch};
@@ -605,54 +603,64 @@ impl ArrowWriterOptions {
/// # ArrowWriter, ArrowWriterOptions, PageKey, PageStore,
PageStoreArgs, PageStoreFactory,
/// # };
/// # use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
- /// # use parquet::errors::{ParquetError, Result};
- /// #[derive(Default)]
- /// struct MapPageStore {
- /// blobs: HashMap<u64, Bytes>,
- /// next: u64,
+ /// # use parquet::errors::Result;
+ /// struct TempFilePageStore {
+ /// file: File,
+ /// /// Total size of the file
+ /// end: u64,
+ /// /// Location of pages: (offset, len)
+ /// locs: Vec<(u64, usize)>,
/// }
///
- /// impl PageStore for MapPageStore {
+ /// impl PageStore for TempFilePageStore {
/// fn put(&mut self, value: Bytes) -> Result<PageKey> {
- /// // Mint a sparse handle (every other integer) to show the
writer
- /// // never assumes anything about the key's value.
- /// let key = PageKey::new(self.next);
- /// self.next += 2;
- /// self.blobs.insert(key.get(), value);
+ /// // Append to the end of the file
+ /// self.file.seek(SeekFrom::Start(self.end))?;
+ /// self.file.write_all(&value)?;
+ /// let key = PageKey::new(self.locs.len() as u64);
+ /// self.locs.push((self.end, value.len()));
+ /// self.end += value.len() as u64;
/// Ok(key)
/// }
///
/// fn take(&mut self, key: PageKey) -> Result<Bytes> {
- /// self.blobs
- /// .remove(&key.get())
- /// .ok_or_else(|| ParquetError::General(format!("invalid key
{}", key.get())))
+ /// let (offset, len) = self.locs[key.get() as usize];
+ /// let mut buf = vec![0u8; len];
+ /// self.file.seek(SeekFrom::Start(offset))?;
+ /// self.file.read_exact(&mut buf)?;
+ /// Ok(Bytes::from(buf))
/// }
/// }
///
+ /// /// Factory for creating [`TempFilePageStore`]
/// #[derive(Debug)]
- /// struct MapPageStoreFactory;
+ /// struct TempFilePageStoreFactory;
///
- /// impl PageStoreFactory for MapPageStoreFactory {
+ /// impl PageStoreFactory for TempFilePageStoreFactory {
/// fn create(&self, args: &PageStoreArgs<'_>) -> Result<Box<dyn
PageStore>> {
/// // `args` exposes the column index and descriptor
(physical/logical
- /// // type, path), so a real backend could spill only large
columns.
+ /// // type, path), so a real backend might choose to spill only
large columns.
/// let _ = (args.column_index(), args.column_descriptor());
- /// Ok(Box::new(MapPageStore::default()))
+ /// Ok(Box::new(TempFilePageStore {
+ /// file: tempfile::tempfile()?, // temp file is cleaned on
drop
+ /// end: 0,
+ /// locs: Vec::new(),
+ /// }))
/// }
/// }
- ///
+ /// // write 1000 integers
/// let col = Arc::new(Int64Array::from_iter_values(0..1000)) as ArrayRef;
/// let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap();
///
/// let options =
- ///
ArrowWriterOptions::new().with_page_store_factory(Arc::new(MapPageStoreFactory));
+ ///
ArrowWriterOptions::new().with_page_store_factory(Arc::new(TempFilePageStoreFactory));
/// let mut buffer = Vec::new();
/// let mut writer =
/// ArrowWriter::try_new_with_options(&mut buffer, to_write.schema(),
options).unwrap();
/// writer.write(&to_write).unwrap();
/// writer.close().unwrap();
///
- /// // The file is byte-identical to one written with the default store.
+ /// // buffer now holds valid Parquet data, which can be read as normal:
/// let mut reader =
ParquetRecordBatchReader::try_new(Bytes::from(buffer), 1024).unwrap();
/// assert_eq!(to_write, reader.next().unwrap().unwrap());
/// ```
diff --git a/parquet/src/column/page_store.rs b/parquet/src/column/page_store.rs
index 4f821c0e5c..14bb2b9e8b 100644
--- a/parquet/src/column/page_store.rs
+++ b/parquet/src/column/page_store.rs
@@ -70,9 +70,12 @@ impl PageKey {
/// thread at a time (both methods take `&mut self`), so it needs no internal
/// synchronization — hence only `Send`, not `Sync`.
///
-/// The default ([`InMemoryPageStore`]) keeps blobs on the heap. Configure a
-/// different backend via
-///
[`ArrowWriterOptions::with_page_store_factory`](crate::arrow::arrow_writer::ArrowWriterOptions::with_page_store_factory).
+/// The default ([`InMemoryPageStore`]) keeps blobs in memory on the heap.
+///
+/// For an example of configuring the Parquet writer to use an alternate
+/// `PageStore` see the [`ArrowWriterOptions::with_page_store_factory`] API.
+///
+/// [`ArrowWriterOptions::with_page_store_factory`]:
crate::arrow::arrow_writer::ArrowWriterOptions::with_page_store_factory
pub trait PageStore: Send {
/// Store `value`, returning a handle that can later be passed to
/// [`take`](Self::take).