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..c2582161397 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,343 @@ 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(()) +} + +/// 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 8de7c30fe47..8e49851abe5 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,74 @@ 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()), + ctx.layout_ctx(), + segment_sink, + eof.advance(), + ) + .await?, + ); + } + Ok(pages) } #[async_trait] @@ -53,11 +118,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 +158,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..3f8b9914090 --- /dev/null +++ b/vortex-layout/src/layouts/paged/mod.rs @@ -0,0 +1,651 @@ +// 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::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 { + // 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 { + 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" + ), + }; + Ok(PagedData { + segment_id: args.segment_ids[0], + layout_ctx: args.layout_read_ctx.clone(), + 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 { + /// 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 vortex_mask::Mask; + + 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}" + ); + }) + } + + /// 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] + 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 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() { + 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..958a31ec937 --- /dev/null +++ b/vortex-layout/src/layouts/paged/reader.rs @@ -0,0 +1,200 @@ +// 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 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_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(); + 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_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( + 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..7b1517d09dd --- /dev/null +++ b/vortex-layout/src/layouts/paged/writer.rs @@ -0,0 +1,49 @@ +// 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, 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, +/// and are kept symbolic when uniform. +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 = layout + .flatbuffer_writer(layout_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(layout_ctx.to_ids()), + array_ctx, + ChunkBoundaries::from_offsets(row_offsets, layout.row_count()), + ) + .into_layout()) +} 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/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 } } 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,