Skip to main content

vortex_layout/layouts/
table.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! A configurable writer strategy for tabular data.
5//!
6//! [`TableStrategy`] is a *dispatcher*: it inspects the dtype of the stream it is handed and
7//! routes struct columns to [`StructStrategy`], list columns to [`ListLayoutStrategy`], and
8//! everything else to the configured leaf strategy. Because it hands *itself* (suitably descended)
9//! to those structural writers as the strategy for their children, arbitrarily nested struct/list
10//! trees are written with no manual wiring.
11//!
12//! The dispatcher also owns field-path overrides, letting callers force a specific leaf field —
13//! at any depth — onto a custom strategy.
14
15use std::env;
16use std::sync::Arc;
17use std::sync::LazyLock;
18
19use async_trait::async_trait;
20use vortex_array::dtype::Field;
21use vortex_array::dtype::FieldName;
22use vortex_array::dtype::FieldPath;
23use vortex_error::VortexResult;
24use vortex_session::VortexSession;
25use vortex_utils::aliases::hash_map::HashMap;
26use vortex_utils::aliases::hash_set::HashSet;
27
28use crate::LayoutRef;
29use crate::LayoutStrategy;
30use crate::LayoutWriterContext;
31use crate::layouts::list::writer::ListLayoutStrategy;
32use crate::layouts::struct_::StructStrategy;
33use crate::segments::SegmentSinkRef;
34use crate::sequence::SendableSequentialStream;
35use crate::sequence::SequencePointer;
36
37/// Whether [`TableStrategy`] writes list fields using a [`ListLayoutStrategy`] by
38/// default. Disabled unless the environment variable `VORTEX_EXPERIMENTAL_LIST_LAYOUT`
39/// is set to `1`.
40///
41/// [`ListLayoutStrategy`]: ListLayoutStrategy
42pub fn use_experimental_list_layout() -> bool {
43    static USE_EXPERIMENTAL_LIST_LAYOUT: LazyLock<bool> =
44        LazyLock::new(|| env::var("VORTEX_EXPERIMENTAL_LIST_LAYOUT").is_ok_and(|v| v == "1"));
45    *USE_EXPERIMENTAL_LIST_LAYOUT
46}
47
48type ListLayoutFactory = Arc<dyn Fn(ListLayoutStrategy) -> Arc<dyn LayoutStrategy> + Send + Sync>;
49
50/// A configurable strategy for writing nested tabular data, dispatching each (sub)stream to the
51/// structural writer for its dtype.
52///
53/// Dispatch rules, applied to the dtype of the stream handed to [`write_stream`]:
54/// - **struct** → [`StructStrategy`], with each field written by its override (if any) or by a
55///   descended copy of this dispatcher.
56/// - **list** → [`ListLayoutStrategy`], with `elements` written by a descended copy of this
57///   dispatcher (so nested structs/lists recurse) and `offsets`/`validity` by the leaf/validity
58///   strategies. Gated: only when list decomposition is enabled via
59///   [`with_list_layout`][Self::with_list_layout] (off by default); otherwise a list falls through
60///   to the leaf strategy.
61/// - **anything else** → the leaf strategy.
62///
63/// [`write_stream`]: LayoutStrategy::write_stream
64pub struct TableStrategy {
65    /// A set of field-path overrides, e.g. to force one column to be compact-compressed. Keys are
66    /// paths relative to the level this dispatcher sits at.
67    leaf_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
68    /// The writer for any validity arrays that may be present, at any level of the tree.
69    validity: Arc<dyn LayoutStrategy>,
70    /// The writer for leaf fields, i.e. anything that is not a struct.
71    leaf: Arc<dyn LayoutStrategy>,
72    /// Optional factory applied to each dynamically constructed [`ListLayoutStrategy`].
73    /// Its presence also enables list decomposition.
74    ///
75    /// [`ListLayoutStrategy`]: ListLayoutStrategy
76    list_layout_factory: Option<ListLayoutFactory>,
77}
78
79impl TableStrategy {
80    /// Create a new dispatcher with the given `validity` strategy and `fallback` leaf strategy and
81    /// no overrides.
82    ///
83    /// Additional per-field overrides can be configured with
84    /// [`with_field_writer`][Self::with_field_writer].
85    ///
86    /// ## Example
87    ///
88    /// ```ignore
89    /// # use std::sync::Arc;
90    /// # use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
91    /// # use vortex_layout::layouts::table::TableStrategy;
92    ///
93    /// // Build a write strategy that does not compress validity or any leaf fields.
94    /// let flat = Arc::new(FlatLayoutStrategy::default());
95    ///
96    /// let strategy = TableStrategy::new(Arc::<FlatLayoutStrategy>::clone(&flat), Arc::<FlatLayoutStrategy>::clone(&flat));
97    /// ```
98    pub fn new(validity: Arc<dyn LayoutStrategy>, fallback: Arc<dyn LayoutStrategy>) -> Self {
99        Self {
100            leaf_writers: Default::default(),
101            validity,
102            leaf: fallback,
103            list_layout_factory: None,
104        }
105    }
106
107    /// Add a custom write strategy for the given leaf field.
108    ///
109    /// ## Example
110    ///
111    /// ```ignore
112    /// # use std::sync::Arc;
113    /// # use vortex_array::dtype::{field_path, Field, FieldPath};
114    /// # use vortex_btrblocks::BtrBlocksCompressor;
115    /// # use vortex_layout::layouts::compressed::CompressingStrategy;
116    /// # use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
117    /// # use vortex_layout::layouts::table::TableStrategy;
118    ///
119    /// // A strategy for compressing data using the balanced BtrBlocks compressor.
120    /// let compress =
121    ///     CompressingStrategy::new(FlatLayoutStrategy::default(), BtrBlocksCompressor::default());
122    ///
123    /// // Our combined strategy uses no compression for validity buffers, BtrBlocks compression
124    /// // for most columns, and stores a nested binary column uncompressed (flat) because it
125    /// // is pre-compressed or never filtered on.
126    /// let strategy = TableStrategy::new(
127    ///         Arc::new(FlatLayoutStrategy::default()),
128    ///         Arc::new(compress),
129    ///     )
130    ///     .with_field_writer(
131    ///         field_path!(request.body.bytes),
132    ///         Arc::new(FlatLayoutStrategy::default()),
133    ///     );
134    /// ```
135    pub fn with_field_writer(
136        mut self,
137        field_path: impl Into<FieldPath>,
138        writer: Arc<dyn LayoutStrategy>,
139    ) -> Self {
140        self.leaf_writers
141            .insert(self.validate_path(field_path.into()), writer);
142        self
143    }
144
145    /// Set writers for several fields at once.
146    ///
147    /// See also: [`with_field_writer`][Self::with_field_writer].
148    pub fn with_field_writers(
149        mut self,
150        writers: impl IntoIterator<Item = (FieldPath, Arc<dyn LayoutStrategy>)>,
151    ) -> Self {
152        for (field_path, strategy) in writers {
153            self.leaf_writers
154                .insert(self.validate_path(field_path), strategy);
155        }
156        self
157    }
158
159    /// Override the default strategy for leaf columns that don't have overrides.
160    pub fn with_default_strategy(mut self, default: Arc<dyn LayoutStrategy>) -> Self {
161        self.leaf = default;
162        self
163    }
164
165    /// Override the strategy for compressing struct validity at all levels of the schema tree.
166    pub fn with_validity_strategy(mut self, validity: Arc<dyn LayoutStrategy>) -> Self {
167        self.validity = validity;
168        self
169    }
170
171    /// Enable writing list fields with [`ListLayoutStrategy`].
172    ///
173    /// **Note**: this is an unstable and experimental layout that is expected to change.
174    /// Using it may lead to unreadable files in the future.
175    pub fn with_list_layout(self) -> Self {
176        self.with_list_layout_factory(|strategy| Arc::new(strategy))
177    }
178
179    /// Enable writing list fields with [`ListLayoutStrategy`] and wrap each list writer. This
180    /// allows repartitioning or zoning to operate in the list's outer-row space before shredding.
181    ///
182    /// **Note**: this is an unstable and experimental layout that is expected to change.
183    /// Using it may lead to unreadable files in the future.
184    pub fn with_list_layout_factory(
185        mut self,
186        factory: impl Fn(ListLayoutStrategy) -> Arc<dyn LayoutStrategy> + Send + Sync + 'static,
187    ) -> Self {
188        self.list_layout_factory = Some(Arc::new(factory));
189        self
190    }
191}
192
193impl TableStrategy {
194    /// Build the [`StructStrategy`] used to write a struct-typed stream at this level.
195    ///
196    /// Each field that has an override (or a deeper override beneath it) is resolved up front;
197    /// every other field falls through to a clean descended dispatcher.
198    fn struct_strategy(&self) -> StructStrategy {
199        let mut field_writers: HashMap<FieldName, Arc<dyn LayoutStrategy>> = HashMap::default();
200
201        // The distinct named first-segments of our override paths are the only fields that need
202        // anything other than the default dispatcher.
203        let mut named_first: HashSet<FieldName> = HashSet::default();
204        for path in self.leaf_writers.keys() {
205            if let Some(Field::Name(name)) = path.parts().first() {
206                named_first.insert(name.clone());
207            }
208        }
209
210        for name in named_first {
211            // `validate_path` forbids overlapping overrides, so a name has *either* an exact
212            // single-segment override *or* deeper overrides, never both.
213            let writer = match self.leaf_writers.get(&FieldPath::from_name(name.clone())) {
214                Some(exact) => Arc::clone(exact),
215                None => {
216                    Arc::new(self.descend(&Field::Name(name.clone()))) as Arc<dyn LayoutStrategy>
217                }
218            };
219            field_writers.insert(name, writer);
220        }
221
222        StructStrategy::new(Arc::clone(&self.validity), Arc::new(self.descend_clean()))
223            .with_field_writers(field_writers)
224    }
225
226    /// Build the [`ListLayoutStrategy`] used to write a list field stream at this level.
227    ///
228    /// The `elements` sub-column is routed back through a clean descended dispatcher so nested
229    /// structs/lists recurse; `offsets` go straight to the leaf (they are always a primitive
230    /// column); and `validity` uses the shared validity strategy.
231    fn list_strategy(&self) -> Option<Arc<dyn LayoutStrategy>> {
232        let factory = self.list_layout_factory.as_ref()?;
233        let list_layout = ListLayoutStrategy::default()
234            .with_elements(Arc::new(self.descend_clean()))
235            .with_offsets(Arc::clone(&self.leaf))
236            .with_validity(Arc::clone(&self.validity))
237            .with_fallback(Arc::clone(&self.leaf));
238        Some(factory(list_layout))
239    }
240
241    /// Descend into a subfield, retaining only the overrides that apply beneath it (rebased to be
242    /// relative to the child).
243    fn descend(&self, field: &Field) -> Self {
244        let mut new_writers = HashMap::with_capacity(self.leaf_writers.len());
245
246        for (field_path, strategy) in &self.leaf_writers {
247            if field_path.parts().first() == Some(field)
248                && let Some(subpath) = field_path.clone().step_into()
249                && !subpath.is_root()
250            {
251                new_writers.insert(subpath, Arc::clone(strategy));
252            }
253        }
254
255        Self {
256            leaf_writers: new_writers,
257            validity: Arc::clone(&self.validity),
258            leaf: Arc::clone(&self.leaf),
259            list_layout_factory: self.list_layout_factory.clone(),
260        }
261    }
262
263    /// A copy of this dispatcher with no overrides, used as the default child strategy for fields
264    /// that carry no override.
265    fn descend_clean(&self) -> Self {
266        Self {
267            leaf_writers: HashMap::default(),
268            validity: Arc::clone(&self.validity),
269            leaf: Arc::clone(&self.leaf),
270            list_layout_factory: self.list_layout_factory.clone(),
271        }
272    }
273
274    fn validate_path(&self, path: FieldPath) -> FieldPath {
275        assert!(
276            !path.is_root(),
277            "Do not set override as a root strategy, instead set the default strategy"
278        );
279
280        // Validate that the field path does not conflict with any overrides
281        // that we've added by overlapping.
282        for field_path in self.leaf_writers.keys() {
283            assert!(
284                !path.overlap(field_path),
285                "Override for field_path {path} conflicts with existing override for {field_path}"
286            );
287        }
288
289        path
290    }
291}
292
293/// Dispatches each stream to the structural writer for its dtype.
294#[async_trait]
295impl LayoutStrategy for TableStrategy {
296    async fn write_stream(
297        &self,
298        ctx: LayoutWriterContext,
299        segment_sink: SegmentSinkRef,
300        stream: SendableSequentialStream,
301        eof: SequencePointer,
302        session: &VortexSession,
303    ) -> VortexResult<LayoutRef> {
304        let dtype = stream.dtype().clone();
305
306        if dtype.is_struct() {
307            return self
308                .struct_strategy()
309                .write_stream(ctx, segment_sink, stream, eof, session)
310                .await;
311        }
312
313        if dtype.is_list()
314            && let Some(list_strategy) = self.list_strategy()
315        {
316            return list_strategy
317                .write_stream(ctx, segment_sink, stream, eof, session)
318                .await;
319        }
320
321        // Leaf: hand off to the leaf strategy.
322        self.leaf
323            .write_stream(ctx, segment_sink, stream, eof, session)
324            .await
325    }
326}
327
328#[cfg(test)]
329mod tests {
330    use std::num::NonZeroUsize;
331    use std::sync::Arc;
332    use std::task::Poll;
333
334    use vortex_array::ArrayContext;
335    use vortex_array::ArrayRef;
336    use vortex_array::IntoArray;
337    use vortex_array::arrays::BoolArray;
338    use vortex_array::arrays::ChunkedArray;
339    use vortex_array::arrays::ListArray;
340    use vortex_array::arrays::PrimitiveArray;
341    use vortex_array::arrays::StructArray;
342    use vortex_array::dtype::DType;
343    use vortex_array::dtype::FieldPath;
344    use vortex_array::dtype::Nullability;
345    use vortex_array::dtype::PType;
346    use vortex_array::dtype::StructFields;
347    use vortex_array::field_path;
348    use vortex_array::validity::Validity;
349    use vortex_buffer::buffer;
350    use vortex_error::VortexExpect;
351    use vortex_error::VortexResult;
352    use vortex_io::session::RuntimeSessionExt;
353
354    use crate::LayoutRef;
355    use crate::LayoutStrategy;
356    use crate::layouts::chunked::writer::ChunkedLayoutStrategy;
357    use crate::layouts::flat::writer::FlatLayoutStrategy;
358    use crate::layouts::list::List;
359    use crate::layouts::repartition::RepartitionStrategy;
360    use crate::layouts::repartition::RepartitionWriterOptions;
361    use crate::layouts::table::TableStrategy;
362    use crate::layouts::zoned::Zoned;
363    use crate::layouts::zoned::writer::ZonedLayoutOptions;
364    use crate::layouts::zoned::writer::ZonedStrategy;
365    use crate::segments::TestSegments;
366    use crate::sequence::SequenceId;
367    use crate::sequence::SequentialArrayStreamExt;
368    use crate::sequence::SequentialStreamAdapter;
369    use crate::sequence::SequentialStreamExt;
370    use crate::test::SESSION;
371    use crate::test::new_session;
372
373    async fn write<S: LayoutStrategy>(strategy: &S, array: ArrayRef) -> VortexResult<LayoutRef> {
374        let segments = Arc::new(TestSegments::default());
375        let (ptr, eof) = SequenceId::root().split();
376        let stream = array.to_array_stream().sequenced(ptr);
377        let session = new_session().with_tokio();
378        strategy
379            .write_stream(
380                ArrayContext::empty().into(),
381                segments,
382                stream,
383                eof,
384                &session,
385            )
386            .await
387    }
388
389    /// A plain table dispatcher with no overrides. `flat` here is both the validity and leaf
390    /// strategy.
391    fn flat_table() -> TableStrategy {
392        let flat: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
393        TableStrategy::new(Arc::clone(&flat), flat)
394    }
395
396    /// The dispatcher shreds a top-level struct into one child per field.
397    #[tokio::test]
398    async fn dispatches_struct() -> VortexResult<()> {
399        let struct_array = StructArray::from_fields(
400            [
401                ("a", buffer![1i32, 2, 3].into_array()),
402                ("b", buffer![10i32, 20, 30].into_array()),
403            ]
404            .as_slice(),
405        )?
406        .into_array();
407
408        let layout = write(&flat_table(), struct_array).await?;
409        insta::assert_snapshot!(layout.display_tree(), @r"
410        vortex.struct, dtype: {a=i32, b=i32}, children: 2
411        ├── a: vortex.flat, dtype: i32, segment: 0
412        └── b: vortex.flat, dtype: i32, segment: 1
413        ");
414        Ok(())
415    }
416
417    /// A `list<list<i32>>` column: the dispatcher recurses into itself so the outer list's
418    /// `elements` are decomposed as a nested `ListLayout`.
419    #[tokio::test]
420    async fn dispatches_nested_list() -> VortexResult<()> {
421        let inner = ListArray::try_new(
422            buffer![1i32, 2, 3, 4, 5, 6].into_array(),
423            buffer![0u32, 2, 5, 5, 6].into_array(),
424            Validity::NonNullable,
425        )?
426        .into_array();
427        let outer = ListArray::try_new(
428            inner,
429            buffer![0u32, 2, 4].into_array(),
430            Validity::NonNullable,
431        )?
432        .into_array();
433
434        let layout = write(&flat_table().with_list_layout(), outer).await?;
435        insta::assert_snapshot!(layout.display_tree(), @r"
436        vortex.list, dtype: list(list(i32)), children: 2
437        ├── elements: vortex.list, dtype: list(i32), children: 2
438        │   ├── elements: vortex.flat, dtype: i32, segment: 1
439        │   └── offsets: vortex.flat, dtype: u64, segment: 2
440        └── offsets: vortex.flat, dtype: u64, segment: 0
441        ");
442        Ok(())
443    }
444
445    /// A `struct<{ items: list<struct<{a,b}>>? }>` column: list decomposition recurses into struct
446    /// decomposition for the elements, and a nullable list writes a validity child.
447    #[tokio::test]
448    async fn dispatches_struct_list_struct() -> VortexResult<()> {
449        let inner_struct = StructArray::from_fields(
450            [
451                ("a", buffer![1i32, 2, 3, 4, 5].into_array()),
452                ("b", buffer![10i32, 20, 30, 40, 50].into_array()),
453            ]
454            .as_slice(),
455        )?
456        .into_array();
457        let items = ListArray::try_new(
458            inner_struct,
459            buffer![0u32, 2, 5, 5].into_array(),
460            Validity::Array(BoolArray::from_iter([true, false, true]).into_array()),
461        )?
462        .into_array();
463        let st = StructArray::from_fields([("items", items)].as_slice())?.into_array();
464
465        let layout = write(&flat_table().with_list_layout(), st).await?;
466        insta::assert_snapshot!(layout.display_tree(), @r"
467        vortex.struct, dtype: {items=list({a=i32, b=i32})?}, children: 1
468        └── items: vortex.list, dtype: list({a=i32, b=i32})?, children: 3
469            ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2
470            │   ├── a: vortex.flat, dtype: i32, segment: 2
471            │   └── b: vortex.flat, dtype: i32, segment: 3
472            ├── offsets: vortex.flat, dtype: u64, segment: 0
473            └── validity: vortex.flat, dtype: bool, segment: 1
474        ");
475        Ok(())
476    }
477
478    /// A multi-chunk `list<i32>` written with a chunked leaf: each sub-column (`elements`,
479    /// `offsets`) becomes its own `ChunkedLayout`, so elements are chunked independently of rows.
480    /// This is the "list-of-chunkeds" topology top-level decomposition unlocks.
481    #[tokio::test]
482    async fn dispatches_chunked_list() -> VortexResult<()> {
483        let chunk0 = ListArray::try_new(
484            buffer![1i32, 2, 3].into_array(),
485            buffer![0u32, 2, 3].into_array(),
486            Validity::NonNullable,
487        )?
488        .into_array();
489        let chunk1 = ListArray::try_new(
490            buffer![4i32, 5, 6, 7].into_array(),
491            buffer![0u32, 1, 4].into_array(),
492            Validity::NonNullable,
493        )?
494        .into_array();
495        let dtype = chunk0.dtype().clone();
496        let chunked = ChunkedArray::try_new(vec![chunk0, chunk1], dtype)?.into_array();
497
498        let flat: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
499        let dispatcher = TableStrategy::new(
500            Arc::clone(&flat),
501            Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())),
502        )
503        .with_list_layout();
504        let layout = write(&dispatcher, chunked).await?;
505        insta::assert_snapshot!(layout.display_tree(), @r"
506        vortex.list, dtype: list(i32), children: 2
507        ├── elements: vortex.chunked, dtype: i32, children: 2
508        │   ├── [0]: vortex.flat, dtype: i32, segment: 0
509        │   └── [1]: vortex.flat, dtype: i32, segment: 1
510        └── offsets: vortex.chunked, dtype: u64, children: 2
511            ├── [0]: vortex.flat, dtype: u64, segment: 2
512            └── [1]: vortex.flat, dtype: u64, segment: 3
513        ");
514        Ok(())
515    }
516
517    /// A wrapper can repartition and zone lists in outer-row space before decomposition.
518    #[tokio::test]
519    async fn wraps_list_strategy_before_decomposition() -> VortexResult<()> {
520        let list = ListArray::try_new(
521            PrimitiveArray::from_iter(0..9_i32).into_array(),
522            PrimitiveArray::from_iter(0..=9_u32).into_array(),
523            Validity::NonNullable,
524        )?
525        .into_array();
526
527        let row_block_size = NonZeroUsize::new(4).vortex_expect("4 is non-zero");
528        let flat: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
529        let stats = Arc::clone(&flat);
530        let chunked: Arc<dyn LayoutStrategy> =
531            Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()));
532        let dispatcher = TableStrategy::new(Arc::clone(&flat), chunked).with_list_layout_factory(
533            move |list_layout| {
534                let zoned = ZonedStrategy::new(
535                    list_layout,
536                    Arc::clone(&stats),
537                    ZonedLayoutOptions {
538                        block_size: row_block_size,
539                        ..Default::default()
540                    },
541                );
542                Arc::new(RepartitionStrategy::new(
543                    zoned,
544                    RepartitionWriterOptions {
545                        block_size_minimum: 0,
546                        block_len_multiple: row_block_size.get(),
547                        block_size_target: None,
548                        canonicalize: false,
549                    },
550                )) as Arc<dyn LayoutStrategy>
551            },
552        );
553
554        let layout = write(&dispatcher, list).await?;
555        let zoned = layout.as_::<Zoned>();
556        assert_eq!(zoned.zone_len(), 4);
557        assert_eq!(zoned.nzones(), 3);
558
559        let data = layout
560            .slot(0)?
561            .vortex_expect("ZonedLayout always has a data child");
562        assert!(data.is::<List>());
563        assert_eq!(data.row_count(), 9);
564        Ok(())
565    }
566
567    /// A non-struct stream is not shredded; it is handed straight to the leaf strategy.
568    #[tokio::test]
569    async fn non_struct_input_uses_leaf() -> VortexResult<()> {
570        let primitive = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
571        let layout = write(&flat_table(), primitive).await?;
572        insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0");
573        Ok(())
574    }
575
576    /// A multi-chunk struct is transposed per field; each column is written by a chunked leaf.
577    #[tokio::test]
578    async fn chunked_struct() -> VortexResult<()> {
579        let validity: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
580        let chunked_flat: Arc<dyn LayoutStrategy> =
581            Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()));
582        let dispatcher = TableStrategy::new(validity, chunked_flat);
583
584        let c0 = StructArray::from_fields(
585            [
586                ("a", buffer![1i32, 2].into_array()),
587                ("b", buffer![10i32, 20].into_array()),
588            ]
589            .as_slice(),
590        )?
591        .into_array();
592        let c1 = StructArray::from_fields(
593            [
594                ("a", buffer![3i32].into_array()),
595                ("b", buffer![30i32].into_array()),
596            ]
597            .as_slice(),
598        )?
599        .into_array();
600        let dtype = c0.dtype().clone();
601        let chunked = ChunkedArray::try_new(vec![c0, c1], dtype)?.into_array();
602
603        let layout = write(&dispatcher, chunked).await?;
604        insta::assert_snapshot!(layout.display_tree(), @r"
605        vortex.struct, dtype: {a=i32, b=i32}, children: 2
606        ├── a: vortex.chunked, dtype: i32, children: 2
607        │   ├── [0]: vortex.flat, dtype: i32, segment: 0
608        │   └── [1]: vortex.flat, dtype: i32, segment: 1
609        └── b: vortex.chunked, dtype: i32, children: 2
610            ├── [0]: vortex.flat, dtype: i32, segment: 2
611            └── [1]: vortex.flat, dtype: i32, segment: 3
612        ");
613        Ok(())
614    }
615
616    /// A field override on a struct field is honored ahead of the default leaf strategy.
617    #[tokio::test]
618    async fn field_override_is_used() -> VortexResult<()> {
619        let struct_array = StructArray::from_fields(
620            [
621                ("a", buffer![1i32, 2, 3].into_array()),
622                ("b", buffer![10i32, 20, 30].into_array()),
623            ]
624            .as_slice(),
625        )?
626        .into_array();
627
628        let strategy =
629            flat_table().with_field_writer(field_path!(a), Arc::new(FlatLayoutStrategy::default()));
630        let layout = write(&strategy, struct_array).await?;
631        insta::assert_snapshot!(layout.display_tree(), @r"
632        vortex.struct, dtype: {a=i32, b=i32}, children: 2
633        ├── a: vortex.flat, dtype: i32, segment: 0
634        └── b: vortex.flat, dtype: i32, segment: 1
635        ");
636        Ok(())
637    }
638
639    #[test]
640    #[should_panic(
641        expected = "Override for field_path $a.$b conflicts with existing override for $a.$b.$c"
642    )]
643    fn test_overlapping_paths_fail() {
644        let flat = Arc::new(FlatLayoutStrategy::default());
645
646        // Success
647        let path = TableStrategy::new(
648            Arc::<FlatLayoutStrategy>::clone(&flat),
649            Arc::<FlatLayoutStrategy>::clone(&flat),
650        )
651        .with_field_writer(field_path!(a.b.c), Arc::<FlatLayoutStrategy>::clone(&flat));
652
653        // Should panic right here.
654        let _path = path.with_field_writer(field_path!(a.b), flat);
655    }
656
657    #[test]
658    #[should_panic(
659        expected = "Do not set override as a root strategy, instead set the default strategy"
660    )]
661    fn test_root_override() {
662        let flat = Arc::new(FlatLayoutStrategy::default());
663        let _strategy = TableStrategy::new(
664            Arc::<FlatLayoutStrategy>::clone(&flat),
665            Arc::<FlatLayoutStrategy>::clone(&flat),
666        )
667        .with_field_writer(FieldPath::root(), flat);
668    }
669
670    #[tokio::test]
671    #[should_panic(expected = "panic while transposing table stream")]
672    async fn table_fanout_panic_propagates() {
673        let ctx = ArrayContext::empty();
674        let segments = Arc::new(TestSegments::default());
675        let (_, eof) = SequenceId::root().split();
676        let dtype = DType::Struct(
677            StructFields::from_iter([(
678                "a",
679                DType::Primitive(PType::I32, Nullability::NonNullable),
680            )]),
681            Nullability::NonNullable,
682        );
683        let stream =
684            futures::stream::poll_fn(|_| -> Poll<Option<VortexResult<(SequenceId, ArrayRef)>>> {
685                panic!("panic while transposing table stream");
686            });
687        let strategy = TableStrategy::new(
688            Arc::new(FlatLayoutStrategy::default()),
689            Arc::new(FlatLayoutStrategy::default()),
690        );
691
692        strategy
693            .write_stream(
694                ctx.into(),
695                segments,
696                SequentialStreamAdapter::new(dtype, stream).sendable(),
697                eof,
698                &SESSION,
699            )
700            .await
701            .unwrap();
702    }
703}