Add list-aware physical scan split boundaries

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<dyn LayoutStrategy> {
+        let data_block_target_bytes = self.data_block_target_bytes;
         let flat: Arc<dyn LayoutStrategy> = if let Some(flat) = self.flat_strategy {
             flat
         } else {
@@ -223,17 +226,17 @@ 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
                 // enough to both (1) allow fine-grained random access reads and (2) allow
                 // 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<dyn LayoutStrategy> {
-                    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<dyn LayoutStrategy> =
+                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<dyn LayoutStrategy> {
+                        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)

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<Vec<MapEntryFixture<'a>>>;
 

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<LayoutRef> = 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"))

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::<Vec<_>>()
             .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));
         }