From 10a9ce802e79843971d4eba0e4bebe726fb9ba31 Mon Sep 17 00:00:00 2001 From: Wael Sulais Date: Mon, 17 Aug 2026 20:59:09 +0300 Subject: [PATCH 1/3] feat(layout): add a paged layout encoding to bound flatbuffer tables A layout is a single recursive flatbuffer, verified on open against `max_tables` (1,000,000 by default). Each chunk layout costs one table, so past the limit the writer produces a file the default reader cannot open: the write reports success and the error appears at open. Add a `vortex.paged` encoding that reports no inline children and holds its subtree in a segment as a nested `Layout` flatbuffer, resolved on descent. Each page is verified in its own right, so the table and depth limits apply per page rather than per file. Opt in via `ChunkedLayoutStrategy::with_page_size` and `WriteStrategyBuilder::with_page_size`; off by default, since a reader without the encoding registered can traverse a paged layout tree but cannot scan it. No flatbuffer schema change. `LayoutChildren` stays synchronous: the descent happens in the reader, whose evaluation methods already return futures. Each page carries its subtree's chunk boundaries so scan planning, which is synchronous, needs no page reads and split boundaries stay identical to the inline layout. The uniform case is kept symbolic, since one integer per chunk would otherwise dominate the footer. Measured on a 2,714,867,981-point file that 0.83.0 writes and then cannot open: 1,215,246 inline layout nodes become 1,264, the footer shrinks from 53.5 MB to 19.6 MB, the file grows 0.001%, and it opens in 54 ms. Assisted-by: Claude Opus 5 (claude-opus-5) via Claude Code 2.1.232 Signed-off-by: Wael Sulais --- vortex-file/src/strategy.rs | 22 +- vortex-file/src/tests.rs | 169 ++++++ vortex-layout/src/layouts/chunked/writer.rs | 191 +++++- vortex-layout/src/layouts/mod.rs | 1 + vortex-layout/src/layouts/paged/mod.rs | 610 ++++++++++++++++++++ vortex-layout/src/layouts/paged/reader.rs | 197 +++++++ vortex-layout/src/layouts/paged/writer.rs | 48 ++ vortex-layout/src/session.rs | 2 + 8 files changed, 1230 insertions(+), 10 deletions(-) create mode 100644 vortex-layout/src/layouts/paged/mod.rs create mode 100644 vortex-layout/src/layouts/paged/reader.rs create mode 100644 vortex-layout/src/layouts/paged/writer.rs diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 9d4dbb90610..3500eee482c 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -62,6 +62,7 @@ pub struct WriteStrategyBuilder { allow_encodings: Option>, flat_strategy: Option>, probe_compressor: Option>, + page_size: Option, /// Whether to write list fields using [`ListLayoutStrategy`]. /// /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy @@ -80,6 +81,9 @@ impl Default for WriteStrategyBuilder { allow_encodings: None, flat_strategy: None, probe_compressor: None, + // Off by default: a reader without the paged layout encoding cannot descend into a + // page, so paging changes who can read the file. + page_size: None, use_list_layout: use_experimental_list_layout(), } } @@ -95,6 +99,21 @@ impl WriteStrategyBuilder { self } + /// Group each chunked layout's children into pages of at most `page_size` chunks, each + /// serialized into its own segment. + /// + /// A layout is a single recursive flatbuffer verified against a table limit, so a file with + /// enough chunks becomes one its own reader rejects. Paging bounds the tables in every + /// flatbuffer to roughly `page_size`, at the cost of one extra segment read per page touched. + /// Zero, or leaving this unset, keeps chunks inline. + /// + /// Note that a reader without the `vortex.paged` layout encoding registered can inspect a + /// paged file's layout tree but cannot scan it. + pub fn with_page_size(mut self, page_size: usize) -> Self { + self.page_size = Some(page_size); + self + } + /// Override the target uncompressed byte size used to coalesce data blocks. /// /// Passing `None` disables byte-size coalescing, so blocks retain the row granularity set by @@ -202,7 +221,8 @@ impl WriteStrategyBuilder { }; // 7. for each chunk create a flat layout - let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); + let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)) + .with_page_size(self.page_size.unwrap_or(0)); // 6. buffer chunks so they end up with closer segment ids physically let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..29fc2ea1a20 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -75,6 +75,7 @@ use vortex_buffer::buffer; use vortex_edition::EditionSession; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_flatbuffers::WriteFlatBufferExt; use vortex_flatbuffers::footer as fb; use vortex_io::session::RuntimeSession; use vortex_layout::DynLayout; @@ -2800,3 +2801,171 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { assert_eq!(result.len(), 1); Ok(()) } + +/// Count the paged layouts anywhere in a layout tree. +fn count_paged(layout: &dyn DynLayout) -> usize { + let here = usize::from(layout.encoding_id().as_ref() == "vortex.paged"); + layout + .children() + .unwrap() + .iter() + .map(|child| count_paged(child.as_ref())) + .sum::() + + here +} + +/// A file whose chunks are grouped into pages must reopen and read back identically. +/// +/// This is the whole claim end to end: the root layout no longer carries a table per chunk, and +/// nothing about opening or scanning the file has to know that. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn paged_chunks_round_trip_through_a_file() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let expected: PrimitiveArray = (0..4096i32).collect(); + let expected = expected.into_array(); + + // Small row blocks and no byte coalescing, so the chunked layer has many chunks to page. + let strategy = crate::strategy::WriteStrategyBuilder::default() + .with_row_block_size(64) + .with_data_block_target_bytes(None) + .with_page_size(4) + .build(); + + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, expected.to_array_stream()) + .await?; + + let pages = count_paged(summary.footer().layout().as_ref()); + assert!( + pages > 1, + "expected the chunks to be grouped into several pages, found {pages}" + ); + + let file = SESSION.open_options().open_buffer(buf)?; + assert_eq!(file.row_count(), 4096); + + let result = file.scan()?.into_array_stream()?.read_all().await?; + assert_arrays_eq!(result, expected, &mut ctx); + + Ok(()) +} + +/// Whether a serialized layout flatbuffer verifies under a table budget. +fn verifies_within(bytes: &[u8], max_tables: usize) -> bool { + let opts = flatbuffers::VerifierOptions { + max_tables, + ..Default::default() + }; + flatbuffers::root_with_opts::(&opts, bytes).is_ok() +} + +/// The smallest flatbuffer table budget these layout bytes verify under. +/// +/// Doubling then bisecting, because the counts of interest run into the tens of thousands. +fn min_tables(bytes: &[u8]) -> usize { + let mut hi = 1; + while !verifies_within(bytes, hi) { + hi *= 2; + } + let mut lo = hi / 2; + // Invariant: `lo` fails (or is 0), `hi` succeeds. + while hi - lo > 1 { + let mid = lo + (hi - lo) / 2; + if verifies_within(bytes, mid) { + hi = mid; + } else { + lo = mid; + } + } + hi +} + +/// The smallest flatbuffer table budget this layout's serialized form verifies under. +fn min_layout_tables(layout: &vortex_layout::LayoutRef) -> usize { + let bytes = layout + .flatbuffer_writer(&vortex_layout::LayoutContext::default()) + .write_flatbuffer_bytes() + .unwrap(); + min_tables(&bytes) +} + +/// The ceiling this exists to remove: a layout is verified against a table limit, and inline +/// chunks spend a table each. The same data paged spends one per page instead. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn paging_lowers_the_root_layout_table_count() -> VortexResult<()> { + let array: PrimitiveArray = (0..4096i32).collect(); + let array = array.into_array(); + + async fn write_with_page_size(array: &ArrayRef, page_size: usize) -> VortexResult { + let strategy = crate::strategy::WriteStrategyBuilder::default() + .with_row_block_size(64) + .with_data_block_target_bytes(None) + .with_page_size(page_size) + .build(); + + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, array.to_array_stream()) + .await?; + + Ok(min_layout_tables(summary.footer().layout())) + } + + let inline = write_with_page_size(&array, 0).await?; + let paged = write_with_page_size(&array, 8).await?; + + assert!( + paged * 2 < inline, + "paging into groups of eight should cut the root layout's tables by more than half, \ + but it needs {paged} against {inline} inline" + ); + + Ok(()) +} + +/// A selective filter must still push down through a page into the zone maps below it. +/// +/// The paged reader intercepts pruning and filter evaluation to await its subtree first. If that +/// forwarding were wrong, pruning would either return wrong rows or stop pruning silently. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn filter_pushes_down_through_pages() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array: PrimitiveArray = (0..4096i32).collect(); + let array = array.into_array(); + + // 64-row zones, so a filter selecting the tail prunes almost every zone — through the pages. + let strategy = crate::strategy::WriteStrategyBuilder::default() + .with_row_block_size(64) + .with_data_block_target_bytes(None) + .with_page_size(8) + .build(); + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, array.to_array_stream()) + .await?; + + let file = SESSION.open_options().open_buffer(buf)?; + let filter = bind_scan_expr(&file, gt(root(), lit(4000i32))); + let result = file + .scan()? + .with_filter(filter) + .into_array_stream()? + .read_all() + .await?; + + let expected: PrimitiveArray = (4001..4096i32).collect(); + assert_arrays_eq!(result, expected.into_array(), &mut ctx); + + Ok(()) +} diff --git a/vortex-layout/src/layouts/chunked/writer.rs b/vortex-layout/src/layouts/chunked/writer.rs index 8de7c30fe47..911290e0843 100644 --- a/vortex-layout/src/layouts/chunked/writer.rs +++ b/vortex-layout/src/layouts/chunked/writer.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::num::NonZeroUsize; use std::sync::Arc; use async_stream::stream; @@ -8,16 +9,20 @@ use async_trait::async_trait; use futures::StreamExt; use futures::TryStreamExt; use futures::stream; +use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; use crate::LayoutRef; use crate::LayoutStrategy; use crate::LayoutWriterContext; use crate::children::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::paged::writer::write_page; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; @@ -28,14 +33,73 @@ use crate::sequence::SequentialStreamExt as _; pub struct ChunkedLayoutStrategy { /// The layout strategy for each chunk. pub chunk_strategy: Arc, + /// If set, chunks are grouped into pages of at most this many, each serialized into its own + /// segment rather than inline in this layout's flatbuffer. See [`crate::layouts::paged`]. + pub page_size: Option, } impl ChunkedLayoutStrategy { pub fn new(chunk_strategy: S) -> Self { Self { chunk_strategy: Arc::new(chunk_strategy), + page_size: None, } } + + /// Group chunks into pages of at most `page_size` chunks each. + /// + /// A layout is a single recursive flatbuffer verified against a table limit, so a file with + /// enough chunks becomes one its own reader rejects. Paging bounds the tables in any one + /// flatbuffer to roughly `page_size`. A page size of zero leaves the chunks inline. + pub fn with_page_size(mut self, page_size: usize) -> Self { + self.page_size = NonZeroUsize::new(page_size); + self + } +} + +/// Replace groups of `page_size` consecutive chunks with pages holding them in segments. +async fn paginate( + children: Vec, + page_size: usize, + dtype: &DType, + ctx: &LayoutWriterContext, + segment_sink: &SegmentSinkRef, + eof: &mut SequencePointer, +) -> VortexResult> { + let mut pages = Vec::with_capacity(children.len().div_ceil(page_size)); + for group in children.chunks(page_size) { + // The chunk boundaries within the page, which travel with it so scans can still be + // planned at chunk granularity without reading it. + let mut row_count = 0u64; + let row_offsets = group + .iter() + .map(|layout| { + row_count = row_count + .checked_add(layout.row_count()) + .ok_or_else(|| vortex_err!("Paged chunk row counts overflow"))?; + Ok(row_count) + }) + .collect::>>()?; + + let page = ChunkedLayout::new( + row_count, + dtype.clone(), + OwnedLayoutChildren::layout_children(group.to_vec()), + ) + .into_layout(); + + pages.push( + write_page( + &page, + &row_offsets, + ReadContext::new(ctx.array_ctx().to_ids()), + segment_sink, + eof.advance(), + ) + .await?, + ); + } + Ok(pages) } #[async_trait] @@ -53,11 +117,16 @@ impl LayoutStrategy for ChunkedLayoutStrategy { let chunk_strategy = Arc::clone(&self.chunk_strategy); let handle = session.handle(); + // The eofs used for the chunks should appear _before_ the pages that reference them. + let mut chunks_eof = eof.split_off(); + let page_ctx = ctx.clone(); + let page_sink = Arc::clone(&segment_sink); + // We spawn each child to allow parallelism when processing chunks. let stream = stream! { let mut stream = stream; while let Some(chunk) = stream.next().await { - let chunk_eof = eof.split_off(); + let chunk_eof = chunks_eof.split_off(); let chunk_strategy = Arc::clone(&chunk_strategy); let ctx = ctx.clone(); @@ -88,15 +157,119 @@ impl LayoutStrategy for ChunkedLayoutStrategy { let mut child_layouts: Vec = stream.buffered(usize::MAX).try_collect().await?; if child_layouts.len() == 1 { - Ok(child_layouts.pop().vortex_expect("must have one child")) - } else { - let row_count = child_layouts.iter().map(|layout| layout.row_count()).sum(); - Ok(ChunkedLayout::new( - row_count, - dtype, - OwnedLayoutChildren::layout_children(child_layouts), + return Ok(child_layouts.pop().vortex_expect("must have one child")); + } + + let row_count = child_layouts.iter().map(|layout| layout.row_count()).sum(); + + if let Some(page_size) = self.page_size { + child_layouts = paginate( + child_layouts, + page_size.get(), + &dtype, + &page_ctx, + &page_sink, + &mut eof, ) - .into_layout()) + .await?; } + + Ok(ChunkedLayout::new( + row_count, + dtype, + OwnedLayoutChildren::layout_children(child_layouts), + ) + .into_layout()) + } +} + +#[cfg(test)] +mod test { + use std::sync::Arc; + + use futures::stream; + use vortex_array::ArrayContext; + use vortex_array::IntoArray; + use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::PType; + use vortex_array::expr::root; + use vortex_io::runtime::single::block_on; + use vortex_io::session::RuntimeSessionExt; + + use crate::LayoutStrategy; + use crate::LayoutWriterContext; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; + use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::segments::TestSegments; + use crate::sequence::SequenceId; + use crate::sequence::SequentialStreamAdapter; + use crate::sequence::SequentialStreamExt as _; + use crate::test::new_session; + + /// Six chunks of three rows written with a page size of two: the root should hold three + /// paged children rather than six inline chunks, and still read back as the same 18 rows. + #[test] + fn paged_chunks_round_trip() { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut exec = session.create_execution_ctx(); + let segments = Arc::new(TestSegments::default()); + let (mut sequence_id, eof) = SequenceId::root().split(); + + let chunks = (0..6i32) + .map(|chunk| { + let base = chunk * 3; + let chunk: PrimitiveArray = (base..base + 3).collect(); + Ok((sequence_id.advance(), chunk.into_array())) + }) + .collect::>(); + + let layout = ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()) + .with_page_size(2) + .write_stream( + LayoutWriterContext::new(ArrayContext::empty()), + Arc::::clone(&segments), + SequentialStreamAdapter::new( + DType::Primitive(PType::I32, NonNullable), + stream::iter(chunks), + ) + .sendable(), + eof, + &session, + ) + .await + .unwrap(); + + assert_eq!(layout.encoding_id().as_ref(), "vortex.chunked"); + assert_eq!(layout.row_count(), 18); + assert_eq!(layout.nchildren(), 3); + + let page = layout.slot(0).unwrap().unwrap(); + assert_eq!(page.encoding_id().as_ref(), "vortex.paged"); + assert_eq!(page.row_count(), 6); + assert_eq!( + page.nchildren(), + 0, + "a page's subtree lives in its segment, not inline" + ); + + let reader = layout + .new_reader("".into(), segments, &session, &Default::default()) + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let result = reader + .projection_evaluation(&(0..18), &expr, MaskFuture::new_true(18)) + .unwrap() + .await + .unwrap(); + + let expected: PrimitiveArray = (0i32..18).collect(); + assert_arrays_eq!(result, expected.into_array(), &mut exec); + }) } } diff --git a/vortex-layout/src/layouts/mod.rs b/vortex-layout/src/layouts/mod.rs index 47fa31aa3d9..2ce01907b37 100644 --- a/vortex-layout/src/layouts/mod.rs +++ b/vortex-layout/src/layouts/mod.rs @@ -17,6 +17,7 @@ pub mod file_stats; pub mod flat; pub(crate) mod foreign; pub mod list; +pub mod paged; pub(crate) mod partitioned; pub mod repartition; pub mod row_idx; diff --git a/vortex-layout/src/layouts/paged/mod.rs b/vortex-layout/src/layouts/paged/mod.rs new file mode 100644 index 00000000000..f62b19800a4 --- /dev/null +++ b/vortex-layout/src/layouts/paged/mod.rs @@ -0,0 +1,610 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A layout whose subtree is stored in a segment rather than inline in its parent's flatbuffer. +//! +//! A Vortex layout is a single recursive flatbuffer, parsed whole before any data is read, and +//! verified against a table limit. A paged layout breaks that recursion: it reports no inline +//! children and instead holds one segment containing its subtree as a nested `Layout` flatbuffer. +//! Each page is therefore verified independently, so the table and depth limits apply per page +//! rather than to the whole file. +//! +//! Because the segment fetch is asynchronous but [`LayoutChildren`](crate::LayoutChildren) is not, +//! the descent happens in the reader, whose evaluation methods already return futures. + +mod reader; +pub mod writer; + +use std::sync::Arc; + +use itertools::Either; +use vortex_array::ProstMetadata; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; +use vortex_session::registry::ReadContext; + +use crate::Layout; +use crate::LayoutChildType; +use crate::LayoutDeserializeArgs; +use crate::LayoutEncodingId; +use crate::LayoutId; +use crate::LayoutParts; +use crate::LayoutReaderContext; +use crate::LayoutReaderRef; +use crate::VTable; +use crate::children::OwnedLayoutChildren; +use crate::layouts::paged::reader::PagedReader; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +/// Paged layout vtable. +#[derive(Clone, Debug)] +pub struct Paged; + +/// Backwards-compatible name for the paged layout plugin. +pub use Paged as PagedLayoutEncoding; + +/// Where a page's subtree splits into chunks, as recorded at its boundary. +/// +/// Scan planning is synchronous and must not read the page, so these boundaries travel with it. +/// Storing one integer per chunk would make them the dominant cost of the footer at any useful +/// page size, so the uniform case — which is what a repartitioner emitting fixed-size row blocks +/// produces — is kept symbolic. +#[derive(Clone, Debug)] +pub enum ChunkBoundaries { + /// Every chunk holds `len` rows, except possibly a shorter last one. + Uniform { + /// Rows per chunk. + len: u64, + /// Number of chunks. + count: usize, + /// Total rows, which the last chunk is truncated to. + row_count: u64, + }, + /// Exclusive row boundaries, for subtrees whose chunks differ in length. + Explicit(Arc<[u64]>), +} + +impl ChunkBoundaries { + /// Choose a representation for `offsets`, the exclusive row boundaries of a subtree. + pub fn from_offsets(offsets: &[u64], row_count: u64) -> Self { + let uniform = offsets.split_last().is_some_and(|(last, rest)| { + let len = offsets[0]; + *last == row_count + && len > 0 + && rest + .iter() + .enumerate() + .all(|(idx, offset)| *offset == (idx as u64 + 1) * len) + && row_count > (offsets.len() as u64 - 1) * len + && row_count <= offsets.len() as u64 * len + }); + + if uniform { + Self::Uniform { + len: offsets[0], + count: offsets.len(), + row_count, + } + } else { + Self::Explicit(offsets.into()) + } + } + + /// The exclusive row boundaries, relative to the subtree's first row. + pub fn offsets(&self) -> impl Iterator + '_ { + match self { + Self::Uniform { + len, + count, + row_count, + } => Either::Left((1..=*count).map(move |idx| (idx as u64 * len).min(*row_count))), + Self::Explicit(offsets) => Either::Right(offsets.iter().copied()), + } + } + + /// The number of chunks. + pub fn len(&self) -> usize { + match self { + Self::Uniform { count, .. } => *count, + Self::Explicit(offsets) => offsets.len(), + } + } + + /// Returns `true` if the subtree has no chunks. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +/// Paged-layout-specific data. +#[derive(Clone, Debug)] +pub struct PagedData { + segment_id: SegmentId, + layout_ctx: ReadContext, + array_ctx: ReadContext, + boundaries: ChunkBoundaries, +} + +/// A layout standing in for a subtree serialized into its own segment. +pub type PagedLayout = Layout; + +impl VTable for Paged { + type LayoutData = PagedData; + type Metadata = ProstMetadata; + + fn id(&self) -> LayoutId { + static ID: CachedId = CachedId::new("vortex.paged"); + *ID + } + + fn metadata(layout: &Layout) -> Self::Metadata { + // The page's own encoding dictionary. The file footer's dictionary is not usable here: it + // is not built until the footer is serialized, long after the page segment was written. + ProstMetadata(PagedLayoutMetadata { + layout_encodings: layout + .layout_ctx + .ids() + .iter() + .map(|id| id.as_str().to_string()) + .collect(), + uniform_chunk_len: match &layout.boundaries { + ChunkBoundaries::Uniform { len, .. } => Some(*len), + ChunkBoundaries::Explicit(_) => None, + }, + chunk_count: match &layout.boundaries { + ChunkBoundaries::Uniform { count, .. } => Some(*count as u64), + ChunkBoundaries::Explicit(_) => None, + }, + row_offsets: match &layout.boundaries { + ChunkBoundaries::Uniform { .. } => Vec::new(), + ChunkBoundaries::Explicit(offsets) => offsets.to_vec(), + }, + }) + } + + fn deserialize( + &self, + args: &LayoutDeserializeArgs<'_>, + metadata: &PagedLayoutMetadata, + ) -> VortexResult { + if args.segment_ids.len() != 1 { + vortex_bail!("Paged layout must have exactly one segment ID"); + } + if args.children.nchildren() != 0 { + vortex_bail!("Paged layout must not have inline children"); + } + let boundaries = match (metadata.uniform_chunk_len, metadata.chunk_count) { + (Some(len), Some(count)) => { + let count = usize::try_from(count)?; + if len == 0 || count == 0 { + vortex_bail!("Paged layout uniform chunk length and count must be non-zero"); + } + // The last chunk may be short, but the rest must exactly cover the rows before it. + if args.row_count <= (count as u64 - 1) * len || args.row_count > count as u64 * len + { + vortex_bail!( + "Paged layout {count} chunks of {len} rows do not cover {} rows", + args.row_count + ); + } + ChunkBoundaries::Uniform { + len, + count, + row_count: args.row_count, + } + } + (None, None) => { + if metadata + .row_offsets + .last() + .is_some_and(|last| *last != args.row_count) + { + vortex_bail!("Paged layout row offsets do not add up to its row count"); + } + ChunkBoundaries::Explicit(metadata.row_offsets.as_slice().into()) + } + _ => vortex_bail!( + "Paged layout must set both a uniform chunk length and a chunk count, or neither" + ), + }; + let layout_ids = metadata + .layout_encodings + .iter() + .map(|id| LayoutEncodingId::from(id.as_str())) + .collect::>(); + Ok(PagedData { + segment_id: args.segment_ids[0], + layout_ctx: ReadContext::new(layout_ids), + array_ctx: args.array_read_ctx.clone(), + boundaries, + }) + } + + fn child_dtype(_layout: &Layout, idx: usize) -> VortexResult { + vortex_bail!("Paged layout has no inline child {idx}; its subtree is in its segment") + } + + fn child_type(_layout: &Layout, idx: usize) -> LayoutChildType { + vortex_panic!("Paged layout has no inline child {idx}; its subtree is in its segment") + } + + fn new_reader( + layout: &Layout, + name: Arc, + segment_source: Arc, + session: &VortexSession, + ctx: &LayoutReaderContext, + ) -> VortexResult { + Ok(Arc::new(PagedReader::new( + layout.clone(), + name, + segment_source, + session.clone(), + ctx.clone(), + ))) + } +} + +impl Layout { + /// Construct a paged layout over an already-written page segment. + /// + /// `boundaries` are the subtree's chunk boundaries, used to plan scans without reading the + /// page. They must cover exactly `row_count` rows. + pub fn new( + row_count: u64, + dtype: DType, + segment_id: SegmentId, + layout_ctx: ReadContext, + array_ctx: ReadContext, + boundaries: ChunkBoundaries, + ) -> Self { + LayoutParts::new( + Paged, + dtype, + row_count, + vec![segment_id], + OwnedLayoutChildren::layout_children(Vec::new()), + PagedData { + segment_id, + layout_ctx, + array_ctx, + boundaries, + }, + ) + .into_typed() + } + + /// Returns the segment holding this page's subtree. + pub fn segment_id(&self) -> SegmentId { + self.segment_id + } + + /// Returns the layout encoding dictionary the page's flatbuffer is written against. + pub fn layout_ctx(&self) -> &ReadContext { + &self.layout_ctx + } + + /// Returns the array read context used to decode arrays within the page. + pub fn array_ctx(&self) -> &ReadContext { + &self.array_ctx + } + + /// Returns the subtree's chunk boundaries, relative to this page's first row. + pub fn boundaries(&self) -> &ChunkBoundaries { + &self.boundaries + } +} + +#[derive(prost::Message)] +pub struct PagedLayoutMetadata { + #[prost(repeated, string, tag = "1")] + pub layout_encodings: Vec, + /// Exclusive row boundaries of the subtree's chunks, relative to this page's first row. + /// + /// Packed varints, so promoting them to the page boundary costs no flatbuffer tables — the + /// tables are what the verifier limits. Empty when the boundaries are uniform, since one + /// integer per chunk would otherwise dominate the footer at any useful page size. + #[prost(repeated, uint64, tag = "2")] + pub row_offsets: Vec, + /// Rows per chunk, when every chunk but the last holds the same number. + #[prost(optional, uint64, tag = "3")] + pub uniform_chunk_len: Option, + /// Number of chunks, set together with [`Self::uniform_chunk_len`]. + #[prost(optional, uint64, tag = "4")] + pub chunk_count: Option, +} + +#[cfg(test)] +mod test { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use flatbuffers::VerifierOptions; + use flatbuffers::root_with_opts; + use futures::stream; + use vortex_array::ArrayContext; + use vortex_array::IntoArray; + use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::FieldMask; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::PType; + use vortex_array::expr::root; + use vortex_flatbuffers::WriteFlatBufferExt; + use vortex_flatbuffers::layout as fbl; + use vortex_io::runtime::single::block_on; + use vortex_io::session::RuntimeSessionExt; + + use super::*; + use crate::LayoutContext; + use crate::LayoutRef; + use crate::LayoutStrategy; + use crate::LayoutWriterContext; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; + use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::scan::split_by::SplitBy; + use crate::segments::TestSegments; + use crate::sequence::SequenceId; + use crate::sequence::SequentialStreamAdapter; + use crate::sequence::SequentialStreamExt as _; + use crate::test::new_session; + + /// Write `chunk_count` single-row chunks, optionally grouped into pages. + async fn write_chunks( + chunk_count: usize, + page_size: usize, + session: &VortexSession, + ) -> (Arc, LayoutRef) { + write_chunk_lens(&vec![1i32; chunk_count], page_size, session).await + } + + /// Write one chunk per entry in `chunk_lens`, optionally grouped into pages. + async fn write_chunk_lens( + chunk_lens: &[i32], + page_size: usize, + session: &VortexSession, + ) -> (Arc, LayoutRef) { + let segments = Arc::new(TestSegments::default()); + let (mut sequence_id, eof) = SequenceId::root().split(); + + let mut next = 0i32; + let chunks = chunk_lens + .iter() + .map(|len| { + let chunk: PrimitiveArray = (next..next + *len).collect(); + next += *len; + Ok((sequence_id.advance(), chunk.into_array())) + }) + .collect::>(); + + let layout = ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()) + .with_page_size(page_size) + .write_stream( + LayoutWriterContext::new(ArrayContext::empty()), + Arc::::clone(&segments), + SequentialStreamAdapter::new( + DType::Primitive(PType::I32, NonNullable), + stream::iter(chunks), + ) + .sendable(), + eof, + session, + ) + .await + .unwrap(); + + (segments, layout) + } + + fn to_flatbuffer(layout: &LayoutRef) -> Vec { + layout + .flatbuffer_writer(&LayoutContext::default()) + .write_flatbuffer_bytes() + .unwrap() + .to_vec() + } + + fn verifies_within(bytes: &[u8], max_tables: usize) -> bool { + let opts = VerifierOptions { + max_tables, + ..Default::default() + }; + root_with_opts::(&opts, bytes).is_ok() + } + + /// A [`SegmentSource`] that counts requests, so IO can be asserted about. + struct CountingSource { + inner: Arc, + requests: Arc, + } + + impl SegmentSource for CountingSource { + fn request(&self, id: SegmentId) -> crate::segments::SegmentFuture { + self.requests.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } + } + + fn splits_of( + session: &VortexSession, + segments: Arc, + layout: &LayoutRef, + ) -> Vec { + let reader = layout + .new_reader("".into(), segments, session, &Default::default()) + .unwrap(); + SplitBy::Layout + .splits(reader.as_ref(), &(0..layout.row_count()), &[FieldMask::All]) + .unwrap() + } + + /// Fixed-size row blocks are what the repartitioner emits, so a page must not spend a varint + /// per chunk recording boundaries it could derive from two integers. Otherwise the offsets + /// become the dominant cost of the footer at any useful page size. + #[test] + fn uniform_pages_record_boundaries_in_constant_space() { + block_on(|handle| async { + let session = new_session().with_handle(handle); + + // One page in each case, holding 8 and 64 equal-length chunks respectively. + let (_, eight) = write_chunks(8, 8, &session).await; + let (_, sixty_four) = write_chunks(64, 64, &session).await; + + let eight = eight.slot(0).unwrap().unwrap().metadata().len(); + let sixty_four = sixty_four.slot(0).unwrap().unwrap().metadata().len(); + + assert!( + sixty_four <= eight + 1, + "a page of 64 uniform chunks should cost no more metadata than one of 8, \ + but it took {sixty_four} bytes against {eight}" + ); + }) + } + + /// Chunks of differing lengths cannot be derived, so those pages fall back to explicit + /// offsets. The split set must stay exact either way. + #[test] + fn non_uniform_pages_keep_exact_splits() { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let lens = [3i32, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8]; + + let (inline_segments, inline) = write_chunk_lens(&lens, 0, &session).await; + let (paged_segments, paged) = write_chunk_lens(&lens, 3, &session).await; + + assert_eq!( + splits_of(&session, paged_segments, &paged), + splits_of(&session, inline_segments, &inline) + ); + }) + } + + /// Paging must not coarsen the scan's batch boundaries. The page carries its subtree's row + /// offsets in its metadata — a flat integer vector, costing no flatbuffer tables — so planning + /// reaches the same answer it would have inline, without reading a page. + #[test] + fn splits_match_the_inline_layout() { + block_on(|handle| async { + let session = new_session().with_handle(handle); + + let (inline_segments, inline) = write_chunks(12, 0, &session).await; + let (paged_segments, paged) = write_chunks(12, 3, &session).await; + + assert_eq!( + splits_of(&session, paged_segments, &paged), + splits_of(&session, inline_segments, &inline) + ); + }) + } + + /// Scan planning is synchronous, so it must reach its answer from the pages' own row counts. + /// If it ever fetched a page to plan, opening a file would cost every page in it — which is + /// the cost paging exists to avoid. + #[test] + fn planning_does_not_fetch_pages() { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let (segments, layout) = write_chunks(12, 3, &session).await; + + let requests = Arc::new(AtomicUsize::new(0)); + let counting = Arc::new(CountingSource { + inner: segments, + requests: Arc::clone(&requests), + }); + + let reader = layout + .new_reader("".into(), counting, &session, &Default::default()) + .unwrap(); + let splits = SplitBy::Layout + .splits(reader.as_ref(), &(0..12), &[FieldMask::All]) + .unwrap(); + + assert_eq!(splits.first(), Some(&0)); + assert_eq!(splits.last(), Some(&12)); + assert_eq!( + requests.load(Ordering::Relaxed), + 0, + "planning must not fetch any page segment" + ); + }) + } + + /// A range that starts and ends inside different pages must still return exactly its rows. + #[test] + fn row_range_spanning_pages() { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut exec = session.create_execution_ctx(); + + // Pages cover 0..3, 3..6, 6..9 and 9..12, so 2..10 clips both end pages. + let (segments, layout) = write_chunks(12, 3, &session).await; + + let reader = layout + .new_reader("".into(), segments, &session, &Default::default()) + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let result = reader + .projection_evaluation(&(2..10), &expr, MaskFuture::new_true(8)) + .unwrap() + .await + .unwrap(); + + let expected: PrimitiveArray = (2i32..10).collect(); + assert_arrays_eq!(result, expected.into_array(), &mut exec); + }) + } + + /// Twelve chunks need thirteen tables inline. Grouped into pages of three, no single + /// flatbuffer in the file needs more than five — which is the whole point: the verifier's + /// table limit applies per flatbuffer, and paging bounds every one of them. + #[test] + fn paging_keeps_every_flatbuffer_under_a_table_budget() { + const BUDGET: usize = 8; + + block_on(|handle| async { + let session = new_session().with_handle(handle); + + let (_, inline) = write_chunks(12, 0, &session).await; + let inline = to_flatbuffer(&inline); + assert!( + verifies_within(&inline, 4 * BUDGET), + "the inline layout must be a valid flatbuffer given enough tables, \ + otherwise the budget below proves nothing" + ); + assert!( + !verifies_within(&inline, BUDGET), + "twelve inline chunks should exceed a {BUDGET}-table budget" + ); + + let (segments, paged) = write_chunks(12, 3, &session).await; + assert!( + verifies_within(&to_flatbuffer(&paged), BUDGET), + "a root of four pages should fit a {BUDGET}-table budget" + ); + + // Each page is a flatbuffer root in its own right, verified independently. + assert_eq!(paged.nchildren(), 4); + for idx in 0..paged.nchildren() { + let page = paged.slot(idx).unwrap().unwrap(); + let page = page.as_::(); + let bytes = (segments.as_ref() as &dyn SegmentSource) + .request(page.segment_id()) + .await + .unwrap() + .to_host_sync(); + assert!( + verifies_within(bytes.as_ref(), BUDGET), + "page {idx} should fit a {BUDGET}-table budget" + ); + } + }) + } +} diff --git a/vortex-layout/src/layouts/paged/reader.rs b/vortex-layout/src/layouts/paged/reader.rs new file mode 100644 index 00000000000..c43ae62dc76 --- /dev/null +++ b/vortex-layout/src/layouts/paged/reader.rs @@ -0,0 +1,197 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; +use once_cell::sync::OnceCell; +use vortex_array::MaskFuture; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldMask; +use vortex_array::expr::BoundExpression; +use vortex_error::SharedVortexResult; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_flatbuffers::FlatBuffer; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::ArrayFuture; +use crate::LayoutReaderContext; +use crate::LayoutReaderRef; +use crate::layout_from_flatbuffer; +use crate::layouts::paged::PagedLayout; +use crate::reader::LayoutReader; +use crate::reader::RowSplits; +use crate::reader::SplitRange; +use crate::segments::SegmentSource; + +/// A reader for the subtree behind a page, resolved once and shared by every evaluation. +type SharedReaderFuture = Shared>>; + +/// A [`LayoutReader`] that fetches and parses its subtree on first use, then delegates to it. +pub struct PagedReader { + layout: PagedLayout, + name: Arc, + segment_source: Arc, + session: VortexSession, + ctx: LayoutReaderContext, + child: OnceCell, +} + +impl PagedReader { + pub(crate) fn new( + layout: PagedLayout, + name: Arc, + segment_source: Arc, + session: VortexSession, + ctx: LayoutReaderContext, + ) -> Self { + Self { + layout, + name, + segment_source, + session, + ctx, + child: OnceCell::new(), + } + } + + /// Return a future resolving to the reader for this page's subtree. + /// + /// The segment request is only issued once the page is actually evaluated, so a page that is + /// pruned away, or never intersected by any scanned row range, costs no IO. + fn child_reader(&self) -> SharedReaderFuture { + self.child + .get_or_init(|| { + let segment_fut = self.segment_source.request(self.layout.segment_id()); + let dtype = self.layout.dtype().clone(); + let layout_ctx = self.layout.layout_ctx().clone(); + let array_ctx = self.layout.array_ctx().clone(); + let segment_source = Arc::clone(&self.segment_source); + let session = self.session.clone(); + let reader_ctx = self.ctx.clone(); + let name = Arc::clone(&self.name); + + async move { + let reader = async { + let segment = segment_fut.await?; + // The page is a `Layout` flatbuffer root, verified in its own right. + let page = FlatBuffer::align_from(segment.to_host_sync()); + let layout = layout_from_flatbuffer( + page, + &dtype, + &layout_ctx, + &array_ctx, + &session, + )?; + layout.new_reader(name, segment_source, &session, &reader_ctx) + } + .await; + reader.map_err(Arc::new) + } + .boxed() + .shared() + }) + .clone() + } +} + +impl LayoutReader for PagedReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn dtype(&self) -> &DType { + self.layout.dtype() + } + + fn row_count(&self) -> u64 { + self.layout.row_count() + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + // Planning is synchronous and must not read the page, so the boundaries come from the + // row offsets the writer promoted into the page's metadata. + split_range.check_bounds(self.layout.row_count())?; + + let row_range = split_range.row_range(); + let boundaries = self.layout.boundaries(); + splits.reserve(boundaries.len()); + for offset in boundaries.offsets() { + // The range's own end is registered below; anything outside it is another split's. + if offset > row_range.start && offset < row_range.end { + splits.push( + split_range + .row_offset() + .checked_add(offset) + .vortex_expect("Paged layout split offset overflow"), + ); + } + } + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + row_range: &Range, + expr: &BoundExpression, + mask: Mask, + ) -> VortexResult { + let child = self.child_reader(); + let row_range = row_range.clone(); + let expr = expr.clone(); + + Ok(MaskFuture::new(mask.len(), async move { + let child = child.await?; + child.pruning_evaluation(&row_range, &expr, mask)?.await + })) + } + + fn filter_evaluation( + &self, + row_range: &Range, + expr: &BoundExpression, + mask: MaskFuture, + ) -> VortexResult { + let child = self.child_reader(); + let row_range = row_range.clone(); + let expr = expr.clone(); + + Ok(MaskFuture::new(mask.len(), async move { + let child = child.await?; + child.filter_evaluation(&row_range, &expr, mask)?.await + })) + } + + fn projection_evaluation( + &self, + row_range: &Range, + expr: &BoundExpression, + mask: MaskFuture, + ) -> VortexResult { + let child = self.child_reader(); + let row_range = row_range.clone(); + let expr = expr.clone(); + + Ok(async move { + let child = child.await?; + child.projection_evaluation(&row_range, &expr, mask)?.await + } + .boxed()) + } +} diff --git a/vortex-layout/src/layouts/paged/writer.rs b/vortex-layout/src/layouts/paged/writer.rs new file mode 100644 index 00000000000..dca20273df0 --- /dev/null +++ b/vortex-layout/src/layouts/paged/writer.rs @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_flatbuffers::WriteFlatBufferExt; +use vortex_session::registry::ReadContext; + +use crate::LayoutContext; +use crate::LayoutRef; +use crate::layouts::paged::ChunkBoundaries; +use crate::layouts::paged::PagedLayout; +use crate::segments::SegmentSinkRef; +use crate::sequence::SequenceId; + +/// Serialize `layout` into its own segment, returning the page that stands in for it. +/// +/// The subtree is written as a nested `Layout` flatbuffer against a fresh encoding dictionary, +/// which travels in the page's metadata. The parent therefore spends one table on the whole +/// subtree, and the page is verified against the table and depth limits on its own. +/// +/// `row_offsets` are the subtree's exclusive chunk boundaries, relative to its first row. They are +/// promoted to the page so scan planning, which is synchronous, can proceed without reading it, +/// and are kept symbolic when uniform. +pub async fn write_page( + layout: &LayoutRef, + row_offsets: &[u64], + array_ctx: ReadContext, + segment_sink: &SegmentSinkRef, + sequence_id: SequenceId, +) -> VortexResult { + let page_ctx = LayoutContext::default(); + let page = layout + .flatbuffer_writer(&page_ctx) + .write_flatbuffer_bytes()?; + let segment_id = segment_sink + .write(sequence_id, vec![page.into_inner()]) + .await?; + + Ok(PagedLayout::new( + layout.row_count(), + layout.dtype().clone(), + segment_id, + ReadContext::new(page_ctx.to_ids()), + array_ctx, + ChunkBoundaries::from_offsets(row_offsets, layout.row_count()), + ) + .into_layout()) +} diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index 0cf2234c09c..eea973e957e 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -15,6 +15,7 @@ use crate::layouts::chunked::Chunked; use crate::layouts::dict::Dict; use crate::layouts::flat::Flat; use crate::layouts::list::List; +use crate::layouts::paged::Paged; use crate::layouts::struct_::Struct; use crate::layouts::zoned::LegacyStats; use crate::layouts::zoned::Zoned; @@ -62,6 +63,7 @@ impl Default for LayoutSession { this.register(&LegacyStats as &dyn LayoutEncoding); this.register(&Dict as &dyn LayoutEncoding); this.register(&List as &dyn LayoutEncoding); + this.register(&Paged as &dyn LayoutEncoding); this } } From 54716686af4947eee2aabda8f94c6353ecbef90b Mon Sep 17 00:00:00 2001 From: Wael Sulais Date: Mon, 17 Aug 2026 21:17:42 +0300 Subject: [PATCH 2/3] feat(layout): intern page layouts into the file's encoding dictionary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A page's nested `Layout` flatbuffer was interned into a dictionary of its own, carried in the page's metadata, because `FooterSerializer` builds the file's dictionary while serializing the root layout — after the page segments have been written. Thread one `LayoutContext` through the write instead. `LayoutWriterContext` carries it, `FooterSerializer` already accepted one, and the file writer now hands the same context to both, so a page indexes into the file's single dictionary and stores no copy. On the read side `LayoutBuildContext` and `LayoutDeserializeArgs` carry the layout read context, so a paged layout resolves its subtree against the footer's dictionary rather than against its own metadata. Worth about 32 bytes per page. For 13,312 chunks the serialized layout drops from 125.9 kB to 73.6 kB at page_size 8, 10.1 kB to 6.7 kB at 128, and 3.4 kB to 2.8 kB at 1024; the file-size overhead of paging at page_size 8 falls from 11.3% to 9.0%, and at 1024 to 0.08%. Also ports the page-size sweep used for those numbers as an ignored test. Assisted-by: Claude Opus 5 (claude-opus-5) via Claude Code 2.1.232 Signed-off-by: Wael Sulais --- vortex-file/src/tests.rs | 172 ++++++++++++++++++++ vortex-file/src/writer.rs | 7 +- vortex-layout/src/children.rs | 1 + vortex-layout/src/encoding.rs | 4 + vortex-layout/src/flatbuffers.rs | 1 + vortex-layout/src/layouts/chunked/writer.rs | 1 + vortex-layout/src/layouts/paged/mod.rs | 39 +++-- vortex-layout/src/layouts/paged/writer.rs | 13 +- vortex-layout/src/layouts/zoned/mod.rs | 8 + vortex-layout/src/layouts/zoned/reader.rs | 2 + vortex-layout/src/strategy.rs | 17 ++ vortex-layout/src/vtable.rs | 1 + 12 files changed, 242 insertions(+), 24 deletions(-) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 29fc2ea1a20..c2582161397 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2969,3 +2969,175 @@ async fn filter_pushes_down_through_pages() -> VortexResult<()> { Ok(()) } + +/// Collect every page's segment id from a layout tree, in tree order. +fn page_segments(layout: &dyn DynLayout, into: &mut Vec) { + if layout.encoding_id().as_ref() == "vortex.paged" { + into.extend(layout.segment_ids()); + } + for child in layout.children().unwrap() { + page_segments(child.as_ref(), into); + } +} + +/// What paging costs and buys, at a scale close to the corpus that motivated it. +/// +/// Not an assertion: a measurement, printed. Run with +/// `cargo test --release -p vortex-file --lib measure_paging -- --ignored --nocapture`. + +#[tokio::test] +#[ignore = "measurement, not an assertion"] +async fn measure_paging() -> VortexResult<()> { + const ROWS: i32 = 65_536; + const COLUMNS: usize = 13; + const BLOCK: usize = 64; + + // Thirteen columns, as in the survey schema, at 64-row blocks: 1,024 blocks per column. + let names: Vec = (0..COLUMNS).map(|c| format!("c{c}")).collect(); + let fields: Vec<(&str, ArrayRef)> = names + .iter() + .enumerate() + .map(|(c, name)| { + let column: PrimitiveArray = (0..ROWS).map(|r| r + (c as i32 * ROWS)).collect(); + (name.as_str(), column.into_array()) + }) + .collect(); + let array = StructArray::from_fields(&fields)?.into_array(); + + println!( + "\n{ROWS} rows x {COLUMNS} i32 columns, {BLOCK}-row blocks \ + ({} blocks per column)\n", + ROWS as usize / BLOCK + ); + println!( + "{:>10} {:>12} {:>11} {:>10} {:>10} {:>9} {:>8} {:>9} {:>10} {:>12}", + "page_size", + "root tables", + "max/page", + "layout kB", + "file kB", + "segments", + "batches", + "scan ms", + "filter ms", + "filter/rowct" + ); + + for page_size in [0usize, 8, 32, 128, 1024] { + let strategy = crate::strategy::WriteStrategyBuilder::default() + .with_row_block_size(BLOCK) + .with_data_block_target_bytes(None) + .with_page_size(page_size) + .build(); + + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, array.to_array_stream()) + .await?; + let buf = buf.freeze(); + + let footer = summary.footer(); + let root_layout = footer.layout(); + let root_tables = min_layout_tables(root_layout); + let layout_bytes = root_layout + .flatbuffer_writer(&vortex_layout::LayoutContext::default()) + .write_flatbuffer_bytes()? + .len(); + + // Every page is a flatbuffer root of its own, so measure the worst one. + let mut pages = vec![]; + page_segments(root_layout.as_ref(), &mut pages); + let max_page_tables = pages + .iter() + .map(|id| { + let spec = &footer.segment_map()[**id as usize]; + let start = spec.offset as usize; + let end = start + spec.length as usize; + min_tables(&buf.as_ref()[start..end]) + }) + .max(); + + let file = SESSION.open_options().open_buffer(buf.clone())?; + + // Best of three, after a warm-up, so allocator and runtime ramp cannot be mistaken + // for a difference between the arms. + let mut batch_count = 0; + let mut scan_ms = f64::INFINITY; + for rep in 0..4 { + let started = std::time::Instant::now(); + let batches: Vec = file.scan()?.into_array_stream()?.try_collect().await?; + let elapsed = started.elapsed().as_secs_f64() * 1e3; + let rows: usize = batches.iter().map(|batch| batch.len()).sum(); + assert_eq!(rows, ROWS as usize, "every arm must read every row"); + if rep > 0 { + scan_ms = scan_ms.min(elapsed); + batch_count = batches.len(); + } + } + + // Selective: the tail of one column, so most zones prune away. + let filter = bind_scan_expr(&file, gt(get_item("c0", root()), lit(ROWS - 500))); + let mut filter_ms = f64::INFINITY; + let mut filter_rowcount_ms = f64::INFINITY; + for rep in 0..4 { + let started = std::time::Instant::now(); + let filtered: Vec = file + .scan()? + .with_filter(filter.clone()) + .into_array_stream()? + .try_collect() + .await?; + let elapsed = started.elapsed().as_secs_f64() * 1e3; + let filtered_rows: usize = filtered.iter().map(|batch| batch.len()).sum(); + assert_eq!(filtered_rows, 499, "every arm must select the same rows"); + + // The same filter, but with splits computed arithmetically instead of by asking the + // layout tree. This bypasses `register_splits`, which is where the inline layout + // builds a reader per chunk. + let started_rc = std::time::Instant::now(); + let filtered: Vec = file + .scan()? + .with_filter(filter.clone()) + .with_split_by(SplitBy::RowCount(BLOCK)) + .into_array_stream()? + .try_collect() + .await?; + let elapsed_rc = started_rc.elapsed().as_secs_f64() * 1e3; + let filtered_rows: usize = filtered.iter().map(|batch| batch.len()).sum(); + assert_eq!( + filtered_rows, 499, + "row-count splits must select the same rows" + ); + + if rep > 0 { + filter_ms = filter_ms.min(elapsed); + filter_rowcount_ms = filter_rowcount_ms.min(elapsed_rc); + } + } + + println!( + "{:>10} {:>12} {:>11} {:>10.1} {:>10.1} {:>9} {:>8} {:>9.1} {:>10.1} {:>12.1}", + if page_size == 0 { + "inline".to_string() + } else { + page_size.to_string() + }, + root_tables, + max_page_tables + .map(|t| t.to_string()) + .unwrap_or_else(|| "-".to_string()), + layout_bytes as f64 / 1024.0, + buf.len() as f64 / 1024.0, + footer.segment_map().len(), + batch_count, + scan_ms, + filter_ms, + filter_rowcount_ms, + ); + } + println!(); + + Ok(()) +} diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 26230146f80..19bd1073f24 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -226,7 +226,12 @@ impl VortexWriteOptions { // The array context is built here, rather than when the options were constructed, so that // encodings registered on the session in between are still eligible for the file. + // One layout context for the whole write: strategies that serialize a layout themselves, + // such as a page, intern into the same dictionary the footer is serialized through, so a + // nested layout flatbuffer indexes into the file's one dictionary. + let layout_ctx = new_layout_context(&self.session); let ctx = LayoutWriterContext::new(new_array_context(&self.session)) + .with_layout_context(layout_ctx.clone()) .with_buffered_bytes_tracker(self.buffered_bytes.clone()) .with_allowed_aggregates(edition_filter(&self.session, ComponentKind::Aggregate)); let dtype = stream.dtype().clone(); @@ -309,7 +314,7 @@ impl VortexWriteOptions { let (footer_buffers, metadata, approx_byte_size) = footer .clone() .into_serializer() - .with_layout_context(new_layout_context(&self.session)) + .with_layout_context(layout_ctx) .with_metadata_segments(self.metadata) .with_offset(position) .with_exclude_dtype(self.exclude_dtype) diff --git a/vortex-layout/src/children.rs b/vortex-layout/src/children.rs index 274ef124a60..c02be3a54fd 100644 --- a/vortex-layout/src/children.rs +++ b/vortex-layout/src/children.rs @@ -270,6 +270,7 @@ impl LayoutChildren for ViewedLayoutChildren { let build_ctx = LayoutBuildContext { session: &self.session, array_read_ctx: &self.array_read_ctx, + layout_read_ctx: &self.layout_read_ctx, }; encoding.build( dtype, diff --git a/vortex-layout/src/encoding.rs b/vortex-layout/src/encoding.rs index e2f06aa443f..ab3be1d6901 100644 --- a/vortex-layout/src/encoding.rs +++ b/vortex-layout/src/encoding.rs @@ -33,6 +33,8 @@ pub struct LayoutDeserializeArgs<'a> { pub session: &'a VortexSession, /// Array read context referenced by serialized array metadata. pub array_read_ctx: &'a ReadContext, + /// Layout read context referenced by nested layout flatbuffers, such as a page's. + pub layout_read_ctx: &'a ReadContext, /// Logical dtype of this layout. pub dtype: &'a DType, /// Number of rows in this layout. @@ -49,6 +51,8 @@ pub struct LayoutBuildContext<'a> { pub session: &'a VortexSession, /// Array read context referenced by serialized array metadata. pub array_read_ctx: &'a ReadContext, + /// Layout read context referenced by nested layout flatbuffers, such as a page's. + pub layout_read_ctx: &'a ReadContext, } /// Object-safe plugin registered for a layout ID. diff --git a/vortex-layout/src/flatbuffers.rs b/vortex-layout/src/flatbuffers.rs index a807db578af..8fe27d507e8 100644 --- a/vortex-layout/src/flatbuffers.rs +++ b/vortex-layout/src/flatbuffers.rs @@ -94,6 +94,7 @@ pub fn layout_from_flatbuffer_with_options( let build_ctx = LayoutBuildContext { session, array_read_ctx: ctx, + layout_read_ctx: layout_ctx, }; let layout = encoding.build( dtype, diff --git a/vortex-layout/src/layouts/chunked/writer.rs b/vortex-layout/src/layouts/chunked/writer.rs index 911290e0843..8e49851abe5 100644 --- a/vortex-layout/src/layouts/chunked/writer.rs +++ b/vortex-layout/src/layouts/chunked/writer.rs @@ -93,6 +93,7 @@ async fn paginate( &page, &row_offsets, ReadContext::new(ctx.array_ctx().to_ids()), + ctx.layout_ctx(), segment_sink, eof.advance(), ) diff --git a/vortex-layout/src/layouts/paged/mod.rs b/vortex-layout/src/layouts/paged/mod.rs index f62b19800a4..24a32fc40c0 100644 --- a/vortex-layout/src/layouts/paged/mod.rs +++ b/vortex-layout/src/layouts/paged/mod.rs @@ -30,7 +30,6 @@ use vortex_session::registry::ReadContext; use crate::Layout; use crate::LayoutChildType; use crate::LayoutDeserializeArgs; -use crate::LayoutEncodingId; use crate::LayoutId; use crate::LayoutParts; use crate::LayoutReaderContext; @@ -143,15 +142,9 @@ impl VTable for Paged { } fn metadata(layout: &Layout) -> Self::Metadata { - // The page's own encoding dictionary. The file footer's dictionary is not usable here: it - // is not built until the footer is serialized, long after the page segment was written. + // No encoding dictionary: the page's flatbuffer is interned into the same context the + // footer serializes through, so it indexes into the file's one dictionary. ProstMetadata(PagedLayoutMetadata { - layout_encodings: layout - .layout_ctx - .ids() - .iter() - .map(|id| id.as_str().to_string()) - .collect(), uniform_chunk_len: match &layout.boundaries { ChunkBoundaries::Uniform { len, .. } => Some(*len), ChunkBoundaries::Explicit(_) => None, @@ -212,14 +205,9 @@ impl VTable for Paged { "Paged layout must set both a uniform chunk length and a chunk count, or neither" ), }; - let layout_ids = metadata - .layout_encodings - .iter() - .map(|id| LayoutEncodingId::from(id.as_str())) - .collect::>(); Ok(PagedData { segment_id: args.segment_ids[0], - layout_ctx: ReadContext::new(layout_ids), + layout_ctx: args.layout_read_ctx.clone(), array_ctx: args.array_read_ctx.clone(), boundaries, }) @@ -302,8 +290,6 @@ impl Layout { #[derive(prost::Message)] pub struct PagedLayoutMetadata { - #[prost(repeated, string, tag = "1")] - pub layout_encodings: Vec, /// Exclusive row boundaries of the subtree's chunks, relative to this page's first row. /// /// Packed varints, so promoting them to the page boundary costs no flatbuffer tables — the @@ -469,6 +455,25 @@ mod test { }) } + /// A page must not repeat the file's layout encoding dictionary. Interning into the same + /// context the footer uses means a page's metadata carries only its chunk boundaries, which + /// for a uniform subtree is two integers. + #[test] + fn pages_do_not_repeat_the_encoding_dictionary() { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let (_, paged) = write_chunks(64, 64, &session).await; + + let page = paged.slot(0).unwrap().unwrap(); + let metadata = page.metadata().len(); + assert!( + metadata <= 16, + "a uniform page should carry only its boundaries, but its metadata is \ + {metadata} bytes" + ); + }) + } + /// Chunks of differing lengths cannot be derived, so those pages fall back to explicit /// offsets. The split set must stay exact either way. #[test] diff --git a/vortex-layout/src/layouts/paged/writer.rs b/vortex-layout/src/layouts/paged/writer.rs index dca20273df0..7b1517d09dd 100644 --- a/vortex-layout/src/layouts/paged/writer.rs +++ b/vortex-layout/src/layouts/paged/writer.rs @@ -14,9 +14,10 @@ use crate::sequence::SequenceId; /// Serialize `layout` into its own segment, returning the page that stands in for it. /// -/// The subtree is written as a nested `Layout` flatbuffer against a fresh encoding dictionary, -/// which travels in the page's metadata. The parent therefore spends one table on the whole -/// subtree, and the page is verified against the table and depth limits on its own. +/// The subtree is written as a nested `Layout` flatbuffer, interned into the same encoding +/// dictionary the footer serializes through, so the page indexes into the file's one dictionary +/// rather than carrying a copy. The parent spends one table on the whole subtree, and the page is +/// verified against the table and depth limits on its own. /// /// `row_offsets` are the subtree's exclusive chunk boundaries, relative to its first row. They are /// promoted to the page so scan planning, which is synchronous, can proceed without reading it, @@ -25,12 +26,12 @@ pub async fn write_page( layout: &LayoutRef, row_offsets: &[u64], array_ctx: ReadContext, + layout_ctx: &LayoutContext, segment_sink: &SegmentSinkRef, sequence_id: SequenceId, ) -> VortexResult { - let page_ctx = LayoutContext::default(); let page = layout - .flatbuffer_writer(&page_ctx) + .flatbuffer_writer(layout_ctx) .write_flatbuffer_bytes()?; let segment_id = segment_sink .write(sequence_id, vec![page.into_inner()]) @@ -40,7 +41,7 @@ pub async fn write_page( layout.row_count(), layout.dtype().clone(), segment_id, - ReadContext::new(page_ctx.to_ids()), + ReadContext::new(layout_ctx.to_ids()), array_ctx, ChunkBoundaries::from_offsets(row_offsets, layout.row_count()), ) diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index bc3d7d0626e..1e6aacb0f8d 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -624,6 +624,8 @@ mod tests { let build_ctx = LayoutBuildContext { session: &session, array_read_ctx: &build_read_ctx, + // No page in a zoned or legacy-stats subtree, so no nested layout to resolve. + layout_read_ctx: &ReadContext::new([]), }; let layout = ::build( @@ -655,6 +657,8 @@ mod tests { let build_ctx = LayoutBuildContext { session: &session, array_read_ctx: &build_read_ctx, + // No page in a zoned or legacy-stats subtree, so no nested layout to resolve. + layout_read_ctx: &ReadContext::new([]), }; let result = ::build( @@ -684,6 +688,8 @@ mod tests { let build_ctx = LayoutBuildContext { session: &session, array_read_ctx: &build_read_ctx, + // No page in a zoned or legacy-stats subtree, so no nested layout to resolve. + layout_read_ctx: &ReadContext::new([]), }; let metadata = ZonedMetadata { @@ -720,6 +726,8 @@ mod tests { let build_ctx = LayoutBuildContext { session: &session, array_read_ctx: &build_read_ctx, + // No page in a zoned or legacy-stats subtree, so no nested layout to resolve. + layout_read_ctx: &ReadContext::new([]), }; let metadata = ZonedMetadata { diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 52a57c1b92b..bb4fced4d6d 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -465,6 +465,8 @@ mod test { let build_ctx = LayoutBuildContext { session: &session, array_read_ctx: &read_ctx, + // No page in a zoned or legacy-stats subtree, so no nested layout to resolve. + layout_read_ctx: &ReadContext::new([]), }; let legacy_layout = ::build( &LegacyStatsLayoutEncoding, diff --git a/vortex-layout/src/strategy.rs b/vortex-layout/src/strategy.rs index 5a0b1025e4a..9ae2ae635c3 100644 --- a/vortex-layout/src/strategy.rs +++ b/vortex-layout/src/strategy.rs @@ -16,6 +16,7 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; +use crate::LayoutContext; use crate::LayoutRef; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; @@ -80,6 +81,7 @@ impl Drop for BufferedBytesReservation { #[derive(Clone)] pub struct LayoutWriterContext { array_ctx: ArrayContext, + layout_ctx: LayoutContext, allowed_aggregates: Option>>, buffered_bytes: BufferedBytesTracker, } @@ -89,6 +91,7 @@ impl LayoutWriterContext { pub fn new(array_ctx: ArrayContext) -> Self { Self { array_ctx, + layout_ctx: LayoutContext::default(), allowed_aggregates: None, buffered_bytes: BufferedBytesTracker::new(), } @@ -121,6 +124,20 @@ impl LayoutWriterContext { } /// Returns the array serialization context. + /// Intern layout encoding ids for any layout this write serializes itself, such as a page. + /// + /// Share the context the footer serializes through, so a nested layout flatbuffer indexes into + /// the file's one dictionary rather than carrying a copy. + pub fn with_layout_context(mut self, layout_ctx: LayoutContext) -> Self { + self.layout_ctx = layout_ctx; + self + } + + /// Returns the layout encoding context shared with the footer. + pub fn layout_ctx(&self) -> &LayoutContext { + &self.layout_ctx + } + pub fn array_ctx(&self) -> &ArrayContext { &self.array_ctx } diff --git a/vortex-layout/src/vtable.rs b/vortex-layout/src/vtable.rs index 7ec27b67046..24eee5c4380 100644 --- a/vortex-layout/src/vtable.rs +++ b/vortex-layout/src/vtable.rs @@ -61,6 +61,7 @@ pub trait VTable: 'static + Clone + Send + Sync + Debug { let args = LayoutDeserializeArgs { session: build_ctx.session, array_read_ctx: build_ctx.array_read_ctx, + layout_read_ctx: build_ctx.layout_read_ctx, dtype, row_count, segment_ids, From d7555910385e48b2fc1943f2c9799ce754ffd740 Mon Sep 17 00:00:00 2001 From: Wael Sulais Date: Mon, 17 Aug 2026 21:22:34 +0300 Subject: [PATCH 3/3] fix(layout): request a page's segment on first poll, not on construction `PagedReader` issued its segment request while constructing an evaluation future rather than when that future is polled. Callers build those futures without knowing whether they will await them: `ZonedReader::pruning_evaluation` builds its data child's pruning future eagerly and drops it unawaited when its own zone map already pruned the range. Pages were therefore read for ranges that were then discarded. Move the request inside the future. A pruning future that is dropped unawaited now issues no segment request, where it previously requested every page intersecting the row range. In-memory timings do not move, because `open_buffer` makes a segment request nearly free. What this avoids is a round trip per pruned page against remote storage. Assisted-by: Claude Opus 5 (claude-opus-5) via Claude Code 2.1.232 Signed-off-by: Wael Sulais --- vortex-layout/src/layouts/paged/mod.rs | 36 +++++++++++++++++++++++ vortex-layout/src/layouts/paged/reader.rs | 11 ++++--- 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/vortex-layout/src/layouts/paged/mod.rs b/vortex-layout/src/layouts/paged/mod.rs index 24a32fc40c0..3f8b9914090 100644 --- a/vortex-layout/src/layouts/paged/mod.rs +++ b/vortex-layout/src/layouts/paged/mod.rs @@ -328,6 +328,7 @@ mod test { use vortex_flatbuffers::layout as fbl; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; + use vortex_mask::Mask; use super::*; use crate::LayoutContext; @@ -542,6 +543,41 @@ mod test { }) } + /// A pruning evaluation that is never awaited must not fetch the page. + /// + /// `ZonedReader::pruning_evaluation` builds its data child's pruning future eagerly and only + /// awaits it when its own zone map did not already prune the range, so a page that issues its + /// segment request while the future is being constructed is read for ranges that are then + /// discarded. + #[test] + fn pruning_that_is_not_awaited_does_not_fetch_the_page() { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let (segments, layout) = write_chunks(12, 3, &session).await; + + let requests = Arc::new(AtomicUsize::new(0)); + let counting = Arc::new(CountingSource { + inner: segments, + requests: Arc::clone(&requests), + }); + let reader = layout + .new_reader("".into(), counting, &session, &Default::default()) + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + + // Built and dropped without ever being polled. + let _pruning = reader + .pruning_evaluation(&(0..12), &expr, Mask::new_true(12)) + .unwrap(); + + assert_eq!( + requests.load(Ordering::Relaxed), + 0, + "a pruning future that was never awaited fetched page segments" + ); + }) + } + /// A range that starts and ends inside different pages must still return exactly its rows. #[test] fn row_range_spanning_pages() { diff --git a/vortex-layout/src/layouts/paged/reader.rs b/vortex-layout/src/layouts/paged/reader.rs index c43ae62dc76..958a31ec937 100644 --- a/vortex-layout/src/layouts/paged/reader.rs +++ b/vortex-layout/src/layouts/paged/reader.rs @@ -63,12 +63,15 @@ impl PagedReader { /// Return a future resolving to the reader for this page's subtree. /// - /// The segment request is only issued once the page is actually evaluated, so a page that is - /// pruned away, or never intersected by any scanned row range, costs no IO. + /// The segment is requested inside the future, on its first poll, rather than when it is + /// constructed. Callers build a page's evaluation future without knowing whether they will + /// await it — `ZonedReader::pruning_evaluation` builds its data child's eagerly and drops it + /// unawaited when its own zone map already pruned the range — so requesting here would read + /// pages for ranges that are then discarded. fn child_reader(&self) -> SharedReaderFuture { self.child .get_or_init(|| { - let segment_fut = self.segment_source.request(self.layout.segment_id()); + let segment_id = self.layout.segment_id(); let dtype = self.layout.dtype().clone(); let layout_ctx = self.layout.layout_ctx().clone(); let array_ctx = self.layout.array_ctx().clone(); @@ -79,7 +82,7 @@ impl PagedReader { async move { let reader = async { - let segment = segment_fut.await?; + let segment = segment_source.request(segment_id).await?; // The page is a `Layout` flatbuffer root, verified in its own right. let page = FlatBuffer::align_from(segment.to_host_sync()); let layout = layout_from_flatbuffer(