From dd326c3ff14339036c80394e3a59e1fdb161eee1 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 1 Jul 2026 11:04:19 -0700 Subject: [PATCH] first pass Signed-off-by: Matt Katz --- Cargo.lock | 1 + vortex-layout/Cargo.toml | 1 + vortex-layout/src/layouts/struct_/mod.rs | 2 + vortex-layout/src/layouts/struct_/writer.rs | 257 +++++++++++++ vortex-layout/src/layouts/table.rs | 406 ++++++++++---------- 5 files changed, 468 insertions(+), 199 deletions(-) create mode 100644 vortex-layout/src/layouts/struct_/writer.rs diff --git a/Cargo.lock b/Cargo.lock index 173843bbaad..fa49b2d6491 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10498,6 +10498,7 @@ dependencies = [ "bit-vec", "flatbuffers", "futures", + "insta", "itertools 0.14.0", "kanal", "moka", diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 61b1253ef43..daa96d4e138 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -55,6 +55,7 @@ vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] futures = { workspace = true, features = ["executor"] } +insta = { workspace = true } rstest = { workspace = true } temp-env = { workspace = true } tokio = { workspace = true, features = ["rt", "macros"] } diff --git a/vortex-layout/src/layouts/struct_/mod.rs b/vortex-layout/src/layouts/struct_/mod.rs index db056f4c57b..853a6793f4c 100644 --- a/vortex-layout/src/layouts/struct_/mod.rs +++ b/vortex-layout/src/layouts/struct_/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod reader; +pub mod writer; use std::sync::Arc; @@ -21,6 +22,7 @@ use vortex_error::vortex_err; use vortex_session::SessionExt; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +pub use writer::StructStrategy; use crate::LayoutBuildContext; use crate::LayoutChildType; diff --git a/vortex-layout/src/layouts/struct_/writer.rs b/vortex-layout/src/layouts/struct_/writer.rs new file mode 100644 index 00000000000..dc2b0a9749b --- /dev/null +++ b/vortex-layout/src/layouts/struct_/writer.rs @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A writer strategy for struct-typed arrays. +//! +//! [`StructStrategy`] transposes a stream of struct chunks into one ordered stream per field +//! (plus a validity stream when the struct is nullable) and writes each through a configurable +//! child strategy, producing a single [`StructLayout`]. It is a *structural* writer: it does not +//! inspect child dtypes or resolve field-path overrides itself. Dispatching a child to the right +//! layout kind is the job of the caller (see [`TableStrategy`]). +//! +//! [`TableStrategy`]: crate::layouts::table::TableStrategy + +use std::sync::Arc; + +use async_trait::async_trait; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::future::try_join_all; +use futures::pin_mut; +use itertools::Itertools; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::Nullability; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_io::kanal_ext::KanalExt; +use vortex_io::session::RuntimeSessionExt; +use vortex_session::VortexSession; +use vortex_utils::aliases::DefaultHashBuilder; +use vortex_utils::aliases::hash_map::HashMap; +use vortex_utils::aliases::hash_set::HashSet; + +use crate::IntoLayout; +use crate::LayoutRef; +use crate::LayoutStrategy; +use crate::layouts::struct_::StructLayout; +use crate::segments::SegmentSinkRef; +use crate::sequence::SendableSequentialStream; +use crate::sequence::SequenceId; +use crate::sequence::SequencePointer; +use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequentialStreamExt; + +/// Writes struct-typed arrays into a [`StructLayout`], one child layout per field. +/// +/// Each field is written through a strategy resolved by direct field name: an entry in +/// `field_writers` if present, otherwise `default`. When the struct is nullable, its validity +/// bitmap is written through `validity`. +/// +/// `StructStrategy` is intentionally unaware of nested dtypes and field-path overrides. To write +/// arbitrarily nested struct trees with per-path overrides, drive it from +/// [`TableStrategy`][crate::layouts::table::TableStrategy], which dispatches on dtype and resolves +/// the per-field child strategies before handing them here. +#[derive(Clone)] +pub struct StructStrategy { + /// Per-field child strategies, keyed by direct field name. Fields without an entry use + /// `default`. + field_writers: HashMap>, + /// Strategy for fields that have no entry in `field_writers`. + default: Arc, + /// Strategy for the struct's own validity bitmap, used only when the struct is nullable. + validity: Arc, +} + +impl StructStrategy { + /// Create a new struct writer that writes every field through `default` and the validity + /// bitmap (when present) through `validity`. + pub fn new(validity: Arc, default: Arc) -> Self { + Self { + field_writers: HashMap::default(), + default, + validity, + } + } + + /// Override the strategy for a single field by name. + pub fn with_field_writer( + mut self, + name: impl Into, + writer: Arc, + ) -> Self { + self.field_writers.insert(name.into(), writer); + self + } + + /// Override the strategy for several fields by name at once. + pub fn with_field_writers( + mut self, + writers: impl IntoIterator)>, + ) -> Self { + self.field_writers.extend(writers); + self + } +} + +#[async_trait] +impl LayoutStrategy for StructStrategy { + async fn write_stream( + &self, + ctx: ArrayContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + mut eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + let dtype = stream.dtype().clone(); + + let Some(struct_dtype) = dtype.as_struct_fields_opt() else { + vortex_bail!("StructStrategy can only write struct-typed streams, got {dtype}"); + }; + + // Check for unique field names at write time. + if HashSet::<_, DefaultHashBuilder>::from_iter(struct_dtype.names().iter()).len() + != struct_dtype.names().len() + { + vortex_bail!("StructLayout must have unique field names"); + } + let is_nullable = dtype.is_nullable(); + + // Optimization: when there are no fields, don't spawn any work and just write a trivial + // StructLayout. + if struct_dtype.nfields() == 0 && !is_nullable { + let row_count = stream + .try_fold( + 0u64, + |acc, (_, arr)| async move { Ok(acc + arr.len() as u64) }, + ) + .await?; + return Ok(StructLayout::new(row_count, dtype, vec![]).into_layout()); + } + + // stream -> stream> + let columns_session = session.clone(); + let columns_vec_stream = stream.map(move |chunk| { + let (sequence_id, chunk) = chunk?; + let mut sequence_pointer = sequence_id.descend(); + let mut ctx = columns_session.create_execution_ctx(); + let struct_chunk = chunk.clone().execute::(&mut ctx)?; + let mut columns: Vec<(SequenceId, ArrayRef)> = Vec::new(); + if is_nullable { + columns.push(( + sequence_pointer.advance(), + chunk + .validity()? + .execute_mask(chunk.len(), &mut ctx)? + .into_array(), + )); + } + + columns.extend( + struct_chunk + .iter_unmasked_fields() + .map(|field| (sequence_pointer.advance(), field.clone())), + ); + + Ok(columns) + }); + + let mut stream_count = struct_dtype.nfields(); + if is_nullable { + stream_count += 1; + } + + let (column_streams_tx, column_streams_rx): (Vec<_>, Vec<_>) = + (0..stream_count).map(|_| kanal::bounded_async(1)).unzip(); + + // Spawn a task to fan out column chunks to their respective transposed streams + let handle = session.handle(); + handle + .spawn(async move { + pin_mut!(columns_vec_stream); + while let Some(result) = columns_vec_stream.next().await { + match result { + Ok(columns) => { + for (tx, column) in column_streams_tx.iter().zip_eq(columns.into_iter()) + { + let _ = tx.send(Ok(column)).await; + } + } + Err(e) => { + let e: Arc = Arc::new(e); + for tx in column_streams_tx.iter() { + let _ = tx.send(Err(VortexError::from(Arc::clone(&e)))).await; + } + break; + } + } + } + }) + .detach(); + + // First child column is the validity, subsequent children are the individual struct fields + let column_dtypes: Vec = if is_nullable { + std::iter::once(DType::Bool(Nullability::NonNullable)) + .chain(struct_dtype.fields()) + .collect() + } else { + struct_dtype.fields().collect() + }; + + let column_names: Vec = if is_nullable { + std::iter::once(FieldName::from("__validity")) + .chain(struct_dtype.names().iter().cloned()) + .collect() + } else { + struct_dtype.names().iter().cloned().collect() + }; + + let layout_futures: Vec<_> = column_dtypes + .into_iter() + .zip_eq(column_streams_rx) + .zip_eq(column_names) + .enumerate() + .map(move |(index, ((dtype, recv), name))| { + let column_stream = + SequentialStreamAdapter::new(dtype, recv.into_stream().boxed()).sendable(); + let child_eof = eof.split_off(); + let session = session.clone(); + let ctx = ctx.clone(); + let segment_sink = Arc::clone(&segment_sink); + handle.spawn_nested(move |h| { + // Validity is written through the validity strategy; every other field + // resolves to its named override or the default strategy. + let writer = if index == 0 && is_nullable { + Arc::clone(&self.validity) + } else { + self.field_writers + .get(&name) + .cloned() + .unwrap_or_else(|| Arc::clone(&self.default)) + }; + let session = session.with_handle(h); + + async move { + writer + .write_stream(ctx, segment_sink, column_stream, child_eof, &session) + .await + } + }) + }) + .collect(); + + let column_layouts = try_join_all(layout_futures).await?; + // TODO(os): transposed stream could count row counts as well, + // This must hold though, all columns must have the same row count of the struct layout + let row_count = column_layouts.first().map(|l| l.row_count()).unwrap_or(0); + Ok(StructLayout::new(row_count, dtype, column_layouts).into_layout()) + } +} diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 28b3af52c87..faf190bc3f3 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -3,64 +3,58 @@ //! A configurable writer strategy for tabular data. //! -//! Allows the caller to override specific leaf fields with custom layout strategies. +//! [`TableStrategy`] is a *dispatcher*: it inspects the dtype of the stream it is handed and +//! routes struct columns to [`StructStrategy`] and everything else to the configured leaf +//! strategy. Because it hands *itself* (suitably descended) to [`StructStrategy`] as the strategy +//! for the struct's children, arbitrarily nested struct trees are written with no manual wiring. +//! +//! The dispatcher also owns field-path overrides, letting callers force a specific leaf field — +//! at any depth — onto a custom strategy. use std::sync::Arc; use async_trait::async_trait; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::future::try_join_all; -use futures::pin_mut; -use itertools::Itertools; use vortex_array::ArrayContext; -use vortex_array::ArrayRef; -use vortex_array::IntoArray; -use vortex_array::VortexSessionExecute; -use vortex_array::arrays::StructArray; -use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::dtype::DType; use vortex_array::dtype::Field; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldPath; -use vortex_array::dtype::Nullability; -use vortex_error::VortexError; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_io::kanal_ext::KanalExt; -use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; -use vortex_utils::aliases::DefaultHashBuilder; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; -use crate::IntoLayout; use crate::LayoutRef; use crate::LayoutStrategy; -use crate::layouts::struct_::StructLayout; +use crate::layouts::struct_::StructStrategy; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; -use crate::sequence::SequenceId; use crate::sequence::SequencePointer; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt; -/// A configurable strategy for writing tables with nested field columns, allowing -/// overrides for specific leaf columns. +/// A configurable strategy for writing nested tabular data, dispatching each (sub)stream to the +/// structural writer for its dtype. +/// +/// Dispatch rules, applied to the dtype of the stream handed to [`write_stream`]: +/// - **struct** → [`StructStrategy`], with each field written by its override (if any) or by a +/// descended copy of this dispatcher. +/// - **anything else** → the leaf strategy. +/// +/// [`write_stream`]: LayoutStrategy::write_stream pub struct TableStrategy { - /// A set of leaf field overrides, e.g. to force one column to be compact-compressed. + /// A set of field-path overrides, e.g. to force one column to be compact-compressed. Keys are + /// paths relative to the level this dispatcher sits at. leaf_writers: HashMap>, - /// The writer for any validity arrays that may be present + /// The writer for any validity arrays that may be present, at any level of the tree. validity: Arc, - /// The fallback writer for any fields that do not have an explicit writer set in `leaf_writers` - fallback: Arc, + /// The writer for leaf fields, i.e. anything that is not a struct. + leaf: Arc, } impl TableStrategy { - /// Create a new writer with the specified write strategies for validity, and for all leaf - /// fields, with no overrides. + /// Create a new dispatcher with the given `validity` strategy and `fallback` leaf strategy and + /// no overrides. /// - /// Additional overrides can be configured using the `with_leaf_strategy` method. + /// Additional per-field overrides can be configured with + /// [`with_field_writer`][Self::with_field_writer]. /// /// ## Example /// @@ -78,7 +72,7 @@ impl TableStrategy { Self { leaf_writers: Default::default(), validity, - fallback, + leaf: fallback, } } @@ -136,7 +130,7 @@ impl TableStrategy { /// Override the default strategy for leaf columns that don't have overrides. pub fn with_default_strategy(mut self, default: Arc) -> Self { - self.fallback = default; + self.leaf = default; self } @@ -148,15 +142,47 @@ impl TableStrategy { } impl TableStrategy { - /// Descend into a subfield for the writer. + /// Build the [`StructStrategy`] used to write a struct-typed stream at this level. + /// + /// Each field that has an override (or a deeper override beneath it) is resolved up front; + /// every other field falls through to a clean descended dispatcher. + fn struct_strategy(&self) -> StructStrategy { + let mut field_writers: HashMap> = HashMap::default(); + + // The distinct named first-segments of our override paths are the only fields that need + // anything other than the default dispatcher. + let mut named_first: HashSet = HashSet::default(); + for path in self.leaf_writers.keys() { + if let Some(Field::Name(name)) = path.parts().first() { + named_first.insert(name.clone()); + } + } + + for name in named_first { + // `validate_path` forbids overlapping overrides, so a name has *either* an exact + // single-segment override *or* deeper overrides, never both. + let writer = match self.leaf_writers.get(&FieldPath::from_name(name.clone())) { + Some(exact) => Arc::clone(exact), + None => { + Arc::new(self.descend(&Field::Name(name.clone()))) as Arc + } + }; + field_writers.insert(name, writer); + } + + StructStrategy::new(Arc::clone(&self.validity), Arc::new(self.descend_clean())) + .with_field_writers(field_writers) + } + + /// Descend into a subfield, retaining only the overrides that apply beneath it (rebased to be + /// relative to the child). fn descend(&self, field: &Field) -> Self { - // Start with the existing set of overrides, then only retain the ones that contain - // the current field let mut new_writers = HashMap::with_capacity(self.leaf_writers.len()); for (field_path, strategy) in &self.leaf_writers { - if field_path.starts_with_field(field) + if field_path.parts().first() == Some(field) && let Some(subpath) = field_path.clone().step_into() + && !subpath.is_root() { new_writers.insert(subpath, Arc::clone(strategy)); } @@ -165,7 +191,17 @@ impl TableStrategy { Self { leaf_writers: new_writers, validity: Arc::clone(&self.validity), - fallback: Arc::clone(&self.fallback), + leaf: Arc::clone(&self.leaf), + } + } + + /// A copy of this dispatcher with no overrides, used as the default child strategy for fields + /// that carry no override. + fn descend_clean(&self) -> Self { + Self { + leaf_writers: HashMap::default(), + validity: Arc::clone(&self.validity), + leaf: Arc::clone(&self.leaf), } } @@ -188,7 +224,7 @@ impl TableStrategy { } } -/// Specialized strategy for when we exactly know the input schema. +/// Dispatches each stream to the structural writer for its dtype. #[async_trait] impl LayoutStrategy for TableStrategy { async fn write_stream( @@ -196,174 +232,22 @@ impl LayoutStrategy for TableStrategy { ctx: ArrayContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, - mut eof: SequencePointer, + eof: SequencePointer, session: &VortexSession, ) -> VortexResult { let dtype = stream.dtype().clone(); - // Fallback: if the array is not a struct, fallback to writing a single array. - if !dtype.is_struct() { + if dtype.is_struct() { return self - .fallback + .struct_strategy() .write_stream(ctx, segment_sink, stream, eof, session) .await; } - let struct_dtype = dtype.as_struct_fields(); - - // Check for unique field names at write time. - if HashSet::<_, DefaultHashBuilder>::from_iter(struct_dtype.names().iter()).len() - != struct_dtype.names().len() - { - vortex_bail!("StructLayout must have unique field names"); - } - let is_nullable = dtype.is_nullable(); - - // Optimization: when there are no fields, don't spawn any work and just write a trivial - // StructLayout. - if struct_dtype.nfields() == 0 && !is_nullable { - let row_count = stream - .try_fold( - 0u64, - |acc, (_, arr)| async move { Ok(acc + arr.len() as u64) }, - ) - .await?; - return Ok(StructLayout::new(row_count, dtype, vec![]).into_layout()); - } - - // stream -> stream> - let columns_session = session.clone(); - let columns_vec_stream = stream.map(move |chunk| { - let (sequence_id, chunk) = chunk?; - let mut sequence_pointer = sequence_id.descend(); - let mut ctx = columns_session.create_execution_ctx(); - let struct_chunk = chunk.clone().execute::(&mut ctx)?; - let mut columns: Vec<(SequenceId, ArrayRef)> = Vec::new(); - if is_nullable { - columns.push(( - sequence_pointer.advance(), - chunk - .validity()? - .execute_mask(chunk.len(), &mut ctx)? - .into_array(), - )); - } - - columns.extend( - struct_chunk - .iter_unmasked_fields() - .map(|field| (sequence_pointer.advance(), field.clone())), - ); - - Ok(columns) - }); - - let mut stream_count = struct_dtype.nfields(); - if is_nullable { - stream_count += 1; - } - - let (column_streams_tx, column_streams_rx): (Vec<_>, Vec<_>) = - (0..stream_count).map(|_| kanal::bounded_async(1)).unzip(); - - // Spawn a task to fan out column chunks to their respective transposed streams - let handle = session.handle(); - handle - .spawn(async move { - pin_mut!(columns_vec_stream); - while let Some(result) = columns_vec_stream.next().await { - match result { - Ok(columns) => { - for (tx, column) in column_streams_tx.iter().zip_eq(columns.into_iter()) - { - let _ = tx.send(Ok(column)).await; - } - } - Err(e) => { - let e: Arc = Arc::new(e); - for tx in column_streams_tx.iter() { - let _ = tx.send(Err(VortexError::from(Arc::clone(&e)))).await; - } - break; - } - } - } - }) - .detach(); - - // First child column is the validity, subsequence children are the individual struct fields - let column_dtypes: Vec = if is_nullable { - std::iter::once(DType::Bool(Nullability::NonNullable)) - .chain(struct_dtype.fields()) - .collect() - } else { - struct_dtype.fields().collect() - }; - - let column_names: Vec = if is_nullable { - std::iter::once(FieldName::from("__validity")) - .chain(struct_dtype.names().iter().cloned()) - .collect() - } else { - struct_dtype.names().iter().cloned().collect() - }; - - let layout_futures: Vec<_> = column_dtypes - .into_iter() - .zip_eq(column_streams_rx) - .zip_eq(column_names) - .enumerate() - .map(move |(index, ((dtype, recv), name))| { - let column_stream = - SequentialStreamAdapter::new(dtype.clone(), recv.into_stream().boxed()) - .sendable(); - let child_eof = eof.split_off(); - let field = Field::Name(name.clone()); - let session = session.clone(); - let ctx = ctx.clone(); - let segment_sink = Arc::clone(&segment_sink); - handle.spawn_nested(move |h| { - let validity = Arc::clone(&self.validity); - // descend further and try with new fields - let writer = self - .leaf_writers - .get(&FieldPath::from_name(name)) - .cloned() - .unwrap_or_else(|| { - if dtype.is_struct() { - // Step into the field path for struct columns - Arc::new(self.descend(&field)) - } else { - // Use fallback for leaf columns - Arc::clone(&self.fallback) - } - }); - let session = session.with_handle(h); - - async move { - // If we have a matching writer, we use it. - // Otherwise, we descend into a new modified one. - // Write validity stream - if index == 0 && is_nullable { - validity - .write_stream(ctx, segment_sink, column_stream, child_eof, &session) - .await - } else { - // Use the underlying writer, otherwise use the fallback writer. - writer - .write_stream(ctx, segment_sink, column_stream, child_eof, &session) - .await - } - } - }) - }) - .collect(); - - let column_layouts = try_join_all(layout_futures).await?; - // TODO(os): transposed stream could count row counts as well, - // This must hold though, all columns must have the same row count of the struct layout - let row_count = column_layouts.first().map(|l| l.row_count()).unwrap_or(0); - Ok(StructLayout::new(row_count, dtype, column_layouts).into_layout()) + // Leaf: hand off to the leaf strategy. + self.leaf + .write_stream(ctx, segment_sink, stream, eof, session) + .await } } @@ -371,11 +255,135 @@ impl LayoutStrategy for TableStrategy { mod tests { use std::sync::Arc; + use vortex_array::ArrayContext; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; use vortex_array::dtype::FieldPath; use vortex_array::field_path; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use crate::LayoutRef; + use crate::LayoutStrategy; + 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::test::SESSION; + + async fn write(strategy: &S, array: ArrayRef) -> VortexResult { + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + strategy + .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) + .await + } + + /// A plain table dispatcher with no overrides. `flat` here is both the validity and leaf + /// strategy. + fn flat_table() -> TableStrategy { + let flat: Arc = Arc::new(FlatLayoutStrategy::default()); + TableStrategy::new(Arc::clone(&flat), flat) + } + + /// The dispatcher shreds a top-level struct into one child per field. + #[tokio::test] + async fn dispatches_struct() -> VortexResult<()> { + let struct_array = StructArray::from_fields( + [ + ("a", buffer![1i32, 2, 3].into_array()), + ("b", buffer![10i32, 20, 30].into_array()), + ] + .as_slice(), + )? + .into_array(); + + let layout = write(&flat_table(), struct_array).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.struct, dtype: {a=i32, b=i32}, children: 2 + ├── a: vortex.flat, dtype: i32, segment: 0 + └── b: vortex.flat, dtype: i32, segment: 1 + "); + 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<()> { + let primitive = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let layout = write(&flat_table(), primitive).await?; + insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0"); + Ok(()) + } + + /// A multi-chunk struct is transposed per field; each column is written by a chunked leaf. + #[tokio::test] + async fn chunked_struct() -> VortexResult<()> { + let validity: Arc = Arc::new(FlatLayoutStrategy::default()); + let chunked_flat: Arc = + Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())); + let dispatcher = TableStrategy::new(validity, chunked_flat); + + let c0 = StructArray::from_fields( + [ + ("a", buffer![1i32, 2].into_array()), + ("b", buffer![10i32, 20].into_array()), + ] + .as_slice(), + )? + .into_array(); + let c1 = StructArray::from_fields( + [ + ("a", buffer![3i32].into_array()), + ("b", buffer![30i32].into_array()), + ] + .as_slice(), + )? + .into_array(); + let dtype = c0.dtype().clone(); + let chunked = ChunkedArray::try_new(vec![c0, c1], dtype)?.into_array(); + + let layout = write(&dispatcher, chunked).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.struct, dtype: {a=i32, b=i32}, children: 2 + ├── a: vortex.chunked, dtype: i32, children: 2 + │ ├── [0]: vortex.flat, dtype: i32, segment: 0 + │ └── [1]: vortex.flat, dtype: i32, segment: 1 + └── b: vortex.chunked, dtype: i32, children: 2 + ├── [0]: vortex.flat, dtype: i32, segment: 2 + └── [1]: vortex.flat, dtype: i32, segment: 3 + "); + Ok(()) + } + + /// A field override on a struct field is honored ahead of the default leaf strategy. + #[tokio::test] + async fn field_override_is_used() -> VortexResult<()> { + let struct_array = StructArray::from_fields( + [ + ("a", buffer![1i32, 2, 3].into_array()), + ("b", buffer![10i32, 20, 30].into_array()), + ] + .as_slice(), + )? + .into_array(); + + let strategy = + flat_table().with_field_writer(field_path!(a), Arc::new(FlatLayoutStrategy::default())); + let layout = write(&strategy, struct_array).await?; + insta::assert_snapshot!(layout.display_tree(), @r" + vortex.struct, dtype: {a=i32, b=i32}, children: 2 + ├── a: vortex.flat, dtype: i32, segment: 0 + └── b: vortex.flat, dtype: i32, segment: 1 + "); + Ok(()) + } #[test] #[should_panic(