feat: write nested struct layouts (#4942)

Part of #4889 

To better support writing nested data, here we update our builtin
StructLayout in a backwards-compatible fashion.

- `StructStrategy` will shred struct fields into their own StructLayout
recurseively
- StructLayout can support nullable structs now. It does this by writing
a new child layout containing the validity buffer for nullable arrays

I use the `RealNest` dataset to evaluate, which contains a copy of ~200k
github pull request webhook events. Nested struct layout reduces file
size over the previous strategy by about ~10%, and also makes pushdown
into the nested columns possible.



Some open questions

* The validity child requires some extra handling. It seems like the
validity handling is very dependent on the expression being pushed down.
For example if I'm doing a simple project of a child field, then adding
the validity to the result is a simple masking operation. If I'm pushing
down an `UNNEST` or something else that increases the result size, it is
hard to map the validity buffer onto the projection_eval result
* I collect all validity chunks into a single buffer at write time. The
idea being that it's better to access the struct validity as a single
unit since it is much smaller than the data size. Assuming an 8MB target
segment size, this lets us comfortably fit ~64mm rows into a single
segment. Another alternative is to bring back roaring, or enable some
other boolean compressors.

---------

Signed-off-by: Andrew Duffy <[email protected]>

Author: a10y

bench-vortex/src/realnest/gharchive.rs

@@ -230,6 +230,10 @@ impl Benchmark for GithubArchive {
     fn data_url(&self) -> &Url {
         &self.data_url
     }
+
+    fn expected_row_counts(&self) -> Option<&[usize]> {
+        Some(&[1, 2, 100, 10, 82468])
+    }
 }
 
 pub async fn register_table(

vortex-cxx/cpp/tests/basic_test.cpp

@@ -242,7 +242,7 @@ TEST_F(VortexTest, ScanBuilderWithRowRangeWithIncludeByIndex) {
 
     RunScanBuilderTest(
         [&include_by_index](vortex::ScanBuilder &scan_builder) {
-            return std::move(scan_builder.WithRowRange(2, 6).WithIncludeByIndex(include_by_index.data(),
+            return std::move(scan_builder.WithRowRange(2, 5).WithIncludeByIndex(include_by_index.data(),
                                                                                 include_by_index.size()))
                 .IntoStream();
         },

vortex-file/src/file.rs

@@ -137,8 +137,9 @@ impl VortexFile {
     }
 
     pub fn splits(&self) -> VortexResult<Vec<Range<u64>>> {
+        let reader = self.layout_reader()?;
         Ok(SplitBy::Layout
-            .splits(self.layout_reader()?.as_ref(), &[FieldMask::All])?
+            .splits(reader.as_ref(), &(0..reader.row_count()), &[FieldMask::All])?
             .into_iter()
             .tuple_windows()
             .map(|(start, end)| start..end)

vortex-file/src/strategy.rs

@@ -8,6 +8,7 @@ use std::sync::Arc;
 use vortex_layout::LayoutStrategy;
 use vortex_layout::layouts::buffered::BufferedStrategy;
 use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy;
+use vortex_layout::layouts::collect::CollectStrategy;
 use vortex_layout::layouts::compressed::{CompressingStrategy, CompressorPlugin};
 use vortex_layout::layouts::dict::writer::DictStrategy;
 use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
@@ -100,7 +101,7 @@ impl WriteStrategyBuilder {
         // 2. calculate stats for each row group
         let stats = ZonedStrategy::new(
             dict,
-            compress_then_flat,
+            compress_then_flat.clone(),
             ZonedLayoutOptions {
                 block_size: self.row_block_size,
                 ..Default::default()
@@ -120,6 +121,8 @@ impl WriteStrategyBuilder {
         );
 
         // 0. start with splitting columns
-        Arc::new(StructStrategy::new(repartition))
+        let validity_strategy = CollectStrategy::new(compress_then_flat);
+
+        Arc::new(StructStrategy::new(repartition, validity_strategy))
     }
 }

vortex-layout/src/layouts/chunked/reader.rs

@@ -8,6 +8,7 @@ use std::sync::Arc;
 use futures::future::BoxFuture;
 use futures::stream::FuturesOrdered;
 use futures::{FutureExt, TryStreamExt};
+use itertools::Itertools;
 use vortex_array::arrays::ChunkedArray;
 use vortex_array::{ArrayRef, MaskFuture};
 use vortex_dtype::{DType, FieldMask};
@@ -43,7 +44,12 @@ impl ChunkedReader {
         }
         chunk_offsets[nchildren] = layout.row_count();
 
-        let lazy_children = LazyReaderChildren::new(layout.children.clone(), segment_source);
+        let dtypes = vec![layout.dtype.clone(); nchildren];
+        let names = (0..nchildren)
+            .map(|idx| Arc::from(format!("{name}.[{idx}]")))
+            .collect();
+        let lazy_children =
+            LazyReaderChildren::new(layout.children.clone(), dtypes, names, segment_source);
 
         Self {
             layout,
@@ -55,11 +61,7 @@ impl ChunkedReader {
 
     /// Return the [`LayoutReader`] for the given chunk.
     fn chunk_reader(&self, idx: usize) -> VortexResult<&LayoutReaderRef> {
-        self.lazy_children.get(
-            idx,
-            self.layout.dtype(),
-            &format!("{}.[{}]", self.name, idx).into(),
-        )
+        self.lazy_children.get(idx)
     }
 
     fn chunk_offset(&self, idx: usize) -> u64 {
@@ -148,16 +150,35 @@ impl LayoutReader for ChunkedReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        let mut offset = row_offset;
-        for i in 0..self.layout.nchildren() {
-            let child = self.chunk_reader(i)?;
-            child.register_splits(field_mask, offset, splits)?;
-            offset += self.layout.child(i)?.row_count();
-            splits.insert(offset);
+        for (index, (&start, &end)) in self
+            .chunk_offsets
+            .iter()
+            .tuple_windows::<(_, _)>()
+            .enumerate()
+        {
+            if end < row_range.start {
+                continue;
+            }
+
+            if start >= row_range.end {
+                break;
+            }
+
+            // Child overlaps in whole or in part with split
+            let child = self.chunk_reader(index)?;
+            let child_range =
+                std::cmp::max(row_range.start, start)..std::cmp::min(row_range.end, end);
+
+            // Register any splits from the child
+            child.register_splits(field_mask, &child_range, splits)?;
+
+            // Register the split indicating the end of this chunk
+            splits.insert(child_range.end);
         }
+
         Ok(())
     }
 

vortex-layout/src/layouts/collect.rs

@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::sync::Arc;
+
+use async_stream::try_stream;
+use async_trait::async_trait;
+use futures::{StreamExt, pin_mut};
+use vortex_array::ArrayContext;
+use vortex_array::arrays::ChunkedArray;
+use vortex_error::{VortexExpect, VortexResult};
+use vortex_io::runtime::Handle;
+
+use crate::segments::SegmentSinkRef;
+use crate::sequence::{
+    SendableSequentialStream, SequencePointer, SequentialStream, SequentialStreamAdapter,
+};
+use crate::{LayoutRef, LayoutStrategy};
+
+/// A strategy that collects all chunks and turns them into a single array chunk to pass into
+/// a child strategy.
+pub struct CollectStrategy {
+    child: Arc<dyn LayoutStrategy>,
+}
+
+impl CollectStrategy {
+    pub fn new<S: LayoutStrategy>(child: S) -> CollectStrategy {
+        CollectStrategy {
+            child: Arc::new(child),
+        }
+    }
+}
+
+#[async_trait]
+impl LayoutStrategy for CollectStrategy {
+    async fn write_stream(
+        &self,
+        ctx: ArrayContext,
+        segment_sink: SegmentSinkRef,
+        stream: SendableSequentialStream,
+        eof: SequencePointer,
+        handle: Handle,
+    ) -> VortexResult<LayoutRef> {
+        // Read the whole stream, then write one Chunked stream to the inner thing
+        let dtype = stream.dtype().clone();
+
+        let _dtype = dtype.clone();
+        let collected_stream = try_stream! {
+            pin_mut!(stream);
+
+            let mut chunks = Vec::new();
+            let mut latest_sequence_id = None;
+            while let Some(chunk) = stream.next().await {
+                let (sequence_id, chunk) = chunk?;
+                latest_sequence_id = Some(sequence_id);
+                chunks.push(chunk);
+            }
+
+            let collected = ChunkedArray::try_new(chunks, _dtype)?.to_array();
+            yield (latest_sequence_id.vortex_expect("must have visited at least one chunk"), collected);
+        };
+
+        let adapted = Box::pin(SequentialStreamAdapter::new(dtype, collected_stream));
+
+        self.child
+            .write_stream(ctx, segment_sink, adapted, eof, handle)
+            .await
+    }
+
+    fn buffered_bytes(&self) -> u64 {
+        todo!()
+    }
+}

vortex-layout/src/layouts/dict/reader.rs

@@ -111,10 +111,10 @@ impl LayoutReader for DictReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        self.codes.register_splits(field_mask, row_offset, splits)
+        self.codes.register_splits(field_mask, row_range, splits)
     }
 
     fn pruning_evaluation(

vortex-layout/src/layouts/flat/reader.rs

@@ -84,10 +84,10 @@ impl LayoutReader for FlatReader {
     fn register_splits(
         &self,
         _field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        splits.insert(row_offset + self.layout.row_count());
+        splits.insert(row_range.start + self.layout.row_count);
         Ok(())
     }
 

vortex-layout/src/layouts/mod.rs

@@ -9,6 +9,7 @@ use vortex_error::SharedVortexResult;
 
 pub mod buffered;
 pub mod chunked;
+pub mod collect;
 #[cfg(feature = "zstd")]
 pub mod compact;
 pub mod compressed;

vortex-layout/src/layouts/row_idx/mod.rs

@@ -136,10 +136,10 @@ impl LayoutReader for RowIdxLayoutReader {
     fn register_splits(
         &self,
         field_mask: &[FieldMask],
-        row_offset: u64,
+        row_range: &Range<u64>,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        self.child.register_splits(field_mask, row_offset, splits)
+        self.child.register_splits(field_mask, row_range, splits)
     }
 
     fn pruning_evaluation(