From d9725b8a5e4ac962448032c40d448d5e883fbb80 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 26 May 2026 10:13:52 -0500 Subject: [PATCH 01/13] parquet: add pluggable PageStore for ArrowWriter page buffering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce a "dumb" key/value page store that the ArrowWriter uses to buffer completed, serialized pages while a row group is being written. The store maps an opaque, store-allocated PageKey to a blob of bytes and knows nothing about pages, dictionaries, ordering, or offsets — the caller keeps the handles and decides what they mean. The default InMemoryPageStore keeps blobs in a Vec, byte-for-byte equivalent to the previous buffering with zero overhead. A PageStoreFactory is threaded through ArrowWriterOptions -> ArrowRowGroupWriterFactory -> ArrowColumnWriterFactory so users can plug in a backend (temp file, object storage) to bound peak write memory independently of row group size. ArrowColumnChunkData now holds (store, keys) and materializes blobs in write order at splice time, preserving the existing append_column path. Tests: - column::page_store unit tests for the in-memory backend contract. - A byte-identical round-trip test using a custom HashMap-backed store with sparse, non-contiguous handles, proving the writer relies only on the opaque-handle contract. - An always-on dhat integration test capturing the in-memory peak-heap baseline (memory grows with the row group), against which a spilling backend will be measured. Co-Authored-By: Claude Opus 4.7 (1M context) --- Cargo.lock | 104 ++++++++++-- parquet/Cargo.toml | 1 + parquet/src/arrow/arrow_writer/mod.rs | 230 ++++++++++++++++++++++++-- parquet/src/column/mod.rs | 1 + parquet/src/column/page_store.rs | 158 ++++++++++++++++++ parquet/tests/page_spill_memory.rs | 158 ++++++++++++++++++ 6 files changed, 625 insertions(+), 27 deletions(-) create mode 100644 parquet/src/column/page_store.rs create mode 100644 parquet/tests/page_spill_memory.rs diff --git a/Cargo.lock b/Cargo.lock index 28fe6a43a5c5..a67498568b5b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "addr2line" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b5d307320b3181d6d7954e663bd7c774a838b8220fe0593c86d9fb09f498b4b" +dependencies = [ + "gimli", +] + [[package]] name = "adler2" version = "2.0.1" @@ -106,7 +115,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -117,7 +126,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -663,6 +672,21 @@ dependencies = [ "tower-service", ] +[[package]] +name = "backtrace" +version = "0.3.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb531853791a215d7c62a30daf0dde835f381ab5de4589cfe7c649d2cbe92bd6" +dependencies = [ + "addr2line", + "cfg-if", + "libc", + "miniz_oxide", + "object", + "rustc-demangle", + "windows-link", +] + [[package]] name = "base64" version = "0.22.1" @@ -1167,6 +1191,22 @@ dependencies = [ "syn", ] +[[package]] +name = "dhat" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98cd11d84628e233de0ce467de10b8633f4ddaecafadefc86e13b84b8739b827" +dependencies = [ + "backtrace", + "lazy_static", + "mintex", + "parking_lot", + "rustc-hash 1.1.0", + "serde", + "serde_json", + "thousands", +] + [[package]] name = "difflib" version = "0.4.0" @@ -1248,7 +1288,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -1484,6 +1524,12 @@ dependencies = [ "wasip3", ] +[[package]] +name = "gimli" +version = "0.32.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" + [[package]] name = "h2" version = "0.4.14" @@ -2100,6 +2146,12 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "mintex" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c505b3e17ed6b70a7ed2e67fbb2c560ee327353556120d6e72f5232b6880d536" + [[package]] name = "mio" version = "1.2.0" @@ -2132,7 +2184,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -2215,6 +2267,15 @@ dependencies = [ "objc2-core-foundation", ] +[[package]] +name = "object" +version = "0.37.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" +dependencies = [ + "memchr", +] + [[package]] name = "object_store" version = "0.13.2" @@ -2331,6 +2392,7 @@ dependencies = [ "clap", "crc32fast", "criterion", + "dhat", "flate2", "futures", "half", @@ -2712,7 +2774,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash", + "rustc-hash 2.1.2", "rustls", "socket2", "thiserror 2.0.18", @@ -2732,7 +2794,7 @@ dependencies = [ "lru-slab", "rand 0.9.4", "ring", - "rustc-hash", + "rustc-hash 2.1.2", "rustls", "rustls-pki-types", "slab", @@ -2933,6 +2995,18 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rustc-demangle" +version = "0.1.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b50b8869d9fc858ce7266cce0194bd74df58b9d0e3f6df3a9fc8eb470d95c09d" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc-hash" version = "2.1.2" @@ -2958,7 +3032,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3222,7 +3296,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -3331,10 +3405,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3344,7 +3418,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" dependencies = [ "rustix", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -3393,6 +3467,12 @@ dependencies = [ "syn", ] +[[package]] +name = "thousands" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3bf63baf9f5039dadc247375c29eb13706706cfde997d0330d05aa63a77d8820" + [[package]] name = "thread_local" version = "1.1.9" @@ -4018,7 +4098,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index dd2c872ede50..b5430d6fe620 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -94,6 +94,7 @@ tokio = { version = "1.0", default-features = false, features = ["macros", "rt-m rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] } object_store = { workspace = true, features = ["azure", "fs"] } sysinfo = { version = "0.38.1", default-features = false, features = ["system"] } +dhat = { version = "0.3", default-features = false } [package.metadata.docs.rs] all-features = true diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 79542caed9b7..3e0cef4f0ef7 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -58,6 +58,11 @@ use levels::{ArrayLevels, calculate_array_levels}; mod byte_array; mod levels; +#[doc(inline)] +pub use crate::column::page_store::{ + InMemoryPageStore, InMemoryPageStoreFactory, PageKey, PageStore, PageStoreFactory, +}; + /// Encodes [`RecordBatch`] to parquet /// /// Writes Arrow `RecordBatch`es to a Parquet writer. Multiple [`RecordBatch`] will be encoded @@ -263,8 +268,12 @@ impl ArrowWriter { let file_writer = SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::clone(&props_ptr))?; - let row_group_writer_factory = + let mut row_group_writer_factory = ArrowRowGroupWriterFactory::new(&file_writer, arrow_schema.clone()); + if let Some(page_store_factory) = options.page_store_factory { + row_group_writer_factory = + row_group_writer_factory.with_page_store_factory(page_store_factory); + } let cdc_chunkers = props_ptr .content_defined_chunking() @@ -556,6 +565,7 @@ pub struct ArrowWriterOptions { skip_arrow_metadata: bool, schema_root: Option, schema_descr: Option, + page_store_factory: Option>, } impl ArrowWriterOptions { @@ -569,6 +579,21 @@ impl ArrowWriterOptions { Self { properties, ..self } } + /// 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. + pub fn with_page_store_factory(self, page_store_factory: Arc) -> Self { + Self { + page_store_factory: Some(page_store_factory), + ..self + } + } + /// Skip encoding the embedded arrow metadata (defaults to `false`) /// /// Parquet files generated by the [`ArrowWriter`] contain embedded arrow schema @@ -603,20 +628,67 @@ impl ArrowWriterOptions { } } -/// A single column chunk produced by [`ArrowColumnWriter`] -#[derive(Default)] +/// A single column chunk produced by [`ArrowColumnWriter`]. +/// +/// Holds the serialized page blobs (each page's header ‖ compressed data, in +/// write order) in a [`PageStore`], plus the handles needed to read them back, +/// in order, when the chunk is spliced into the output file. struct ArrowColumnChunkData { + length: usize, + store: Box, + keys: Vec, +} + +impl ArrowColumnChunkData { + fn new(store: Box) -> Self { + Self { + length: 0, + store, + keys: Vec::new(), + } + } + + /// Append a serialized blob to the store, recording its handle in write + /// order. + fn push(&mut self, value: Bytes) -> Result<()> { + let key = self.store.put(value)?; + self.keys.push(key); + Ok(()) + } + + /// Take every buffered blob back out of the store, in write order, ready to + /// splice into the output file. + /// + /// TODO(spill): stream blobs one at a time into the splice instead of + /// collecting them, so the splice phase is also bounded for spilling + /// backends (it currently materializes one column chunk at a time). + fn drain(&mut self) -> Result { + let mut data = Vec::with_capacity(self.keys.len()); + for key in std::mem::take(&mut self.keys) { + data.push(self.store.take(key)?); + } + Ok(MaterializedColumnChunk { + length: self.length, + data, + }) + } +} + +/// The buffered pages of one column chunk, taken back out of the [`PageStore`] +/// in write order and ready to be spliced into the output file via +/// [`SerializedRowGroupWriter::append_column`]. +struct MaterializedColumnChunk { length: usize, data: Vec, } -impl Length for ArrowColumnChunkData { +impl Length for MaterializedColumnChunk { fn len(&self) -> u64 { self.length as _ } } -impl ChunkReader for ArrowColumnChunkData { +impl ChunkReader for MaterializedColumnChunk { type T = ArrowColumnChunkReader; fn get_read(&self, start: u64) -> Result { @@ -631,7 +703,7 @@ impl ChunkReader for ArrowColumnChunkData { } } -/// A [`Read`] for [`ArrowColumnChunkData`] +/// A [`Read`] for [`MaterializedColumnChunk`] struct ArrowColumnChunkReader(Peekable>); impl Read for ArrowColumnChunkReader { @@ -660,7 +732,6 @@ impl Read for ArrowColumnChunkReader { /// [`ArrowRowGroupWriter`] on flush, without requiring self-referential borrows type SharedColumnChunk = Arc>; -#[derive(Default)] struct ArrowPageWriter { buffer: SharedColumnChunk, #[cfg(feature = "encryption")] @@ -668,6 +739,15 @@ struct ArrowPageWriter { } impl ArrowPageWriter { + /// Create a page writer that buffers completed pages in `store`. + fn new(store: Box) -> Self { + Self { + buffer: Arc::new(Mutex::new(ArrowColumnChunkData::new(store))), + #[cfg(feature = "encryption")] + page_encryptor: None, + } + } + #[cfg(feature = "encryption")] pub fn with_encryptor(mut self, page_encryptor: Option) -> Self { self.page_encryptor = page_encryptor; @@ -726,8 +806,8 @@ impl PageWriter for ArrowPageWriter { spec.bytes_written = compressed_size as u64; buf.length += compressed_size; - buf.data.push(header); - buf.data.push(data); + buf.push(header)?; + buf.push(data)?; Ok(spec) } @@ -787,10 +867,11 @@ impl ArrowColumnChunk { /// Calls [`SerializedRowGroupWriter::append_column`] with this column's data pub fn append_to_row_group( - self, + mut self, writer: &mut SerializedRowGroupWriter<'_, W>, ) -> Result<()> { - writer.append_column(&self.data, self.close) + let materialized = self.data.drain()?; + writer.append_column(&materialized, self.close) } } @@ -1082,6 +1163,7 @@ pub struct ArrowRowGroupWriterFactory { schema: SchemaDescPtr, arrow_schema: SchemaRef, props: WriterPropertiesPtr, + page_store_factory: Arc, #[cfg(feature = "encryption")] file_encryptor: Option>, } @@ -1098,11 +1180,20 @@ impl ArrowRowGroupWriterFactory { schema, arrow_schema, props, + page_store_factory: Arc::new(InMemoryPageStoreFactory), #[cfg(feature = "encryption")] file_encryptor: file_writer.file_encryptor(), } } + /// Set the [`PageStoreFactory`] used to allocate the buffer for each column + /// chunk, e.g. to spill completed pages to a temp file or object storage + /// instead of the heap. Defaults to [`InMemoryPageStoreFactory`]. + pub fn with_page_store_factory(mut self, page_store_factory: Arc) -> Self { + self.page_store_factory = page_store_factory; + self + } + fn create_row_group_writer(&self, row_group_index: usize) -> Result { let writers = self.create_column_writers(row_group_index)?; Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema)) @@ -1127,12 +1218,13 @@ impl ArrowRowGroupWriterFactory { #[cfg(feature = "encryption")] fn column_writer_factory(&self, row_group_idx: usize) -> ArrowColumnWriterFactory { ArrowColumnWriterFactory::new() + .with_page_store_factory(self.page_store_factory.clone()) .with_file_encryptor(row_group_idx, self.file_encryptor.clone()) } #[cfg(not(feature = "encryption"))] fn column_writer_factory(&self, _row_group_idx: usize) -> ArrowColumnWriterFactory { - ArrowColumnWriterFactory::new() + ArrowColumnWriterFactory::new().with_page_store_factory(self.page_store_factory.clone()) } } @@ -1159,6 +1251,8 @@ pub fn get_column_writers( /// Creates [`ArrowColumnWriter`] instances struct ArrowColumnWriterFactory { + /// Allocates the per-column-chunk [`PageStore`] backing each page writer. + page_store_factory: Arc, #[cfg(feature = "encryption")] row_group_index: usize, #[cfg(feature = "encryption")] @@ -1168,6 +1262,7 @@ struct ArrowColumnWriterFactory { impl ArrowColumnWriterFactory { pub fn new() -> Self { Self { + page_store_factory: Arc::new(InMemoryPageStoreFactory), #[cfg(feature = "encryption")] row_group_index: 0, #[cfg(feature = "encryption")] @@ -1175,6 +1270,12 @@ impl ArrowColumnWriterFactory { } } + /// Use `page_store_factory` to allocate the buffer for each column chunk. + pub fn with_page_store_factory(mut self, page_store_factory: Arc) -> Self { + self.page_store_factory = page_store_factory; + self + } + #[cfg(feature = "encryption")] pub fn with_file_encryptor( mut self, @@ -1199,8 +1300,9 @@ impl ArrowColumnWriterFactory { column_index, &column_path, )?; + let store = self.page_store_factory.create(column_index)?; Ok(Box::new( - ArrowPageWriter::default().with_encryptor(page_encryptor), + ArrowPageWriter::new(store).with_encryptor(page_encryptor), )) } @@ -1208,9 +1310,10 @@ impl ArrowColumnWriterFactory { fn create_page_writer( &self, _column_descriptor: &ColumnDescPtr, - _column_index: usize, + column_index: usize, ) -> Result> { - Ok(Box::::default()) + let store = self.page_store_factory.create(column_index)?; + Ok(Box::new(ArrowPageWriter::new(store))) } /// Gets an [`ArrowColumnWriter`] for the given `data_type`, appending the @@ -1738,6 +1841,103 @@ mod tests { statistics::Statistics, }; + /// A [`PageStore`] that allocates *sparse, non-contiguous* handles and keeps + /// blobs in a `HashMap` — nothing like the default `Vec`. Used to + /// prove the writer relies only on the opaque-handle contract and never on + /// handles being dense `Vec` indices. Records how many blobs were stored. + #[derive(Debug, Default)] + struct RecordingPageStore { + next: u64, + blobs: HashMap, + puts: Arc, + } + + impl PageStore for RecordingPageStore { + fn put(&mut self, value: Bytes) -> Result { + // Deliberately non-sequential, never-zero handles. + let id = 100 + self.next * 7; + self.next += 1; + self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.blobs.insert(id, value); + Ok(PageKey(id)) + } + + fn take(&mut self, key: PageKey) -> Result { + self.blobs + .remove(&key.0) + .ok_or_else(|| ParquetError::General(format!("missing key {}", key.0))) + } + } + + #[derive(Debug)] + struct RecordingPageStoreFactory { + puts: Arc, + } + + impl PageStoreFactory for RecordingPageStoreFactory { + fn create(&self, _column_index: usize) -> Result> { + Ok(Box::new(RecordingPageStore { + puts: self.puts.clone(), + ..Default::default() + })) + } + } + + /// A custom [`PageStore`] must produce byte-identical files to the in-memory + /// default, across dictionary and non-dictionary columns and multiple row + /// groups (so multiple store instances are exercised). + #[test] + fn custom_page_store_is_byte_identical_to_default() { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, true), + // A low-cardinality string column to exercise the dictionary path. + Field::new("s", DataType::Utf8, true), + ])); + let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]); + let s = StringArray::from(vec![ + Some("a"), + Some("bb"), + Some("a"), + None, + Some("bb"), + Some("ccc"), + ]); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap(); + + // Small row groups so multiple column chunks (hence multiple store + // instances) are produced. + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .build(); + + let write = |factory: Option>| { + let mut buffer = Vec::new(); + let mut opts = ArrowWriterOptions::new().with_properties(props.clone()); + if let Some(factory) = factory { + opts = opts.with_page_store_factory(factory); + } + let mut writer = + ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + buffer + }; + + let default_bytes = write(None); + + let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory { puts: puts.clone() }))); + + assert!( + puts.load(std::sync::atomic::Ordering::Relaxed) > 0, + "custom PageStore was never written to" + ); + assert_eq!( + default_bytes, custom_bytes, + "a custom PageStore must produce byte-identical output to the default" + ); + } + #[test] fn arrow_writer() { // define schema diff --git a/parquet/src/column/mod.rs b/parquet/src/column/mod.rs index 115c8dd01b80..69a87cb631fa 100644 --- a/parquet/src/column/mod.rs +++ b/parquet/src/column/mod.rs @@ -120,6 +120,7 @@ #[cfg(feature = "arrow")] pub(crate) mod chunker; pub mod page; +pub mod page_store; #[cfg(feature = "encryption")] pub(crate) mod page_encryption; #[cfg(not(feature = "encryption"))] diff --git a/parquet/src/column/page_store.rs b/parquet/src/column/page_store.rs new file mode 100644 index 000000000000..78181de427a2 --- /dev/null +++ b/parquet/src/column/page_store.rs @@ -0,0 +1,158 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Pluggable storage for completed, serialized page blobs. +//! +//! While a row group is being written the [`ArrowWriter`] must buffer every +//! column's encoded pages, because Parquet requires each column chunk to be +//! contiguous on disk while record batches arrive with all columns interleaved. +//! By default that buffer lives on the heap, so the writer's peak memory grows +//! with the row group size. A [`PageStore`] lets the buffer live somewhere else +//! — a local temp file, object storage, etc. — bounding peak write memory +//! independently of the row group size. +//! +//! [`ArrowWriter`]: crate::arrow::arrow_writer::ArrowWriter + +use std::fmt::Debug; + +use bytes::Bytes; + +use crate::errors::{ParquetError, Result}; + +/// An opaque, store-allocated handle to a blob held by a [`PageStore`]. +/// +/// Handles are allocated by the store — densely and sequentially — and are only +/// meaningful to the store that produced them. The caller treats them as opaque +/// tokens and decides what they *mean* (ordering, which one is the dictionary +/// page, etc.). +/// +/// Letting the store allocate the handle (rather than the caller choosing keys) +/// lets each backend pick the cheapest possible locator with no hashing: an +/// in-memory backend uses the handle as an index into a `Vec`, a temp-file +/// backend as an index into a `Vec<(offset, len)>`, and so on. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PageKey(pub(crate) u64); + +/// A pluggable store for completed, serialized page blobs. +/// +/// The store is intentionally "dumb": it only maps an opaque [`PageKey`] to a +/// blob of bytes. It knows nothing about pages, dictionaries, ordering, or +/// offsets. The caller keeps the handles it gets back from [`put`](Self::put) +/// and decides what they mean. +/// +/// Each store instance is owned by a single column writer and mutated by one +/// 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). +pub trait PageStore: Send { + /// Store `value`, returning a handle that can later be passed to + /// [`take`](Self::take). + fn put(&mut self, value: Bytes) -> Result; + + /// Take back the blob previously stored under `key`. + /// + /// The caller takes ownership of the returned bytes and will **not** request + /// `key` again, so the store may release any resources backing it — eagerly + /// here, or when the store is dropped. + fn take(&mut self, key: PageKey) -> Result; +} + +/// Creates a fresh [`PageStore`] for each column chunk. +/// +/// See +/// [`ArrowWriterOptions::with_page_store_factory`](crate::arrow::arrow_writer::ArrowWriterOptions::with_page_store_factory). +pub trait PageStoreFactory: Send + Sync + Debug { + /// Create a new, empty [`PageStore`] for the leaf column at `column_index`. + /// + /// `column_index` is a hint a backend may use to e.g. name spill files or + /// shard across a bounded pool; it carries no ordering or coordination + /// requirement. + fn create(&self, column_index: usize) -> Result>; +} + +/// The default [`PageStore`], holding blobs on the heap in a `Vec`. +/// +/// This is byte-for-byte equivalent to the writer's historical buffering +/// behavior and adds no overhead: peak memory still grows with the row group +/// size. Use a spilling backend to bound it. +#[derive(Debug, Default)] +pub struct InMemoryPageStore { + blobs: Vec, +} + +impl PageStore for InMemoryPageStore { + fn put(&mut self, value: Bytes) -> Result { + let key = PageKey(self.blobs.len() as u64); + self.blobs.push(value); + Ok(key) + } + + fn take(&mut self, key: PageKey) -> Result { + // Replace the slot with an empty `Bytes` so the stored blob is released + // as soon as it is taken, keeping memory bounded while the chunk is + // streamed into the output file. + self.blobs + .get_mut(key.0 as usize) + .map(std::mem::take) + .ok_or_else(|| ParquetError::General(format!("invalid page key {}", key.0))) + } +} + +/// Factory for [`InMemoryPageStore`] — the default used by +/// [`ArrowWriter`](crate::arrow::arrow_writer::ArrowWriter). +#[derive(Debug, Default)] +pub struct InMemoryPageStoreFactory; + +impl PageStoreFactory for InMemoryPageStoreFactory { + fn create(&self, _column_index: usize) -> Result> { + Ok(Box::new(InMemoryPageStore::default())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn in_memory_round_trips_blobs_in_handle_order() { + let mut store = InMemoryPageStore::default(); + let k0 = store.put(Bytes::from_static(b"hello")).unwrap(); + let k1 = store.put(Bytes::from_static(b"world")).unwrap(); + assert_ne!(k0, k1); + assert_eq!(&store.take(k0).unwrap()[..], b"hello"); + assert_eq!(&store.take(k1).unwrap()[..], b"world"); + } + + #[test] + fn in_memory_take_releases_the_slot() { + let mut store = InMemoryPageStore::default(); + let k = store.put(Bytes::from_static(b"abc")).unwrap(); + assert_eq!(&store.take(k).unwrap()[..], b"abc"); + // A second take yields the emptied placeholder rather than the blob, + // confirming the bytes were released on the first take. + assert!(store.take(k).unwrap().is_empty()); + } + + #[test] + fn in_memory_invalid_key_errors() { + let mut store = InMemoryPageStore::default(); + assert!(store.take(PageKey(99)).is_err()); + } +} diff --git a/parquet/tests/page_spill_memory.rs b/parquet/tests/page_spill_memory.rs new file mode 100644 index 000000000000..3fef222eab88 --- /dev/null +++ b/parquet/tests/page_spill_memory.rs @@ -0,0 +1,158 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Heap-memory regression tests for the Parquet writer's page buffering, +//! measured with [`dhat`]. +//! +//! These prove the headline invariant of the pluggable [`PageStore`]: while a +//! row group is being written, the heap used to buffer completed pages grows +//! with the row group size for the default in-memory store, but stays bounded +//! (≈ a few pages per leaf column) once a spilling backend is plugged in. +//! +//! The whole test binary uses dhat's allocator, so every test here observes +//! precise peak-heap statistics. dhat's profiler is process-global and only one +//! may be live at a time, so all measurements run in a single, serialized test. +//! +//! [`PageStore`]: parquet::arrow::arrow_writer::PageStore + +use std::sync::Arc; + +use arrow::array::{ArrayRef, BinaryArray, Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_writer::ArrowWriterOptions; +use parquet::file::properties::WriterProperties; + +#[global_allocator] +static ALLOC: dhat::Alloc = dhat::Alloc; + +/// Width of each value in the one "fat" column, in bytes. +const FAT_VALUE_LEN: usize = 4096; +/// Rows per input batch fed to the writer. Kept small so each batch is dropped +/// promptly — only the writer's *buffering* should accumulate, not the input. +const ROWS_PER_BATCH: usize = 64; +/// Number of batches, all funnelled into a single large row group. +const NUM_BATCHES: usize = 64; +/// Total bytes of fat-column payload written (≈ 16 MiB). +const TOTAL_FAT_BYTES: usize = FAT_VALUE_LEN * ROWS_PER_BATCH * NUM_BATCHES; + +/// A wide schema: one fat, high-cardinality binary column (the spill target) +/// plus several tiny integer columns. +fn skewed_schema() -> SchemaRef { + let mut fields = vec![Field::new("fat", DataType::Binary, false)]; + for i in 0..8 { + fields.push(Field::new(format!("small_{i}"), DataType::Int32, false)); + } + Arc::new(Schema::new(fields)) +} + +/// Build one batch of `ROWS_PER_BATCH` rows. The fat column holds unique, +/// high-entropy values (so they neither dictionary-encode nor compress away), +/// derived deterministically from `batch_index`. +fn make_batch(schema: &SchemaRef, batch_index: usize) -> RecordBatch { + let mut fat: Vec = vec![0u8; FAT_VALUE_LEN * ROWS_PER_BATCH]; + // A cheap xorshift fill keyed by the batch index → distinct, incompressible. + let mut state = (batch_index as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + for byte in fat.iter_mut() { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + *byte = (state >> 24) as u8; + } + let offsets: Vec = (0..=ROWS_PER_BATCH) + .map(|i| (i * FAT_VALUE_LEN) as i32) + .collect(); + let fat_array = BinaryArray::try_new( + arrow::buffer::OffsetBuffer::new(offsets.into()), + arrow::buffer::Buffer::from_vec(fat), + None, + ) + .unwrap(); + + let mut columns: Vec = vec![Arc::new(fat_array)]; + for c in 0..8 { + let vals: Vec = (0..ROWS_PER_BATCH) + .map(|r| (batch_index * ROWS_PER_BATCH + r + c) as i32) + .collect(); + columns.push(Arc::new(Int32Array::from(vals))); + } + RecordBatch::try_new(schema.clone(), columns).unwrap() +} + +/// Writer properties forcing the whole dataset into a single, uncompressed row +/// group (so the page buffer is the only thing that grows). +fn single_row_group_props() -> WriterProperties { + WriterProperties::builder() + .set_compression(parquet::basic::Compression::UNCOMPRESSED) + // One row group for everything: never auto-flush on row count. + .set_max_row_group_row_count(Some(ROWS_PER_BATCH * NUM_BATCHES * 2)) + .build() +} + +/// Write the full skewed dataset with the given writer options, feeding small +/// batches (each dropped immediately) into one row group. +/// +/// The output is sent to [`io::sink`] so the produced file bytes never live on +/// the heap — the measured peak then reflects only the writer's internal page +/// *buffering*, which is exactly what a [`PageStore`] governs. +fn write_skewed_dataset(options: ArrowWriterOptions) { + let schema = skewed_schema(); + let mut writer = + ArrowWriter::try_new_with_options(std::io::sink(), schema.clone(), options).unwrap(); + for b in 0..NUM_BATCHES { + let batch = make_batch(&schema, b); + writer.write(&batch).unwrap(); + // `batch` dropped here — only the writer's internal buffering persists. + } + writer.close().unwrap(); +} + +/// Run `f` under a fresh dhat profiler and return the peak live heap (bytes) +/// observed during it. +fn peak_heap_bytes(f: impl FnOnce()) -> usize { + let profiler = dhat::Profiler::builder().testing().build(); + f(); + let stats = dhat::HeapStats::get(); + drop(profiler); + stats.max_bytes +} + +#[test] +fn in_memory_store_buffers_whole_row_group() { + // Baseline: with the default in-memory page store, peak heap while writing a + // single large row group is at least the size of the buffered column data — + // memory grows with the row group, unbounded. A spilling backend (added in a + // later commit) is measured against this. + let props = single_row_group_props(); + let peak = peak_heap_bytes(|| { + let opts = ArrowWriterOptions::new().with_properties(props.clone()); + write_skewed_dataset(opts); + }); + + eprintln!( + "in-memory peak heap: {peak} bytes ({:.1} MiB); total fat payload {TOTAL_FAT_BYTES} bytes", + peak as f64 / (1024.0 * 1024.0) + ); + + // The fat column alone is ~16 MiB and must be fully resident in the page + // buffer at flush. Allow generous headroom below the total to stay robust. + let floor = TOTAL_FAT_BYTES * 3 / 4; + assert!( + peak >= floor, + "expected in-memory peak heap >= {floor} bytes (3/4 of buffered data), got {peak}" + ); +} From 3fd51c43bb62705e702356017781e5ac21c4626f Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 26 May 2026 10:17:48 -0500 Subject: [PATCH 02/13] parquet: stream column chunks out of the PageStore at splice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the splice materialized an entire column chunk back into a Vec before copying it into the output file, so peak memory during the splice phase was bounded by the largest column chunk — defeating a spilling backend for skewed schemas. Replace the materialize-then-copy path with StreamingColumnChunkReader, a Read that takes each page blob back out of the store in write order as it is consumed and releases it immediately, so the splice holds at most one page in memory at a time. SerializedRowGroupWriter::append_column is refactored to delegate to a new append_column_from_read that consumes an owned Read (append_column itself is unchanged for external ChunkReader callers). For the default in-memory store this is behavior-preserving (it already holds the bytes); for a spilling store it keeps the splice within the memory bound. Co-Authored-By: Claude Opus 4.7 (1M context) --- parquet/src/arrow/arrow_writer/mod.rs | 94 ++++++++++----------------- parquet/src/file/writer.rs | 26 +++++++- 2 files changed, 59 insertions(+), 61 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 3e0cef4f0ef7..e3d311c760bf 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -21,7 +21,6 @@ use crate::column::chunker::ContentDefinedChunker; use bytes::Bytes; use std::io::{Read, Write}; -use std::iter::Peekable; use std::slice::Iter; use std::sync::{Arc, Mutex}; use std::vec::IntoIter; @@ -49,7 +48,6 @@ use crate::encryption::encrypt::FileEncryptor; use crate::errors::{ParquetError, Result}; use crate::file::metadata::{KeyValue, ParquetMetaData, RowGroupMetaData}; use crate::file::properties::{WriterProperties, WriterPropertiesPtr}; -use crate::file::reader::{ChunkReader, Length}; use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter}; use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift}; use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, SchemaDescriptor}; @@ -655,72 +653,46 @@ impl ArrowColumnChunkData { self.keys.push(key); Ok(()) } - - /// Take every buffered blob back out of the store, in write order, ready to - /// splice into the output file. - /// - /// TODO(spill): stream blobs one at a time into the splice instead of - /// collecting them, so the splice phase is also bounded for spilling - /// backends (it currently materializes one column chunk at a time). - fn drain(&mut self) -> Result { - let mut data = Vec::with_capacity(self.keys.len()); - for key in std::mem::take(&mut self.keys) { - data.push(self.store.take(key)?); - } - Ok(MaterializedColumnChunk { - length: self.length, - data, - }) - } } -/// The buffered pages of one column chunk, taken back out of the [`PageStore`] -/// in write order and ready to be spliced into the output file via -/// [`SerializedRowGroupWriter::append_column`]. -struct MaterializedColumnChunk { - length: usize, - data: Vec, -} - -impl Length for MaterializedColumnChunk { - fn len(&self) -> u64 { - self.length as _ - } +/// A streaming [`Read`] over one column chunk's buffered pages. +/// +/// Takes each blob back out of the [`PageStore`] in write order *as it is +/// consumed*, releasing it 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 { + store: Box, + keys: IntoIter, + /// The blob currently being drained into the output; emptied as it is read. + current: Bytes, } -impl ChunkReader for MaterializedColumnChunk { - type T = ArrowColumnChunkReader; - - fn get_read(&self, start: u64) -> Result { - assert_eq!(start, 0); // Assume append_column writes all data in one-shot - Ok(ArrowColumnChunkReader( - self.data.clone().into_iter().peekable(), - )) - } - - fn get_bytes(&self, _start: u64, _length: usize) -> Result { - unimplemented!() +impl StreamingColumnChunkReader { + fn new(data: ArrowColumnChunkData) -> Self { + Self { + store: data.store, + keys: data.keys.into_iter(), + current: Bytes::new(), + } } } -/// A [`Read`] for [`MaterializedColumnChunk`] -struct ArrowColumnChunkReader(Peekable>); - -impl Read for ArrowColumnChunkReader { +impl Read for StreamingColumnChunkReader { fn read(&mut self, out: &mut [u8]) -> std::io::Result { - let buffer = loop { - match self.0.peek_mut() { - Some(b) if b.is_empty() => { - self.0.next(); - continue; + // Refill from the next stored blob whenever the current one is drained. + while self.current.is_empty() { + match self.keys.next() { + Some(key) => { + self.current = self.store.take(key).map_err(std::io::Error::other)?; } - Some(b) => break b, None => return Ok(0), } - }; + } - let len = buffer.len().min(out.len()); - let b = buffer.split_to(len); + let len = self.current.len().min(out.len()); + let b = self.current.split_to(len); out[..len].copy_from_slice(&b); Ok(len) } @@ -865,13 +837,15 @@ impl ArrowColumnChunk { &mut self.close } - /// Calls [`SerializedRowGroupWriter::append_column`] with this column's data + /// Splices this column's buffered pages into the row group, streaming them + /// back out of the [`PageStore`] one page at a time. pub fn append_to_row_group( - mut self, + self, writer: &mut SerializedRowGroupWriter<'_, W>, ) -> Result<()> { - let materialized = self.data.drain()?; - writer.append_column(&materialized, self.close) + let ArrowColumnChunk { data, close } = self; + let reader = StreamingColumnChunkReader::new(data); + writer.append_column_from_read(reader, close) } } diff --git a/parquet/src/file/writer.rs b/parquet/src/file/writer.rs index 942013ea6238..8ec16ba36739 100644 --- a/parquet/src/file/writer.rs +++ b/parquet/src/file/writer.rs @@ -684,6 +684,30 @@ impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> { pub fn append_column( &mut self, reader: &R, + close: ColumnCloseResult, + ) -> Result<()> { + // Position a reader at the start of the buffered chunk, then splice the + // bytes through the shared streaming path. + let metadata = &close.metadata; + let src_offset = metadata + .dictionary_page_offset() + .unwrap_or_else(|| metadata.data_page_offset()); + let read = reader.get_read(src_offset as _)?; + self.append_column_from_read(read, close) + } + + /// Splice an already-encoded column chunk into the row group, reading its + /// bytes sequentially from `read`. + /// + /// `read` must be positioned at the start of the chunk (the dictionary page + /// if present, otherwise the first data page — i.e. `src_offset` below) and + /// yield exactly the chunk's compressed bytes. Unlike [`Self::append_column`] + /// this consumes an owned [`Read`], which lets the caller stream the bytes + /// back from a [`PageStore`](crate::column::page_store::PageStore) one page + /// at a time without materializing the whole chunk in memory. + pub(crate) fn append_column_from_read( + &mut self, + read: R, mut close: ColumnCloseResult, ) -> Result<()> { self.assert_previous_writer_closed()?; @@ -707,7 +731,7 @@ impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> { let src_length = metadata.compressed_size(); let write_offset = self.buf.bytes_written(); - let mut read = reader.get_read(src_offset as _)?.take(src_length as _); + 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 { From c5de71eca972cdabdd92b76ced2d4ac4e5e56f04 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 26 May 2026 10:21:21 -0500 Subject: [PATCH 03/13] parquet: make PageKey constructible + prove the memory bound with a spilling backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PageKey's field was private, so an external PageStore implementor could not mint the handle it must return from put() — the trait was unusable outside the crate. Add public PageKey::new/get so any backend can allocate its own opaque, dense handles. Extend the dhat integration test with a temp-file PageStore backend (one unlinked temp file per column chunk; put appends, take seeks+reads) and assert the headline invariant: writing a skewed ~16 MiB single row group, peak heap drops from ~18 MiB with the in-memory store to ~3 MiB with the spilling store — bounded by the in-flight encoder/dictionary buffers rather than the row group size. Co-Authored-By: Claude Opus 4.7 (1M context) --- parquet/src/arrow/arrow_writer/mod.rs | 20 +++-- parquet/src/column/mod.rs | 2 +- parquet/src/column/page_store.rs | 19 ++++- parquet/tests/page_spill_memory.rs | 115 ++++++++++++++++++++++---- 4 files changed, 133 insertions(+), 23 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index e3d311c760bf..4b4530cafa83 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -1163,7 +1163,10 @@ impl ArrowRowGroupWriterFactory { /// Set the [`PageStoreFactory`] used to allocate the buffer for each column /// chunk, e.g. to spill completed pages to a temp file or object storage /// instead of the heap. Defaults to [`InMemoryPageStoreFactory`]. - pub fn with_page_store_factory(mut self, page_store_factory: Arc) -> Self { + pub fn with_page_store_factory( + mut self, + page_store_factory: Arc, + ) -> Self { self.page_store_factory = page_store_factory; self } @@ -1245,7 +1248,10 @@ impl ArrowColumnWriterFactory { } /// Use `page_store_factory` to allocate the buffer for each column chunk. - pub fn with_page_store_factory(mut self, page_store_factory: Arc) -> Self { + pub fn with_page_store_factory( + mut self, + page_store_factory: Arc, + ) -> Self { self.page_store_factory = page_store_factory; self } @@ -1833,13 +1839,13 @@ mod tests { self.next += 1; self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed); self.blobs.insert(id, value); - Ok(PageKey(id)) + Ok(PageKey::new(id)) } fn take(&mut self, key: PageKey) -> Result { self.blobs - .remove(&key.0) - .ok_or_else(|| ParquetError::General(format!("missing key {}", key.0))) + .remove(&key.get()) + .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get()))) } } @@ -1900,7 +1906,9 @@ mod tests { let default_bytes = write(None); let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory { puts: puts.clone() }))); + let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory { + puts: puts.clone(), + }))); assert!( puts.load(std::sync::atomic::Ordering::Relaxed) > 0, diff --git a/parquet/src/column/mod.rs b/parquet/src/column/mod.rs index 69a87cb631fa..9c7e77d29cba 100644 --- a/parquet/src/column/mod.rs +++ b/parquet/src/column/mod.rs @@ -120,11 +120,11 @@ #[cfg(feature = "arrow")] pub(crate) mod chunker; pub mod page; -pub mod page_store; #[cfg(feature = "encryption")] pub(crate) mod page_encryption; #[cfg(not(feature = "encryption"))] #[path = "page_encryption_disabled.rs"] pub(crate) mod page_encryption; +pub mod page_store; pub mod reader; pub mod writer; diff --git a/parquet/src/column/page_store.rs b/parquet/src/column/page_store.rs index 78181de427a2..5825e07ca35c 100644 --- a/parquet/src/column/page_store.rs +++ b/parquet/src/column/page_store.rs @@ -45,7 +45,24 @@ use crate::errors::{ParquetError, Result}; /// in-memory backend uses the handle as an index into a `Vec`, a temp-file /// backend as an index into a `Vec<(offset, len)>`, and so on. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct PageKey(pub(crate) u64); +pub struct PageKey(u64); + +impl PageKey { + /// Create a handle wrapping `raw`. + /// + /// A [`PageStore`] implementation calls this to mint the handle it returns + /// from [`put`](PageStore::put). The value is opaque to the caller, so a + /// store is free to use a dense counter, a packed locator, or anything else + /// it can later resolve in [`take`](PageStore::take). + pub const fn new(raw: u64) -> Self { + Self(raw) + } + + /// The raw value passed to [`new`](Self::new). + pub const fn get(self) -> u64 { + self.0 + } +} /// A pluggable store for completed, serialized page blobs. /// diff --git a/parquet/tests/page_spill_memory.rs b/parquet/tests/page_spill_memory.rs index 3fef222eab88..8a3455d783bc 100644 --- a/parquet/tests/page_spill_memory.rs +++ b/parquet/tests/page_spill_memory.rs @@ -29,12 +29,16 @@ //! //! [`PageStore`]: parquet::arrow::arrow_writer::PageStore +use std::fs::File; +use std::io::{Read as _, Seek, SeekFrom, Write as _}; use std::sync::Arc; use arrow::array::{ArrayRef, BinaryArray, Int32Array, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; use parquet::arrow::ArrowWriter; -use parquet::arrow::arrow_writer::ArrowWriterOptions; +use parquet::arrow::arrow_writer::{ArrowWriterOptions, PageKey, PageStore, PageStoreFactory}; +use parquet::errors::Result; use parquet::file::properties::WriterProperties; #[global_allocator] @@ -121,8 +125,60 @@ fn write_skewed_dataset(options: ArrowWriterOptions) { writer.close().unwrap(); } +/// A spilling [`PageStore`]: one temp file per column chunk. `put` appends the +/// blob and records its `(offset, len)`; `take` seeks and reads it back. The +/// file is unlinked on creation (via [`tempfile::tempfile`]) so it is cleaned up +/// when the store is dropped. This is the canonical "spill completed pages off +/// the heap" backend the design targets. +struct TempFilePageStore { + file: File, + end: u64, + locs: Vec<(u64, usize)>, +} + +impl TempFilePageStore { + fn new() -> Result { + Ok(Self { + file: tempfile::tempfile()?, + end: 0, + locs: Vec::new(), + }) + } +} + +impl PageStore for TempFilePageStore { + fn put(&mut self, value: Bytes) -> Result { + // Always append at the logical end (a prior `take` may have moved the + // OS file cursor). + 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 { + 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)) + } +} + +#[derive(Debug, Default)] +struct TempFilePageStoreFactory; + +impl PageStoreFactory for TempFilePageStoreFactory { + fn create(&self, _column_index: usize) -> Result> { + Ok(Box::new(TempFilePageStore::new()?)) + } +} + /// Run `f` under a fresh dhat profiler and return the peak live heap (bytes) -/// observed during it. +/// observed during it. dhat's profiler is process-global, so callers must run +/// sequentially (a single `#[test]`, not parallel tests). fn peak_heap_bytes(f: impl FnOnce()) -> usize { let profiler = dhat::Profiler::builder().testing().build(); f(); @@ -131,28 +187,57 @@ fn peak_heap_bytes(f: impl FnOnce()) -> usize { stats.max_bytes } +/// The whole test runs in one function because dhat allows only one live +/// profiler at a time; running the two measurements as separate parallel tests +/// would race on the global profiler. #[test] -fn in_memory_store_buffers_whole_row_group() { - // Baseline: with the default in-memory page store, peak heap while writing a - // single large row group is at least the size of the buffered column data — - // memory grows with the row group, unbounded. A spilling backend (added in a - // later commit) is measured against this. +fn page_store_bounds_write_memory() { let props = single_row_group_props(); - let peak = peak_heap_bytes(|| { + + // Baseline: the default in-memory store buffers the whole row group, so peak + // heap is at least the size of the buffered column data. + let in_memory_peak = peak_heap_bytes(|| { let opts = ArrowWriterOptions::new().with_properties(props.clone()); write_skewed_dataset(opts); }); + // Spilling: the temp-file store keeps completed pages off the heap, so peak + // heap stays bounded by the in-flight encoder/dictionary buffers plus a page + // or two in flight — independent of the row group size. + let spill_peak = peak_heap_bytes(|| { + let opts = ArrowWriterOptions::new() + .with_properties(props.clone()) + .with_page_store_factory(Arc::new(TempFilePageStoreFactory)); + write_skewed_dataset(opts); + }); + eprintln!( - "in-memory peak heap: {peak} bytes ({:.1} MiB); total fat payload {TOTAL_FAT_BYTES} bytes", - peak as f64 / (1024.0 * 1024.0) + "peak heap — in-memory: {:.1} MiB, temp-file spill: {:.1} MiB (total fat payload {:.1} MiB)", + in_memory_peak as f64 / (1024.0 * 1024.0), + spill_peak as f64 / (1024.0 * 1024.0), + TOTAL_FAT_BYTES as f64 / (1024.0 * 1024.0), ); - // The fat column alone is ~16 MiB and must be fully resident in the page - // buffer at flush. Allow generous headroom below the total to stay robust. - let floor = TOTAL_FAT_BYTES * 3 / 4; + // The in-memory store must hold most of the ~16 MiB of buffered data. + let in_memory_floor = TOTAL_FAT_BYTES * 3 / 4; + assert!( + in_memory_peak >= in_memory_floor, + "expected in-memory peak >= {in_memory_floor} bytes, got {in_memory_peak}" + ); + + // The spilling store must stay near the per-column bound — roughly + // (data_page_size + dict_page_size) per leaf column, ~2 MiB × 9 columns — + // and far below the in-memory baseline. We assert a generous 8 MiB ceiling + // (well under the ~16 MiB row group) to stay robust across platforms. + const SPILL_CEILING: usize = 8 * 1024 * 1024; + assert!( + spill_peak < SPILL_CEILING, + "expected spilling peak < {SPILL_CEILING} bytes (bounded by page/dict size × columns), \ + got {spill_peak}" + ); assert!( - peak >= floor, - "expected in-memory peak heap >= {floor} bytes (3/4 of buffered data), got {peak}" + spill_peak * 2 < in_memory_peak, + "expected spilling peak ({spill_peak}) to be far below the in-memory baseline \ + ({in_memory_peak})" ); } From 4e3685b4682bf5894a4268b1f1d2339cade3ac50 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 26 May 2026 11:39:42 -0500 Subject: [PATCH 04/13] parquet: spill dictionary-column data pages too (close accumulation point #2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dictionary-encoded columns buffered every completed data page in GenericColumnWriter.data_pages until close(), because the dictionary page must be written first but isn't final until all values are seen. Those pages never reached the PageStore, so spilling couldn't bound them — a low-cardinality 4.2M-row column peaked ~2.5 MiB regardless of backend. Add PageWriter::defers_dictionary_ordering(): a writer that buffers the whole chunk and splices it later (the Arrow path) can accept data pages before the dictionary page and order them itself. When set, the column writer streams dictionary-column data pages straight through instead of buffering them. ArrowPageWriter returns true, holds the (bounded) dictionary page in memory since it now arrives last, and at splice emits it first; the buffer-relative page offsets recorded in production order are rewritten to the dictionary-first layout there. The column-at-a-time SerializedFileWriter path is unchanged (defaults to false). Also fix memory_size() accounting: instead of counting bytes written (which over-reports once pages are spilled off-heap), ask the page writer how much it actually holds resident via PageWriter::buffered_memory_size() and PageStore::memory_size(). For the in-memory store this is unchanged; for a spilling store it drops to ~0 plus the retained dictionary page. Result: the dict-column case drops from ~2.69 MiB to ~0.48 MiB peak heap with a spilling backend. Adds an offset-index-disabled dictionary round-trip test and store memory-size unit tests; extends the dhat test with the dictionary-column scenario. Co-Authored-By: Claude Opus 4.7 (1M context) --- parquet/src/arrow/arrow_writer/mod.rs | 145 +++++++++++++++++++++++--- parquet/src/column/page.rs | 34 ++++++ parquet/src/column/page_store.rs | 55 +++++++++- parquet/src/column/writer/mod.rs | 22 +++- parquet/tests/page_spill_memory.rs | 48 +++++++++ 5 files changed, 287 insertions(+), 17 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 4b4530cafa83..a5e0bfc3ae64 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -36,6 +36,7 @@ use super::schema::{add_encoded_arrow_schema_to_metadata, decimal_length_from_pr use crate::arrow::ArrowSchemaConverter; use crate::arrow::arrow_writer::byte_array::ByteArrayEncoder; +use crate::basic::PageType; use crate::column::page::{CompressedPage, PageWriteSpec, PageWriter}; use crate::column::page_encryption::PageEncryptor; use crate::column::writer::encoder::ColumnValueEncoder; @@ -635,6 +636,17 @@ struct ArrowColumnChunkData { length: usize, store: Box, keys: Vec, + /// The dictionary page's serialized blobs (header ‖ data), held in memory + /// rather than the store. + /// + /// A dictionary page is produced at most once and bounded by + /// `dict_page_size_limit`, but it must be written *first* in the chunk even + /// though the data pages reach the writer before it (see + /// [`PageWriter::defers_dictionary_ordering`]). Spilling it would only + /// round-trip ~1 page to the backend and straight back, so it is kept here + /// and emitted ahead of the data pages at splice. Empty for non-dictionary + /// columns. + dictionary: Vec, } impl ArrowColumnChunkData { @@ -643,26 +655,47 @@ impl ArrowColumnChunkData { length: 0, store, keys: Vec::new(), + dictionary: Vec::new(), } } - /// Append a serialized blob to the store, recording its handle in write + /// Append a data-page blob to the store, recording its handle in write /// order. fn push(&mut self, value: Bytes) -> Result<()> { let key = self.store.put(value)?; self.keys.push(key); Ok(()) } + + /// Retain a dictionary-page blob in memory (emitted first at splice). + fn push_dictionary(&mut self, value: Bytes) { + self.dictionary.push(value); + } + + /// Total serialized size of the in-memory dictionary page, in bytes. + fn dictionary_len(&self) -> usize { + self.dictionary.iter().map(Bytes::len).sum() + } + + /// Bytes this chunk currently holds on the heap: whatever the store keeps + /// resident (zero for a spilling backend) plus the in-memory dictionary + /// page. + fn memory_size(&self) -> usize { + self.store.memory_size() + self.dictionary_len() + } } -/// A streaming [`Read`] over one column chunk's buffered pages. +/// A streaming [`Read`] over one column chunk's buffered pages, in final file +/// order: the in-memory dictionary page (if any) first, then the data pages. /// -/// Takes each blob back out of the [`PageStore`] in write order *as it is -/// consumed*, releasing it immediately afterwards, so splicing a chunk into the +/// Each data-page 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 { + /// Dictionary-page blobs, emitted before any data page. + dictionary: IntoIter, store: Box, keys: IntoIter, /// The blob currently being drained into the output; emptied as it is read. @@ -672,6 +705,7 @@ struct StreamingColumnChunkReader { impl StreamingColumnChunkReader { fn new(data: ArrowColumnChunkData) -> Self { Self { + dictionary: data.dictionary.into_iter(), store: data.store, keys: data.keys.into_iter(), current: Bytes::new(), @@ -681,13 +715,15 @@ impl StreamingColumnChunkReader { impl Read for StreamingColumnChunkReader { fn read(&mut self, out: &mut [u8]) -> std::io::Result { - // Refill from the next stored blob whenever the current one is drained. + // Refill from the next blob whenever the current one is drained: the + // dictionary page first, then each data page from the store. while self.current.is_empty() { - match self.keys.next() { - Some(key) => { - self.current = self.store.take(key).map_err(std::io::Error::other)?; - } - None => return Ok(0), + if let Some(blob) = self.dictionary.next() { + self.current = blob; + } else if let Some(key) = self.keys.next() { + self.current = self.store.take(key).map_err(std::io::Error::other)?; + } else { + return Ok(0); } } @@ -778,12 +814,35 @@ impl PageWriter for ArrowPageWriter { spec.bytes_written = compressed_size as u64; buf.length += compressed_size; - buf.push(header)?; - buf.push(data)?; + if spec.page_type == PageType::DICTIONARY_PAGE { + // Held in memory and emitted first at splice — see + // `ArrowColumnChunkData::dictionary`. The buffer-relative offset in + // `spec` (the dictionary arrives after the data pages on this path) + // is rewritten to its true, dictionary-first position at splice. + buf.push_dictionary(header); + buf.push_dictionary(data); + } else { + buf.push(header)?; + buf.push(data)?; + } Ok(spec) } + fn defers_dictionary_ordering(&self) -> bool { + // The Arrow chunk is buffered in full and spliced at row-group flush, so + // data pages may be accepted before the dictionary page and reordered + // then. This lets `GenericColumnWriter` stream dictionary-column data + // pages straight through instead of buffering them in memory. + true + } + + fn buffered_memory_size(&self) -> usize { + // Only what is actually resident: a spilling store reports ~0 here even + // though the chunk's bytes have all passed through it. + self.buffer.try_lock().unwrap().memory_size() + } + fn close(&mut self) -> Result<()> { Ok(()) } @@ -843,7 +902,31 @@ impl ArrowColumnChunk { self, writer: &mut SerializedRowGroupWriter<'_, W>, ) -> Result<()> { - let ArrowColumnChunk { data, close } = self; + let ArrowColumnChunk { data, mut close } = self; + + // On the Arrow path the dictionary page is produced *after* the data + // pages (so the data pages can stream straight through rather than + // accumulating in memory), but it must be written *first*. The encoder + // therefore recorded buffer-relative page offsets in production order; + // rewrite them to the final dictionary-first layout (dict at 0, data + // pages following) so the splice's offset remap lands them correctly. + let dictionary_len = data.dictionary_len(); + if dictionary_len > 0 { + close.metadata = close + .metadata + .into_builder() + .set_dictionary_page_offset(Some(0)) + .set_data_page_offset(dictionary_len as i64) + .build()?; + if let Some(offset_index) = close.offset_index.as_mut() { + let mut offset = dictionary_len as i64; + for location in offset_index.page_locations.iter_mut() { + location.offset = offset; + offset += location.compressed_page_size as i64; + } + } + } + let reader = StreamingColumnChunkReader::new(data); writer.append_column_from_read(reader, close) } @@ -1920,6 +2003,42 @@ mod tests { ); } + /// A dictionary-encoded column written through the deferred-ordering Arrow + /// path must round-trip correctly even with the offset index disabled, when + /// only the chunk-level dictionary/data page offsets are rewritten (there is + /// no offset index to rebuild). Spans multiple data pages so the + /// dictionary-first reordering is exercised. + #[test] + fn dictionary_column_round_trips_with_offset_index_disabled() { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)])); + + // Low cardinality so the column stays dictionary-encoded; enough rows to + // span several data pages within a single row group. + let values: Vec> = (0..50_000).map(|i| Some(i % 8)).collect(); + let array = Int32Array::from(values.clone()); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap(); + + let props = WriterProperties::builder() + .set_offset_index_disabled(true) + .set_data_page_row_count_limit(4096) + .build(); + let opts = ArrowWriterOptions::new().with_properties(props); + + let mut buffer = Vec::new(); + let mut writer = + ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap(); + let read: Vec = reader.collect::>().unwrap(); + let read_values: Vec> = read + .iter() + .flat_map(|b| b.column(0).as_primitive::().iter()) + .collect(); + assert_eq!(read_values, values); + } + #[test] fn arrow_writer() { // define schema diff --git a/parquet/src/column/page.rs b/parquet/src/column/page.rs index 4cfc07a02883..88cc0d18ef73 100644 --- a/parquet/src/column/page.rs +++ b/parquet/src/column/page.rs @@ -430,6 +430,40 @@ pub trait PageWriter: Send { /// either data page or dictionary page. fn write_page(&mut self, page: CompressedPage) -> Result; + /// Whether this writer resolves the final page layout itself (at flush) + /// rather than committing bytes to their final position as pages arrive. + /// + /// The dictionary page of a column chunk must be written *first*, but it is + /// not finalized until every value has been seen. A writer that commits + /// bytes live (e.g. straight to a file) therefore relies on the column + /// writer buffering the dictionary-encoded data pages in memory until the + /// dictionary page is ready — see [`GenericColumnWriter`]'s `data_pages`. + /// + /// A writer that instead buffers the whole chunk and splices it later (the + /// [`ArrowWriter`] path) can accept data pages *before* the dictionary page + /// and order them itself at flush. Returning `true` tells the column writer + /// to skip that in-memory buffering and stream dictionary-column data pages + /// straight through, bounding the column writer's memory. + /// + /// [`GenericColumnWriter`]: crate::column::writer::GenericColumnWriter + /// [`ArrowWriter`]: crate::arrow::arrow_writer::ArrowWriter + fn defers_dictionary_ordering(&self) -> bool { + false + } + + /// The number of bytes this writer is currently holding **in memory** for + /// pages it has been handed (i.e. completed pages not yet committed to their + /// final destination). + /// + /// Used by the column writer to report its memory footprint. The default is + /// `0`: a writer that streams pages straight to their destination retains + /// nothing. A writer that buffers pages should report what it actually holds + /// on the heap — which, when it spills to a backing store, can be far less + /// than the bytes written. + fn buffered_memory_size(&self) -> usize { + 0 + } + /// Closes resources and flushes underlying sink. /// Page writer should not be used after this method is called. fn close(&mut self) -> Result<()>; diff --git a/parquet/src/column/page_store.rs b/parquet/src/column/page_store.rs index 5825e07ca35c..603e012c013f 100644 --- a/parquet/src/column/page_store.rs +++ b/parquet/src/column/page_store.rs @@ -89,6 +89,18 @@ pub trait PageStore: Send { /// `key` again, so the store may release any resources backing it — eagerly /// here, or when the store is dropped. fn take(&mut self, key: PageKey) -> Result; + + /// The number of bytes this store currently holds **in memory** (resident + /// on the heap), used to report the writer's memory footprint. + /// + /// The default is `0`, which is exactly right for a backend that moves + /// every blob off-heap (a temp file, object storage): the bytes it has been + /// handed no longer occupy heap. The in-memory backend overrides this to + /// report its resident blobs. A backend that keeps a partial in-memory + /// buffer should report that buffer's size. + fn memory_size(&self) -> usize { + 0 + } } /// Creates a fresh [`PageStore`] for each column chunk. @@ -112,11 +124,14 @@ pub trait PageStoreFactory: Send + Sync + Debug { #[derive(Debug, Default)] pub struct InMemoryPageStore { blobs: Vec, + /// Running total of resident blob bytes, kept in step with `put`/`take`. + resident: usize, } impl PageStore for InMemoryPageStore { fn put(&mut self, value: Bytes) -> Result { let key = PageKey(self.blobs.len() as u64); + self.resident += value.len(); self.blobs.push(value); Ok(key) } @@ -125,10 +140,17 @@ impl PageStore for InMemoryPageStore { // Replace the slot with an empty `Bytes` so the stored blob is released // as soon as it is taken, keeping memory bounded while the chunk is // streamed into the output file. - self.blobs + let blob = self + .blobs .get_mut(key.0 as usize) .map(std::mem::take) - .ok_or_else(|| ParquetError::General(format!("invalid page key {}", key.0))) + .ok_or_else(|| ParquetError::General(format!("invalid page key {}", key.0)))?; + self.resident -= blob.len(); + Ok(blob) + } + + fn memory_size(&self) -> usize { + self.resident } } @@ -172,4 +194,33 @@ mod tests { let mut store = InMemoryPageStore::default(); assert!(store.take(PageKey(99)).is_err()); } + + #[test] + fn in_memory_reports_resident_bytes() { + let mut store = InMemoryPageStore::default(); + assert_eq!(store.memory_size(), 0); + let k0 = store.put(Bytes::from_static(b"hello")).unwrap(); + let k1 = store.put(Bytes::from_static(b"!")).unwrap(); + assert_eq!(store.memory_size(), 6); + store.take(k0).unwrap(); + assert_eq!(store.memory_size(), 1); + store.take(k1).unwrap(); + assert_eq!(store.memory_size(), 0); + } + + #[test] + fn default_store_memory_size_is_zero() { + // A spilling backend that does not override `memory_size` reports 0, + // reflecting that its blobs no longer occupy the heap. + struct OffHeap; + impl PageStore for OffHeap { + fn put(&mut self, _value: Bytes) -> Result { + Ok(PageKey::new(0)) + } + fn take(&mut self, _key: PageKey) -> Result { + Ok(Bytes::new()) + } + } + assert_eq!(OffHeap.memory_size(), 0); + } } diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 4e53230bbf89..9e2c9fb9ee6f 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -632,7 +632,18 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { /// of the current memory usage and not the final anticipated encoded size. #[cfg(feature = "arrow")] pub(crate) fn memory_size(&self) -> usize { - self.column_metrics.total_bytes_written as usize + self.encoder.estimated_memory_size() + // In-flight encoder buffers, plus any completed pages still held on the + // heap: the dictionary-column data pages buffered here (column-at-a-time + // path), plus whatever the page writer keeps resident. A page writer + // that spills completed pages off-heap reports far less than the bytes + // it was handed, so this tracks real memory rather than bytes written. + self.encoder.estimated_memory_size() + + self + .data_pages + .iter() + .map(|page| page.data().len()) + .sum::() + + self.page_writer.buffered_memory_size() } /// Returns total number of bytes written by this column writer so far. @@ -1271,7 +1282,14 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { }; // Check if we need to buffer data page or flush it to the sink directly. - if self.encoder.has_dictionary() { + // + // For dictionary-encoded columns the dictionary page must be written + // first, but it is not final until all values are seen, so completed + // data pages are normally buffered here until `close`. A page writer + // that defers final layout (the Arrow path) instead orders pages itself + // at flush, so we stream the data pages straight through and never let + // them accumulate in memory. + if self.encoder.has_dictionary() && !self.page_writer.defers_dictionary_ordering() { self.data_pages.push_back(compressed_page); } else { self.write_data_page(compressed_page)?; diff --git a/parquet/tests/page_spill_memory.rs b/parquet/tests/page_spill_memory.rs index 8a3455d783bc..85969fe85923 100644 --- a/parquet/tests/page_spill_memory.rs +++ b/parquet/tests/page_spill_memory.rs @@ -176,6 +176,34 @@ impl PageStoreFactory for TempFilePageStoreFactory { } } +/// Rows per batch / batches for the dictionary-column scenario (~4.2M rows). +const DICT_ROWS_PER_BATCH: usize = 8192; +const DICT_NUM_BATCHES: usize = 512; + +/// Write a single, low-cardinality (16 distinct values), high-row-count column +/// as one row group. Such a column stays dictionary-encoded, so its completed +/// data pages would historically pile up in `GenericColumnWriter` until close — +/// the second accumulation point that plain page-buffer spilling does not reach. +fn write_dict_dataset(options: ArrowWriterOptions) { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); + let props = WriterProperties::builder() + .set_compression(parquet::basic::Compression::UNCOMPRESSED) + .set_max_row_group_row_count(Some(DICT_ROWS_PER_BATCH * DICT_NUM_BATCHES * 2)) + .build(); + let options = options.with_properties(props); + let mut writer = + ArrowWriter::try_new_with_options(std::io::sink(), schema.clone(), options).unwrap(); + for b in 0..DICT_NUM_BATCHES { + let vals: Vec = (0..DICT_ROWS_PER_BATCH) + .map(|r| ((b + r) % 16) as i32) + .collect(); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vals))]).unwrap(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); +} + /// Run `f` under a fresh dhat profiler and return the peak live heap (bytes) /// observed during it. dhat's profiler is process-global, so callers must run /// sequentially (a single `#[test]`, not parallel tests). @@ -240,4 +268,24 @@ fn page_store_bounds_write_memory() { "expected spilling peak ({spill_peak}) to be far below the in-memory baseline \ ({in_memory_peak})" ); + + // Dictionary-encoded column: completed data pages reach the page writer (and + // thus the store) as they are produced, so spilling bounds them too. + let dict_in_memory = peak_heap_bytes(|| write_dict_dataset(ArrowWriterOptions::new())); + let dict_spill = peak_heap_bytes(|| { + write_dict_dataset( + ArrowWriterOptions::new().with_page_store_factory(Arc::new(TempFilePageStoreFactory)), + ) + }); + eprintln!( + "dict column ({} rows) peak heap — in-memory: {:.2} MiB, temp-file spill: {:.2} MiB", + DICT_ROWS_PER_BATCH * DICT_NUM_BATCHES, + dict_in_memory as f64 / (1024.0 * 1024.0), + dict_spill as f64 / (1024.0 * 1024.0), + ); + assert!( + dict_spill * 2 < dict_in_memory, + "expected dict-column spilling peak ({dict_spill}) to be far below the in-memory \ + baseline ({dict_in_memory}) — dictionary data pages should spill, not accumulate" + ); } From c2c70b9d0f54607194dc062be81b7ba8394dda94 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 26 May 2026 11:46:12 -0500 Subject: [PATCH 05/13] parquet: gate page_spill_memory integration test on the arrow feature The test uses parquet::arrow, so without a required-features entry it was auto-discovered and compiled under --all-targets --no-default-features, breaking that CI compilation check. Mirror the other arrow integration tests with required-features = ["arrow"]. Co-Authored-By: Claude Opus 4.7 (1M context) --- parquet/Cargo.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index b5430d6fe620..fdd679f3e6ca 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -160,6 +160,10 @@ path = "./examples/read_with_rowgroup.rs" name = "arrow_writer_layout" required-features = ["arrow"] +[[test]] +name = "page_spill_memory" +required-features = ["arrow"] + [[test]] name = "arrow_reader" required-features = ["arrow"] From abd6b863be25ea44015bb05a1e39db2f5b4c499c Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 2 Jun 2026 15:58:54 -0500 Subject: [PATCH 06/13] parquet: drop dhat from the page-spill memory test Replace the experimental, single-maintainer dhat crate with the in-tree thread-local tracking-allocator pattern already used by parquet/benches/arrow_reader_peak_memory.rs, and fold the test into the existing arrow_writer test binary instead of a dedicated one (saving a compile/link). The measurement still observes real process peak heap and the assertions are unchanged. Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 84 +------ parquet/Cargo.toml | 5 - parquet/tests/arrow_writer.rs | 341 ++++++++++++++++++++++++++++- parquet/tests/page_spill_memory.rs | 291 ------------------------ 4 files changed, 339 insertions(+), 382 deletions(-) delete mode 100644 parquet/tests/page_spill_memory.rs diff --git a/Cargo.lock b/Cargo.lock index a67498568b5b..7033b5da4431 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,15 +2,6 @@ # It is not intended for manual editing. version = 4 -[[package]] -name = "addr2line" -version = "0.25.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b5d307320b3181d6d7954e663bd7c774a838b8220fe0593c86d9fb09f498b4b" -dependencies = [ - "gimli", -] - [[package]] name = "adler2" version = "2.0.1" @@ -672,21 +663,6 @@ dependencies = [ "tower-service", ] -[[package]] -name = "backtrace" -version = "0.3.76" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb531853791a215d7c62a30daf0dde835f381ab5de4589cfe7c649d2cbe92bd6" -dependencies = [ - "addr2line", - "cfg-if", - "libc", - "miniz_oxide", - "object", - "rustc-demangle", - "windows-link", -] - [[package]] name = "base64" version = "0.22.1" @@ -1191,22 +1167,6 @@ dependencies = [ "syn", ] -[[package]] -name = "dhat" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98cd11d84628e233de0ce467de10b8633f4ddaecafadefc86e13b84b8739b827" -dependencies = [ - "backtrace", - "lazy_static", - "mintex", - "parking_lot", - "rustc-hash 1.1.0", - "serde", - "serde_json", - "thousands", -] - [[package]] name = "difflib" version = "0.4.0" @@ -1524,12 +1484,6 @@ dependencies = [ "wasip3", ] -[[package]] -name = "gimli" -version = "0.32.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" - [[package]] name = "h2" version = "0.4.14" @@ -2146,12 +2100,6 @@ dependencies = [ "simd-adler32", ] -[[package]] -name = "mintex" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c505b3e17ed6b70a7ed2e67fbb2c560ee327353556120d6e72f5232b6880d536" - [[package]] name = "mio" version = "1.2.0" @@ -2267,15 +2215,6 @@ dependencies = [ "objc2-core-foundation", ] -[[package]] -name = "object" -version = "0.37.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" -dependencies = [ - "memchr", -] - [[package]] name = "object_store" version = "0.13.2" @@ -2392,7 +2331,6 @@ dependencies = [ "clap", "crc32fast", "criterion", - "dhat", "flate2", "futures", "half", @@ -2774,7 +2712,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash 2.1.2", + "rustc-hash", "rustls", "socket2", "thiserror 2.0.18", @@ -2794,7 +2732,7 @@ dependencies = [ "lru-slab", "rand 0.9.4", "ring", - "rustc-hash 2.1.2", + "rustc-hash", "rustls", "rustls-pki-types", "slab", @@ -2995,18 +2933,6 @@ dependencies = [ "windows-sys 0.52.0", ] -[[package]] -name = "rustc-demangle" -version = "0.1.27" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b50b8869d9fc858ce7266cce0194bd74df58b9d0e3f6df3a9fc8eb470d95c09d" - -[[package]] -name = "rustc-hash" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" - [[package]] name = "rustc-hash" version = "2.1.2" @@ -3467,12 +3393,6 @@ dependencies = [ "syn", ] -[[package]] -name = "thousands" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3bf63baf9f5039dadc247375c29eb13706706cfde997d0330d05aa63a77d8820" - [[package]] name = "thread_local" version = "1.1.9" diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index fdd679f3e6ca..dd2c872ede50 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -94,7 +94,6 @@ tokio = { version = "1.0", default-features = false, features = ["macros", "rt-m rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] } object_store = { workspace = true, features = ["azure", "fs"] } sysinfo = { version = "0.38.1", default-features = false, features = ["system"] } -dhat = { version = "0.3", default-features = false } [package.metadata.docs.rs] all-features = true @@ -160,10 +159,6 @@ path = "./examples/read_with_rowgroup.rs" name = "arrow_writer_layout" required-features = ["arrow"] -[[test]] -name = "page_spill_memory" -required-features = ["arrow"] - [[test]] name = "arrow_reader" required-features = ["arrow"] diff --git a/parquet/tests/arrow_writer.rs b/parquet/tests/arrow_writer.rs index 020b4c6267e0..bd8dac0e6331 100644 --- a/parquet/tests/arrow_writer.rs +++ b/parquet/tests/arrow_writer.rs @@ -17,13 +17,20 @@ //! Tests for [`ArrowWriter`] -use arrow::array::Float64Array; -use arrow::datatypes::{DataType, Field, Schema}; -use arrow::record_batch::RecordBatch; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::cell::Cell; +use std::fs::File; +use std::io::{Read as _, Seek, SeekFrom, Write as _}; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BinaryArray, Float64Array, Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_writer::{ArrowWriterOptions, PageKey, PageStore, PageStoreFactory}; use parquet::basic::Encoding; +use parquet::errors::Result; use parquet::file::properties::WriterProperties; -use std::sync::Arc; #[test] #[should_panic( @@ -48,3 +55,329 @@ fn test_delta_bit_pack_type() { let mut writer = ArrowWriter::try_new(&mut buffer, record_batch.schema(), Some(props)).unwrap(); let _ = writer.write(&record_batch); } + +// --------------------------------------------------------------------------- +// Heap-memory regression test for the writer's page buffering. +// +// This proves the headline invariant of the pluggable [`PageStore`]: while a +// row group is being written, the heap used to buffer completed pages grows +// with the row group size for the default in-memory store, but stays bounded +// (≈ a few pages per leaf column) once a spilling backend is plugged in. +// +// Peak heap is measured with a thread-local tracking allocator (the same +// pattern used by `parquet/benches/arrow_reader_peak_memory.rs`), so the test +// needs no external profiling dependency. Tracking is thread-local, so the +// measured peak reflects only allocations made on the measuring thread; the +// default `ArrowWriter` is single-threaded, so the writer's buffering all lands +// there. Each measurement resets the peak to the current live baseline and +// reports the delta, so the threads of unrelated tests in this binary do not +// perturb it. +// +// [`PageStore`]: parquet::arrow::arrow_writer::PageStore +// --------------------------------------------------------------------------- + +thread_local! { + static LIVE_BYTES: Cell = const { Cell::new(0) }; + static PEAK_BYTES: Cell = const { Cell::new(0) }; +} + +struct TrackingAllocator { + inner: System, +} + +#[global_allocator] +static GLOBAL: TrackingAllocator = TrackingAllocator { inner: System }; + +fn add_live_bytes(size: usize) { + LIVE_BYTES.with(|live| { + let new = live.get().saturating_add(size); + live.set(new); + PEAK_BYTES.with(|peak| { + if new > peak.get() { + peak.set(new); + } + }); + }); +} + +fn subtract_live_bytes(size: usize) { + LIVE_BYTES.with(|live| { + live.set(live.get().saturating_sub(size)); + }); +} + +#[allow(unsafe_code)] +unsafe impl GlobalAlloc for TrackingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { self.inner.alloc(layout) }; + if !ptr.is_null() { + add_live_bytes(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + subtract_live_bytes(layout.size()); + unsafe { self.inner.dealloc(ptr, layout) }; + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { self.inner.realloc(ptr, layout, new_size) }; + if !new_ptr.is_null() { + let old_size = layout.size(); + if new_size > old_size { + add_live_bytes(new_size - old_size); + } else { + subtract_live_bytes(old_size - new_size); + } + } + new_ptr + } +} + +/// Run `f` and return the peak *additional* live heap (bytes) observed on this +/// thread during it — the delta from the live heap when `f` began. +fn peak_heap_bytes(f: impl FnOnce()) -> usize { + let start = LIVE_BYTES.with(Cell::get); + // Reset the peak to the window's baseline so prior allocations don't count. + PEAK_BYTES.with(|peak| peak.set(start)); + f(); + PEAK_BYTES.with(Cell::get).saturating_sub(start) +} + +/// Width of each value in the one "fat" column, in bytes. +const FAT_VALUE_LEN: usize = 4096; +/// Rows per input batch fed to the writer. Kept small so each batch is dropped +/// promptly — only the writer's *buffering* should accumulate, not the input. +const ROWS_PER_BATCH: usize = 64; +/// Number of batches, all funnelled into a single large row group. +const NUM_BATCHES: usize = 64; +/// Total bytes of fat-column payload written (≈ 16 MiB). +const TOTAL_FAT_BYTES: usize = FAT_VALUE_LEN * ROWS_PER_BATCH * NUM_BATCHES; + +/// A wide schema: one fat, high-cardinality binary column (the spill target) +/// plus several tiny integer columns. +fn skewed_schema() -> SchemaRef { + let mut fields = vec![Field::new("fat", DataType::Binary, false)]; + for i in 0..8 { + fields.push(Field::new(format!("small_{i}"), DataType::Int32, false)); + } + Arc::new(Schema::new(fields)) +} + +/// Build one batch of `ROWS_PER_BATCH` rows. The fat column holds unique, +/// high-entropy values (so they neither dictionary-encode nor compress away), +/// derived deterministically from `batch_index`. +fn make_batch(schema: &SchemaRef, batch_index: usize) -> RecordBatch { + let mut fat: Vec = vec![0u8; FAT_VALUE_LEN * ROWS_PER_BATCH]; + // A cheap xorshift fill keyed by the batch index → distinct, incompressible. + let mut state = (batch_index as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + for byte in fat.iter_mut() { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + *byte = (state >> 24) as u8; + } + let offsets: Vec = (0..=ROWS_PER_BATCH) + .map(|i| (i * FAT_VALUE_LEN) as i32) + .collect(); + let fat_array = BinaryArray::try_new( + arrow::buffer::OffsetBuffer::new(offsets.into()), + arrow::buffer::Buffer::from_vec(fat), + None, + ) + .unwrap(); + + let mut columns: Vec = vec![Arc::new(fat_array)]; + for c in 0..8 { + let vals: Vec = (0..ROWS_PER_BATCH) + .map(|r| (batch_index * ROWS_PER_BATCH + r + c) as i32) + .collect(); + columns.push(Arc::new(Int32Array::from(vals))); + } + RecordBatch::try_new(schema.clone(), columns).unwrap() +} + +/// Writer properties forcing the whole dataset into a single, uncompressed row +/// group (so the page buffer is the only thing that grows). +fn single_row_group_props() -> WriterProperties { + WriterProperties::builder() + .set_compression(parquet::basic::Compression::UNCOMPRESSED) + // One row group for everything: never auto-flush on row count. + .set_max_row_group_row_count(Some(ROWS_PER_BATCH * NUM_BATCHES * 2)) + .build() +} + +/// Write the full skewed dataset with the given writer options, feeding small +/// batches (each dropped immediately) into one row group. +/// +/// The output is sent to [`io::sink`] so the produced file bytes never live on +/// the heap — the measured peak then reflects only the writer's internal page +/// *buffering*, which is exactly what a [`PageStore`] governs. +fn write_skewed_dataset(options: ArrowWriterOptions) { + let schema = skewed_schema(); + let mut writer = + ArrowWriter::try_new_with_options(std::io::sink(), schema.clone(), options).unwrap(); + for b in 0..NUM_BATCHES { + let batch = make_batch(&schema, b); + writer.write(&batch).unwrap(); + // `batch` dropped here — only the writer's internal buffering persists. + } + writer.close().unwrap(); +} + +/// A spilling [`PageStore`]: one temp file per column chunk. `put` appends the +/// blob and records its `(offset, len)`; `take` seeks and reads it back. The +/// file is unlinked on creation (via [`tempfile::tempfile`]) so it is cleaned up +/// when the store is dropped. This is the canonical "spill completed pages off +/// the heap" backend the design targets. +struct TempFilePageStore { + file: File, + end: u64, + locs: Vec<(u64, usize)>, +} + +impl TempFilePageStore { + fn new() -> Result { + Ok(Self { + file: tempfile::tempfile()?, + end: 0, + locs: Vec::new(), + }) + } +} + +impl PageStore for TempFilePageStore { + fn put(&mut self, value: Bytes) -> Result { + // Always append at the logical end (a prior `take` may have moved the + // OS file cursor). + 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 { + 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)) + } +} + +#[derive(Debug, Default)] +struct TempFilePageStoreFactory; + +impl PageStoreFactory for TempFilePageStoreFactory { + fn create(&self, _column_index: usize) -> Result> { + Ok(Box::new(TempFilePageStore::new()?)) + } +} + +/// Rows per batch / batches for the dictionary-column scenario (~4.2M rows). +const DICT_ROWS_PER_BATCH: usize = 8192; +const DICT_NUM_BATCHES: usize = 512; + +/// Write a single, low-cardinality (16 distinct values), high-row-count column +/// as one row group. Such a column stays dictionary-encoded, so its completed +/// data pages would historically pile up in `GenericColumnWriter` until close — +/// the second accumulation point that plain page-buffer spilling does not reach. +fn write_dict_dataset(options: ArrowWriterOptions) { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); + let props = WriterProperties::builder() + .set_compression(parquet::basic::Compression::UNCOMPRESSED) + .set_max_row_group_row_count(Some(DICT_ROWS_PER_BATCH * DICT_NUM_BATCHES * 2)) + .build(); + let options = options.with_properties(props); + let mut writer = + ArrowWriter::try_new_with_options(std::io::sink(), schema.clone(), options).unwrap(); + for b in 0..DICT_NUM_BATCHES { + let vals: Vec = (0..DICT_ROWS_PER_BATCH) + .map(|r| ((b + r) % 16) as i32) + .collect(); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vals))]).unwrap(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); +} + +/// All measurements run in one function so they execute sequentially on a single +/// thread — the tracking allocator is thread-local, so running them as separate +/// parallel tests would each see only their own thread's allocations (which is +/// fine), but keeping them together also keeps the in-memory/spill comparison on +/// one consistent baseline. +#[test] +fn page_store_bounds_write_memory() { + let props = single_row_group_props(); + + // Baseline: the default in-memory store buffers the whole row group, so peak + // heap is at least the size of the buffered column data. + let in_memory_peak = peak_heap_bytes(|| { + let opts = ArrowWriterOptions::new().with_properties(props.clone()); + write_skewed_dataset(opts); + }); + + // Spilling: the temp-file store keeps completed pages off the heap, so peak + // heap stays bounded by the in-flight encoder/dictionary buffers plus a page + // or two in flight — independent of the row group size. + let spill_peak = peak_heap_bytes(|| { + let opts = ArrowWriterOptions::new() + .with_properties(props.clone()) + .with_page_store_factory(Arc::new(TempFilePageStoreFactory)); + write_skewed_dataset(opts); + }); + + eprintln!( + "peak heap — in-memory: {:.1} MiB, temp-file spill: {:.1} MiB (total fat payload {:.1} MiB)", + in_memory_peak as f64 / (1024.0 * 1024.0), + spill_peak as f64 / (1024.0 * 1024.0), + TOTAL_FAT_BYTES as f64 / (1024.0 * 1024.0), + ); + + // The in-memory store must hold most of the ~16 MiB of buffered data. + let in_memory_floor = TOTAL_FAT_BYTES * 3 / 4; + assert!( + in_memory_peak >= in_memory_floor, + "expected in-memory peak >= {in_memory_floor} bytes, got {in_memory_peak}" + ); + + // The spilling store must stay near the per-column bound — roughly + // (data_page_size + dict_page_size) per leaf column, ~2 MiB × 9 columns — + // and far below the in-memory baseline. We assert a generous 8 MiB ceiling + // (well under the ~16 MiB row group) to stay robust across platforms. + const SPILL_CEILING: usize = 8 * 1024 * 1024; + assert!( + spill_peak < SPILL_CEILING, + "expected spilling peak < {SPILL_CEILING} bytes (bounded by page/dict size × columns), \ + got {spill_peak}" + ); + assert!( + spill_peak * 2 < in_memory_peak, + "expected spilling peak ({spill_peak}) to be far below the in-memory baseline \ + ({in_memory_peak})" + ); + + // Dictionary-encoded column: completed data pages reach the page writer (and + // thus the store) as they are produced, so spilling bounds them too. + let dict_in_memory = peak_heap_bytes(|| write_dict_dataset(ArrowWriterOptions::new())); + let dict_spill = peak_heap_bytes(|| { + write_dict_dataset( + ArrowWriterOptions::new().with_page_store_factory(Arc::new(TempFilePageStoreFactory)), + ) + }); + eprintln!( + "dict column ({} rows) peak heap — in-memory: {:.2} MiB, temp-file spill: {:.2} MiB", + DICT_ROWS_PER_BATCH * DICT_NUM_BATCHES, + dict_in_memory as f64 / (1024.0 * 1024.0), + dict_spill as f64 / (1024.0 * 1024.0), + ); + assert!( + dict_spill * 2 < dict_in_memory, + "expected dict-column spilling peak ({dict_spill}) to be far below the in-memory \ + baseline ({dict_in_memory}) — dictionary data pages should spill, not accumulate" + ); +} diff --git a/parquet/tests/page_spill_memory.rs b/parquet/tests/page_spill_memory.rs deleted file mode 100644 index 85969fe85923..000000000000 --- a/parquet/tests/page_spill_memory.rs +++ /dev/null @@ -1,291 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Heap-memory regression tests for the Parquet writer's page buffering, -//! measured with [`dhat`]. -//! -//! These prove the headline invariant of the pluggable [`PageStore`]: while a -//! row group is being written, the heap used to buffer completed pages grows -//! with the row group size for the default in-memory store, but stays bounded -//! (≈ a few pages per leaf column) once a spilling backend is plugged in. -//! -//! The whole test binary uses dhat's allocator, so every test here observes -//! precise peak-heap statistics. dhat's profiler is process-global and only one -//! may be live at a time, so all measurements run in a single, serialized test. -//! -//! [`PageStore`]: parquet::arrow::arrow_writer::PageStore - -use std::fs::File; -use std::io::{Read as _, Seek, SeekFrom, Write as _}; -use std::sync::Arc; - -use arrow::array::{ArrayRef, BinaryArray, Int32Array, RecordBatch}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use bytes::Bytes; -use parquet::arrow::ArrowWriter; -use parquet::arrow::arrow_writer::{ArrowWriterOptions, PageKey, PageStore, PageStoreFactory}; -use parquet::errors::Result; -use parquet::file::properties::WriterProperties; - -#[global_allocator] -static ALLOC: dhat::Alloc = dhat::Alloc; - -/// Width of each value in the one "fat" column, in bytes. -const FAT_VALUE_LEN: usize = 4096; -/// Rows per input batch fed to the writer. Kept small so each batch is dropped -/// promptly — only the writer's *buffering* should accumulate, not the input. -const ROWS_PER_BATCH: usize = 64; -/// Number of batches, all funnelled into a single large row group. -const NUM_BATCHES: usize = 64; -/// Total bytes of fat-column payload written (≈ 16 MiB). -const TOTAL_FAT_BYTES: usize = FAT_VALUE_LEN * ROWS_PER_BATCH * NUM_BATCHES; - -/// A wide schema: one fat, high-cardinality binary column (the spill target) -/// plus several tiny integer columns. -fn skewed_schema() -> SchemaRef { - let mut fields = vec![Field::new("fat", DataType::Binary, false)]; - for i in 0..8 { - fields.push(Field::new(format!("small_{i}"), DataType::Int32, false)); - } - Arc::new(Schema::new(fields)) -} - -/// Build one batch of `ROWS_PER_BATCH` rows. The fat column holds unique, -/// high-entropy values (so they neither dictionary-encode nor compress away), -/// derived deterministically from `batch_index`. -fn make_batch(schema: &SchemaRef, batch_index: usize) -> RecordBatch { - let mut fat: Vec = vec![0u8; FAT_VALUE_LEN * ROWS_PER_BATCH]; - // A cheap xorshift fill keyed by the batch index → distinct, incompressible. - let mut state = (batch_index as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; - for byte in fat.iter_mut() { - state ^= state << 13; - state ^= state >> 7; - state ^= state << 17; - *byte = (state >> 24) as u8; - } - let offsets: Vec = (0..=ROWS_PER_BATCH) - .map(|i| (i * FAT_VALUE_LEN) as i32) - .collect(); - let fat_array = BinaryArray::try_new( - arrow::buffer::OffsetBuffer::new(offsets.into()), - arrow::buffer::Buffer::from_vec(fat), - None, - ) - .unwrap(); - - let mut columns: Vec = vec![Arc::new(fat_array)]; - for c in 0..8 { - let vals: Vec = (0..ROWS_PER_BATCH) - .map(|r| (batch_index * ROWS_PER_BATCH + r + c) as i32) - .collect(); - columns.push(Arc::new(Int32Array::from(vals))); - } - RecordBatch::try_new(schema.clone(), columns).unwrap() -} - -/// Writer properties forcing the whole dataset into a single, uncompressed row -/// group (so the page buffer is the only thing that grows). -fn single_row_group_props() -> WriterProperties { - WriterProperties::builder() - .set_compression(parquet::basic::Compression::UNCOMPRESSED) - // One row group for everything: never auto-flush on row count. - .set_max_row_group_row_count(Some(ROWS_PER_BATCH * NUM_BATCHES * 2)) - .build() -} - -/// Write the full skewed dataset with the given writer options, feeding small -/// batches (each dropped immediately) into one row group. -/// -/// The output is sent to [`io::sink`] so the produced file bytes never live on -/// the heap — the measured peak then reflects only the writer's internal page -/// *buffering*, which is exactly what a [`PageStore`] governs. -fn write_skewed_dataset(options: ArrowWriterOptions) { - let schema = skewed_schema(); - let mut writer = - ArrowWriter::try_new_with_options(std::io::sink(), schema.clone(), options).unwrap(); - for b in 0..NUM_BATCHES { - let batch = make_batch(&schema, b); - writer.write(&batch).unwrap(); - // `batch` dropped here — only the writer's internal buffering persists. - } - writer.close().unwrap(); -} - -/// A spilling [`PageStore`]: one temp file per column chunk. `put` appends the -/// blob and records its `(offset, len)`; `take` seeks and reads it back. The -/// file is unlinked on creation (via [`tempfile::tempfile`]) so it is cleaned up -/// when the store is dropped. This is the canonical "spill completed pages off -/// the heap" backend the design targets. -struct TempFilePageStore { - file: File, - end: u64, - locs: Vec<(u64, usize)>, -} - -impl TempFilePageStore { - fn new() -> Result { - Ok(Self { - file: tempfile::tempfile()?, - end: 0, - locs: Vec::new(), - }) - } -} - -impl PageStore for TempFilePageStore { - fn put(&mut self, value: Bytes) -> Result { - // Always append at the logical end (a prior `take` may have moved the - // OS file cursor). - 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 { - 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)) - } -} - -#[derive(Debug, Default)] -struct TempFilePageStoreFactory; - -impl PageStoreFactory for TempFilePageStoreFactory { - fn create(&self, _column_index: usize) -> Result> { - Ok(Box::new(TempFilePageStore::new()?)) - } -} - -/// Rows per batch / batches for the dictionary-column scenario (~4.2M rows). -const DICT_ROWS_PER_BATCH: usize = 8192; -const DICT_NUM_BATCHES: usize = 512; - -/// Write a single, low-cardinality (16 distinct values), high-row-count column -/// as one row group. Such a column stays dictionary-encoded, so its completed -/// data pages would historically pile up in `GenericColumnWriter` until close — -/// the second accumulation point that plain page-buffer spilling does not reach. -fn write_dict_dataset(options: ArrowWriterOptions) { - let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, false)])); - let props = WriterProperties::builder() - .set_compression(parquet::basic::Compression::UNCOMPRESSED) - .set_max_row_group_row_count(Some(DICT_ROWS_PER_BATCH * DICT_NUM_BATCHES * 2)) - .build(); - let options = options.with_properties(props); - let mut writer = - ArrowWriter::try_new_with_options(std::io::sink(), schema.clone(), options).unwrap(); - for b in 0..DICT_NUM_BATCHES { - let vals: Vec = (0..DICT_ROWS_PER_BATCH) - .map(|r| ((b + r) % 16) as i32) - .collect(); - let batch = - RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vals))]).unwrap(); - writer.write(&batch).unwrap(); - } - writer.close().unwrap(); -} - -/// Run `f` under a fresh dhat profiler and return the peak live heap (bytes) -/// observed during it. dhat's profiler is process-global, so callers must run -/// sequentially (a single `#[test]`, not parallel tests). -fn peak_heap_bytes(f: impl FnOnce()) -> usize { - let profiler = dhat::Profiler::builder().testing().build(); - f(); - let stats = dhat::HeapStats::get(); - drop(profiler); - stats.max_bytes -} - -/// The whole test runs in one function because dhat allows only one live -/// profiler at a time; running the two measurements as separate parallel tests -/// would race on the global profiler. -#[test] -fn page_store_bounds_write_memory() { - let props = single_row_group_props(); - - // Baseline: the default in-memory store buffers the whole row group, so peak - // heap is at least the size of the buffered column data. - let in_memory_peak = peak_heap_bytes(|| { - let opts = ArrowWriterOptions::new().with_properties(props.clone()); - write_skewed_dataset(opts); - }); - - // Spilling: the temp-file store keeps completed pages off the heap, so peak - // heap stays bounded by the in-flight encoder/dictionary buffers plus a page - // or two in flight — independent of the row group size. - let spill_peak = peak_heap_bytes(|| { - let opts = ArrowWriterOptions::new() - .with_properties(props.clone()) - .with_page_store_factory(Arc::new(TempFilePageStoreFactory)); - write_skewed_dataset(opts); - }); - - eprintln!( - "peak heap — in-memory: {:.1} MiB, temp-file spill: {:.1} MiB (total fat payload {:.1} MiB)", - in_memory_peak as f64 / (1024.0 * 1024.0), - spill_peak as f64 / (1024.0 * 1024.0), - TOTAL_FAT_BYTES as f64 / (1024.0 * 1024.0), - ); - - // The in-memory store must hold most of the ~16 MiB of buffered data. - let in_memory_floor = TOTAL_FAT_BYTES * 3 / 4; - assert!( - in_memory_peak >= in_memory_floor, - "expected in-memory peak >= {in_memory_floor} bytes, got {in_memory_peak}" - ); - - // The spilling store must stay near the per-column bound — roughly - // (data_page_size + dict_page_size) per leaf column, ~2 MiB × 9 columns — - // and far below the in-memory baseline. We assert a generous 8 MiB ceiling - // (well under the ~16 MiB row group) to stay robust across platforms. - const SPILL_CEILING: usize = 8 * 1024 * 1024; - assert!( - spill_peak < SPILL_CEILING, - "expected spilling peak < {SPILL_CEILING} bytes (bounded by page/dict size × columns), \ - got {spill_peak}" - ); - assert!( - spill_peak * 2 < in_memory_peak, - "expected spilling peak ({spill_peak}) to be far below the in-memory baseline \ - ({in_memory_peak})" - ); - - // Dictionary-encoded column: completed data pages reach the page writer (and - // thus the store) as they are produced, so spilling bounds them too. - let dict_in_memory = peak_heap_bytes(|| write_dict_dataset(ArrowWriterOptions::new())); - let dict_spill = peak_heap_bytes(|| { - write_dict_dataset( - ArrowWriterOptions::new().with_page_store_factory(Arc::new(TempFilePageStoreFactory)), - ) - }); - eprintln!( - "dict column ({} rows) peak heap — in-memory: {:.2} MiB, temp-file spill: {:.2} MiB", - DICT_ROWS_PER_BATCH * DICT_NUM_BATCHES, - dict_in_memory as f64 / (1024.0 * 1024.0), - dict_spill as f64 / (1024.0 * 1024.0), - ); - assert!( - dict_spill * 2 < dict_in_memory, - "expected dict-column spilling peak ({dict_spill}) to be far below the in-memory \ - baseline ({dict_in_memory}) — dictionary data pages should spill, not accumulate" - ); -} From e06d8b8936a4463b0d3f22f6abae5a55131a576e Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 2 Jun 2026 16:05:29 -0500 Subject: [PATCH 07/13] Update parquet/src/column/page_store.rs Co-authored-by: Andrew Lamb --- parquet/src/column/page_store.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parquet/src/column/page_store.rs b/parquet/src/column/page_store.rs index 603e012c013f..c6d90c4f7685 100644 --- a/parquet/src/column/page_store.rs +++ b/parquet/src/column/page_store.rs @@ -19,7 +19,7 @@ //! //! While a row group is being written the [`ArrowWriter`] must buffer every //! column's encoded pages, because Parquet requires each column chunk to be -//! contiguous on disk while record batches arrive with all columns interleaved. +//! contiguous in the file while record batches arrive with all columns interleaved. //! By default that buffer lives on the heap, so the writer's peak memory grows //! with the row group size. A [`PageStore`] lets the buffer live somewhere else //! — a local temp file, object storage, etc. — bounding peak write memory From 4c379a3f48af0f4f74d8b0a128fe7e3e3b2924f1 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 3 Jun 2026 10:07:12 -0400 Subject: [PATCH 08/13] parquet: add spill_page_store example demonstrating the PageStore API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `parquet/examples/spill_page_store.rs`, a runnable demonstration of the pluggable `PageStore` API. It implements a spilling `TempFilePageStore` (one temp file per column chunk) and writes a wide, skewed Parquet file — a few Int64 columns, some small (~20 byte) string columns, and a configurable number of large (~8 KiB) string columns — into a single row group, reporting peak `ArrowWriter::memory_size()` with and without spilling. --large-string-columns N number of fat ~8 KiB string columns (default 10) --spill use the spilling TempFilePageStore vs the default On the default 10 fat columns the spilling store cuts peak writer memory from ~161 MiB (whole row group buffered) to ~21 MiB (in-flight encoder buffers only), and produces a byte-identical file to the in-memory path. Co-Authored-By: Claude Opus 4.8 (1M context) --- parquet/Cargo.toml | 5 + parquet/examples/spill_page_store.rs | 371 +++++++++++++++++++++++++++ 2 files changed, 376 insertions(+) create mode 100644 parquet/examples/spill_page_store.rs diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index dd2c872ede50..ec447c90d468 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -155,6 +155,11 @@ name = "read_with_rowgroup" required-features = ["arrow", "async"] path = "./examples/read_with_rowgroup.rs" +[[example]] +name = "spill_page_store" +required-features = ["arrow", "cli"] +path = "./examples/spill_page_store.rs" + [[test]] name = "arrow_writer_layout" required-features = ["arrow"] diff --git a/parquet/examples/spill_page_store.rs b/parquet/examples/spill_page_store.rs new file mode 100644 index 000000000000..ed3f61afdfc1 --- /dev/null +++ b/parquet/examples/spill_page_store.rs @@ -0,0 +1,371 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Demonstrates the pluggable [`PageStore`] API by implementing a **spilling** +//! page store that keeps completed Parquet pages in temp files instead of on the +//! heap. +//! +//! # Background +//! +//! Parquet requires every column chunk to be contiguous in the file, but Arrow +//! record batches arrive with all columns interleaved. So while a row group is +//! being written, [`ArrowWriter`] must buffer every column's completed, +//! compressed pages until the row group is flushed. Peak write memory therefore +//! grows with the row group size — painful for wide schemas with large, skewed +//! columns (e.g. a few `id` columns next to a pile of fat string columns). +//! +//! A [`PageStore`] lets that page buffer live somewhere other than the heap. This +//! example plugs in a [`TempFilePageStore`] (one temp file per column chunk) and +//! compares peak writer memory against the default in-memory buffering. +//! +//! # Running +//! +//! ```sh +//! # Default in-memory page buffering (baseline): peak writer memory grows with +//! # the row group. +//! cargo run --release --features cli --example spill_page_store +//! +//! # Spill completed pages to temp files: peak writer memory stays bounded by +//! # the in-flight encoder buffers, independent of the row group size. +//! cargo run --release --features cli --example spill_page_store -- --spill +//! +//! # Make the schema wider / the skew worse: +//! cargo run --release --features cli --example spill_page_store -- --spill --large-string-columns 40 +//! ``` +//! +//! [`ArrowWriter`]: parquet::arrow::ArrowWriter +//! [`PageStore`]: parquet::arrow::arrow_writer::PageStore + +use std::fs::File; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringBuilder}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; +use clap::Parser; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_writer::{ + ArrowWriterOptions, PageKey, PageStore, PageStoreFactory, +}; +use parquet::basic::Compression; +use parquet::errors::Result; +use parquet::file::properties::WriterProperties; +use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System}; + +/// Write a skewed, wide Parquet file and compare peak writer memory with and +/// without a spilling `PageStore`. +#[derive(Parser)] +#[command(version, about, long_about = None)] +struct Args { + /// Number of large (~8 KiB average) string columns to write — the fat + /// columns that make the in-memory page buffer blow up. + #[arg(long, default_value_t = 10)] + large_string_columns: usize, + + /// Number of small (~20 byte average) string columns to write. + #[arg(long, default_value_t = 5)] + small_string_columns: usize, + + /// Number of integer (`Int64`) columns to write. + #[arg(long, default_value_t = 3)] + int_columns: usize, + + /// Total number of rows, all funnelled into a single row group. + #[arg(long, default_value_t = 2048)] + rows: usize, + + /// Rows per input batch fed to the writer. Each batch is dropped right after + /// it is written, so only the writer's internal buffering accumulates. + #[arg(long, default_value_t = 256)] + batch_size: usize, + + /// Use the spilling [`TempFilePageStore`] instead of the default in-memory + /// page buffering. + #[arg(long)] + spill: bool, + + /// Optional path to write the Parquet file to. Defaults to `io::sink()` so + /// the produced file bytes never live on the heap and the reported memory + /// reflects only the writer's page buffering. + #[arg(long)] + output: Option, +} + +/// Average length, in bytes, of values in a "large" string column. +const LARGE_AVG_LEN: usize = 8 * 1024; +/// Average length, in bytes, of values in a "small" string column. +const SMALL_AVG_LEN: usize = 20; + +// --------------------------------------------------------------------------- +// The spilling page store. +// +// A `PageStore` is intentionally "dumb": it maps an opaque, store-allocated +// `PageKey` to a blob of bytes and knows nothing about pages, dictionaries, or +// ordering. The caller (`ArrowWriter`) keeps the handles and decides what they +// mean. That is all a backend has to implement to move the page buffer off the +// heap. +// --------------------------------------------------------------------------- + +/// A spilling [`PageStore`]: one temp file per column chunk. +/// +/// `put` appends the blob to the file and records its `(offset, len)`; `take` +/// seeks and reads it back. The file is unlinked on creation (via +/// [`tempfile::tempfile`]) so the OS reclaims it when the store is dropped. +struct TempFilePageStore { + file: File, + /// Logical end of the file — where the next `put` appends. + end: u64, + /// `(offset, len)` for each stored blob, indexed by the `PageKey` we minted. + locs: Vec<(u64, usize)>, +} + +impl TempFilePageStore { + fn new() -> Result { + Ok(Self { + file: tempfile::tempfile()?, + end: 0, + locs: Vec::new(), + }) + } +} + +impl PageStore for TempFilePageStore { + fn put(&mut self, value: Bytes) -> Result { + // Always append at the logical end (a prior `take` may have moved the + // OS file cursor). + 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 { + 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)) + } + + // `memory_size` keeps its default of 0: once a blob is handed to `put` it + // lives in the temp file, not on the heap. That zero is what makes + // `ArrowWriter::memory_size()` drop to just the in-flight encoder buffers. +} + +/// Creates a fresh [`TempFilePageStore`] for every column chunk the writer opens. +#[derive(Debug, Default)] +struct TempFilePageStoreFactory; + +impl PageStoreFactory for TempFilePageStoreFactory { + fn create(&self, _column_index: usize) -> Result> { + Ok(Box::new(TempFilePageStore::new()?)) + } +} + +// --------------------------------------------------------------------------- +// Schema + data generation. +// --------------------------------------------------------------------------- + +/// Build the wide, skewed schema: a few integer columns, then a bunch of small +/// string columns, then the fat large string columns. +fn build_schema(args: &Args) -> SchemaRef { + let mut fields = Vec::new(); + for i in 0..args.int_columns { + fields.push(Field::new(format!("int_{i}"), DataType::Int64, false)); + } + for i in 0..args.small_string_columns { + fields.push(Field::new(format!("small_str_{i}"), DataType::Utf8, false)); + } + for i in 0..args.large_string_columns { + fields.push(Field::new(format!("large_str_{i}"), DataType::Utf8, false)); + } + Arc::new(Schema::new(fields)) +} + +/// A tiny deterministic xorshift RNG so runs are reproducible without pulling in +/// a `rand` dependency. +struct XorShift(u64); + +impl XorShift { + fn next(&mut self) -> u64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + self.0 + } +} + +/// Build a string column of `rows` rows whose values average `avg_len` bytes. +/// +/// Lengths vary uniformly in `[1, 2 * avg_len)` (so the mean is ≈ `avg_len`) and +/// the content is high-entropy printable ASCII, so the values neither +/// dictionary-encode nor compress away — the page buffer holds real bytes. +fn make_string_array(rows: usize, avg_len: usize, rng: &mut XorShift) -> ArrayRef { + let mut builder = StringBuilder::with_capacity(rows, rows * avg_len); + let mut value = String::new(); + for _ in 0..rows { + let len = 1 + (rng.next() as usize % (2 * avg_len - 1)); + value.clear(); + for _ in 0..len { + // Map to printable ASCII (33..=126). + value.push((33 + (rng.next() % 94) as u8) as char); + } + builder.append_value(&value); + } + Arc::new(builder.finish()) +} + +/// Build one record batch of `rows` rows for `schema`. +fn make_batch(schema: &SchemaRef, args: &Args, rows: usize, rng: &mut XorShift) -> RecordBatch { + let mut columns: Vec = Vec::with_capacity(schema.fields().len()); + for _ in 0..args.int_columns { + let vals: Vec = (0..rows).map(|_| rng.next() as i64).collect(); + columns.push(Arc::new(Int64Array::from(vals))); + } + for _ in 0..args.small_string_columns { + columns.push(make_string_array(rows, SMALL_AVG_LEN, rng)); + } + for _ in 0..args.large_string_columns { + columns.push(make_string_array(rows, LARGE_AVG_LEN, rng)); + } + RecordBatch::try_new(schema.clone(), columns).unwrap() +} + +// --------------------------------------------------------------------------- +// Memory reporting. +// --------------------------------------------------------------------------- + +/// Current process resident set size (RSS), in bytes. +fn rss_bytes(system: &mut System) -> u64 { + let Ok(pid) = sysinfo::get_current_pid() else { + return 0; + }; + system.refresh_processes_specifics( + ProcessesToUpdate::Some(&[pid]), + true, + ProcessRefreshKind::everything(), + ); + system.process(pid).map(|p| p.memory()).unwrap_or(0) +} + +fn mib(bytes: usize) -> f64 { + bytes as f64 / (1024.0 * 1024.0) +} + +fn main() -> Result<()> { + let args = Args::parse(); + let schema = build_schema(&args); + + // One uncompressed row group for the whole dataset, so the page buffer (the + // thing a PageStore governs) is the only thing that grows. Uncompressed keeps + // the reported numbers easy to reason about — the buffer holds the raw page + // bytes. + let props = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_max_row_group_row_count(Some(args.rows * 2)) + .build(); + + let mut options = ArrowWriterOptions::new().with_properties(props); + if args.spill { + options = options.with_page_store_factory(Arc::new(TempFilePageStoreFactory)); + } + + // Total logical payload across the large columns — the part that dominates. + let large_payload = args.large_string_columns * LARGE_AVG_LEN * args.rows; + println!( + "Writing {} rows × {} columns ({} int, {} small-string ~{}B, {} large-string ~{}B)", + args.rows, + args.int_columns + args.small_string_columns + args.large_string_columns, + args.int_columns, + args.small_string_columns, + SMALL_AVG_LEN, + args.large_string_columns, + LARGE_AVG_LEN, + ); + println!( + "Page buffering: {} (large-column payload ≈ {:.1} MiB)", + if args.spill { + "TempFilePageStore (spilling to temp files)" + } else { + "InMemoryPageStore (default, on the heap)" + }, + mib(large_payload), + ); + + let mut system = System::new_with_specifics(RefreshKind::everything()); + let rss_start = rss_bytes(&mut system); + + // The output sink. `io::sink()` discards the file bytes so they never inflate + // the heap — the measured peak then reflects only the writer's buffering. + let writer_sink: Box = match &args.output { + Some(path) => Box::new(File::create(path)?), + None => Box::new(std::io::sink()), + }; + let mut writer = ArrowWriter::try_new_with_options(writer_sink, schema.clone(), options)?; + + let mut rng = XorShift(0x9E37_79B9_7F4A_7C15); + let mut peak_writer_memory = 0usize; + let mut peak_rss = rss_start; + + let mut written = 0; + while written < args.rows { + let rows = args.batch_size.min(args.rows - written); + let batch = make_batch(&schema, &args, rows, &mut rng); + writer.write(&batch)?; + written += rows; + // `batch` is dropped here — only the writer's internal page buffering + // persists. `memory_size()` reports the bytes the writer holds *resident* + // on the heap: with the in-memory store this climbs toward the whole row + // group; with the spilling store it stays flat. + peak_writer_memory = peak_writer_memory.max(writer.memory_size()); + peak_rss = peak_rss.max(rss_bytes(&mut system)); + } + + peak_writer_memory = peak_writer_memory.max(writer.memory_size()); + writer.close()?; + peak_rss = peak_rss.max(rss_bytes(&mut system)); + + println!(); + println!("Done. Wrote {written} rows."); + println!( + "Peak ArrowWriter::memory_size() : {:>8.1} MiB <- bytes the writer held on the heap", + mib(peak_writer_memory), + ); + println!( + "Peak process RSS delta : {:>8.1} MiB", + (peak_rss.saturating_sub(rss_start)) as f64 / (1024.0 * 1024.0), + ); + println!(); + if args.spill { + println!( + "With spilling, peak writer memory is bounded by the in-flight encoder \n\ + buffers (a page or two per column), not the {:.1} MiB row group.", + mib(large_payload), + ); + } else { + println!( + "Re-run with --spill to keep those pages off the heap and watch peak \n\ + writer memory drop well below the {:.1} MiB row group payload.", + mib(large_payload), + ); + } + + Ok(()) +} From d94a9c6c54bb75555f7caf7a7aaa1a8d187afe46 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 3 Jun 2026 13:11:32 -0400 Subject: [PATCH 09/13] parquet: report spill stats and default batch_size to 8192 in example MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The spill_page_store example now threads a shared SpillStats (atomics) through the TempFilePageStoreFactory into each per-column store, and prints the number of temp files created (one per column chunk) plus total bytes spilled — written on `put` and read back on `take`. Also bumps the default `--batch-size` to 8192. Co-Authored-By: Claude Opus 4.8 (1M context) --- parquet/examples/spill_page_store.rs | 58 ++++++++++++++++++++++++---- 1 file changed, 51 insertions(+), 7 deletions(-) diff --git a/parquet/examples/spill_page_store.rs b/parquet/examples/spill_page_store.rs index ed3f61afdfc1..7852f272bd43 100644 --- a/parquet/examples/spill_page_store.rs +++ b/parquet/examples/spill_page_store.rs @@ -53,6 +53,7 @@ use std::fs::File; use std::io::{Read, Seek, SeekFrom, Write}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringBuilder}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; @@ -91,7 +92,7 @@ struct Args { /// Rows per input batch fed to the writer. Each batch is dropped right after /// it is written, so only the writer's internal buffering accumulates. - #[arg(long, default_value_t = 256)] + #[arg(long, default_value_t = 8192)] batch_size: usize, /// Use the spilling [`TempFilePageStore`] instead of the default in-memory @@ -121,6 +122,20 @@ const SMALL_AVG_LEN: usize = 20; // heap. // --------------------------------------------------------------------------- +/// Shared counters describing how much spilling actually happened, aggregated +/// across every per-column [`TempFilePageStore`]. Each store holds an `Arc` to +/// the same instance and bumps the atomics, so the totals survive the stores +/// being dropped at row group flush. +#[derive(Debug, Default)] +struct SpillStats { + /// Number of temp files created — one per column chunk that was opened. + files_created: AtomicU64, + /// Total bytes handed to `put` and written to a temp file. + bytes_written: AtomicU64, + /// Total bytes read back out of a temp file by `take` (at row group flush). + bytes_read: AtomicU64, +} + /// A spilling [`PageStore`]: one temp file per column chunk. /// /// `put` appends the blob to the file and records its `(offset, len)`; `take` @@ -132,14 +147,18 @@ struct TempFilePageStore { end: u64, /// `(offset, len)` for each stored blob, indexed by the `PageKey` we minted. locs: Vec<(u64, usize)>, + /// Shared, cross-store spill counters (see [`SpillStats`]). + stats: Arc, } impl TempFilePageStore { - fn new() -> Result { + fn new(stats: Arc) -> Result { + stats.files_created.fetch_add(1, Ordering::Relaxed); Ok(Self { file: tempfile::tempfile()?, end: 0, locs: Vec::new(), + stats, }) } } @@ -150,6 +169,9 @@ impl PageStore for TempFilePageStore { // OS file cursor). self.file.seek(SeekFrom::Start(self.end))?; self.file.write_all(&value)?; + self.stats + .bytes_written + .fetch_add(value.len() as u64, Ordering::Relaxed); let key = PageKey::new(self.locs.len() as u64); self.locs.push((self.end, value.len())); self.end += value.len() as u64; @@ -161,6 +183,9 @@ impl PageStore for TempFilePageStore { let mut buf = vec![0u8; len]; self.file.seek(SeekFrom::Start(offset))?; self.file.read_exact(&mut buf)?; + self.stats + .bytes_read + .fetch_add(len as u64, Ordering::Relaxed); Ok(Bytes::from(buf)) } @@ -169,13 +194,16 @@ impl PageStore for TempFilePageStore { // `ArrowWriter::memory_size()` drop to just the in-flight encoder buffers. } -/// Creates a fresh [`TempFilePageStore`] for every column chunk the writer opens. -#[derive(Debug, Default)] -struct TempFilePageStoreFactory; +/// Creates a fresh [`TempFilePageStore`] for every column chunk the writer opens, +/// handing each one an `Arc` to the shared [`SpillStats`]. +#[derive(Debug)] +struct TempFilePageStoreFactory { + stats: Arc, +} impl PageStoreFactory for TempFilePageStoreFactory { fn create(&self, _column_index: usize) -> Result> { - Ok(Box::new(TempFilePageStore::new()?)) + Ok(Box::new(TempFilePageStore::new(self.stats.clone())?)) } } @@ -283,8 +311,11 @@ fn main() -> Result<()> { .build(); let mut options = ArrowWriterOptions::new().with_properties(props); + let spill_stats = Arc::new(SpillStats::default()); if args.spill { - options = options.with_page_store_factory(Arc::new(TempFilePageStoreFactory)); + options = options.with_page_store_factory(Arc::new(TempFilePageStoreFactory { + stats: spill_stats.clone(), + })); } // Total logical payload across the large columns — the part that dominates. @@ -354,6 +385,19 @@ fn main() -> Result<()> { ); println!(); if args.spill { + let files = spill_stats.files_created.load(Ordering::Relaxed); + let written_bytes = spill_stats.bytes_written.load(Ordering::Relaxed); + let read_bytes = spill_stats.bytes_read.load(Ordering::Relaxed); + println!("Spill temp files created : {files:>8} <- one per column chunk"); + println!( + "Total bytes spilled (written) : {:>8.1} MiB", + mib(written_bytes as usize), + ); + println!( + "Total bytes read back (take) : {:>8.1} MiB", + mib(read_bytes as usize), + ); + println!(); println!( "With spilling, peak writer memory is bounded by the in-flight encoder \n\ buffers (a page or two per column), not the {:.1} MiB row group.", From 4294548b53ee863eb9513b13d15c89e4f17216a1 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 3 Jun 2026 13:32:30 -0400 Subject: [PATCH 10/13] parquet: list each spill file's path and byte count in example Switches the example's spilling backend from anonymous tempfile() handles to NamedTempFile so each per-column temp file has a reportable path, records (column index, path, bytes) for every file when its store is dropped, and prints one line per spill file alongside the existing totals. Co-Authored-By: Claude Opus 4.8 (1M context) --- parquet/examples/spill_page_store.rs | 83 ++++++++++++++++++++++------ 1 file changed, 65 insertions(+), 18 deletions(-) diff --git a/parquet/examples/spill_page_store.rs b/parquet/examples/spill_page_store.rs index 7852f272bd43..9f04b2007154 100644 --- a/parquet/examples/spill_page_store.rs +++ b/parquet/examples/spill_page_store.rs @@ -52,8 +52,11 @@ use std::fs::File; use std::io::{Read, Seek, SeekFrom, Write}; -use std::sync::Arc; +use std::path::PathBuf; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; + +use tempfile::NamedTempFile; use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringBuilder}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; @@ -122,40 +125,54 @@ const SMALL_AVG_LEN: usize = 20; // heap. // --------------------------------------------------------------------------- -/// Shared counters describing how much spilling actually happened, aggregated -/// across every per-column [`TempFilePageStore`]. Each store holds an `Arc` to -/// the same instance and bumps the atomics, so the totals survive the stores -/// being dropped at row group flush. +/// Shared spill bookkeeping, aggregated across every per-column +/// [`TempFilePageStore`]. Each store holds an `Arc` to the same instance, so the +/// totals and per-file records survive the stores being dropped at row group +/// flush. #[derive(Debug, Default)] struct SpillStats { - /// Number of temp files created — one per column chunk that was opened. - files_created: AtomicU64, /// Total bytes handed to `put` and written to a temp file. bytes_written: AtomicU64, /// Total bytes read back out of a temp file by `take` (at row group flush). bytes_read: AtomicU64, + /// One record per temp file, pushed when each store is dropped. + files: Mutex>, +} + +/// What a single spill temp file ended up holding. +#[derive(Debug)] +struct FileRecord { + /// Leaf column index the store was created for. + column_index: usize, + /// Filesystem path of the temp file (valid while the store was alive). + path: PathBuf, + /// Total bytes written into the file. + bytes: u64, } /// A spilling [`PageStore`]: one temp file per column chunk. /// /// `put` appends the blob to the file and records its `(offset, len)`; `take` -/// seeks and reads it back. The file is unlinked on creation (via -/// [`tempfile::tempfile`]) so the OS reclaims it when the store is dropped. +/// seeks and reads it back. A [`NamedTempFile`] is used (rather than an +/// anonymous one) so the file has a reportable path; the OS reclaims it when the +/// store is dropped at row group flush. struct TempFilePageStore { - file: File, + file: NamedTempFile, + /// Leaf column index this store backs (used only for reporting). + column_index: usize, /// Logical end of the file — where the next `put` appends. end: u64, /// `(offset, len)` for each stored blob, indexed by the `PageKey` we minted. locs: Vec<(u64, usize)>, - /// Shared, cross-store spill counters (see [`SpillStats`]). + /// Shared, cross-store spill bookkeeping (see [`SpillStats`]). stats: Arc, } impl TempFilePageStore { - fn new(stats: Arc) -> Result { - stats.files_created.fetch_add(1, Ordering::Relaxed); + fn new(stats: Arc, column_index: usize) -> Result { Ok(Self { - file: tempfile::tempfile()?, + file: NamedTempFile::new()?, + column_index, end: 0, locs: Vec::new(), stats, @@ -194,6 +211,18 @@ impl PageStore for TempFilePageStore { // `ArrowWriter::memory_size()` drop to just the in-flight encoder buffers. } +impl Drop for TempFilePageStore { + fn drop(&mut self) { + // Record this file's path and final byte count before the NamedTempFile + // is unlinked, so `main` can list it after the writer is closed. + self.stats.files.lock().unwrap().push(FileRecord { + column_index: self.column_index, + path: self.file.path().to_path_buf(), + bytes: self.end, + }); + } +} + /// Creates a fresh [`TempFilePageStore`] for every column chunk the writer opens, /// handing each one an `Arc` to the shared [`SpillStats`]. #[derive(Debug)] @@ -202,8 +231,11 @@ struct TempFilePageStoreFactory { } impl PageStoreFactory for TempFilePageStoreFactory { - fn create(&self, _column_index: usize) -> Result> { - Ok(Box::new(TempFilePageStore::new(self.stats.clone())?)) + fn create(&self, column_index: usize) -> Result> { + Ok(Box::new(TempFilePageStore::new( + self.stats.clone(), + column_index, + )?)) } } @@ -385,10 +417,14 @@ fn main() -> Result<()> { ); println!(); if args.spill { - let files = spill_stats.files_created.load(Ordering::Relaxed); + let mut files = spill_stats.files.lock().unwrap(); + files.sort_by_key(|f| f.column_index); let written_bytes = spill_stats.bytes_written.load(Ordering::Relaxed); let read_bytes = spill_stats.bytes_read.load(Ordering::Relaxed); - println!("Spill temp files created : {files:>8} <- one per column chunk"); + println!( + "Spill temp files created : {:>8} <- one per column chunk", + files.len(), + ); println!( "Total bytes spilled (written) : {:>8.1} MiB", mib(written_bytes as usize), @@ -398,6 +434,17 @@ fn main() -> Result<()> { mib(read_bytes as usize), ); println!(); + println!("Per spill file (column index, path, bytes stored):"); + for f in files.iter() { + println!( + " col {:>3} {} {} bytes ({:.1} MiB)", + f.column_index, + f.path.display(), + f.bytes, + mib(f.bytes as usize), + ); + } + println!(); println!( "With spilling, peak writer memory is bounded by the in-flight encoder \n\ buffers (a page or two per column), not the {:.1} MiB row group.", From d382296ef5463620d4562494c572ce2a6eac0dcb Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 3 Jun 2026 13:40:12 -0400 Subject: [PATCH 11/13] parquet: encode columns concurrently in the spill example + show column type Rewrites the spill_page_store example to use the low-level ArrowColumnWriter API so columns are encoded on multiple threads (one worker thread per column, the documented parallel pattern). Each worker owns its own ArrowColumnWriter backed by its own PageStore, receives ArrowLeafColumns over a bounded channel, tracks its writer's peak resident bytes, and returns the finished ArrowColumnChunk; the main thread splices each chunk into the row group, streaming pages back out of the store. The spilling benefit is preserved (append_to_row_group reads one page at a time from the store), and the output stays byte-identical to the in-memory path. Also adds the column data type to the per-spill-file listing. Co-Authored-By: Claude Opus 4.8 (1M context) --- parquet/examples/spill_page_store.rs | 140 +++++++++++++++++++++------ 1 file changed, 108 insertions(+), 32 deletions(-) diff --git a/parquet/examples/spill_page_store.rs b/parquet/examples/spill_page_store.rs index 9f04b2007154..e31cd6bc4622 100644 --- a/parquet/examples/spill_page_store.rs +++ b/parquet/examples/spill_page_store.rs @@ -23,15 +23,27 @@ //! //! Parquet requires every column chunk to be contiguous in the file, but Arrow //! record batches arrive with all columns interleaved. So while a row group is -//! being written, [`ArrowWriter`] must buffer every column's completed, -//! compressed pages until the row group is flushed. Peak write memory therefore -//! grows with the row group size — painful for wide schemas with large, skewed -//! columns (e.g. a few `id` columns next to a pile of fat string columns). +//! being written, the writer must buffer every column's completed, compressed +//! pages until the row group is flushed. Peak write memory therefore grows with +//! the row group size — painful for wide schemas with large, skewed columns +//! (e.g. a few `id` columns next to a pile of fat string columns). //! //! A [`PageStore`] lets that page buffer live somewhere other than the heap. This //! example plugs in a [`TempFilePageStore`] (one temp file per column chunk) and //! compares peak writer memory against the default in-memory buffering. //! +//! # Concurrency +//! +//! To exercise the store from multiple threads, this example uses the low-level +//! [`ArrowColumnWriter`] API (rather than the single-threaded [`ArrowWriter`]): +//! it spawns one worker thread per column, hands each its own +//! [`ArrowColumnWriter`], and fans each record batch's columns out to the +//! workers via [`compute_leaves`]. Every column encodes on its own thread into +//! its own [`PageStore`], and the finished [`ArrowColumnChunk`]s are spliced into +//! the row group — streaming back out of the store one page at a time — on the +//! main thread. (One thread per column is for clarity; production code would use +//! a bounded pool such as rayon or tokio.) +//! //! # Running //! //! ```sh @@ -48,13 +60,18 @@ //! ``` //! //! [`ArrowWriter`]: parquet::arrow::ArrowWriter +//! [`ArrowColumnWriter`]: parquet::arrow::arrow_writer::ArrowColumnWriter +//! [`ArrowColumnChunk`]: parquet::arrow::arrow_writer::ArrowColumnChunk +//! [`compute_leaves`]: parquet::arrow::arrow_writer::compute_leaves //! [`PageStore`]: parquet::arrow::arrow_writer::PageStore use std::fs::File; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::PathBuf; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::mpsc::sync_channel; use std::sync::{Arc, Mutex}; +use std::thread; use tempfile::NamedTempFile; @@ -62,13 +79,15 @@ use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringBuilder}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use bytes::Bytes; use clap::Parser; -use parquet::arrow::ArrowWriter; +use parquet::arrow::ArrowSchemaConverter; use parquet::arrow::arrow_writer::{ - ArrowWriterOptions, PageKey, PageStore, PageStoreFactory, + ArrowColumnChunk, ArrowLeafColumn, ArrowRowGroupWriterFactory, PageKey, PageStore, + PageStoreFactory, compute_leaves, }; use parquet::basic::Compression; use parquet::errors::Result; use parquet::file::properties::WriterProperties; +use parquet::file::writer::SerializedFileWriter; use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System}; /// Write a skewed, wide Parquet file and compare peak writer memory with and @@ -337,18 +356,14 @@ fn main() -> Result<()> { // thing a PageStore governs) is the only thing that grows. Uncompressed keeps // the reported numbers easy to reason about — the buffer holds the raw page // bytes. - let props = WriterProperties::builder() - .set_compression(Compression::UNCOMPRESSED) - .set_max_row_group_row_count(Some(args.rows * 2)) - .build(); + let props = Arc::new( + WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_max_row_group_row_count(Some(args.rows * 2)) + .build(), + ); - let mut options = ArrowWriterOptions::new().with_properties(props); let spill_stats = Arc::new(SpillStats::default()); - if args.spill { - options = options.with_page_store_factory(Arc::new(TempFilePageStoreFactory { - stats: spill_stats.clone(), - })); - } // Total logical payload across the large columns — the part that dominates. let large_payload = args.large_string_columns * LARGE_AVG_LEN * args.rows; @@ -371,44 +386,102 @@ fn main() -> Result<()> { }, mib(large_payload), ); + println!("Encoding columns concurrently: one worker thread per column."); let mut system = System::new_with_specifics(RefreshKind::everything()); let rss_start = rss_bytes(&mut system); - // The output sink. `io::sink()` discards the file bytes so they never inflate - // the heap — the measured peak then reflects only the writer's buffering. + // Build the lower-level file writer so columns can be encoded in parallel. + // `io::sink()` discards the produced file bytes so they never inflate the + // heap — the measured peak then reflects only the writer's page buffering. + let parquet_schema = ArrowSchemaConverter::new() + .with_coerce_types(props.coerce_types()) + .convert(&schema)?; let writer_sink: Box = match &args.output { Some(path) => Box::new(File::create(path)?), None => Box::new(std::io::sink()), }; - let mut writer = ArrowWriter::try_new_with_options(writer_sink, schema.clone(), options)?; + let mut file_writer = + SerializedFileWriter::new(writer_sink, parquet_schema.root_schema_ptr(), props.clone())?; + // One `ArrowColumnWriter` per leaf column, each backed by its own PageStore. + // With `--spill`, that store is a TempFilePageStore, so each worker thread's + // completed pages land in a temp file instead of the heap. + let mut factory = ArrowRowGroupWriterFactory::new(&file_writer, schema.clone()); + if args.spill { + factory = factory.with_page_store_factory(Arc::new(TempFilePageStoreFactory { + stats: spill_stats.clone(), + })); + } + let col_writers = factory.create_column_writers(0)?; + + // Spawn a worker per column. Each owns its `ArrowColumnWriter`, receives + // `ArrowLeafColumn`s over a small bounded channel (back-pressure keeps in- + // flight input from piling up), tracks the peak bytes its writer held + // resident, and returns the finished chunk plus that peak. The bounded + // channel makes the workers run concurrently with batch generation. + let workers: Vec<_> = col_writers + .into_iter() + .map(|mut col_writer| { + let (send, recv) = sync_channel::(2); + let handle = thread::spawn(move || -> Result<(ArrowColumnChunk, usize)> { + let mut peak_memory = 0usize; + for leaf in recv { + col_writer.write(&leaf)?; + // `memory_size()` is the bytes this column's writer holds + // resident — pages in its PageStore plus in-flight encoder + // buffers. With the in-memory store it climbs toward the whole + // column chunk; with spilling it stays flat. + peak_memory = peak_memory.max(col_writer.memory_size()); + } + peak_memory = peak_memory.max(col_writer.memory_size()); + Ok((col_writer.close()?, peak_memory)) + }); + (handle, send) + }) + .collect(); + + // Generate batches and fan each batch's columns out to the workers. `schema` + // is flat, so leaf order matches field order and each field maps to exactly + // one worker. let mut rng = XorShift(0x9E37_79B9_7F4A_7C15); - let mut peak_writer_memory = 0usize; let mut peak_rss = rss_start; - let mut written = 0; while written < args.rows { let rows = args.batch_size.min(args.rows - written); let batch = make_batch(&schema, &args, rows, &mut rng); - writer.write(&batch)?; + for (col_idx, (field, array)) in schema.fields().iter().zip(batch.columns()).enumerate() { + for leaf in compute_leaves(field, array)? { + // Blocks if this worker is busy — bounding in-flight input. + workers[col_idx].1.send(leaf).unwrap(); + } + } written += rows; - // `batch` is dropped here — only the writer's internal page buffering - // persists. `memory_size()` reports the bytes the writer holds *resident* - // on the heap: with the in-memory store this climbs toward the whole row - // group; with the spilling store it stays flat. - peak_writer_memory = peak_writer_memory.max(writer.memory_size()); + // `batch` is dropped here; its data lives on in whatever leaves are still + // in flight to the workers. peak_rss = peak_rss.max(rss_bytes(&mut system)); } - peak_writer_memory = peak_writer_memory.max(writer.memory_size()); - writer.close()?; + // Signal end-of-input to every worker, then join in column order and splice + // each finished chunk into the row group (streaming pages back out of the + // store). Columns must be appended in schema order. + let mut row_group_writer = file_writer.next_row_group()?; + let mut peak_writer_memory = 0usize; + for (handle, send) in workers { + drop(send); // closes the channel so the worker's `for leaf in recv` ends + let (chunk, col_peak) = handle.join().expect("worker thread panicked")?; + peak_writer_memory += col_peak; + chunk.append_to_row_group(&mut row_group_writer)?; + peak_rss = peak_rss.max(rss_bytes(&mut system)); + } + row_group_writer.close()?; + file_writer.close()?; peak_rss = peak_rss.max(rss_bytes(&mut system)); println!(); println!("Done. Wrote {written} rows."); println!( - "Peak ArrowWriter::memory_size() : {:>8.1} MiB <- bytes the writer held on the heap", + "Peak writer memory (Σ per-column): {:>8.1} MiB <- bytes the column writers held on the heap", mib(peak_writer_memory), ); println!( @@ -434,11 +507,14 @@ fn main() -> Result<()> { mib(read_bytes as usize), ); println!(); - println!("Per spill file (column index, path, bytes stored):"); + println!("Per spill file (column index, type, path, bytes stored):"); for f in files.iter() { + // `schema` is flat, so the leaf column index is also the field index. + let field = schema.field(f.column_index); println!( - " col {:>3} {} {} bytes ({:.1} MiB)", + " col {:>3} {:<8} {} {} bytes ({:.1} MiB)", f.column_index, + format!("{}", field.data_type()), f.path.display(), f.bytes, mib(f.bytes as usize), From 995f80df2b2b7ad4e72d6ca5bf3b9821e317f16d Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 3 Jun 2026 13:46:11 -0400 Subject: [PATCH 12/13] parquet: print total elapsed time in the spill example Adds a wall-clock timer (std::time::Instant) over the whole run and reports "Total elapsed time" in the final summary. Also includes the faster data generation (a single batch is built once and reused across iterations). Co-Authored-By: Claude Opus 4.8 (1M context) --- parquet/examples/spill_page_store.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/parquet/examples/spill_page_store.rs b/parquet/examples/spill_page_store.rs index e31cd6bc4622..3b02d9d0141d 100644 --- a/parquet/examples/spill_page_store.rs +++ b/parquet/examples/spill_page_store.rs @@ -72,6 +72,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::mpsc::sync_channel; use std::sync::{Arc, Mutex}; use std::thread; +use std::time::Instant; use tempfile::NamedTempFile; @@ -349,6 +350,7 @@ fn mib(bytes: usize) -> f64 { } fn main() -> Result<()> { + let start = Instant::now(); let args = Args::parse(); let schema = build_schema(&args); @@ -447,9 +449,10 @@ fn main() -> Result<()> { let mut rng = XorShift(0x9E37_79B9_7F4A_7C15); let mut peak_rss = rss_start; let mut written = 0; + let rows = args.batch_size.min(args.rows - written); + let batch = make_batch(&schema, &args, rows, &mut rng); while written < args.rows { - let rows = args.batch_size.min(args.rows - written); - let batch = make_batch(&schema, &args, rows, &mut rng); + //let batch = make_batch(&schema, &args, rows, &mut rng); for (col_idx, (field, array)) in schema.fields().iter().zip(batch.columns()).enumerate() { for leaf in compute_leaves(field, array)? { // Blocks if this worker is busy — bounding in-flight input. @@ -477,6 +480,7 @@ fn main() -> Result<()> { row_group_writer.close()?; file_writer.close()?; peak_rss = peak_rss.max(rss_bytes(&mut system)); + let elapsed = start.elapsed(); println!(); println!("Done. Wrote {written} rows."); @@ -488,6 +492,10 @@ fn main() -> Result<()> { "Peak process RSS delta : {:>8.1} MiB", (peak_rss.saturating_sub(rss_start)) as f64 / (1024.0 * 1024.0), ); + println!( + "Total elapsed time : {:>8.3} s", + elapsed.as_secs_f64(), + ); println!(); if args.spill { let mut files = spill_stats.files.lock().unwrap(); From ac4f0956fad62fa7d5dc519eef78be69a3a013f2 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 3 Jun 2026 14:14:41 -0400 Subject: [PATCH 13/13] parquet: split spill example into generator + encoder thread pools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Restructures the example to keep all cores busy: N/2 generator threads build record batches and N/2 encoder threads encode them (N = available cores). The main thread re-orders generated batches by index — keeping the output file deterministic regardless of thread timing — and broadcasts each to the encoders, which each own a disjoint subset of columns (and their PageStores). Finished chunks are sorted into schema order and spliced in. Also removes the XorShift RNG in favor of deterministic counter-derived values, so a run is fully reproducible (verified byte-identical between the in-memory and spilling backends, and run-to-run). Channel depths are kept shallow so in-flight input batches don't dominate RSS on wide/huge schemas. Co-Authored-By: Claude Opus 4.8 (1M context) --- parquet/examples/spill_page_store.rs | 276 ++++++++++++++++++--------- 1 file changed, 191 insertions(+), 85 deletions(-) diff --git a/parquet/examples/spill_page_store.rs b/parquet/examples/spill_page_store.rs index 3b02d9d0141d..38e1a40e0494 100644 --- a/parquet/examples/spill_page_store.rs +++ b/parquet/examples/spill_page_store.rs @@ -34,15 +34,26 @@ //! //! # Concurrency //! -//! To exercise the store from multiple threads, this example uses the low-level -//! [`ArrowColumnWriter`] API (rather than the single-threaded [`ArrowWriter`]): -//! it spawns one worker thread per column, hands each its own -//! [`ArrowColumnWriter`], and fans each record batch's columns out to the -//! workers via [`compute_leaves`]. Every column encodes on its own thread into -//! its own [`PageStore`], and the finished [`ArrowColumnChunk`]s are spliced into -//! the row group — streaming back out of the store one page at a time — on the -//! main thread. (One thread per column is for clarity; production code would use -//! a bounded pool such as rayon or tokio.) +//! To keep every core busy this example splits the work across two thread pools, +//! sized to the machine: `N/2` **generator** threads building record batches and +//! `N/2` **encoder** threads encoding them (`N` = available cores). It uses the +//! low-level [`ArrowColumnWriter`] API (rather than the single-threaded +//! [`ArrowWriter`]) so encoding can be parallelized: +//! +//! - Each generator claims the next batch index from a shared counter, builds +//! that batch deterministically, and sends it to the main thread. +//! - The main thread re-orders batches by index (so the output is deterministic +//! regardless of how the generators interleave) and broadcasts each one to all +//! encoders. Batches are cheap to share — the columns are reference-counted. +//! - The columns are distributed across the encoder threads, each owning a +//! disjoint subset of [`ArrowColumnWriter`]s backed by their own [`PageStore`]. +//! Each encoder picks out its columns from every batch via [`compute_leaves`]. +//! - The finished [`ArrowColumnChunk`]s are collected, sorted into schema order, +//! and spliced into the row group — streaming back out of the store one page +//! at a time — on the main thread. +//! +//! (For clarity this hand-rolls the pools with threads and channels; production +//! code would use rayon or tokio.) //! //! # Running //! @@ -65,10 +76,11 @@ //! [`compute_leaves`]: parquet::arrow::arrow_writer::compute_leaves //! [`PageStore`]: parquet::arrow::arrow_writer::PageStore +use std::collections::HashMap; use std::fs::File; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::PathBuf; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::mpsc::sync_channel; use std::sync::{Arc, Mutex}; use std::thread; @@ -82,7 +94,7 @@ use bytes::Bytes; use clap::Parser; use parquet::arrow::ArrowSchemaConverter; use parquet::arrow::arrow_writer::{ - ArrowColumnChunk, ArrowLeafColumn, ArrowRowGroupWriterFactory, PageKey, PageStore, + ArrowColumnChunk, ArrowColumnWriter, ArrowRowGroupWriterFactory, PageKey, PageStore, PageStoreFactory, compute_leaves, }; use parquet::basic::Compression; @@ -279,51 +291,69 @@ fn build_schema(args: &Args) -> SchemaRef { Arc::new(Schema::new(fields)) } -/// A tiny deterministic xorshift RNG so runs are reproducible without pulling in -/// a `rand` dependency. -struct XorShift(u64); +/// The per-batch column counts a generator needs (a small `Copy` view of `Args` +/// so it can be handed to each generator thread). +#[derive(Clone, Copy)] +struct BatchSpec { + int_columns: usize, + small_string_columns: usize, + large_string_columns: usize, +} -impl XorShift { - fn next(&mut self) -> u64 { - self.0 ^= self.0 << 13; - self.0 ^= self.0 >> 7; - self.0 ^= self.0 << 17; - self.0 +/// Fill `buf` with a deterministic value of exactly `len` bytes derived from the +/// counter `n`. +/// +/// The 20-digit zero-padded counter makes every value distinct, so the fat +/// columns stay plain-encoded (high cardinality) rather than dictionary-encoding +/// away; the remainder is padded with a fixed `a`–`z` cycle. No RNG — the content +/// is a pure function of `n`, so a run is fully reproducible. +fn fill_value(buf: &mut String, n: u64, len: usize) { + use std::fmt::Write; + buf.clear(); + let _ = write!(buf, "{n:020}"); + while buf.len() < len { + buf.push((b'a' + (buf.len() % 26) as u8) as char); } + buf.truncate(len); // all bytes are ASCII, so this is a clean char boundary } -/// Build a string column of `rows` rows whose values average `avg_len` bytes. -/// -/// Lengths vary uniformly in `[1, 2 * avg_len)` (so the mean is ≈ `avg_len`) and -/// the content is high-entropy printable ASCII, so the values neither -/// dictionary-encode nor compress away — the page buffer holds real bytes. -fn make_string_array(rows: usize, avg_len: usize, rng: &mut XorShift) -> ArrayRef { - let mut builder = StringBuilder::with_capacity(rows, rows * avg_len); +/// Build a string column of `rows` values, each exactly `len` bytes, keyed by the +/// global row index (`row_offset + r`) and a per-column `salt` so values are +/// distinct within the column. +fn make_string_array(rows: usize, row_offset: u64, salt: u64, len: usize) -> ArrayRef { + let mut builder = StringBuilder::with_capacity(rows, rows * len); let mut value = String::new(); - for _ in 0..rows { - let len = 1 + (rng.next() as usize % (2 * avg_len - 1)); - value.clear(); - for _ in 0..len { - // Map to printable ASCII (33..=126). - value.push((33 + (rng.next() % 94) as u8) as char); - } + for r in 0..rows { + let n = (row_offset + r as u64).wrapping_mul(101).wrapping_add(salt); + fill_value(&mut value, n, len); builder.append_value(&value); } Arc::new(builder.finish()) } -/// Build one record batch of `rows` rows for `schema`. -fn make_batch(schema: &SchemaRef, args: &Args, rows: usize, rng: &mut XorShift) -> RecordBatch { +/// Build the record batch covering rows `[row_offset, row_offset + rows)`. +/// +/// Fully deterministic in `row_offset`: every column's values are a pure function +/// of the global row index, so the batch a generator produces depends only on its +/// claimed index — not on thread timing. +fn make_batch(schema: &SchemaRef, spec: BatchSpec, row_offset: u64, rows: usize) -> RecordBatch { let mut columns: Vec = Vec::with_capacity(schema.fields().len()); - for _ in 0..args.int_columns { - let vals: Vec = (0..rows).map(|_| rng.next() as i64).collect(); + let mut salt = 0u64; // distinguishes columns so they don't all hold equal values + for _ in 0..spec.int_columns { + let s = salt; + salt += 1; + let vals: Vec = (0..rows) + .map(|r| (row_offset + r as u64 + s) as i64) + .collect(); columns.push(Arc::new(Int64Array::from(vals))); } - for _ in 0..args.small_string_columns { - columns.push(make_string_array(rows, SMALL_AVG_LEN, rng)); + for _ in 0..spec.small_string_columns { + columns.push(make_string_array(rows, row_offset, salt, SMALL_AVG_LEN)); + salt += 1; } - for _ in 0..args.large_string_columns { - columns.push(make_string_array(rows, LARGE_AVG_LEN, rng)); + for _ in 0..spec.large_string_columns { + columns.push(make_string_array(rows, row_offset, salt, LARGE_AVG_LEN)); + salt += 1; } RecordBatch::try_new(schema.clone(), columns).unwrap() } @@ -388,7 +418,18 @@ fn main() -> Result<()> { }, mib(large_payload), ); - println!("Encoding columns concurrently: one worker thread per column."); + // Split the cores: half generate batches, half encode them. + let num_cores = thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(2); + let num_generators = (num_cores / 2).max(1); + let num_encoders = (num_cores / 2).max(1); + let num_batches = args.rows.div_ceil(args.batch_size); + println!( + "Cores: {num_cores} ({num_generators} generator threads, {num_encoders} encoder threads, \ + {num_batches} batches of ≤{} rows)", + args.batch_size, + ); let mut system = System::new_with_specifics(RefreshKind::everything()); let rss_start = rss_bytes(&mut system); @@ -417,62 +458,126 @@ fn main() -> Result<()> { } let col_writers = factory.create_column_writers(0)?; - // Spawn a worker per column. Each owns its `ArrowColumnWriter`, receives - // `ArrowLeafColumn`s over a small bounded channel (back-pressure keeps in- - // flight input from piling up), tracks the peak bytes its writer held - // resident, and returns the finished chunk plus that peak. The bounded - // channel makes the workers run concurrently with batch generation. - let workers: Vec<_> = col_writers - .into_iter() - .map(|mut col_writer| { - let (send, recv) = sync_channel::(2); - let handle = thread::spawn(move || -> Result<(ArrowColumnChunk, usize)> { - let mut peak_memory = 0usize; - for leaf in recv { - col_writer.write(&leaf)?; + // Distribute the columns across the encoder threads round-robin, so the fat + // columns (contiguous at the end of the schema) spread evenly. Each encoder + // owns a disjoint set of `(column index, ArrowColumnWriter)` pairs. + let mut encoder_cols: Vec> = + (0..num_encoders).map(|_| Vec::new()).collect(); + for (idx, writer) in col_writers.into_iter().enumerate() { + encoder_cols[idx % num_encoders].push((idx, writer)); + } + + // Spawn the encoder pool. Each encoder receives whole batches over a small + // bounded channel (back-pressure), encodes only its own columns from each, + // tracks the peak bytes each writer held resident, and returns the finished + // chunks plus those peaks. + let mut encoder_txs = Vec::with_capacity(num_encoders); + let mut encoder_handles = Vec::with_capacity(num_encoders); + for mut cols in encoder_cols { + let (tx, rx) = sync_channel::>(1); + let schema = schema.clone(); + let handle = thread::spawn(move || -> Result> { + let mut peaks = vec![0usize; cols.len()]; + for batch in rx { + for (slot, (idx, writer)) in cols.iter_mut().enumerate() { + let field = &schema.fields()[*idx]; + for leaf in compute_leaves(field, batch.column(*idx))? { + writer.write(&leaf)?; + } // `memory_size()` is the bytes this column's writer holds // resident — pages in its PageStore plus in-flight encoder // buffers. With the in-memory store it climbs toward the whole // column chunk; with spilling it stays flat. - peak_memory = peak_memory.max(col_writer.memory_size()); + peaks[slot] = peaks[slot].max(writer.memory_size()); } - peak_memory = peak_memory.max(col_writer.memory_size()); - Ok((col_writer.close()?, peak_memory)) - }); - (handle, send) - }) - .collect(); + } + cols.into_iter() + .zip(peaks) + .map(|((idx, writer), peak)| { + let peak = peak.max(writer.memory_size()); + Ok((idx, writer.close()?, peak)) + }) + .collect() + }); + encoder_txs.push(tx); + encoder_handles.push(handle); + } + + // Spawn the generator pool. Each generator claims the next batch index from a + // shared counter, builds that batch deterministically, and sends `(index, + // batch)` to the main thread. The bounded channel applies back-pressure so + // generators don't race arbitrarily far ahead of the encoders. + // Keep the result channel shallow: with wide/huge schemas each batch can be + // hundreds of MiB, and they pipeline, so the in-flight input — not the + // (spilled) page buffer — dominates process RSS. A depth of 1 bounds it to + // roughly the working set the generators and encoders are actively touching. + let next_batch = Arc::new(AtomicUsize::new(0)); + let (result_tx, result_rx) = sync_channel::<(usize, Arc)>(1); + let mut gen_handles = Vec::with_capacity(num_generators); + let spec = BatchSpec { + int_columns: args.int_columns, + small_string_columns: args.small_string_columns, + large_string_columns: args.large_string_columns, + }; + for _ in 0..num_generators { + let schema = schema.clone(); + let next_batch = next_batch.clone(); + let tx = result_tx.clone(); + let (rows, batch_size) = (args.rows, args.batch_size); + let handle = thread::spawn(move || { + loop { + let i = next_batch.fetch_add(1, Ordering::Relaxed); + if i >= num_batches { + break; + } + let row_offset = (i * batch_size) as u64; + let rows_in = batch_size.min(rows - i * batch_size); + let batch = make_batch(&schema, spec, row_offset, rows_in); + if tx.send((i, Arc::new(batch))).is_err() { + break; // main hung up (shouldn't happen on the happy path) + } + } + }); + gen_handles.push(handle); + } + drop(result_tx); // only the generators hold senders now - // Generate batches and fan each batch's columns out to the workers. `schema` - // is flat, so leaf order matches field order and each field maps to exactly - // one worker. - let mut rng = XorShift(0x9E37_79B9_7F4A_7C15); + // Main: pull generated batches, re-order them by index, and broadcast each in + // index order to every encoder. Re-ordering keeps the written row order — and + // therefore the output file — deterministic regardless of generator timing. let mut peak_rss = rss_start; - let mut written = 0; - let rows = args.batch_size.min(args.rows - written); - let batch = make_batch(&schema, &args, rows, &mut rng); - while written < args.rows { - //let batch = make_batch(&schema, &args, rows, &mut rng); - for (col_idx, (field, array)) in schema.fields().iter().zip(batch.columns()).enumerate() { - for leaf in compute_leaves(field, array)? { - // Blocks if this worker is busy — bounding in-flight input. - workers[col_idx].1.send(leaf).unwrap(); + let mut pending: HashMap> = HashMap::new(); + let mut next_emit = 0usize; + while next_emit < num_batches { + if let Some(batch) = pending.remove(&next_emit) { + for tx in &encoder_txs { + tx.send(batch.clone()).unwrap(); // cheap: clones an Arc } + next_emit += 1; + peak_rss = peak_rss.max(rss_bytes(&mut system)); + } else { + let (i, batch) = result_rx.recv().expect("generators ended early"); + pending.insert(i, batch); } - written += rows; - // `batch` is dropped here; its data lives on in whatever leaves are still - // in flight to the workers. + } + drop(encoder_txs); // signal end-of-input to the encoders + for handle in gen_handles { + handle.join().expect("generator thread panicked"); + } + + // Collect the encoded chunks, sort them back into schema order, and splice + // each into the row group (streaming pages back out of the store). Columns + // must be appended in schema order. + let mut chunks: Vec<(usize, ArrowColumnChunk, usize)> = Vec::new(); + for handle in encoder_handles { + chunks.extend(handle.join().expect("encoder thread panicked")?); peak_rss = peak_rss.max(rss_bytes(&mut system)); } + chunks.sort_by_key(|(idx, _, _)| *idx); - // Signal end-of-input to every worker, then join in column order and splice - // each finished chunk into the row group (streaming pages back out of the - // store). Columns must be appended in schema order. let mut row_group_writer = file_writer.next_row_group()?; let mut peak_writer_memory = 0usize; - for (handle, send) in workers { - drop(send); // closes the channel so the worker's `for leaf in recv` ends - let (chunk, col_peak) = handle.join().expect("worker thread panicked")?; + for (_idx, chunk, col_peak) in chunks { peak_writer_memory += col_peak; chunk.append_to_row_group(&mut row_group_writer)?; peak_rss = peak_rss.max(rss_bytes(&mut system)); @@ -481,6 +586,7 @@ fn main() -> Result<()> { file_writer.close()?; peak_rss = peak_rss.max(rss_bytes(&mut system)); let elapsed = start.elapsed(); + let written = args.rows; println!(); println!("Done. Wrote {written} rows.");