diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 9d4dbb90610..120bcd16d37 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -25,7 +25,6 @@ use vortex_layout::layouts::list::writer::ListLayoutStrategy; use vortex_layout::layouts::repartition::RepartitionStrategy; use vortex_layout::layouts::repartition::RepartitionWriterOptions; use vortex_layout::layouts::table::TableStrategy; -use vortex_layout::layouts::table::use_experimental_list_layout; use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; use vortex_layout::layouts::zoned::writer::ZonedStrategy; use vortex_utils::aliases::hash_map::HashMap; @@ -33,6 +32,11 @@ use vortex_utils::aliases::hash_set::HashSet; const ONE_MEG: u64 = 1 << 20; +#[cfg(feature = "unstable_encodings")] +const USE_LIST_LAYOUT_BY_DEFAULT: bool = true; +#[cfg(not(feature = "unstable_encodings"))] +const USE_LIST_LAYOUT_BY_DEFAULT: bool = false; + /// How the compressor was configured on [`WriteStrategyBuilder`]. enum CompressorConfig { /// A [`BtrBlocksCompressorBuilder`] that [`WriteStrategyBuilder::build`] will finalize. @@ -63,6 +67,7 @@ pub struct WriteStrategyBuilder { flat_strategy: Option>, probe_compressor: Option>, /// Whether to write list fields using [`ListLayoutStrategy`]. + /// Enabled by default with the `unstable_encodings` feature. /// /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy use_list_layout: bool, @@ -80,7 +85,7 @@ impl Default for WriteStrategyBuilder { allow_encodings: None, flat_strategy: None, probe_compressor: None, - use_list_layout: use_experimental_list_layout(), + use_list_layout: USE_LIST_LAYOUT_BY_DEFAULT, } } } @@ -106,6 +111,8 @@ impl WriteStrategyBuilder { /// Enable writing list fields with [`ListLayoutStrategy`]. /// + /// This is already enabled by default when the `unstable_encodings` feature is active. + /// /// **Note**: this is an unstable and experimental layout that is expected to change. /// Using it may lead to unreadable files in the future. /// @@ -201,8 +208,19 @@ impl WriteStrategyBuilder { flat }; - // 7. for each chunk create a flat layout - let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); + // 7. write each compressed chunk as either a shallow list layout or a flat layout. + // Chunking stays outside ListLayoutStrategy so every list layout receives exactly one + // complete page and its elements, offsets, and validity are direct flat leaves. + let terminal: Arc = if self.use_list_layout { + Arc::new( + ListLayoutStrategy::default() + .with_leaf(Arc::clone(&flat)) + .with_fallback(Arc::clone(&flat)), + ) + } else { + Arc::clone(&flat) + }; + let chunked = ChunkedLayoutStrategy::new(terminal); // 6. buffer chunks so they end up with closer segment ids physically let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB @@ -243,7 +261,8 @@ impl WriteStrategyBuilder { CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), CompressorConfig::Opaque(compressor) => compressor, }; - let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); + let compress_then_flat = + CompressingStrategy::new(Arc::clone(&flat), Arc::clone(&stats_compressor)); // 3. apply dict encoding or fallback let probe_compressor = if let Some(probe_compressor) = self.probe_compressor { @@ -285,37 +304,11 @@ impl WriteStrategyBuilder { ); // 0. start with splitting columns - let validity_strategy = CollectStrategy::new(compress_then_flat.clone()); + let validity_strategy = CollectStrategy::new(compress_then_flat); // Take any field overrides from the builder and apply them to the final strategy. - let mut table_strategy = - TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition)) - .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 table_strategy = TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition)) + .with_field_writers(self.field_writers); Arc::new(table_strategy) } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..bbe665da934 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1898,8 +1898,7 @@ async fn write_read_roundtrip_with_layout( .await } -/// A `list>` column round-trips through the `TableStrategy` dispatcher, exercising list -/// decomposition recursing into itself (the outer list's `elements` are themselves lists). +/// A `list>` column round-trips through the shallow list-layout dispatcher. #[tokio::test] #[cfg_attr(miri, ignore)] async fn nested_list_of_list_roundtrip() -> VortexResult<()> { @@ -1987,9 +1986,8 @@ async fn struct_with_map_column_roundtrip() -> VortexResult<()> { Ok(()) } -/// A `struct<{ items: list>? }>` column round-trips, exercising list decomposition -/// recursing into struct decomposition (list `elements` are structs) plus a nullable list validity -/// child. +/// A `struct<{ items: list>? }>` column round-trips with its element struct kept in a +/// flat layout, plus a nullable list validity child. #[tokio::test] #[cfg_attr(miri, ignore)] async fn nested_struct_list_struct_roundtrip() -> VortexResult<()> { diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 53227635de6..33e77b4643d 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -12,7 +12,6 @@ use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ListArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; @@ -23,7 +22,6 @@ use vortex_array::expr::BoundExpression; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -47,9 +45,6 @@ type OptionalArrayFuture = BoxFuture<'static, VortexResult>>; /// and above which we evaluate the expression over all rows and intersect afterward. const EXPR_EVAL_THRESHOLD: f64 = 0.2; -/// Maximum number of outer-row scan ranges contributed by one list layout. -const MAX_LIST_SPLIT_COUNT: u64 = 64; - /// Reader for [`ListLayout`]. #[derive(Clone)] pub struct ListReader { @@ -69,20 +64,22 @@ impl ListReader { session: VortexSession, ctx: &LayoutReaderContext, ) -> VortexResult { - let elements = layout.elements()?.new_reader( + let elements_layout = layout.elements()?; + let offsets_layout = layout.offsets()?; + let validity_layout = layout.validity()?; + let elements = elements_layout.new_reader( format!("{name}.elements").into(), Arc::clone(&segment_source), &session, ctx, )?; - let offsets = layout.offsets()?.new_reader( + let offsets = offsets_layout.new_reader( format!("{name}.offsets").into(), Arc::clone(&segment_source), &session, ctx, )?; - let validity = layout - .validity()? + let validity = validity_layout .map(|v| { v.new_reader( format!("{name}.validity").into(), @@ -144,36 +141,18 @@ impl ListReader { .boxed()) } - /// Projection for [`ListChildrenNeeded::All`] expressions. - /// - /// An all-true mask over the full local range reads every child concurrently. Otherwise, the - /// read is bounded to the first and last selected list. + /// Projection for [`ListChildrenNeeded::All`] expressions. Registers complete child reads + /// eagerly and reconstructs the list page before applying the outer-row mask. fn project_all( &self, row_range: &Range, expr: &BoundExpression, mask: MaskFuture, ) -> VortexResult { - let is_full_range = row_range.start == 0 && row_range.end == self.layout.row_count(); - let reader = self.clone(); - let row_range = row_range.clone(); - let expr = expr.clone(); - Ok(async move { - let mask = mask.await?; - if is_full_range && mask.all_true() { - reader.project_all_full(&expr)?.await - } else { - reader.project_all_bounded(&row_range, &expr, mask)?.await - } - } - .boxed()) - } - - /// Fetch the complete `elements`, `offsets`, and `validity` children concurrently. - fn project_all_full(&self, expr: &BoundExpression) -> VortexResult { let row_count = self.layout.row_count(); let elements_row_count = self.elements.row_count(); let nullability = self.layout.dtype().nullability(); + let row_range = row_range.clone(); let expr = expr.clone(); let offsets_fut = self.fetch_raw_offsets(&(0..row_count))?; @@ -185,68 +164,26 @@ impl ListReader { )?; Ok(async move { + let mask = mask.await?; + if mask.all_false() { + return Ok(Canonical::empty(expr.dtype()).into_array()); + } + let (offsets, elements, validity) = try_join!(offsets_fut, elements_fut, validity_fut)?; // SAFETY: ListLayout is constructed from a valid ListArray and reading its children // without transformation preserves the list invariants. - let list = unsafe { + let mut list = unsafe { ListArray::new_unchecked(elements, offsets, create_validity(validity, nullability)) } .into_array(); - list.apply_bound(&expr) - } - .boxed()) - } - - /// Bounded read for a sub-range or selective mask. - /// - /// Crops leading and trailing unselected lists, reads their offsets, and translates the first - /// and last offset into the element-row range to fetch. Any holes in the selection are filtered - /// after reconstructing the list array. - fn project_all_bounded( - &self, - row_range: &Range, - expr: &BoundExpression, - mask: Mask, - ) -> VortexResult { - // Crop to the smallest contiguous row range containing every selected list. - let Some(selected_rows) = selected_row_range(&mask) else { - let empty = Canonical::empty(expr.dtype()).into_array(); - return Ok(async move { Ok(empty) }.boxed()); - }; - - let selected_mask = mask.slice(selected_rows.clone()); - let selected_row_range = (row_range.start + u64::try_from(selected_rows.start)?) - ..(row_range.start + u64::try_from(selected_rows.end)?); - - let nullability = self.layout.dtype().nullability(); - let expr = expr.clone(); - let reader = self.clone(); - let offsets_fut = self.fetch_raw_offsets(&selected_row_range)?; - - Ok(async move { - let offsets = offsets_fut.await?; - - let elements_range = elements_range_from_offsets(&offsets, &reader.session)?; - let elements_fut = reader.fetch_raw_elements(&elements_range)?; - let validity_fut = fetch_validity( - reader.validity.as_ref(), - &selected_row_range, - MaskFuture::new_true(selected_mask.len()), - )?; - let (elements, validity) = try_join!(elements_fut, validity_fut)?; - - let offsets = rebase_offsets(offsets, elements_range.start)?; - // SAFETY: the selected offsets remain monotonically increasing, rebasing them against - // the selected element range preserves their lengths, and validity covers the same - // cropped list rows. - let list = unsafe { - ListArray::new_unchecked(elements, offsets, create_validity(validity, nullability)) + if row_range.start != 0 || row_range.end != row_count { + list = + list.slice(usize::try_from(row_range.start)?..usize::try_from(row_range.end)?)?; } - .into_array(); - let list = if selected_mask.all_true() { + let list = if mask.all_true() { list } else { - list.filter(selected_mask)? + list.filter(mask)? }; list.apply_bound(&expr) } @@ -319,10 +256,6 @@ impl ListReader { } } -fn selected_row_range(mask: &Mask) -> Option> { - Some(mask.first()?..mask.last()? + 1) -} - fn create_validity(validity_array: Option, nullability: Nullability) -> Validity { match validity_array { Some(arr) => Validity::Array(arr), @@ -357,53 +290,6 @@ impl LayoutReader for ListReader { splits: &mut RowSplits, ) -> VortexResult<()> { split_range.check_bounds(self.layout.row_count())?; - - // 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. - // - // Scan splits must be expressed in the list layout's outer-row space, but the elements child - // reports its natural boundaries in element-row space. So we translate the element splits using a - // heuristic to outer-row space. - - let element_row_count = self.elements.row_count(); - if element_row_count != 0 { - let mut element_splits = RowSplits::new_capacity(128); - self.elements.register_splits( - &[FieldMask::All], - &SplitRange::root(0..element_row_count)?, - &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( - element_split, - element_row_count, - self.layout.row_count(), - MAX_LIST_SPLIT_COUNT, - ) else { - continue; - }; - if split <= row_range.start { - continue; - } - if split >= row_range.end { - break; - } - if last_split == Some(split) { - continue; - } - splits.push( - split_range - .row_offset() - .checked_add(split) - .vortex_expect("List layout split offset overflow"), - ); - last_split = Some(split); - } - } splits.push(split_range.root_row_range().end); Ok(()) } @@ -474,52 +360,6 @@ impl LayoutReader for ListReader { } } -/// 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 -/// reports its natural boundaries in element-row space. Translating a boundary exactly would -/// require consulting the list offsets, so this function instead preserves its relative position: -/// -/// ```text -/// element_split / element_row_count ≈ outer_split / outer_row_count -/// ``` -/// -/// The relative position is first rounded onto a grid containing at most `max_split_count` scan -/// ranges, then mapped into outer-row space. Multiple element boundaries may therefore map to the -/// same outer split and are deduplicated by the caller. With a grid size of 64, at most 63 interior -/// splits—and therefore 64 scan ranges—can be produced. -/// -/// The result is only a task-sizing hint. It is always between outer rows, but it is not guaranteed -/// to correspond exactly to the original physical element boundary. Endpoint boundaries are -/// omitted because they do not subdivide the scan -fn map_element_split_to_outer_grid( - element_split: u64, - element_row_count: u64, - outer_row_count: u64, - max_split_count: u64, -) -> Option { - if element_split == 0 - || element_split >= element_row_count - || outer_row_count == 0 - || max_split_count < 2 - { - return None; - } - debug_assert!(max_split_count.is_power_of_two()); - - let grid_index = (u128::from(element_split) * u128::from(max_split_count) - + u128::from(element_row_count / 2)) - / u128::from(element_row_count); - if grid_index == 0 || grid_index >= u128::from(max_split_count) { - return None; - } - - let outer_split = grid_index * u128::from(outer_row_count) / u128::from(max_split_count); - let outer_split = u64::try_from(outer_split) - .vortex_expect("Outer split is bounded by the list layout row count"); - (outer_split != 0 && outer_split < outer_row_count).then_some(outer_split) -} - /// Fetch the validity child for `row_range` under `mask`, yielding `None` for a non-nullable list /// (which has no validity child). fn fetch_validity( @@ -542,40 +382,6 @@ fn fetch_validity( .boxed()) } -/// Read `offsets[0]` and `offsets[-1]` and return the elements range they bound. -fn elements_range_from_offsets( - offsets: &ArrayRef, - session: &VortexSession, -) -> VortexResult> { - if offsets.is_empty() { - return Ok(0..0); - } - let mut exec_ctx = session.create_execution_ctx(); - let start = offsets - .execute_scalar(0, &mut exec_ctx)? - .as_primitive() - .as_::() - .vortex_expect("offset value fits in u64"); - let end = offsets - .execute_scalar(offsets.len() - 1, &mut exec_ctx)? - .as_primitive() - .as_::() - .vortex_expect("offset value fits in u64"); - Ok(start..end) -} - -/// Subtract `first` from every offset so they index into a sliced `elements[first..]` buffer that -/// starts at zero. -fn rebase_offsets(offsets: ArrayRef, first: u64) -> VortexResult { - if first == 0 { - return Ok(offsets); - } - let constant = ConstantArray::new(first, offsets.len()) - .into_array() - .cast(offsets.dtype().clone())?; - offsets.binary(constant, Operator::Sub) -} - /// Compute `offsets[i + 1] - offsets[i]` as the unmasked list length values. fn list_lengths_from_offsets(offsets: ArrayRef) -> VortexResult { let len = offsets.len().saturating_sub(1); @@ -607,14 +413,17 @@ fn predicate_array_to_mask(array: ArrayRef, session: &VortexSession) -> VortexRe #[cfg(test)] mod tests { use std::ops::Range; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; use rstest::rstest; use vortex_array::ArrayContext; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::Dict; + use vortex_array::arrays::DictArray; use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::assert_arrays_eq; use vortex_array::expr::Expression; use vortex_array::expr::cast; @@ -630,13 +439,8 @@ mod tests { use super::*; use crate::LayoutRef; use crate::LayoutStrategy; - use crate::layouts::chunked::writer::ChunkedLayoutStrategy; - use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::list::writer::ListLayoutStrategy; - use crate::layouts::repartition::RepartitionStrategy; - use crate::layouts::repartition::RepartitionWriterOptions; use crate::scan::split_by::SplitBy; - use crate::segments::SegmentFuture; use crate::segments::SegmentSource; use crate::segments::TestSegments; use crate::sequence::SequenceId; @@ -867,12 +671,12 @@ mod tests { Ok((segments_ref, layout, session)) } - fn materialize_u64_array(array: ArrayRef) -> Vec { + fn materialize_u32_array(array: ArrayRef) -> Vec { let mut ctx = SESSION.create_execution_ctx(); array .execute::(&mut ctx) .unwrap() - .as_slice::() + .as_slice::() .to_vec() } @@ -919,7 +723,7 @@ mod tests { .expect("ListReader"); let offsets = reader.fetch_raw_offsets(&(1..3))?.await?; - assert_eq!(materialize_u64_array(offsets), vec![2u64, 4, 5]); + assert_eq!(materialize_u32_array(offsets), vec![2u32, 4, 5]); Ok(()) } @@ -1020,11 +824,10 @@ mod tests { Ok(()) } - /// A partial range against a single flat elements segment takes the bounded path. The flat - /// reader may still fetch its whole segment, but the list reader reconstructs only the requested - /// element range. + /// A partial range against flat children is reconstructed from the complete encoded page and + /// then sliced in outer-row space. #[tokio::test] - async fn projection_evaluation_partial_range_bounded() -> VortexResult<()> { + async fn projection_evaluation_partial_range_flat_children() -> VortexResult<()> { let list = create_wider_list_array(false); let ctx = LayoutReaderContext::new(); let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; @@ -1041,167 +844,46 @@ mod tests { Ok(()) } - #[test] - fn maps_element_splits_to_outer_grid() { - assert_eq!(map_element_split_to_outer_grid(0, 100, 100, 8), None); - assert_eq!(map_element_split_to_outer_grid(20, 100, 100, 8), Some(25)); - assert_eq!(map_element_split_to_outer_grid(25, 100, 100, 8), Some(25)); - assert_eq!(map_element_split_to_outer_grid(50, 100, 100, 8), Some(50)); - assert_eq!(map_element_split_to_outer_grid(75, 100, 100, 8), Some(75)); - assert_eq!(map_element_split_to_outer_grid(100, 100, 100, 8), None); - - let mut splits = (1..1_000) - .filter_map(|split| { - map_element_split_to_outer_grid(split, 1_000, 100_000, MAX_LIST_SPLIT_COUNT) - }) - .collect::>(); - splits.dedup(); - - let expected = (1..MAX_LIST_SPLIT_COUNT) - .map(|grid_index| grid_index * 100_000 / MAX_LIST_SPLIT_COUNT) - .collect::>(); - assert_eq!(splits, expected); - } - #[tokio::test] - async fn nested_list_propagates_element_splits() -> VortexResult<()> { - let inner = ListArray::try_new( - PrimitiveArray::from_iter(0..128_i32).into_array(), - PrimitiveArray::from_iter((0..=8_u32).map(|idx| idx * 16)).into_array(), - Validity::NonNullable, + async fn projection_preserves_dictionary_encoded_elements() -> VortexResult<()> { + let elements = DictArray::try_new( + buffer![0u8, 1, 0, 2, 1, 2, 0, 1].into_array(), + VarBinArray::from(vec!["alpha", "beta", "gamma"]).into_array(), )? .into_array(); - let outer = ListArray::try_new( - inner, + let list = ListArray::try_new( + elements, buffer![0u32, 2, 4, 6, 8].into_array(), Validity::NonNullable, )? .into_array(); - - let inner_strategy = - ListLayoutStrategy::default().with_elements(chunked_elements_strategy()); - let strategy = ListLayoutStrategy::default().with_elements(Arc::new(inner_strategy)); - let (segments, layout, session) = write_layout(&strategy, outer).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, vec![0, 1, 2, 3, 4]); - Ok(()) - } - - /// A list strategy whose `elements` child is repartitioned into two-element chunks, so the - /// reader takes the bounded (chunk-skipping) path for strict sub-ranges. Offsets stay flat. - fn chunked_elements_list_strategy() -> ListLayoutStrategy { - ListLayoutStrategy::default().with_elements(chunked_elements_strategy()) - } - - fn chunked_elements_strategy() -> Arc { - Arc::new(RepartitionStrategy::new( - ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), - RepartitionWriterOptions { - block_size_minimum: 0, - block_len_multiple: 2, - block_size_target: None, - canonicalize: true, - }, - )) - } - - struct CountingSegmentSource { - inner: Arc, - request_count: Arc, - } - - impl SegmentSource for CountingSegmentSource { - fn request(&self, id: crate::segments::SegmentId) -> SegmentFuture { - self.request_count.fetch_add(1, Ordering::Relaxed); - self.inner.request(id) - } - } - - /// The chunked-elements strategy must actually produce a chunked `elements` layout, otherwise - /// the reader would silently take the whole-chunk path and the bounded read would be untested. - #[tokio::test] - async fn chunked_elements_produces_chunked_layout() -> VortexResult<()> { - let list = create_wider_list_array(false); - let (_segments, layout, _session) = - write_layout(&chunked_elements_list_strategy(), list).await?; - let tree = layout.display_tree().to_string(); - assert!( - tree.contains("elements: vortex.chunked"), - "elements should be chunked:\n{tree}" - ); - Ok(()) - } - - /// A sparse mask must remain selective even when the requested row range covers the complete - /// local list layout. The selected first list touches one element chunk plus the offsets - /// segment; reading all five element chunks would make the count six. - #[tokio::test] - async fn full_range_sparse_mask_crops_element_read() -> VortexResult<()> { - let list = create_wider_list_array(false); let ctx = LayoutReaderContext::new(); - let (segments, layout, session) = - write_layout(&chunked_elements_list_strategy(), list.clone()).await?; - let request_count = Arc::new(AtomicUsize::new(0)); - let source = Arc::new(CountingSegmentSource { - inner: segments, - request_count: Arc::clone(&request_count), - }); - let reader = layout.new_reader("".into(), source, &session, &ctx)?; - - let mask = Mask::from_iter([true, false, false, false, false]); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list.clone()).await?; + let reader = layout.new_reader("".into(), segments, &session, &ctx)?; + + let mask = Mask::from_iter([true, false, true]); let expr = root().bind(reader.dtype())?; let result = reader - .projection_evaluation(&(0..5), &expr, MaskFuture::ready(mask.clone()))? + .projection_evaluation(&(1..4), &expr, MaskFuture::ready(mask.clone()))? .await?; - let expected = list.filter(mask)?; + let expected = list.slice(1..4)?.filter(mask)?; let mut exec_ctx = session.create_execution_ctx(); assert_arrays_eq!(result, expected, &mut exec_ctx); - assert_eq!(request_count.load(Ordering::Relaxed), 2); + let result_view = result.execute::(&mut exec_ctx)?; + assert!(result_view.elements().is::()); Ok(()) } - /// With chunked elements, sub-range projections take the bounded read path. Every - /// range/mask/nullability combination must match the same projection over the ground-truth - /// array (`list.slice(range).filter(mask)`). - #[rstest] - #[case::full_all_true(0..5, Mask::new_true(5), false)] - #[case::subrange_all_true(1..4, Mask::new_true(3), false)] - #[case::subrange_sparse(1..4, Mask::from_iter([true, false, true]), false)] - #[case::partial_start(0..2, Mask::new_true(2), false)] - #[case::partial_end(2..5, Mask::new_true(3), false)] - #[case::single_non_empty(0..1, Mask::new_true(1), false)] - #[case::single_empty_row(2..3, Mask::from_iter([true]), false)] - #[case::empty_range(2..2, Mask::new_true(0), false)] - #[case::subrange_all_false(1..4, Mask::new_false(3), false)] - #[case::subrange_single_interior(0..4, Mask::from_iter([false, true, false, false]), false)] - #[case::subrange_sparse_nullable(1..4, Mask::from_iter([true, false, true]), true)] - #[case::partial_end_nullable(2..5, Mask::new_true(3), true)] #[tokio::test] - async fn chunked_elements_round_trips( - #[case] row_range: Range, - #[case] mask: Mask, - #[case] nullable: bool, - ) -> VortexResult<()> { - let list = create_wider_list_array(nullable); - let ctx = LayoutReaderContext::new(); - let (segments, layout, session) = - write_layout(&chunked_elements_list_strategy(), list.clone()).await?; - let reader = layout.new_reader("".into(), segments, &session, &ctx)?; - - let expr = root().bind(reader.dtype())?; - let result = reader - .projection_evaluation(&row_range, &expr, MaskFuture::ready(mask.clone()))? - .await?; + async fn list_does_not_propagate_element_splits() -> VortexResult<()> { + let list = create_wider_list_array(false); + let (segments, layout, session) = write_layout(&flat_list_strategy(), list).await?; + let reader = + layout.new_reader("".into(), segments, &session, &LayoutReaderContext::new())?; - let sliced = - list.slice(usize::try_from(row_range.start)?..usize::try_from(row_range.end)?)?; - let expected = sliced.filter(mask)?; - let mut exec_ctx = session.create_execution_ctx(); - assert_arrays_eq!(result, expected, &mut exec_ctx); + let splits = SplitBy::Layout.splits(reader.as_ref(), &(0..5), &[FieldMask::All])?; + assert_eq!(splits, vec![0, 5]); Ok(()) } } diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 4f41dc3b20d..1b43274e3ec 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -5,102 +5,73 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt; -use futures::future::try_join; use futures::future::try_join_all; +use futures::stream::once; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -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::ListDataParts; use vortex_array::arrays::listview::list_from_list_view; -use vortex_array::builtins::ArrayBuiltins; -use vortex_array::dtype::DType; -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_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_io::kanal_ext::KanalExt; +use vortex_error::vortex_ensure; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; use crate::LayoutWriterContext; +use crate::layouts::flat::Flat; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::list::ListLayout; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; -use crate::sequence::SequenceId; use crate::sequence::SequencePointer; -use crate::sequence::SequentialStream; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; -/// Item carried on each child sub-stream: a sequenced, materialized chunk. -type ChildChunk = VortexResult<(SequenceId, ArrayRef)>; - /// 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`, /// `offsets`, and (when nullable) `validity` sub-columns, each written through its own downstream /// strategy, producing a single [`ListLayout`]. /// -/// For list-typed input the strategy transposes the whole column stream into three sub-streams: -/// 1. Each chunk is canonicalized to a [`ListArray`] (rebuilding a [`ListView`] via -/// [`list_from_list_view`] when necessary). -/// 2. `offsets` are rebased to global `u64` positions (cumulative across chunks) so the single -/// `offsets` child indexes into the concatenated `elements` child. -/// 3. `elements`, `offsets`, and `validity` are streamed to their child strategies concurrently. +/// For list-typed input the strategy accepts exactly one chunk, canonicalizes its outer list +/// container (rebuilding a [`ListView`] via [`list_from_list_view`] when necessary), and writes +/// its encoded `elements`, `offsets`, and `validity` children concurrently as direct flat leaves. +/// Chunking belongs outside this strategy: callers with a multi-chunk stream should wrap it in a +/// [`ChunkedLayoutStrategy`]. Keeping the input chunk intact preserves the compressor's encoding +/// choices for both elements and offsets. /// -/// For input whose dtype is not [`DType::List`], the stream is forwarded unchanged to the -/// configured `fallback` strategy. +/// For non-list input, the stream is forwarded unchanged to the configured `fallback` strategy. /// -/// [`ListArray`]: vortex_array::arrays::ListArray +/// [`ChunkedLayoutStrategy`]: crate::layouts::chunked::writer::ChunkedLayoutStrategy #[derive(Clone)] pub struct ListLayoutStrategy { - elements: Arc, - offsets: Arc, - validity: Arc, + leaf: Arc, fallback: Arc, } impl Default for ListLayoutStrategy { - /// Routes every child (elements, offsets, validity) and the non-list fallback through - /// [`FlatLayoutStrategy`]. Override individual children with the `with_*` builder methods. + /// Writes every child and the non-list fallback through [`FlatLayoutStrategy`]. fn default() -> Self { let flat: Arc = Arc::new(FlatLayoutStrategy::default()); Self { - elements: Arc::clone(&flat), - offsets: Arc::clone(&flat), - validity: Arc::clone(&flat), + leaf: Arc::clone(&flat), fallback: flat, } } } impl ListLayoutStrategy { - /// Strategy for the `elements` child. - pub fn with_elements(mut self, elements: Arc) -> Self { - self.elements = elements; - self - } - - /// Strategy for the `offsets` child. - pub fn with_offsets(mut self, offsets: Arc) -> Self { - self.offsets = offsets; - self - } - - /// Strategy for the `validity` child (written only when the list is nullable). - pub fn with_validity(mut self, validity: Arc) -> Self { - self.validity = validity; + /// Strategy used for every list child. It must produce a single [`Flat`] leaf. + pub fn with_leaf(mut self, leaf: Arc) -> Self { + self.leaf = leaf; self } @@ -117,7 +88,7 @@ impl LayoutStrategy for ListLayoutStrategy { &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, + mut stream: SendableSequentialStream, mut eof: SequencePointer, session: &VortexSession, ) -> VortexResult { @@ -129,60 +100,46 @@ impl LayoutStrategy for ListLayoutStrategy { .await; } - let is_nullable = dtype.is_nullable(); - let element_dtype = dtype - .as_list_element_opt() - .vortex_expect("DType is List") - .as_ref() - .clone(); - // Global (whole-column) offsets are cumulative and may exceed the input offset width, - // so definsively widen. - 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 (validity_tx, validity_rx) = if is_nullable { - let (tx, rx) = kanal::bounded_async::(1); - (Some(tx), Some(rx)) - } else { - (None, None) + let Some((sequence_id, array)) = stream.next().await.transpose()? else { + vortex_bail!("ListLayoutStrategy needs exactly one chunk"); }; + if stream.next().await.is_some() { + vortex_bail!("ListLayoutStrategy received more than one chunk"); + } - // 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, - ); - - // Spawn a writer per child sub-stream, concurrently. - let handle = session.handle(); - let mut child_specs: Vec<( - DType, - Arc, - kanal::AsyncReceiver, - )> = vec![ - (element_dtype, Arc::clone(&self.elements), elements_rx), - (offsets_dtype, Arc::clone(&self.offsets), offsets_rx), + let mut exec_ctx = session.create_execution_ctx(); + let ListDataParts { + elements, + offsets, + validity, + .. + } = canonicalize_to_list_parts(array, &mut exec_ctx)?; + let row_count = offsets.len().saturating_sub(1); + let mut sequence = sequence_id.descend(); + let mut child_specs = vec![ + (elements, Arc::clone(&self.leaf), sequence.advance()), + (offsets, Arc::clone(&self.leaf), sequence.advance()), ]; - if let Some(validity_rx) = validity_rx { + if dtype.is_nullable() { child_specs.push(( - DType::Bool(Nullability::NonNullable), - Arc::clone(&self.validity), - validity_rx, + validity + .execute_mask(row_count, &mut exec_ctx)? + .into_array(), + Arc::clone(&self.leaf), + sequence.advance(), )); } + let handle = session.handle(); let layout_futures: Vec<_> = child_specs .into_iter() - .map(|(child_dtype, strategy, rx)| { - let child_stream = - SequentialStreamAdapter::new(child_dtype, rx.into_stream().boxed()).sendable(); + .map(|(child, strategy, child_sequence)| { + let child_dtype = child.dtype().clone(); + let child_stream = SequentialStreamAdapter::new( + child_dtype, + once(async move { Ok((child_sequence, child)) }), + ) + .sendable(); let child_eof = eof.split_off(); let ctx = ctx.clone(); let segment_sink = Arc::clone(&segment_sink); @@ -195,77 +152,25 @@ impl LayoutStrategy for ListLayoutStrategy { }) .collect(); - let (_, layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; - let mut layouts = layouts.into_iter(); + let mut layouts = try_join_all(layout_futures).await?.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")); + let validity_layout = dtype + .is_nullable() + .then(|| layouts.next().vortex_expect("validity layout present")); + vortex_ensure!( + elements_layout.is::() + && offsets_layout.is::() + && validity_layout + .as_ref() + .is_none_or(|layout| layout.is::()), + "ListLayout children must be flat leaves" + ); Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).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 -/// `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. -async fn transpose_list_column( - mut stream: SendableSequentialStream, - session: VortexSession, - elements_tx: kanal::AsyncSender, - offsets_tx: kanal::AsyncSender, - validity_tx: Option>, -) -> VortexResult<()> { - let mut exec_ctx = session.create_execution_ctx(); - let mut element_base: u64 = 0; - let mut first = true; - 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 { - elements, - offsets, - 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"); - } - } - } - if !saw_chunk { - vortex_bail!("ListLayoutStrategy needs at least one chunk"); - } - Ok(()) -} - /// Canonicalize a list-dtype array into [`ListDataParts`]. fn canonicalize_to_list_parts( array: ArrayRef, @@ -281,33 +186,6 @@ fn canonicalize_to_list_parts( } } -/// Rebase a chunk's local `offsets` into global `u64` positions for the whole-column `offsets` -/// child. Each chunk's offsets are shifted by `element_base` (the number of elements already -/// emitted) so they index into the concatenated `elements`. The duplicated boundary offset is -/// dropped on every chunk after the first, so the concatenation of all chunks' contributions is a -/// single monotonic `[0, .., total_elements]` array of length `row_count + 1`. -fn global_offsets( - offsets: ArrayRef, - element_base: u64, - first: bool, - exec_ctx: &mut ExecutionCtx, -) -> VortexResult { - let widened = offsets.cast(DType::Primitive(PType::U64, Nullability::NonNullable))?; - let based = if element_base == 0 { - widened - } else { - let base = ConstantArray::new(element_base, widened.len()).into_array(); - widened.binary(base, Operator::Add)? - }; - let based = if first { - based - } else { - based.slice(1..based.len())? - }; - // Materialize so the child sub-stream carries a concrete array rather than a lazy expression. - Ok(based.execute::(exec_ctx)?.into_array()) -} - /// Matcher for `Array` or `Array`. struct AnyList; @@ -323,21 +201,25 @@ impl Matcher for AnyList { mod tests { use futures::stream; use vortex_array::ArrayContext; + use vortex_array::MaskFuture; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::Dict; + use vortex_array::arrays::DictArray; use vortex_array::arrays::ListArray; - use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::expr::root; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_io::session::RuntimeSession; use super::*; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; - use crate::layouts::flat::writer::FlatLayoutStrategy; - use crate::layouts::table::TableStrategy; use crate::segments::TestSegments; + use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::session::LayoutSession; @@ -399,7 +281,7 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(i32), children: 2 ├── elements: vortex.flat, dtype: i32, segment: 0 - └── offsets: vortex.flat, dtype: u64, segment: 1 + └── offsets: vortex.flat, dtype: u32, segment: 1 "); Ok(()) } @@ -416,12 +298,74 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(i32)?, children: 3 ├── elements: vortex.flat, dtype: i32, segment: 0 - ├── offsets: vortex.flat, dtype: u64, segment: 1 + ├── offsets: vortex.flat, dtype: u32, segment: 1 └── validity: vortex.flat, dtype: bool, segment: 2 "); Ok(()) } + #[tokio::test] + async fn preserves_encoded_element_leaf() -> VortexResult<()> { + let elements = DictArray::try_new( + buffer![0u8, 1, 0, 2, 1].into_array(), + VarBinArray::from(vec!["alpha", "beta", "gamma"]).into_array(), + )? + .into_array(); + let list = ListArray::try_new( + elements, + buffer![0u32, 2, 5].into_array(), + Validity::NonNullable, + )? + .into_array(); + + let session = layout_test_session(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let layout = flat_list_strategy() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + list.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + let elements_layout = layout + .slot(0)? + .vortex_expect("ListLayout elements child is present"); + let reader = + elements_layout.new_reader("".into(), segments, &session, &Default::default())?; + let expr = root().bind(reader.dtype())?; + let read = reader + .projection_evaluation( + &(0..elements_layout.row_count()), + &expr, + MaskFuture::new_true(usize::try_from(elements_layout.row_count())?), + )? + .await?; + + assert!(read.is::()); + Ok(()) + } + + #[tokio::test] + async fn rejects_non_flat_child_layouts() { + let inner = create_basic_list(Validity::NonNullable); + let outer = ListArray::try_new( + inner, + buffer![0u32, 2, 3].into_array(), + Validity::NonNullable, + ) + .unwrap() + .into_array(); + let strategy = + ListLayoutStrategy::default().with_leaf(Arc::new(ListLayoutStrategy::default())); + + let result = write(&strategy, outer).await; + assert!(result.is_err()); + } + /// Non-list input dispatches to the fallback strategy unchanged. #[tokio::test] async fn non_list_input_routes_to_fallback() -> VortexResult<()> { @@ -451,101 +395,6 @@ mod tests { assert!(res.is_err()) } - #[tokio::test] - async fn list_of_struct_tree() -> VortexResult<()> { - let struct_array = StructArray::from_fields( - [ - ("a", buffer![1i32, 2, 3, 4, 5].into_array()), - ("b", buffer![10i32, 20, 30, 40, 50].into_array()), - ] - .as_slice(), - )? - .into_array(); - let list = ListArray::try_new( - struct_array, - buffer![0u32, 2, 5, 5].into_array(), - Validity::NonNullable, - )? - .into_array(); - - let flat: Arc = Arc::new(FlatLayoutStrategy::default()); - let table_strategy: Arc = - Arc::new(TableStrategy::new(Arc::clone(&flat), Arc::clone(&flat))); - let writer = ListLayoutStrategy::default().with_elements(table_strategy); - - let layout = write(&writer, list).await?; - insta::assert_snapshot!(layout.display_tree(), @" - vortex.list, dtype: list({a=i32, b=i32}), children: 2 - ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 - │ ├── a: vortex.flat, dtype: i32, segment: 1 - │ └── b: vortex.flat, dtype: i32, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 - "); - Ok(()) - } - - #[tokio::test] - async fn list_of_list_tree() -> VortexResult<()> { - let inner_list = ListArray::try_new( - buffer![1i32, 2, 3, 4, 5, 6].into_array(), - buffer![0u32, 2, 5, 5, 6].into_array(), - Validity::NonNullable, - )? - .into_array(); - let list = ListArray::try_new( - inner_list, - buffer![0u32, 2, 4].into_array(), - Validity::NonNullable, - )? - .into_array(); - - let writer = - ListLayoutStrategy::default().with_elements(Arc::new(ListLayoutStrategy::default())); - let layout = write(&writer, list).await?; - insta::assert_snapshot!(layout.display_tree(), @" - vortex.list, dtype: list(list(i32)), children: 2 - ├── elements: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 1 - │ └── offsets: vortex.flat, dtype: u64, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 - "); - Ok(()) - } - - #[tokio::test] - async fn list_of_list_of_list_tree() -> VortexResult<()> { - let innermost = ListArray::try_new( - buffer![1i32, 2, 3, 4].into_array(), - buffer![0u32, 2, 4].into_array(), - Validity::NonNullable, - )? - .into_array(); - let middle = ListArray::try_new( - innermost, - buffer![0u32, 2].into_array(), - Validity::NonNullable, - )? - .into_array(); - let outer = - ListArray::try_new(middle, buffer![0u32, 1].into_array(), Validity::NonNullable)? - .into_array(); - - let writer = ListLayoutStrategy::default().with_elements(Arc::new( - ListLayoutStrategy::default().with_elements(Arc::new(ListLayoutStrategy::default())), - )); - let layout = write(&writer, outer).await?; - insta::assert_snapshot!(layout.display_tree(), @" - vortex.list, dtype: list(list(list(i32))), children: 2 - ├── elements: vortex.list, dtype: list(list(i32)), children: 2 - │ ├── elements: vortex.list, dtype: list(i32), children: 2 - │ │ ├── elements: vortex.flat, dtype: i32, segment: 2 - │ │ └── offsets: vortex.flat, dtype: u64, segment: 3 - │ └── offsets: vortex.flat, dtype: u64, segment: 1 - └── offsets: vortex.flat, dtype: u64, segment: 0 - "); - Ok(()) - } - #[tokio::test] async fn chunked_list_input_with_chunked_strategy_succeeds() -> VortexResult<()> { let chunk0 = ListArray::try_new( @@ -572,10 +421,10 @@ mod tests { vortex.chunked, dtype: list(i32), children: 2 ├── [0]: vortex.list, dtype: list(i32), children: 2 │ ├── elements: vortex.flat, dtype: i32, segment: 0 - │ └── offsets: vortex.flat, dtype: u64, segment: 1 + │ └── offsets: vortex.flat, dtype: u32, segment: 1 └── [1]: vortex.list, dtype: list(i32), children: 2 ├── elements: vortex.flat, dtype: i32, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 3 + └── offsets: vortex.flat, dtype: u32, segment: 3 "); Ok(()) } diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 1a3c1adc524..37011c182ed 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -5,16 +5,13 @@ //! //! [`TableStrategy`] is a *dispatcher*: it inspects the dtype of the stream it is handed and //! routes struct columns to [`StructStrategy`], list columns to [`ListLayoutStrategy`], and -//! everything else to the configured leaf strategy. Because it hands *itself* (suitably descended) -//! to those structural writers as the strategy for their children, arbitrarily nested struct/list -//! trees are written with no manual wiring. +//! everything else to the configured leaf strategy. Struct fields recurse through the dispatcher, +//! while list elements always go directly to the leaf strategy. //! //! The dispatcher also owns field-path overrides, letting callers force a specific leaf field — //! at any depth — onto a custom strategy. -use std::env; use std::sync::Arc; -use std::sync::LazyLock; use async_trait::async_trait; use vortex_array::dtype::Field; @@ -28,23 +25,13 @@ use vortex_utils::aliases::hash_set::HashSet; use crate::LayoutRef; use crate::LayoutStrategy; use crate::LayoutWriterContext; +use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::list::writer::ListLayoutStrategy; use crate::layouts::struct_::StructStrategy; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; -/// Whether [`TableStrategy`] writes list fields using a [`ListLayoutStrategy`] by -/// default. Disabled unless the environment variable `VORTEX_EXPERIMENTAL_LIST_LAYOUT` -/// is set to `1`. -/// -/// [`ListLayoutStrategy`]: ListLayoutStrategy -pub fn use_experimental_list_layout() -> bool { - static USE_EXPERIMENTAL_LIST_LAYOUT: LazyLock = - LazyLock::new(|| env::var("VORTEX_EXPERIMENTAL_LIST_LAYOUT").is_ok_and(|v| v == "1")); - *USE_EXPERIMENTAL_LIST_LAYOUT -} - type ListLayoutFactory = Arc Arc + Send + Sync>; /// A configurable strategy for writing nested tabular data, dispatching each (sub)stream to the @@ -53,9 +40,9 @@ type ListLayoutFactory = Arc Arc Self { - self.with_list_layout_factory(|strategy| Arc::new(strategy)) + self.with_list_layout_factory(|strategy| Arc::new(ChunkedLayoutStrategy::new(strategy))) } /// Enable writing list fields with [`ListLayoutStrategy`] and wrap each list writer. This @@ -225,16 +212,11 @@ impl TableStrategy { /// Build the [`ListLayoutStrategy`] used to write a list field stream at this level. /// - /// The `elements` sub-column is routed back through a clean descended dispatcher so nested - /// structs/lists recurse; `offsets` go straight to the leaf (they are always a primitive - /// column); and `validity` uses the shared validity strategy. + /// The `elements`, `offsets`, and optional `validity` sub-columns are direct flat leaves. In + /// particular, nested lists and structs in the elements subtree are not decomposed recursively. 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_offsets(Arc::clone(&self.leaf)) - .with_validity(Arc::clone(&self.validity)) - .with_fallback(Arc::clone(&self.leaf)); + let list_layout = ListLayoutStrategy::default().with_fallback(Arc::clone(&self.leaf)); Some(factory(list_layout)) } @@ -353,9 +335,9 @@ mod tests { use crate::LayoutRef; use crate::LayoutStrategy; + use crate::layouts::chunked::Chunked; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; - use crate::layouts::list::List; use crate::layouts::repartition::RepartitionStrategy; use crate::layouts::repartition::RepartitionWriterOptions; use crate::layouts::table::TableStrategy; @@ -414,8 +396,7 @@ mod tests { Ok(()) } - /// A `list>` column: the dispatcher recurses into itself so the outer list's - /// `elements` are decomposed as a nested `ListLayout`. + /// A `list>` column keeps its nested list elements in a flat layout. #[tokio::test] async fn dispatches_nested_list() -> VortexResult<()> { let inner = ListArray::try_new( @@ -434,16 +415,14 @@ mod tests { let layout = write(&flat_table().with_list_layout(), outer).await?; insta::assert_snapshot!(layout.display_tree(), @r" vortex.list, dtype: list(list(i32)), children: 2 - ├── elements: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 1 - │ └── offsets: vortex.flat, dtype: u64, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + ├── elements: vortex.flat, dtype: list(i32), segment: 0 + └── offsets: vortex.flat, dtype: u32, segment: 1 "); Ok(()) } - /// A `struct<{ items: list>? }>` column: list decomposition recurses into struct - /// decomposition for the elements, and a nullable list writes a validity child. + /// A `struct<{ items: list>? }>` column keeps the element struct in a flat + /// layout, and a nullable list writes a validity child. #[tokio::test] async fn dispatches_struct_list_struct() -> VortexResult<()> { let inner_struct = StructArray::from_fields( @@ -466,18 +445,14 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @r" vortex.struct, dtype: {items=list({a=i32, b=i32})?}, children: 1 └── items: vortex.list, dtype: list({a=i32, b=i32})?, children: 3 - ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 - │ ├── a: vortex.flat, dtype: i32, segment: 2 - │ └── b: vortex.flat, dtype: i32, segment: 3 - ├── offsets: vortex.flat, dtype: u64, segment: 0 - └── validity: vortex.flat, dtype: bool, segment: 1 + ├── elements: vortex.flat, dtype: {a=i32, b=i32}, segment: 0 + ├── offsets: vortex.flat, dtype: u32, segment: 1 + └── validity: vortex.flat, dtype: bool, segment: 2 "); Ok(()) } - /// A multi-chunk `list` written with a chunked leaf: each sub-column (`elements`, - /// `offsets`) becomes its own `ChunkedLayout`, so elements are chunked independently of rows. - /// This is the "list-of-chunkeds" topology top-level decomposition unlocks. + /// A multi-chunk `list` is chunked in outer-row space before each page is shredded. #[tokio::test] async fn dispatches_chunked_list() -> VortexResult<()> { let chunk0 = ListArray::try_new( @@ -503,18 +478,18 @@ mod tests { .with_list_layout(); let layout = write(&dispatcher, chunked).await?; insta::assert_snapshot!(layout.display_tree(), @r" - vortex.list, dtype: list(i32), children: 2 - ├── elements: vortex.chunked, dtype: i32, children: 2 - │ ├── [0]: vortex.flat, dtype: i32, segment: 0 - │ └── [1]: vortex.flat, dtype: i32, segment: 1 - └── offsets: vortex.chunked, dtype: u64, children: 2 - ├── [0]: vortex.flat, dtype: u64, segment: 2 - └── [1]: vortex.flat, dtype: u64, segment: 3 + vortex.chunked, dtype: list(i32), children: 2 + ├── [0]: vortex.list, dtype: list(i32), children: 2 + │ ├── elements: vortex.flat, dtype: i32, segment: 0 + │ └── offsets: vortex.flat, dtype: u32, segment: 1 + └── [1]: vortex.list, dtype: list(i32), children: 2 + ├── elements: vortex.flat, dtype: i32, segment: 2 + └── offsets: vortex.flat, dtype: u32, segment: 3 "); Ok(()) } - /// A wrapper can repartition and zone lists in outer-row space before decomposition. + /// A wrapper can repartition lists into physical outer-row chunks before decomposition. #[tokio::test] async fn wraps_list_strategy_before_decomposition() -> VortexResult<()> { let list = ListArray::try_new( @@ -532,7 +507,7 @@ mod tests { let dispatcher = TableStrategy::new(Arc::clone(&flat), chunked).with_list_layout_factory( move |list_layout| { let zoned = ZonedStrategy::new( - list_layout, + ChunkedLayoutStrategy::new(list_layout), Arc::clone(&stats), ZonedLayoutOptions { block_size: row_block_size, @@ -559,7 +534,7 @@ mod tests { let data = layout .slot(0)? .vortex_expect("ZonedLayout always has a data child"); - assert!(data.is::()); + assert!(data.is::()); assert_eq!(data.row_count(), 9); Ok(()) }