feat: write nested struct layouts

Part of the push to fix deeply nested data in Vortex.

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

Author: a10y

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/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/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/struct_/mod.rs

@@ -8,7 +8,7 @@ use std::sync::Arc;
 
 use reader::StructReader;
 use vortex_array::{ArrayContext, DeserializeMetadata, EmptyMetadata};
-use vortex_dtype::{DType, Field, FieldMask, StructFields};
+use vortex_dtype::{DType, Field, FieldMask, Nullability, StructFields};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err, vortex_panic};
 
 use crate::children::{LayoutChildren, OwnedLayoutChildren};
@@ -52,24 +52,43 @@ impl VTable for StructVTable {
         layout.struct_fields().nfields()
     }
 
-    fn child(layout: &Self::Layout, idx: usize) -> VortexResult<LayoutRef> {
-        layout.children.child(
-            idx,
-            &layout
+    fn child(layout: &Self::Layout, index: usize) -> VortexResult<LayoutRef> {
+        let schema_index = if layout.dtype.is_nullable() {
+            index.saturating_sub(1)
+        } else {
+            index
+        };
+
+        let child_dtype = if index == 0 && layout.dtype.is_nullable() {
+            DType::Bool(Nullability::NonNullable)
+        } else {
+            layout
                 .struct_fields()
-                .field_by_index(idx)
-                .ok_or_else(|| vortex_err!("Missing field {idx}"))?,
-        )
+                .field_by_index(schema_index)
+                .ok_or_else(|| vortex_err!("Missing field {schema_index}"))?
+        };
+
+        layout.children.child(index, &child_dtype)
     }
 
     fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType {
-        LayoutChildType::Field(
-            layout
-                .struct_fields()
-                .field_name(idx)
-                .vortex_expect("Field index out of bounds")
-                .clone(),
-        )
+        let schema_index = if layout.dtype.is_nullable() {
+            idx.saturating_sub(1)
+        } else {
+            idx
+        };
+
+        if idx == 0 && layout.dtype.is_nullable() {
+            LayoutChildType::Auxiliary("validity".into())
+        } else {
+            LayoutChildType::Field(
+                layout
+                    .struct_fields()
+                    .field_name(schema_index)
+                    .vortex_expect("Field index out of bounds")
+                    .clone(),
+            )
+        }
     }
 
     fn new_reader(

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

@@ -277,7 +277,8 @@ mod tests {
         let ctx = ArrayContext::empty();
         let segments = Arc::new(TestSegments::default());
         let (ptr, eof) = SequenceId::root().split();
-        let strategy = StructStrategy::new(FlatLayoutStrategy::default());
+        let strategy =
+            StructStrategy::new(FlatLayoutStrategy::default(), FlatLayoutStrategy::default());
         let layout = block_on(|handle| {
             strategy.write_stream(
                 ctx,