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