perf(layout): zone lists before decomposition

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

vortex-file/src/strategy.rs

@@ -50,6 +50,7 @@ use vortex_layout::layouts::compressed::CompressingStrategy;
 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::RepartitionStrategy;
 use vortex_layout::layouts::repartition::RepartitionWriterOptions;
 use vortex_layout::layouts::table::TableStrategy;
@@ -322,12 +323,14 @@ impl WriteStrategyBuilder {
             probe_compressor,
         );
 
+        let row_block_size = NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0");
+
         // 2. calculate stats for each row group
         let stats = ZonedStrategy::new(
             dict,
             compress_then_flat.clone(),
             ZonedLayoutOptions {
-                block_size: NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0"),
+                block_size: row_block_size,
                 ..Default::default()
             },
         );
@@ -346,15 +349,36 @@ impl WriteStrategyBuilder {
         );
 
         // 0. start with splitting columns
-        let validity_strategy = CollectStrategy::new(compress_then_flat);
+        let validity_strategy = CollectStrategy::new(compress_then_flat.clone());
 
         // 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 {
-            table_strategy = table_strategy.with_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,
+                        },
+                    ))
+                },
+            );
         }
 
         Arc::new(table_strategy)

vortex-layout/src/layouts/table.rs

@@ -45,6 +45,8 @@ pub fn use_experimental_list_layout() -> bool {
     *USE_EXPERIMENTAL_LIST_LAYOUT
 }
 
+type ListLayoutFactory = Arc<dyn Fn(ListLayoutStrategy) -> Arc<dyn LayoutStrategy> + Send + Sync>;
+
 /// A configurable strategy for writing nested tabular data, dispatching each (sub)stream to the
 /// structural writer for its dtype.
 ///
@@ -67,10 +69,11 @@ pub struct TableStrategy {
     validity: Arc<dyn LayoutStrategy>,
     /// The writer for leaf fields, i.e. anything that is not a struct.
     leaf: Arc<dyn LayoutStrategy>,
-    /// Whether to write list fields using [`ListLayoutStrategy`].
+    /// Optional factory applied to each dynamically constructed [`ListLayoutStrategy`].
+    /// Its presence also enables list decomposition.
     ///
     /// [`ListLayoutStrategy`]: crate::layouts::list::writer::ListLayoutStrategy
-    use_list_layout: bool,
+    list_layout_factory: Option<ListLayoutFactory>,
 }
 
 impl TableStrategy {
@@ -97,7 +100,7 @@ impl TableStrategy {
             leaf_writers: Default::default(),
             validity,
             leaf: fallback,
-            use_list_layout: false,
+            list_layout_factory: None,
         }
     }
 
@@ -166,8 +169,18 @@ impl TableStrategy {
     }
 
     /// Enable writing list fields with [`ListLayoutStrategy`].
-    pub fn with_list_layout(mut self) -> Self {
-        self.use_list_layout = true;
+    pub fn with_list_layout(self) -> Self {
+        self.with_list_layout_factory(|strategy| Arc::new(strategy))
+    }
+
+    /// Enable list layouts and transform each structural list writer into a physical strategy.
+    /// The factory receives a fully configured writer and is invoked independently for every list,
+    /// including nested lists.
+    pub fn with_list_layout_factory(
+        mut self,
+        factory: impl Fn(ListLayoutStrategy) -> Arc<dyn LayoutStrategy> + Send + Sync + 'static,
+    ) -> Self {
+        self.list_layout_factory = Some(Arc::new(factory));
         self
     }
 }
@@ -210,12 +223,14 @@ impl TableStrategy {
     /// 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.
-    fn list_strategy(&self) -> ListLayoutStrategy {
-        ListLayoutStrategy::default()
+    fn list_strategy(&self) -> Option<Arc<dyn LayoutStrategy>> {
+        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))
+            .with_fallback(Arc::clone(&self.leaf));
+        Some(factory(list_layout))
     }
 
     /// Descend into a subfield, retaining only the overrides that apply beneath it (rebased to be
