From 241de0311a545bc171b998ff887ef4bd8ba3e107 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 24 Aug 2026 09:51:43 -0400 Subject: [PATCH] Add list-aware physical scan split boundaries Signed-off-by: Matt Katz --- vortex-file/src/strategy.rs | 79 ++- vortex-file/src/tests.rs | 67 +++ vortex-layout/src/layouts/chunked/writer.rs | 39 +- vortex-layout/src/layouts/dict/writer.rs | 49 +- vortex-layout/src/layouts/list/mod.rs | 162 +++++- vortex-layout/src/layouts/list/reader.rs | 72 ++- vortex-layout/src/layouts/list/repartition.rs | 296 ++++++++++ vortex-layout/src/layouts/list/writer.rs | 510 +++++++++++++++--- vortex-layout/src/layouts/repartition.rs | 133 +++++ vortex-layout/src/layouts/struct_/writer.rs | 34 +- vortex-layout/src/layouts/table.rs | 73 ++- vortex-layout/src/layouts/zoned/writer.rs | 14 +- vortex-layout/src/strategy.rs | 81 ++- 13 files changed, 1471 insertions(+), 138 deletions(-) create mode 100644 vortex-layout/src/layouts/list/repartition.rs diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 9d4dbb90610..58174f81fba 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -3,6 +3,7 @@ //! This module defines the default layout strategy for a Vortex file. +use std::num::NonZeroU64; use std::num::NonZeroUsize; use std::sync::Arc; @@ -22,6 +23,7 @@ use vortex_layout::layouts::compressed::CompressorPlugin; use vortex_layout::layouts::dict::writer::DictStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::list::writer::ListLayoutStrategy; +use vortex_layout::layouts::repartition::CoalescingStrategy; use vortex_layout::layouts::repartition::RepartitionStrategy; use vortex_layout::layouts::repartition::RepartitionWriterOptions; use vortex_layout::layouts::table::TableStrategy; @@ -176,6 +178,7 @@ impl WriteStrategyBuilder { /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides /// applied. pub fn build(self) -> Arc { + let data_block_target_bytes = self.data_block_target_bytes; let flat: Arc = if let Some(flat) = self.flat_strategy { flat } else { @@ -223,7 +226,7 @@ impl WriteStrategyBuilder { // 4. prior to compression, coalesce up to a minimum size let coalescing = RepartitionStrategy::new( - compressing, + compressing.clone(), RepartitionWriterOptions { // Write stream partitions roughly become segments. Because Vortex never reads less // than one segment, the size of segments and, therefore, partitions, must be small @@ -231,9 +234,9 @@ impl WriteStrategyBuilder { // sufficient read concurrency for the desired throughput. One megabyte is small // enough to achieve this for S3 (Durner et al., "Exploiting Cloud Object Storage for // High-Performance Analytics", VLDB Vol 16, Iss 11). - block_size_minimum: self.data_block_target_bytes.unwrap_or(0), + block_size_minimum: data_block_target_bytes.unwrap_or(0), block_len_multiple: self.row_block_size, - block_size_target: self.data_block_target_bytes, + block_size_target: data_block_target_bytes, canonicalize: true, }, ); @@ -256,7 +259,7 @@ impl WriteStrategyBuilder { compress_then_flat.clone(), coalescing, Default::default(), - probe_compressor, + Arc::clone(&probe_compressor), ); let row_block_size = NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0"); @@ -293,28 +296,52 @@ impl WriteStrategyBuilder { .with_field_writers(self.field_writers); if self.use_list_layout { - // We need a closure here to enable recursive application of list layout. - table_strategy = table_strategy.with_list_layout_factory( - move |list_layout: ListLayoutStrategy| -> Arc { - let zoned = ZonedStrategy::new( - list_layout, - compress_then_flat.clone(), - ZonedLayoutOptions { - block_size: row_block_size, - ..Default::default() - }, - ); - Arc::new(RepartitionStrategy::new( - zoned, - RepartitionWriterOptions { - block_size_minimum: 0, - block_len_multiple: row_block_size.get(), - block_size_target: None, - canonicalize: false, - }, - )) - }, - ); + let list_repartition_target = data_block_target_bytes.and_then(NonZeroU64::new); + let list_coalescing: Arc = + if let Some(target) = list_repartition_target { + Arc::new(CoalescingStrategy::new(compressing, target)) + } else { + Arc::new(compressing) + }; + // The list writer chooses these element chunk boundaries before decomposition. Keep + // them through dictionary encoding and coalesce whole chunks only; never split them. + let list_leaf = Arc::new(DictStrategy::new( + Arc::clone(&list_coalescing), + compress_then_flat.clone(), + list_coalescing, + Default::default(), + probe_compressor, + )); + + // The factory is applied recursively to nested lists. + table_strategy = table_strategy + .with_list_elements_strategy(list_leaf) + .with_list_layout_factory( + move |list_layout: ListLayoutStrategy| -> Arc { + let list_layout = if let Some(target) = list_repartition_target { + list_layout.with_list_aware_repartition(target) + } else { + list_layout + }; + let zoned = ZonedStrategy::new( + list_layout, + compress_then_flat.clone(), + ZonedLayoutOptions { + block_size: row_block_size, + ..Default::default() + }, + ); + Arc::new(RepartitionStrategy::new( + zoned, + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: row_block_size.get(), + block_size_target: None, + canonicalize: false, + }, + )) + }, + ); } Arc::new(table_strategy) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..685577eef1d 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1922,6 +1922,73 @@ async fn nested_list_of_list_roundtrip() -> VortexResult<()> { Ok(()) } +/// List-element chunks are chosen in element space, then translated through offsets so file +/// splits remain at complete outer list rows rather than the unrelated input row-block boundary. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn list_layout_uses_offset_aligned_element_chunks_for_splits() -> VortexResult<()> { + let items = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(), + buffer![0u32, 2, 4, 9, 10].into_array(), + Validity::NonNullable, + )? + .into_array(); + let table = StructArray::from_fields(&[("items", items)])?.into_array(); + + let strategy = crate::strategy::WriteStrategyBuilder::default() + .with_row_block_size(4) + .with_data_block_target_bytes(Some(16)) + .with_list_layout() + .build(); + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, table.to_array_stream()) + .await?; + + let file = SESSION.open_options().open_buffer(buf)?; + assert_eq!(file.splits()?, [0..2, 2..3, 3..4]); + Ok(()) +} + +/// Child list boundaries climb through a struct's element-row space, then the outer list retains +/// only the one that lands on an outer list offset. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn list_of_struct_of_list_uses_composed_chunk_splits() -> VortexResult<()> { + let nested = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6, 7].into_array(), + buffer![0u32, 2, 4, 6, 8].into_array(), + Validity::NonNullable, + )? + .into_array(); + let elements = StructArray::from_fields(&[("nested", nested)])?.into_array(); + let items = ListArray::try_new( + elements, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let table = StructArray::from_fields(&[("items", items)])?.into_array(); + + let strategy = crate::strategy::WriteStrategyBuilder::default() + .with_row_block_size(4) + .with_data_block_target_bytes(Some(8)) + .with_list_layout() + .build(); + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, table.to_array_stream()) + .await?; + + let file = SESSION.open_options().open_buffer(buf)?; + assert_eq!(file.splits()?, [0..1, 1..2]); + Ok(()) +} + type MapEntryFixture<'a> = (i32, Option<&'a str>); type MapRowFixture<'a> = Option>>; diff --git a/vortex-layout/src/layouts/chunked/writer.rs b/vortex-layout/src/layouts/chunked/writer.rs index 8de7c30fe47..5a2e74d60ac 100644 --- a/vortex-layout/src/layouts/chunked/writer.rs +++ b/vortex-layout/src/layouts/chunked/writer.rs @@ -52,6 +52,7 @@ impl LayoutStrategy for ChunkedLayoutStrategy { let dtype2 = dtype.clone(); let chunk_strategy = Arc::clone(&self.chunk_strategy); let handle = session.handle(); + let child_context = ctx.clone(); // We spawn each child to allow parallelism when processing chunks. let stream = stream! { @@ -60,16 +61,16 @@ impl LayoutStrategy for ChunkedLayoutStrategy { let chunk_eof = eof.split_off(); let chunk_strategy = Arc::clone(&chunk_strategy); - let ctx = ctx.clone(); + let (child_ctx, child_info) = child_context.child_context(); let segment_sink = Arc::clone(&segment_sink); let dtype = dtype2.clone(); let session = session.clone(); yield handle.spawn_nested(move |handle| async move { let session = session.with_handle(handle); - chunk_strategy + let layout = chunk_strategy .write_stream( - ctx, + child_ctx, segment_sink, SequentialStreamAdapter::new( dtype, @@ -79,13 +80,41 @@ impl LayoutStrategy for ChunkedLayoutStrategy { chunk_eof, &session, ) - .await + .await?; + Ok::<_, vortex_error::VortexError>((layout, child_info)) }) } }; // Poll all of our children concurrently to accumulate their layouts. - let mut child_layouts: Vec = stream.buffered(usize::MAX).try_collect().await?; + let child_results: Vec<_> = stream.buffered(usize::MAX).try_collect().await?; + let mut child_layouts = Vec::with_capacity(child_results.len()); + let mut chunk_boundaries = Vec::new(); + let mut row_offset = 0; + + for (layout, child_info) in child_results { + let row_count = layout.row_count(); + chunk_boundaries.extend( + child_info + .chunk_boundaries() + .into_iter() + .filter(|&boundary| boundary != 0 && boundary < row_count) + .map(|boundary| row_offset + boundary), + ); + row_offset += row_count; + child_layouts.push(layout); + } + chunk_boundaries.extend( + child_layouts + .iter() + .map(|layout| layout.row_count()) + .scan(0, |row_offset, row_count| { + *row_offset += row_count; + Some(*row_offset) + }) + .take(child_layouts.len().saturating_sub(1)), + ); + ctx.report_chunk_boundaries(chunk_boundaries); if child_layouts.len() == 1 { Ok(child_layouts.pop().vortex_expect("must have one child")) diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 878a83f54c9..92dea533a24 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -187,36 +187,38 @@ impl LayoutStrategy for DictStrategy { let handle = session.handle(); let dtype2 = dtype.clone(); + let child_context = ctx.clone(); let child_layouts = stream! { pin_mut!(runs); while let Some((codes_stream, values_fut)) = runs.next().await { let codes = Arc::clone(&self.codes); let codes_eof = eof.split_off(); - let ctx2 = ctx.clone(); + let (codes_ctx, codes_info) = child_context.child_context(); let segment_sink2 = Arc::clone(&segment_sink); let session2 = session.clone(); let codes_fut = handle.spawn_nested(move |h| async move { let session2 = session2.with_handle(h); - codes.write_stream( - ctx2, + let layout = codes.write_stream( + codes_ctx, segment_sink2, codes_stream.sendable(), codes_eof, &session2, - ).await + ).await?; + Ok((layout, codes_info)) }); let values = Arc::clone(&self.values); let values_eof = eof.split_off(); - let ctx2 = ctx.clone(); + let (values_ctx, _) = child_context.child_context(); let segment_sink2 = Arc::clone(&segment_sink); let dtype2 = dtype2.clone(); let session2 = session.clone(); let values_layout = handle.spawn_nested(move |h| async move { let session2 = session2.with_handle(h); values.write_stream( - ctx2, + values_ctx, segment_sink2, SequentialStreamAdapter::new(dtype2, once(values_fut)).sendable(), values_eof, @@ -233,13 +235,44 @@ impl LayoutStrategy for DictStrategy { let mut child_layouts = child_layouts .buffered(usize::MAX) .map(|result| { - let (codes_layout, values_layout) = result?; + let ((codes_layout, codes_info), values_layout) = result?; // All values are referenced when created via dictionary encoding - Ok::<_, VortexError>(DictLayout::new(values_layout, codes_layout).into_layout()) + Ok::<_, VortexError>(( + DictLayout::new(values_layout, codes_layout).into_layout(), + codes_info, + )) }) .try_collect::>() .await?; + let mut chunk_boundaries = Vec::new(); + let mut row_offset = 0; + let mut layouts = Vec::with_capacity(child_layouts.len()); + for (layout, codes_info) in child_layouts.drain(..) { + let row_count = layout.row_count(); + chunk_boundaries.extend( + codes_info + .chunk_boundaries() + .into_iter() + .filter(|&boundary| boundary != 0 && boundary < row_count) + .map(|boundary| row_offset + boundary), + ); + row_offset += row_count; + layouts.push(layout); + } + chunk_boundaries.extend( + layouts + .iter() + .map(|layout| layout.row_count()) + .scan(0, |row_offset, row_count| { + *row_offset += row_count; + Some(*row_offset) + }) + .take(layouts.len().saturating_sub(1)), + ); + ctx.report_chunk_boundaries(chunk_boundaries); + let mut child_layouts = layouts; + if child_layouts.len() == 1 { return Ok(child_layouts.remove(0)); } diff --git a/vortex-layout/src/layouts/list/mod.rs b/vortex-layout/src/layouts/list/mod.rs index 0ba5be17dbc..17b0762ee4b 100644 --- a/vortex-layout/src/layouts/list/mod.rs +++ b/vortex-layout/src/layouts/list/mod.rs @@ -5,6 +5,7 @@ mod expr; mod reader; +mod repartition; pub mod writer; use std::sync::Arc; @@ -55,6 +56,35 @@ pub use List as ListLayoutEncoding; #[derive(Clone, Debug)] pub struct ListData { offsets_ptype: PType, + chunk_boundaries: Arc<[ListChunkBoundary]>, +} + +/// A physical elements-child chunk boundary expressed in both list row spaces. +/// +/// The boundary is stored only when an elements chunk end is also an outer-list row boundary. +/// This lets scan planning split list rows without putting either side of the split across an +/// elements chunk. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct ListChunkBoundary { + outer_row_end: u64, + element_row_end: u64, +} + +impl ListChunkBoundary { + pub(super) fn new(outer_row_end: u64, element_row_end: u64) -> Self { + Self { + outer_row_end, + element_row_end, + } + } + + pub(super) fn outer_row_end(self) -> u64 { + self.outer_row_end + } + + pub(super) fn element_row_end(self) -> u64 { + self.element_row_end + } } /// A list layout shredded into elements, offsets, and optional validity children. @@ -70,7 +100,10 @@ impl VTable for List { } fn metadata(layout: &Layout) -> Self::Metadata { - ProstMetadata(ListLayoutMetadata::new(layout.offsets_ptype)) + ProstMetadata(ListLayoutMetadata::new_with_chunk_boundaries( + layout.offsets_ptype, + &layout.chunk_boundaries, + )) } fn deserialize( @@ -83,7 +116,7 @@ impl VTable for List { .dtype .as_list_element_opt() .ok_or_else(|| vortex_err!("ListLayout requires a List dtype, got {}", args.dtype))?; - args.children.child(ELEMENTS_CHILD_INDEX, elements_dtype)?; + let elements = args.children.child(ELEMENTS_CHILD_INDEX, elements_dtype)?; let offsets_dtype = DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable); let offsets = args.children.child(OFFSETS_CHILD_INDEX, &offsets_dtype)?; vortex_error::vortex_ensure!( @@ -99,8 +132,17 @@ impl VTable for List { "List validity row count does not match parent" ); } + let chunk_boundaries = metadata + .chunk_boundaries + .iter() + .map(|boundary| { + ListChunkBoundary::new(boundary.outer_row_end, boundary.element_row_end) + }) + .collect::>(); + validate_chunk_boundaries(args.row_count, elements.row_count(), &chunk_boundaries)?; Ok(ListData { offsets_ptype: metadata.offsets_ptype(), + chunk_boundaries: chunk_boundaries.into(), }) } @@ -170,9 +212,21 @@ impl Layout { elements: LayoutRef, offsets: LayoutRef, validity: Option, + ) -> Self { + Self::new_with_chunk_boundaries(dtype, elements, offsets, validity, Vec::new()) + } + + pub(super) fn new_with_chunk_boundaries( + dtype: DType, + elements: LayoutRef, + offsets: LayoutRef, + validity: Option, + chunk_boundaries: Vec, ) -> Self { let row_count = offsets.row_count().saturating_sub(1); let offsets_ptype = offsets.dtype().as_ptype(); + validate_chunk_boundaries(row_count, elements.row_count(), &chunk_boundaries) + .vortex_expect("invalid list chunk boundaries"); let mut children = vec![elements, offsets]; children.extend(validity); Self::validate_children(&dtype, children.len()).vortex_expect("invalid list children"); @@ -182,7 +236,10 @@ impl Layout { row_count, Vec::new(), OwnedLayoutChildren::layout_children(children), - ListData { offsets_ptype }, + ListData { + offsets_ptype, + chunk_boundaries: chunk_boundaries.into(), + }, ) .into_typed() } @@ -209,6 +266,10 @@ impl Layout { self.offsets_ptype } + pub(super) fn chunk_boundaries(&self) -> &[ListChunkBoundary] { + &self.chunk_boundaries + } + /// Returns the list element dtype. pub fn elements_dtype(&self) -> &DType { self.dtype() @@ -227,12 +288,107 @@ impl Layout { pub struct ListLayoutMetadata { #[prost(enumeration = "PType", tag = "1")] offsets_ptype: i32, + #[prost(message, repeated, tag = "2")] + chunk_boundaries: Vec, +} + +#[derive(Clone, PartialEq, Eq, prost::Message)] +struct ListChunkBoundaryMetadata { + #[prost(uint64, tag = "1")] + outer_row_end: u64, + #[prost(uint64, tag = "2")] + element_row_end: u64, } impl ListLayoutMetadata { pub fn new(offsets_ptype: PType) -> Self { + Self::new_with_chunk_boundaries(offsets_ptype, &[]) + } + + fn new_with_chunk_boundaries( + offsets_ptype: PType, + chunk_boundaries: &[ListChunkBoundary], + ) -> Self { let mut metadata = Self::default(); metadata.set_offsets_ptype(offsets_ptype); + metadata.chunk_boundaries = chunk_boundaries + .iter() + .map(|boundary| ListChunkBoundaryMetadata { + outer_row_end: boundary.outer_row_end(), + element_row_end: boundary.element_row_end(), + }) + .collect(); metadata } } + +fn validate_chunk_boundaries( + outer_row_count: u64, + element_row_count: u64, + chunk_boundaries: &[ListChunkBoundary], +) -> VortexResult<()> { + vortex_error::vortex_ensure!( + chunk_boundaries + .iter() + .all(|boundary| boundary.outer_row_end() != 0 + && boundary.outer_row_end() < outer_row_count + && boundary.element_row_end() != 0 + && boundary.element_row_end() < element_row_count), + "List chunk boundaries must be interior to their row spaces" + ); + vortex_error::vortex_ensure!( + chunk_boundaries + .windows(2) + .all(|pair| pair[0].outer_row_end() < pair[1].outer_row_end() + && pair[0].element_row_end() < pair[1].element_row_end()), + "List chunk boundaries must be strictly increasing" + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use vortex_array::DeserializeMetadata; + use vortex_array::SerializeMetadata; + + use super::*; + + #[test] + fn chunk_boundaries_round_trip_through_metadata() -> VortexResult<()> { + let boundaries = [ + ListChunkBoundary::new(8, 21), + ListChunkBoundary::new(16, 50), + ]; + let encoded = ProstMetadata(ListLayoutMetadata::new_with_chunk_boundaries( + PType::U64, + &boundaries, + )) + .serialize(); + let decoded = + as DeserializeMetadata>::deserialize(&encoded)?; + + assert_eq!(decoded.offsets_ptype(), PType::U64); + assert_eq!( + decoded + .chunk_boundaries + .iter() + .map(|boundary| { + ListChunkBoundary::new(boundary.outer_row_end, boundary.element_row_end) + }) + .collect::>(), + boundaries + ); + Ok(()) + } + + #[test] + fn legacy_metadata_has_no_chunk_boundaries() -> VortexResult<()> { + let encoded = ProstMetadata(ListLayoutMetadata::new(PType::U32)).serialize(); + let decoded = + as DeserializeMetadata>::deserialize(&encoded)?; + + assert_eq!(decoded.offsets_ptype(), PType::U32); + assert!(decoded.chunk_boundaries.is_empty()); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 53227635de6..01daf2c3fc8 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -34,6 +34,7 @@ use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::RowSplits; use crate::SplitRange; +use crate::layouts::list::ListChunkBoundary; use crate::layouts::list::ListLayout; use crate::layouts::list::expr::ListChildrenNeeded; use crate::layouts::list::expr::get_necessary_bound_list_children; @@ -358,6 +359,14 @@ impl LayoutReader for ListReader { ) -> VortexResult<()> { split_range.check_bounds(self.layout.row_count())?; + let row_range = split_range.row_range(); + let chunk_boundaries = self.layout.chunk_boundaries(); + if !chunk_boundaries.is_empty() { + register_chunk_boundaries(chunk_boundaries, split_range, splits); + splits.push(split_range.root_row_range().end); + return Ok(()); + } + // Splits are difficult to calculate because all children live in different row coordinate spaces. // List elements typically comprise the majority of the data in a list, and validity/offsets can be treated // as metadata. We therefore want to parallelize the scan based on element work. @@ -375,7 +384,6 @@ impl LayoutReader for ListReader { &mut element_splits, )?; - let row_range = split_range.row_range(); let mut last_split = None; for element_split in element_splits.into_sorted_deduped() { let Some(split) = map_element_split_to_outer_grid( @@ -474,6 +482,30 @@ impl LayoutReader for ListReader { } } +/// Adds persisted list chunk boundaries that fall strictly inside `split_range`. +fn register_chunk_boundaries( + chunk_boundaries: &[ListChunkBoundary], + split_range: &SplitRange, + splits: &mut RowSplits, +) { + let row_range = split_range.row_range(); + for boundary in chunk_boundaries { + let split = boundary.outer_row_end(); + if split <= row_range.start { + continue; + } + if split >= row_range.end { + break; + } + splits.push( + split_range + .row_offset() + .checked_add(split) + .vortex_expect("List layout split offset overflow"), + ); + } +} + /// Converts a natural boundary from element-row space into an approximate outer-row scan split. /// /// Scan splits must be expressed in the list layout's outer-row space, but the elements child @@ -613,6 +645,7 @@ mod tests { use rstest::rstest; use vortex_array::ArrayContext; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -840,6 +873,16 @@ mod tests { ListLayoutStrategy::default() } + fn chunk_preserving_list_strategy() -> ListLayoutStrategy { + ListLayoutStrategy::default() + .with_elements(Arc::new(ChunkedLayoutStrategy::new( + FlatLayoutStrategy::default(), + ))) + .with_offsets(Arc::new(ChunkedLayoutStrategy::new( + FlatLayoutStrategy::default(), + ))) + } + fn layout_test_session() -> VortexSession { vortex_array::array_session() .with::() @@ -1063,6 +1106,33 @@ mod tests { assert_eq!(splits, expected); } + #[tokio::test] + async fn persisted_chunk_boundaries_drive_exact_list_splits() -> VortexResult<()> { + let chunk0 = ListArray::try_new( + PrimitiveArray::from_iter(0..3_i32).into_array(), + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + )? + .into_array(); + let chunk1 = ListArray::try_new( + PrimitiveArray::from_iter(3..7_i32).into_array(), + buffer![0u32, 1, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + let dtype = chunk0.dtype().clone(); + let list = ChunkedArray::try_new(vec![chunk0, chunk1], dtype)?.into_array(); + + let (segments, layout, session) = + write_layout(&chunk_preserving_list_strategy(), list).await?; + let reader = + layout.new_reader("".into(), segments, &session, &LayoutReaderContext::new())?; + + let splits = SplitBy::Layout.splits(reader.as_ref(), &(0..4), &[FieldMask::All])?; + assert_eq!(splits, [0, 2, 4]); + Ok(()) + } + #[tokio::test] async fn nested_list_propagates_element_splits() -> VortexResult<()> { let inner = ListArray::try_new( diff --git a/vortex-layout/src/layouts/list/repartition.rs b/vortex-layout/src/layouts/list/repartition.rs new file mode 100644 index 00000000000..34a2128d7d6 --- /dev/null +++ b/vortex-layout/src/layouts/list/repartition.rs @@ -0,0 +1,296 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::num::NonZeroU64; + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +/// Element arrays and their interior ends, chosen on the enclosing list-offset grid. +pub(super) struct RepartitionedListElements { + pub(super) arrays: Vec, + pub(super) boundaries: Vec, +} + +/// Repartition list elements at leaf fields, snapping every chunk end to a list boundary. +pub(super) fn repartition_list_elements( + elements: ArrayRef, + offsets: &[u64], + target_element_bytes: NonZeroU64, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let elements = elements.execute::(exec_ctx)?.into_array(); + let offset_base = offsets.first().copied().unwrap_or(0); + let element_end = offsets.last().copied().unwrap_or(offset_base); + let elements = elements.slice( + usize::try_from(offset_base).vortex_expect("list offset must fit usize") + ..usize::try_from(element_end).vortex_expect("list offset must fit usize"), + )?; + + if elements.dtype().is_struct() { + let (elements, mut boundaries) = + chunk_struct_fields(elements, offsets, target_element_bytes.get(), exec_ctx)?; + let elements = elements.into_array(); + if boundaries.is_empty() { + boundaries = + chunk_boundaries_at_list_offsets(&elements, offsets, target_element_bytes.get()); + } + return Ok(RepartitionedListElements { + arrays: split_at_boundaries(elements, &boundaries)?, + boundaries, + }); + } + if elements.dtype().is_list() { + // Retain outer list fences even when a nested list refines the chunks inside them. + let boundaries = + chunk_boundaries_at_list_offsets(&elements, offsets, target_element_bytes.get()); + return Ok(RepartitionedListElements { + arrays: split_at_boundaries(elements, &boundaries)?, + boundaries, + }); + } + + let elements = chunk_leaf_field(elements, offsets, target_element_bytes.get(), exec_ctx)?; + let arrays = if let Some(chunked) = elements.as_opt::() { + chunked.chunks() + } else { + vec![elements] + }; + let boundaries = arrays + .iter() + .map(|array| array.len() as u64) + .scan(0, |row_end, len| { + *row_end += len; + Some(*row_end) + }) + .take(arrays.len().saturating_sub(1)) + .collect(); + Ok(RepartitionedListElements { arrays, boundaries }) +} + +/// Recursively descend through structs and chunk each non-struct field independently. +fn chunk_struct_fields( + array: ArrayRef, + offsets: &[u64], + target_element_bytes: u64, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult<(StructArray, Vec)> { + let struct_array = array.execute::(exec_ctx)?; + let mut fields = Vec::with_capacity(struct_array.struct_fields().nfields()); + let mut boundaries = Vec::new(); + for field in struct_array.iter_unmasked_fields() { + let (field, field_boundaries) = if field.dtype().is_struct() { + let (field, boundaries) = + chunk_struct_fields(field.clone(), offsets, target_element_bytes, exec_ctx)?; + (field.into_array(), boundaries) + } else if field.dtype().is_list() { + // A nested list has its own finer fences. The enclosing struct gets a fallback fence + // schedule below if it has no non-list leaves. + (field.clone(), Vec::new()) + } else { + let field = chunk_leaf_field(field.clone(), offsets, target_element_bytes, exec_ctx)?; + let boundaries = chunk_boundaries(&field); + (field, boundaries) + }; + fields.push(field); + boundaries.extend(field_boundaries); + } + boundaries.sort_unstable(); + boundaries.dedup(); + + Ok(( + StructArray::try_new_with_dtype( + fields, + struct_array.struct_fields().clone(), + struct_array.len(), + struct_array.validity()?, + )?, + boundaries, + )) +} + +fn chunk_boundaries(array: &ArrayRef) -> Vec { + let Some(chunked) = array.as_opt::() else { + return Vec::new(); + }; + chunked + .iter_chunks() + .map(|chunk| chunk.len() as u64) + .scan(0, |row_end, len| { + *row_end += len; + Some(*row_end) + }) + .take(chunked.nchunks().saturating_sub(1)) + .collect() +} + +/// Canonicalize one leaf field and split it only at enclosing list boundaries. +fn chunk_leaf_field( + field: ArrayRef, + offsets: &[u64], + target_element_bytes: u64, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let field = field.execute::(exec_ctx)?.into_array(); + let row_count = offsets.len().saturating_sub(1); + if row_count == 0 { + return Ok(field); + } + + let boundaries = chunk_boundaries_at_list_offsets(&field, offsets, target_element_bytes); + let chunks = split_at_boundaries(field.clone(), &boundaries)?; + if chunks.len() == 1 { + Ok(chunks.into_iter().next().vortex_expect("one leaf chunk")) + } else { + ChunkedArray::try_new(chunks, field.dtype().clone()).map(IntoArray::into_array) + } +} + +/// Choose interior element ends at enclosing list offsets using an array's proportional byte +/// contribution. The caller turns these ends into physical chunks appropriate for its dtype. +fn chunk_boundaries_at_list_offsets( + elements: &ArrayRef, + offsets: &[u64], + target_element_bytes: u64, +) -> Vec { + let row_count = offsets.len().saturating_sub(1); + if row_count == 0 { + return Vec::new(); + } + + let offset_base = offsets[0]; + let element_count = offsets[row_count] - offset_base; + let element_bytes = estimated_element_bytes(elements); + let mut boundaries = Vec::new(); + let mut range_bytes = 0u64; + for row in 0..row_count { + range_bytes = range_bytes.saturating_add(estimated_range_bytes( + offsets[row] - offset_base, + offsets[row + 1] - offset_base, + element_count, + element_bytes, + )); + let boundary = offsets[row + 1] - offset_base; + if range_bytes >= target_element_bytes && boundary < element_count { + boundaries.push(boundary); + range_bytes = 0; + } + } + boundaries +} + +fn split_at_boundaries(array: ArrayRef, boundaries: &[u64]) -> VortexResult> { + let mut chunks = Vec::with_capacity(boundaries.len() + 1); + let mut start = 0; + for &end in boundaries { + let end = usize::try_from(end).vortex_expect("list offset must fit usize"); + chunks.push(array.slice(start..end)?); + start = end; + } + chunks.push(array.slice(start..array.len())?); + Ok(chunks) +} + +/// Estimate the flattened elements' contribution to the repartition target. +fn estimated_element_bytes(elements: &ArrayRef) -> u64 { + elements + .dtype() + .element_size() + .and_then(|element_size| { + u64::try_from(element_size) + .ok()? + .checked_mul(elements.len() as u64) + }) + .unwrap_or_else(|| elements.nbytes()) +} + +/// Estimate a range's bytes by its proportional position in the flattened element array. +/// Taking the difference between two prefix estimates preserves the total byte count exactly. +fn estimated_range_bytes(start: u64, end: u64, element_count: u64, nbytes: u64) -> u64 { + if element_count == 0 { + return 0; + } + let prefix = |offset: u64| { + u64::try_from(u128::from(offset) * u128::from(nbytes) / u128::from(element_count)) + .vortex_expect("estimated prefix bytes cannot exceed the element array size") + }; + prefix(end) - prefix(start) +} + +#[cfg(test)] +mod tests { + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::StructArray; + use vortex_buffer::buffer; + + use super::*; + + fn target(target_element_bytes: u64) -> NonZeroU64 { + NonZeroU64::new(target_element_bytes).vortex_expect("test target is non-zero") + } + + #[test] + fn keeps_sublists_whole_at_chunk_boundaries() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let chunks = repartition_list_elements( + buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(), + &[0, 2, 4, 9, 10], + target(16), + &mut ctx, + )?; + + assert_eq!(chunks.arrays.len(), 3); + assert_eq!(chunks.arrays[0].len(), 4); + assert_eq!(chunks.arrays[1].len(), 5); + assert_eq!(chunks.arrays[2].len(), 1); + assert_eq!(chunk_boundaries_from_arrays(&chunks.arrays), [4, 9]); + Ok(()) + } + + #[test] + fn chunks_struct_fields_independently() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let elements = StructArray::from_fields( + [ + ( + "wide", + buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(), + ), + ( + "narrow", + buffer![0u8, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(), + ), + ] + .as_slice(), + )? + .into_array(); + + let chunks = repartition_list_elements(elements, &[0, 2, 4, 9, 10], target(16), &mut ctx)?; + + assert_eq!(chunks.boundaries, [4, 9]); + assert_eq!(chunk_boundaries_from_arrays(&chunks.arrays), [4, 9]); + Ok(()) + } + + fn chunk_boundaries_from_arrays(arrays: &[ArrayRef]) -> Vec { + arrays + .iter() + .map(|array| array.len() as u64) + .scan(0, |row_end, len| { + *row_end += len; + Some(*row_end) + }) + .take(arrays.len().saturating_sub(1)) + .collect() + } +} diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 4d8565fdd10..359ee82f211 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::num::NonZeroU64; use std::sync::Arc; use async_trait::async_trait; @@ -15,6 +16,7 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::List; use vortex_array::arrays::ListView; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::list::ListArrayExt; use vortex_array::arrays::list::ListDataParts; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::builtins::ArrayBuiltins; @@ -23,6 +25,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::matcher::Matcher; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -30,10 +33,13 @@ use vortex_io::kanal_ext::KanalExt; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; +use super::repartition::repartition_list_elements; use crate::LayoutRef; use crate::LayoutStrategy; use crate::LayoutWriterContext; +use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; +use crate::layouts::list::ListChunkBoundary; use crate::layouts::list::ListLayout; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; @@ -46,6 +52,20 @@ use crate::sequence::SequentialStreamExt; /// Item carried on each child sub-stream: a sequenced, materialized chunk. type ChildChunk = VortexResult<(SequenceId, ArrayRef)>; +struct ListChildSenders { + elements: kanal::AsyncSender, + offsets: kanal::AsyncSender, + validity: Option>, +} + +#[derive(Default)] +struct ListTransposeState { + element_base: u64, + outer_base: u64, + chunk_boundaries: Vec, + emitted_first_chunk: bool, +} + /// Strategy for writing list-typed arrays, with a fallback for non-list dtypes. /// /// This is a *structural* writer that decomposes a list column into independent `elements`, @@ -69,6 +89,7 @@ pub struct ListLayoutStrategy { offsets: Arc, validity: Arc, fallback: Arc, + element_repartition_target: Option, } impl Default for ListLayoutStrategy { @@ -81,6 +102,7 @@ impl Default for ListLayoutStrategy { offsets: Arc::clone(&flat), validity: Arc::clone(&flat), fallback: flat, + element_repartition_target: None, } } } @@ -92,6 +114,13 @@ impl ListLayoutStrategy { self } + /// Repartition the elements stream toward the requested byte size without splitting a sublist + /// across chunks. + pub fn with_list_aware_repartition(mut self, target_element_bytes: NonZeroU64) -> Self { + self.element_repartition_target = Some(target_element_bytes); + self + } + /// Strategy for the `offsets` child. pub fn with_offsets(mut self, offsets: Arc) -> Self { self.offsets = offsets; @@ -135,150 +164,344 @@ impl LayoutStrategy for ListLayoutStrategy { .vortex_expect("DType is List") .as_ref() .clone(); - // Global (whole-column) offsets are cumulative and may exceed the input offset width, - // so definsively widen. + // Global offsets are cumulative and may exceed the input offset width. let offsets_dtype = DType::Primitive(PType::U64, Nullability::NonNullable); - // One bounded sub-stream per child: elements, offsets, and (when nullable) validity. - let (elements_tx, elements_rx) = kanal::bounded_async::(1); - let (offsets_tx, offsets_rx) = kanal::bounded_async::(1); + let (elements_tx, elements_rx) = kanal::bounded_async(1); + let (offsets_tx, offsets_rx) = kanal::bounded_async(1); let (validity_tx, validity_rx) = if is_nullable { - let (tx, rx) = kanal::bounded_async::(1); + let (tx, rx) = kanal::bounded_async(1); (Some(tx), Some(rx)) } else { (None, None) }; + let child_senders = ListChildSenders { + elements: elements_tx, + offsets: offsets_tx, + validity: validity_tx, + }; // Transpose the list column into its child sub-streams and rebase offsets to global // positions. Kept joined with the child writers below so producer errors surface rather // than being hidden as an early channel close. - let fanout_fut = transpose_list_column( - stream, - session.clone(), - elements_tx, - offsets_tx, - validity_tx, - ); + let transpose_session = session.clone(); + let element_repartition_target = self.element_repartition_target; + let fanout_fut = async move { + if let Some(target) = element_repartition_target { + transpose_repartitioned_list_column( + stream, + transpose_session, + child_senders, + target, + ) + .await + } else { + transpose_list_column(stream, transpose_session, child_senders).await + } + }; - // Spawn a writer per child sub-stream, concurrently. - let handle = session.handle(); + let (elements_ctx, elements_info) = ctx.child_context(); + let (offsets_ctx, _) = ctx.child_context(); + // Repartitioning creates a parent-owned fence schedule in element-row space. Make each + // selected fence a physical layout boundary before handing it to an arbitrary structural + // elements writer: otherwise a nested List/Struct writer could merge the stream back + // together and the enclosing list could not safely persist the fence. + let elements_strategy: Arc = if element_repartition_target.is_some() { + Arc::new(ChunkedLayoutStrategy::new(Arc::clone(&self.elements))) + } else { + Arc::clone(&self.elements) + }; let mut child_specs: Vec<( DType, Arc, kanal::AsyncReceiver, + LayoutWriterContext, )> = vec![ - (element_dtype, Arc::clone(&self.elements), elements_rx), - (offsets_dtype, Arc::clone(&self.offsets), offsets_rx), + (element_dtype, elements_strategy, elements_rx, elements_ctx), + ( + offsets_dtype, + Arc::clone(&self.offsets), + offsets_rx, + offsets_ctx, + ), ]; if let Some(validity_rx) = validity_rx { + let (validity_ctx, _) = ctx.child_context(); child_specs.push(( DType::Bool(Nullability::NonNullable), Arc::clone(&self.validity), validity_rx, + validity_ctx, )); } + let handle = session.handle(); let layout_futures: Vec<_> = child_specs .into_iter() - .map(|(child_dtype, strategy, rx)| { + .map(|(child_dtype, strategy, rx, child_ctx)| { let child_stream = SequentialStreamAdapter::new(child_dtype, rx.into_stream().boxed()).sendable(); let child_eof = eof.split_off(); - let ctx = ctx.clone(); let segment_sink = Arc::clone(&segment_sink); let session = session.clone(); handle.spawn_nested(move |h| async move { let session = session.with_handle(h); strategy - .write_stream(ctx, segment_sink, child_stream, child_eof, &session) + .write_stream(child_ctx, segment_sink, child_stream, child_eof, &session) .await }) }) .collect(); - let (_, layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; + let (planned_chunk_boundaries, layouts) = + try_join(fanout_fut, try_join_all(layout_futures)).await?; let mut layouts = layouts.into_iter(); let elements_layout = layouts.next().vortex_expect("elements layout present"); let offsets_layout = layouts.next().vortex_expect("offsets layout present"); let validity_layout = is_nullable.then(|| layouts.next().vortex_expect("validity layout present")); - Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) + let element_chunk_boundaries = elements_info.chunk_boundaries(); + let mut chunk_boundaries = planned_chunk_boundaries + .into_iter() + .filter(|boundary| { + element_chunk_boundaries + .binary_search(&boundary.element_row_end()) + .is_ok() + }) + .filter(|boundary| { + boundary.outer_row_end() != 0 + && boundary.outer_row_end() < offsets_layout.row_count().saturating_sub(1) + && boundary.element_row_end() != 0 + && boundary.element_row_end() < elements_layout.row_count() + }) + .collect::>(); + chunk_boundaries.sort_unstable_by_key(|boundary| boundary.outer_row_end()); + chunk_boundaries.dedup_by_key(|boundary| boundary.outer_row_end()); + ctx.report_chunk_boundaries( + chunk_boundaries + .iter() + .map(|boundary| boundary.outer_row_end()), + ); + + Ok(ListLayout::new_with_chunk_boundaries( + dtype, + elements_layout, + offsets_layout, + validity_layout, + chunk_boundaries, + ) + .into_layout()) } } /// Transpose a list column into its `elements`, `offsets`, and (when present) `validity` child -/// sub-streams, rebasing each chunk's local `offsets` to global `u64` positions so the single +/// sub-streams. Rebases each chunk's local `offsets` to global `u64` positions so the single /// `offsets` child indexes into the concatenated `elements` child. /// -/// `validity_tx` is `Some` exactly when the list is nullable. Errors surface to the caller, which -/// joins this against the child writers, rather than being hidden as an early channel close. +/// The validity sender is present only when the list is nullable. Errors surface to the caller, +/// which joins this against the child writers, rather than being hidden as an early channel close. async fn transpose_list_column( mut stream: SendableSequentialStream, session: VortexSession, - elements_tx: kanal::AsyncSender, - offsets_tx: kanal::AsyncSender, - validity_tx: Option>, -) -> VortexResult<()> { + child_senders: ListChildSenders, +) -> VortexResult> { let mut exec_ctx = session.create_execution_ctx(); - let mut element_base: u64 = 0; - let mut first = true; + let mut state = ListTransposeState::default(); let mut saw_chunk = false; + while let Some(chunk) = stream.next().await { let (sequence_id, array) = chunk?; saw_chunk = true; + let mut sp = sequence_id.descend(); - let ListDataParts { + let (elements, offsets, validity) = canonicalize_list_chunk(array, &mut exec_ctx)?; + // An input list chunk end is an exact outer-list offset boundary. Keep it as a candidate + // and persist it only if the elements writer confirms that it remained physical. + state.chunk_boundaries.push(ListChunkBoundary::new( + state.outer_base + offsets.len().saturating_sub(1) as u64, + state.element_base + elements.len() as u64, + )); + emit_list_parts( + vec![elements], + offsets.into_array(), + validity, + &mut sp, + &child_senders, + &mut state, + &mut exec_ctx, + ) + .await?; + } + + if !saw_chunk { + vortex_bail!("ListLayoutStrategy needs at least one chunk"); + } + + Ok(state.chunk_boundaries) +} + +/// Transpose a list column while snapping independently sized element chunks to list boundaries. +async fn transpose_repartitioned_list_column( + mut stream: SendableSequentialStream, + session: VortexSession, + child_senders: ListChildSenders, + element_repartition_target: NonZeroU64, +) -> VortexResult> { + let mut exec_ctx = session.create_execution_ctx(); + let mut state = ListTransposeState::default(); + let mut saw_chunk = false; + + while let Some(chunk) = stream.next().await { + let (sequence_id, array) = chunk?; + saw_chunk = true; + + let mut sp = sequence_id.descend(); + let (elements, offsets, validity) = canonicalize_list_chunk(array, &mut exec_ctx)?; + let element_chunks = repartition_list_elements( elements, - offsets, + offsets.as_slice::(), + element_repartition_target, + &mut exec_ctx, + )?; + state + .chunk_boundaries + .extend(map_element_boundaries_to_list_rows( + &element_chunks.boundaries, + offsets.as_slice::(), + state.outer_base, + state.element_base, + )); + emit_list_parts( + element_chunks.arrays, + offsets.into_array(), validity, - .. - } = canonicalize_to_list_parts(array, &mut exec_ctx)?; - let n_elements = elements.len() as u64; - let row_count = offsets.len().saturating_sub(1); - let offsets = global_offsets(offsets, element_base, first, &mut exec_ctx)?; - element_base += n_elements; - first = false; - - if elements_tx - .send(Ok((sp.advance(), elements))) - .await - .is_err() - || offsets_tx.send(Ok((sp.advance(), offsets))).await.is_err() - { - vortex_bail!("list child writer finished before all chunks were sent"); - } - if let Some(validity_tx) = &validity_tx { - let validity = validity - .execute_mask(row_count, &mut exec_ctx)? - .into_array(); - if validity_tx - .send(Ok((sp.advance(), validity))) - .await - .is_err() - { - vortex_bail!("list validity writer finished before all chunks were sent"); - } - } + &mut sp, + &child_senders, + &mut state, + &mut exec_ctx, + ) + .await?; } + if !saw_chunk { vortex_bail!("ListLayoutStrategy needs at least one chunk"); } - Ok(()) + + Ok(state.chunk_boundaries) } -/// Canonicalize a list-dtype array into [`ListDataParts`]. -fn canonicalize_to_list_parts( +/// Translate producer-selected element ends into the enclosing list's row space. A boundary that +/// falls inside a list value is intentionally dropped: it cannot be a list scan split. +fn map_element_boundaries_to_list_rows( + element_boundaries: &[u64], + offsets: &[u64], + outer_base: u64, + element_base: u64, +) -> Vec { + let offset_base = offsets.first().copied().unwrap_or_default(); + element_boundaries + .iter() + .filter_map(|&element_row_end| { + let offset = offset_base.checked_add(element_row_end)?; + let offset_index = offsets.partition_point(|&candidate| candidate <= offset); + (offset_index != 0 && offsets[offset_index - 1] == offset).then_some(())?; + Some(ListChunkBoundary::new( + outer_base + (offset_index - 1) as u64, + element_base + element_row_end, + )) + }) + .collect() +} + +/// Canonicalize a list chunk into elements, `u64` offsets, and validity. +fn canonicalize_list_chunk( array: ArrayRef, exec_ctx: &mut ExecutionCtx, -) -> VortexResult { +) -> VortexResult<(ArrayRef, PrimitiveArray, Validity)> { let canonical = array.execute_until::(exec_ctx)?; - if let Some(list) = canonical.as_opt::() { - Ok(list.into_owned().into_data_parts()) + let ListDataParts { + elements, + offsets, + validity, + .. + } = if let Some(list) = canonical.as_opt::() { + list.reset_offsets(false, exec_ctx)?.into_data_parts() } else if let Some(view) = canonical.as_opt::() { - Ok(list_from_list_view(view.into_owned(), exec_ctx)?.into_data_parts()) + list_from_list_view(view.into_owned(), exec_ctx)?.into_data_parts() } else { unreachable!("AnyList matcher guarantees List or ListView") + }; + let offsets = offsets + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::(exec_ctx)?; + Ok((elements, offsets, validity)) +} + +/// Emit one set of list parts to the child writers, rebasing its offsets to the global element +/// stream. +async fn emit_list_parts( + elements: Vec, + offsets: ArrayRef, + validity: Validity, + sp: &mut SequencePointer, + child_senders: &ListChildSenders, + state: &mut ListTransposeState, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let n_elements: u64 = elements.iter().map(|elements| elements.len() as u64).sum(); + let row_count = offsets.len().saturating_sub(1); + + for elements in elements { + if child_senders + .elements + .send(Ok((sp.advance(), elements))) + .await + .is_err() + { + vortex_bail!("list elements writer finished before all chunks were sent"); + } + } + + let offsets = global_offsets( + offsets, + state.element_base, + !state.emitted_first_chunk, + exec_ctx, + )?; + state.element_base += n_elements; + state.outer_base += row_count as u64; + state.emitted_first_chunk = true; + + if child_senders + .offsets + .send(Ok((sp.advance(), offsets))) + .await + .is_err() + { + vortex_bail!("list offsets writer finished before all chunks were sent"); + } + if let Some(validity_tx) = &child_senders.validity { + let validity = validity.execute_mask(row_count, exec_ctx)?.into_array(); + if validity_tx + .send(Ok((sp.advance(), validity))) + .await + .is_err() + { + vortex_bail!("list validity writer finished before all chunks were sent"); + } + } + Ok(()) +} + +/// Matcher for `Array` or `Array`. +struct AnyList; + +impl Matcher for AnyList { + type Match<'a> = (); + + fn try_match(array: &ArrayRef) -> Option> { + (array.as_opt::().is_some() || array.as_opt::().is_some()).then_some(()) } } @@ -293,12 +516,11 @@ fn global_offsets( first: bool, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let widened = offsets.cast(DType::Primitive(PType::U64, Nullability::NonNullable))?; let based = if element_base == 0 { - widened + offsets } else { - let base = ConstantArray::new(element_base, widened.len()).into_array(); - widened.binary(base, Operator::Add)? + let base = ConstantArray::new(element_base, offsets.len()).into_array(); + offsets.binary(base, Operator::Add)? }; let based = if first { based @@ -309,17 +531,6 @@ fn global_offsets( Ok(based.execute::(exec_ctx)?.into_array()) } -/// Matcher for `Array` or `Array`. -struct AnyList; - -impl Matcher for AnyList { - type Match<'a> = (); - - fn try_match(array: &ArrayRef) -> Option> { - (array.as_opt::().is_some() || array.as_opt::().is_some()).then_some(()) - } -} - #[cfg(test)] mod tests { use futures::stream; @@ -337,6 +548,8 @@ mod tests { use super::*; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::layouts::list::List; + use crate::layouts::struct_::StructStrategy; use crate::layouts::table::TableStrategy; use crate::segments::TestSegments; use crate::sequence::SequentialArrayStreamExt; @@ -423,6 +636,137 @@ mod tests { Ok(()) } + #[tokio::test] + async fn maps_independent_element_chunks_to_outer_row_boundaries() -> VortexResult<()> { + let list = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(), + buffer![0u32, 2, 4, 9, 10].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let writer = ListLayoutStrategy::default() + .with_list_aware_repartition(NonZeroU64::new(16).vortex_expect("non-zero target")) + .with_elements(Arc::new(ChunkedLayoutStrategy::new( + FlatLayoutStrategy::default(), + ))); + let session = layout_test_session(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let (ctx, info) = LayoutWriterContext::new(ArrayContext::empty()).child_context(); + let layout = writer + .write_stream( + ctx, + segments, + list.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + assert_eq!( + layout.as_::().chunk_boundaries(), + [ListChunkBoundary::new(2, 4), ListChunkBoundary::new(3, 9),] + ); + assert_eq!(info.chunk_boundaries(), [2, 3]); + Ok(()) + } + + #[tokio::test] + async fn nested_list_boundaries_are_mapped_compositionally() -> VortexResult<()> { + let inner = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6, 7].into_array(), + buffer![0u32, 2, 4, 6, 8].into_array(), + Validity::NonNullable, + )? + .into_array(); + let outer = ListArray::try_new( + inner, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let inner_strategy = ListLayoutStrategy::default() + .with_list_aware_repartition(NonZeroU64::new(8).vortex_expect("non-zero target")) + .with_elements(Arc::new(ChunkedLayoutStrategy::new( + FlatLayoutStrategy::default(), + ))); + let strategy = ListLayoutStrategy::default() + .with_list_aware_repartition(NonZeroU64::new(8).vortex_expect("non-zero target")) + .with_elements(Arc::new(inner_strategy)); + let session = layout_test_session(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let (ctx, info) = LayoutWriterContext::new(ArrayContext::empty()).child_context(); + let layout = strategy + .write_stream( + ctx, + segments, + outer.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + assert_eq!( + layout.as_::().chunk_boundaries(), + [ListChunkBoundary::new(1, 2)] + ); + assert_eq!(info.chunk_boundaries(), [1]); + Ok(()) + } + + #[tokio::test] + async fn list_of_struct_of_list_maps_only_outer_aligned_boundaries() -> VortexResult<()> { + let nested = ListArray::try_new( + buffer![0i32, 1, 2, 3, 4, 5, 6, 7].into_array(), + buffer![0u32, 2, 4, 6, 8].into_array(), + Validity::NonNullable, + )? + .into_array(); + let elements = StructArray::from_fields([("nested", nested)].as_slice())?.into_array(); + let outer = ListArray::try_new( + elements, + buffer![0u32, 2, 4].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + let nested_strategy = ListLayoutStrategy::default() + .with_list_aware_repartition(NonZeroU64::new(8).vortex_expect("non-zero target")) + .with_elements(Arc::new(ChunkedLayoutStrategy::new( + FlatLayoutStrategy::default(), + ))); + let struct_strategy = StructStrategy::new(Arc::clone(&flat), flat) + .with_field_writer("nested", Arc::new(nested_strategy)); + let strategy = ListLayoutStrategy::default() + .with_list_aware_repartition(NonZeroU64::new(8).vortex_expect("non-zero target")) + .with_elements(Arc::new(struct_strategy)); + + let session = layout_test_session(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let (ctx, info) = LayoutWriterContext::new(ArrayContext::empty()).child_context(); + let layout = strategy + .write_stream( + ctx, + segments, + outer.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + assert_eq!( + layout.as_::().chunk_boundaries(), + [ListChunkBoundary::new(1, 2)] + ); + assert_eq!(info.chunk_boundaries(), [1]); + Ok(()) + } + /// Non-list input dispatches to the fallback strategy unchanged. #[tokio::test] async fn non_list_input_routes_to_fallback() -> VortexResult<()> { diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 7344ee7be51..86dd5fe555c 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::collections::VecDeque; +use std::num::NonZeroU64; use std::sync::Arc; use async_stream::try_stream; @@ -23,6 +24,7 @@ use crate::LayoutStrategy; use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; +use crate::sequence::SequenceId; use crate::sequence::SequencePointer; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; @@ -82,6 +84,97 @@ pub struct RepartitionStrategy { options: RepartitionWriterOptions, } +/// Coalesce adjacent input chunks toward a target size without splitting any input chunk. +/// +/// This is used below list elements, where input chunk boundaries have already been snapped to +/// list offsets and must remain valid boundaries for the eventual physical chunks. +#[derive(Clone)] +pub struct CoalescingStrategy { + child: Arc, + target_bytes: NonZeroU64, +} + +impl CoalescingStrategy { + /// Create a strategy that merges whole chunks until their logical size reaches `target_bytes`. + pub fn new(child: S, target_bytes: NonZeroU64) -> Self { + Self { + child: Arc::new(child), + target_bytes, + } + } +} + +#[async_trait] +impl LayoutStrategy for CoalescingStrategy { + async fn write_stream( + &self, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + let dtype = stream.dtype().clone(); + let dtype_clone = dtype.clone(); + let target_bytes = self.target_bytes.get(); + let coalescing_session = session.clone(); + let coalesced = try_stream! { + pin_mut!(stream); + + let mut exec_ctx = coalescing_session.create_execution_ctx(); + let mut chunks = Vec::new(); + let mut nbytes = 0u64; + let mut last_sequence_id: Option = None; + + while let Some(chunk) = stream.next().await { + let (sequence_id, array) = chunk?; + nbytes = nbytes.saturating_add(array.nbytes()); + chunks.push(array); + last_sequence_id = Some(sequence_id); + + if nbytes >= target_bytes { + let array = canonicalize_chunks(&mut chunks, &dtype_clone, &mut exec_ctx)?; + let mut sequence_pointer = last_sequence_id + .take() + .vortex_expect("coalesced chunks have a sequence id") + .descend(); + yield (sequence_pointer.advance(), array); + nbytes = 0; + } + } + + if !chunks.is_empty() { + let array = canonicalize_chunks(&mut chunks, &dtype_clone, &mut exec_ctx)?; + let mut sequence_pointer = last_sequence_id + .vortex_expect("coalesced chunks have a sequence id") + .descend(); + yield (sequence_pointer.advance(), array); + } + }; + + self.child + .write_stream( + ctx, + segment_sink, + SequentialStreamAdapter::new(dtype, coalesced).sendable(), + eof, + session, + ) + .await + } +} + +fn canonicalize_chunks( + chunks: &mut Vec, + dtype: &DType, + exec_ctx: &mut vortex_array::ExecutionCtx, +) -> VortexResult { + ChunkedArray::try_new(std::mem::take(chunks), dtype.clone())? + .into_array() + .execute::(exec_ctx) + .map(IntoArray::into_array) +} + impl RepartitionStrategy { pub fn new(child: S, options: RepartitionWriterOptions) -> Self { Self { @@ -355,6 +448,46 @@ mod tests { assert_eq!(options.effective_block_len(&dtype), 1); } + #[test] + fn coalescing_strategy_merges_whole_input_chunks() -> VortexResult<()> { + let chunks = vec![ + PrimitiveArray::from_iter([0i32, 1, 2]).into_array(), + PrimitiveArray::from_iter([3i32, 4, 5]).into_array(), + PrimitiveArray::from_iter([6i32, 7, 8]).into_array(), + ]; + let dtype = chunks[0].dtype().clone(); + let array = ChunkedArray::try_new(chunks, dtype)?.into_array(); + + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let target_bytes = NonZeroU64::new(16).vortex_expect("16 is non-zero"); + let strategy = CoalescingStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + target_bytes, + ); + let (ctx, info) = LayoutWriterContext::new(ArrayContext::empty()).child_context(); + let stream = array.to_array_stream().sequenced(ptr); + let layout = block_on(|handle| async move { + let session = new_session().with_handle(handle); + strategy + .write_stream( + ctx, + Arc::::clone(&segments), + stream, + eof, + &session, + ) + .await + })?; + + assert_eq!(layout.nchildren(), 2); + let children = layout.children()?; + assert_eq!(children[0].row_count(), 6); + assert_eq!(children[1].row_count(), 3); + assert_eq!(info.chunk_boundaries(), [6]); + Ok(()) + } + #[test] fn repartition_large_element_type_produces_small_blocks() -> VortexResult<()> { // Create a FixedSizeList(f64, 1000) array with 1000 lists. diff --git a/vortex-layout/src/layouts/struct_/writer.rs b/vortex-layout/src/layouts/struct_/writer.rs index fee71596fee..3b9151f9a82 100644 --- a/vortex-layout/src/layouts/struct_/writer.rs +++ b/vortex-layout/src/layouts/struct_/writer.rs @@ -217,6 +217,7 @@ impl LayoutStrategy for StructStrategy { struct_dtype.names().iter().cloned().collect() }; + let child_context = ctx.clone(); let layout_futures: Vec<_> = column_dtypes .into_iter() .zip_eq(column_streams_rx) @@ -227,7 +228,7 @@ impl LayoutStrategy for StructStrategy { SequentialStreamAdapter::new(dtype, recv.into_stream().boxed()).sendable(); let child_eof = eof.split_off(); let session = session.clone(); - let ctx = ctx.clone(); + let (child_ctx, child_info) = child_context.child_context(); let segment_sink = Arc::clone(&segment_sink); handle.spawn_nested(move |h| { // Validity is written through the validity strategy; every other field @@ -243,15 +244,38 @@ impl LayoutStrategy for StructStrategy { let session = session.with_handle(h); async move { - writer - .write_stream(ctx, segment_sink, column_stream, child_eof, &session) - .await + let layout = writer + .write_stream( + child_ctx, + segment_sink, + column_stream, + child_eof, + &session, + ) + .await?; + Ok((layout, child_info)) } }) }) .collect(); - let (_success, column_layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; + let (_success, column_results) = try_join(fanout_fut, try_join_all(layout_futures)).await?; + let mut column_layouts = Vec::with_capacity(column_results.len()); + let mut chunk_boundaries: Option> = None; + for (layout, child_info) in column_results { + let child_boundaries = child_info.chunk_boundaries(); + if let Some(chunk_boundaries) = &mut chunk_boundaries { + chunk_boundaries + .retain(|boundary| child_boundaries.binary_search(boundary).is_ok()); + } else { + chunk_boundaries = Some(child_boundaries); + } + column_layouts.push(layout); + } + // A struct boundary is usable only when no field (including validity) crosses it. + // Reporting the union would let a parent treat a boundary from one field as though the + // whole struct had been physically partitioned there. + ctx.report_chunk_boundaries(chunk_boundaries.unwrap_or_default()); // TODO(os): transposed stream could count row counts as well, // This must hold though, all columns must have the same row count of the struct layout let row_count = column_layouts.first().map(|l| l.row_count()).unwrap_or(0); diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 1a3c1adc524..2bec5ac3a93 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -16,7 +16,12 @@ use std::env; use std::sync::Arc; use std::sync::LazyLock; +use async_stream::try_stream; use async_trait::async_trait; +use futures::StreamExt; +use futures::pin_mut; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::dtype::Field; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldPath; @@ -33,6 +38,8 @@ use crate::layouts::struct_::StructStrategy; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt; /// Whether [`TableStrategy`] writes list fields using a [`ListLayoutStrategy`] by /// default. Disabled unless the environment variable `VORTEX_EXPERIMENTAL_LIST_LAYOUT` @@ -47,6 +54,12 @@ pub fn use_experimental_list_layout() -> bool { type ListLayoutFactory = Arc Arc + Send + Sync>; +#[derive(Clone, Copy)] +enum LeafRole { + Default, + ListElements, +} + /// A configurable strategy for writing nested tabular data, dispatching each (sub)stream to the /// structural writer for its dtype. /// @@ -69,6 +82,10 @@ pub struct TableStrategy { validity: Arc, /// The writer for leaf fields, i.e. anything that is not a struct. leaf: Arc, + /// Optional leaf strategy used below a list's elements child. + list_elements: Option>, + /// Selects which leaf strategy this descended dispatcher uses. + leaf_role: LeafRole, /// Optional factory applied to each dynamically constructed [`ListLayoutStrategy`]. /// Its presence also enables list decomposition. /// @@ -100,6 +117,8 @@ impl TableStrategy { leaf_writers: Default::default(), validity, leaf: fallback, + list_elements: None, + leaf_role: LeafRole::Default, list_layout_factory: None, } } @@ -162,6 +181,12 @@ impl TableStrategy { self } + /// Override the leaf strategy used below a list's elements child. + pub fn with_list_elements_strategy(mut self, strategy: Arc) -> Self { + self.list_elements = Some(strategy); + self + } + /// Override the strategy for compressing struct validity at all levels of the schema tree. pub fn with_validity_strategy(mut self, validity: Arc) -> Self { self.validity = validity; @@ -231,13 +256,26 @@ impl TableStrategy { fn list_strategy(&self) -> Option> { let factory = self.list_layout_factory.as_ref()?; let list_layout = ListLayoutStrategy::default() - .with_elements(Arc::new(self.descend_clean())) + .with_elements(Arc::new(self.descend_list_elements())) .with_offsets(Arc::clone(&self.leaf)) .with_validity(Arc::clone(&self.validity)) - .with_fallback(Arc::clone(&self.leaf)); + .with_fallback(self.leaf_strategy()); Some(factory(list_layout)) } + fn leaf_strategy(&self) -> Arc { + match (self.leaf_role, &self.list_elements) { + (LeafRole::ListElements, Some(strategy)) => Arc::clone(strategy), + _ => Arc::clone(&self.leaf), + } + } + + fn descend_list_elements(&self) -> Self { + let mut elements = self.descend_clean(); + elements.leaf_role = LeafRole::ListElements; + elements + } + /// Descend into a subfield, retaining only the overrides that apply beneath it (rebased to be /// relative to the child). fn descend(&self, field: &Field) -> Self { @@ -256,6 +294,8 @@ impl TableStrategy { leaf_writers: new_writers, validity: Arc::clone(&self.validity), leaf: Arc::clone(&self.leaf), + list_elements: self.list_elements.clone(), + leaf_role: self.leaf_role, list_layout_factory: self.list_layout_factory.clone(), } } @@ -267,6 +307,8 @@ impl TableStrategy { leaf_writers: HashMap::default(), validity: Arc::clone(&self.validity), leaf: Arc::clone(&self.leaf), + list_elements: self.list_elements.clone(), + leaf_role: self.leaf_role, list_layout_factory: self.list_layout_factory.clone(), } } @@ -290,6 +332,26 @@ impl TableStrategy { } } +/// Expand chunked leaf arrays created by list-aware structural repartitioning. +fn flatten_chunked_stream(stream: SendableSequentialStream) -> SendableSequentialStream { + let dtype = stream.dtype().clone(); + let flattened = try_stream! { + pin_mut!(stream); + while let Some(chunk) = stream.next().await { + let (sequence_id, array) = chunk?; + if let Some(chunked) = array.as_opt::() { + let mut sequence_pointer = sequence_id.descend(); + for chunk in chunked.chunks() { + yield (sequence_pointer.advance(), chunk); + } + } else { + yield (sequence_id, array); + } + } + }; + SequentialStreamAdapter::new(dtype, flattened).sendable() +} + /// Dispatches each stream to the structural writer for its dtype. #[async_trait] impl LayoutStrategy for TableStrategy { @@ -319,7 +381,12 @@ impl LayoutStrategy for TableStrategy { } // Leaf: hand off to the leaf strategy. - self.leaf + let stream = if matches!(self.leaf_role, LeafRole::ListElements) { + flatten_chunked_stream(stream) + } else { + stream + }; + self.leaf_strategy() .write_stream(ctx, segment_sink, stream, eof, session) .await } diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 4151679e6c3..94ffab9c3c2 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -161,10 +161,11 @@ impl LayoutStrategy for ZonedStrategy { // The eof used for the data child should appear _before_ our own stats tables. let data_eof = eof.split_off(); + let (data_ctx, data_info) = ctx.child_context(); let data_layout = self .child .write_stream( - ctx.clone(), + data_ctx, Arc::clone(&segment_sink), stream, data_eof, @@ -178,6 +179,7 @@ impl LayoutStrategy for ZonedStrategy { else { // If we have no stats (e.g. the DType doesn't support them), then we just return the // child layout. + ctx.report_chunk_boundaries(data_info.chunk_boundaries()); return Ok(data_layout); }; @@ -187,11 +189,19 @@ impl LayoutStrategy for ZonedStrategy { .into_array() .to_array_stream() .sequenced(eof.split_off()); + let (stats_ctx, _) = ctx.child_context(); let zones_layout = self .stats - .write_stream(ctx, Arc::clone(&segment_sink), stats_stream, eof, session) + .write_stream( + stats_ctx, + Arc::clone(&segment_sink), + stats_stream, + eof, + session, + ) .await?; + ctx.report_chunk_boundaries(data_info.chunk_boundaries()); Ok( ZonedLayout::try_new(data_layout, zones_layout, block_size, aggregate_fns)? .into_layout(), diff --git a/vortex-layout/src/strategy.rs b/vortex-layout/src/strategy.rs index 5a0b1025e4a..e1ed2278a6f 100644 --- a/vortex-layout/src/strategy.rs +++ b/vortex-layout/src/strategy.rs @@ -7,6 +7,7 @@ use std::sync::atomic::Ordering; use async_trait::async_trait; use futures::StreamExt; +use parking_lot::Mutex; use vortex_array::ArrayContext; use vortex_array::ArrayId; use vortex_array::aggregate_fn::AggregateFnId; @@ -72,6 +73,35 @@ impl Drop for BufferedBytesReservation { } } +/// Write-time information reported by a child layout strategy to its parent. +/// +/// A parent creates an isolated report with +/// [`LayoutWriterContext::child_context`] before invoking a child strategy. The child strategy +/// writes facts about its output into the context and the parent reads the completed report after +/// the child finishes. This keeps writer-side communication out of serialized layouts until the +/// parent chooses the information it needs to persist. +#[derive(Clone, Debug, Default)] +pub struct LayoutWriterInfo(Arc>); + +#[derive(Debug, Default)] +struct LayoutWriterInfoData { + chunk_boundaries: Vec, +} + +impl LayoutWriterInfo { + /// Returns the strictly increasing output-row boundaries reported by the child strategy. + pub fn chunk_boundaries(&self) -> Vec { + self.0.lock().chunk_boundaries.clone() + } + + fn set_chunk_boundaries(&self, chunk_boundaries: impl IntoIterator) { + let mut chunk_boundaries = chunk_boundaries.into_iter().collect::>(); + chunk_boundaries.sort_unstable(); + chunk_boundaries.dedup(); + self.0.lock().chunk_boundaries = chunk_boundaries; + } +} + /// State shared by every strategy participating in a single layout write. /// /// Clones share the [`BufferedBytesTracker`] while retaining the array serialization context. @@ -82,6 +112,7 @@ pub struct LayoutWriterContext { array_ctx: ArrayContext, allowed_aggregates: Option>>, buffered_bytes: BufferedBytesTracker, + writer_info: LayoutWriterInfo, } impl LayoutWriterContext { @@ -91,6 +122,7 @@ impl LayoutWriterContext { array_ctx, allowed_aggregates: None, buffered_bytes: BufferedBytesTracker::new(), + writer_info: LayoutWriterInfo::default(), } } @@ -139,6 +171,34 @@ impl LayoutWriterContext { pub fn reserve_buffered_bytes(&self, bytes: u64) -> BufferedBytesReservation { self.buffered_bytes.reserve(bytes) } + + /// Creates a context and report for a child strategy with a distinct output row space. + /// + /// The returned context retains this write's array context, aggregate policy, and buffered + /// bytes tracker, but reports its chunk boundaries to the returned [`LayoutWriterInfo`] + /// rather than this context. Parents can use the report to translate child boundaries into + /// their own row space once the child write completes. + pub fn child_context(&self) -> (Self, LayoutWriterInfo) { + let writer_info = LayoutWriterInfo::default(); + ( + Self { + array_ctx: self.array_ctx.clone(), + allowed_aggregates: self.allowed_aggregates.clone(), + buffered_bytes: self.buffered_bytes.clone(), + writer_info: writer_info.clone(), + }, + writer_info, + ) + } + + /// Reports the output-row boundaries between physical chunks written by this strategy. + /// + /// Boundaries must use the strategy's output row coordinate space. The report is normalized + /// to sorted, unique values; structural parents are responsible for translating it when they + /// change coordinate spaces. + pub fn report_chunk_boundaries(&self, chunk_boundaries: impl IntoIterator) { + self.writer_info.set_chunk_boundaries(chunk_boundaries) + } } impl From for LayoutWriterContext { @@ -170,8 +230,10 @@ pub trait LayoutStrategy: 'static + Send + Sync { /// with a sequence pointer that indicates its position in the overall array. By passing /// around these pointers (essentially vector clocks), the writer can support concurrent /// and parallel processing while maintaining a deterministic order of data in the file. - /// The `ctx` parameter carries both array serialization state and writer-scoped accounting - /// through every child strategy. + /// The `ctx` parameter carries array serialization state, writer-scoped accounting, and the + /// current output report through every child strategy. Structural strategies that change row + /// coordinates create a [`LayoutWriterContext::child_context`] for each child, then translate + /// the completed child report before publishing their own. /// /// The `eof` parameter is a guaranteed to be greater than all sequence pointers in the stream. /// @@ -270,7 +332,10 @@ impl LayoutStrategy for Arc { #[cfg(test)] mod tests { + use vortex_array::ArrayContext; + use crate::strategy::BufferedBytesTracker; + use crate::strategy::LayoutWriterContext; #[test] fn reservations_accumulate_and_release() { @@ -300,4 +365,16 @@ mod tests { drop(reservation); assert_eq!(observer.buffered_bytes(), 0); } + + #[test] + fn child_context_has_an_isolated_writer_report() { + let ctx = LayoutWriterContext::new(ArrayContext::empty()); + let (child_ctx, child_info) = ctx.child_context(); + + child_ctx.report_chunk_boundaries([8, 3, 8]); + + assert_eq!(child_info.chunk_boundaries(), [3, 8]); + let (_, parent_info) = ctx.child_context(); + assert!(parent_info.chunk_boundaries().is_empty()); + } }