@@ -236,7 +251,7 @@ impl TableStrategy {
             leaf_writers: new_writers,
             validity: Arc::clone(&self.validity),
             leaf: Arc::clone(&self.leaf),
-            use_list_layout: self.use_list_layout,
+            list_layout_factory: self.list_layout_factory.clone(),
         }
     }
 
@@ -247,7 +262,7 @@ impl TableStrategy {
             leaf_writers: HashMap::default(),
             validity: Arc::clone(&self.validity),
             leaf: Arc::clone(&self.leaf),
-            use_list_layout: self.use_list_layout,
+            list_layout_factory: self.list_layout_factory.clone(),
         }
     }
 
@@ -290,9 +305,10 @@ impl LayoutStrategy for TableStrategy {
                 .await;
         }
 
-        if dtype.is_list() && self.use_list_layout {
-            return self
-                .list_strategy()
+        if dtype.is_list()
+            && let Some(list_strategy) = self.list_strategy()
+        {
+            return list_strategy
                 .write_stream(ctx, segment_sink, stream, eof, session)
                 .await;
         }
@@ -306,6 +322,7 @@ impl LayoutStrategy for TableStrategy {
 
 #[cfg(test)]
 mod tests {
+    use std::num::NonZeroUsize;
     use std::sync::Arc;
     use std::task::Poll;
 
@@ -325,6 +342,7 @@ mod tests {
     use vortex_array::field_path;
     use vortex_array::validity::Validity;
     use vortex_buffer::buffer;
+    use vortex_error::VortexExpect;
     use vortex_error::VortexResult;
     use vortex_io::runtime::single::block_on;
     use vortex_io::session::RuntimeSessionExt;
@@ -333,7 +351,13 @@ mod tests {
     use crate::LayoutStrategy;
     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;
+    use crate::layouts::zoned::Zoned;
+    use crate::layouts::zoned::writer::ZonedLayoutOptions;
+    use crate::layouts::zoned::writer::ZonedStrategy;
     use crate::segments::TestSegments;
     use crate::sequence::SequenceId;
     use crate::sequence::SequentialArrayStreamExt;
@@ -478,6 +502,54 @@ mod tests {
         Ok(())
     }
 
+    /// A wrapper can repartition and zone lists in outer-row space before decomposition.
+    #[tokio::test]
+    async fn wraps_list_strategy_before_decomposition() -> VortexResult<()> {
+        let list = ListArray::try_new(
+            PrimitiveArray::from_iter(0..9_i32).into_array(),
+            PrimitiveArray::from_iter(0..=9_u32).into_array(),
+            Validity::NonNullable,
+        )?
+        .into_array();
+
+        let row_block_size = NonZeroUsize::new(4).vortex_expect("4 is non-zero");
+        let flat: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
+        let stats = Arc::clone(&flat);
+        let chunked: Arc<dyn LayoutStrategy> =
+            Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()));
+        let dispatcher = TableStrategy::new(Arc::clone(&flat), chunked).with_list_layout_factory(
+            move |list_layout| {
+                let zoned = ZonedStrategy::new(
+                    list_layout,
+                    Arc::clone(&stats),
+                    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,
+                    },
+                )) as Arc<dyn LayoutStrategy>
+            },
+        );
+
+        let layout = write(&dispatcher, list).await?;
+        let zoned = layout.as_::<Zoned>();
+        assert_eq!(zoned.zone_len(), 4);
+        assert_eq!(zoned.nzones(), 3);
+
+        let data = layout.child(0)?;
+        assert!(data.is::<List>());
+        assert_eq!(data.row_count(), 9);
+        Ok(())
+    }
+
     /// A non-struct stream is not shredded; it is handed straight to the leaf strategy.
     #[tokio::test]
     async fn non_struct_input_uses_leaf() -> VortexResult<()> {