vortex_layout/
strategy.rs

1//! Each [`LayoutWriter`] is passed horizontal chunks of a Vortex array one-by-one, and is
2//! eventually asked to return a [`crate::Layout`]. The writers can buffer, re-chunk, flush, or
3//! otherwise manipulate the chunks of data enabling experimentation with different strategies
4//! all while remaining independent of the read code.
5
6use vortex_array::ArrayContext;
7use vortex_dtype::DType;
8use vortex_error::VortexResult;
9
10use crate::layouts::flat::FlatLayout;
11use crate::layouts::flat::writer::FlatLayoutWriter;
12use crate::layouts::struct_::writer::StructLayoutWriter;
13use crate::writer::{LayoutWriter, LayoutWriterExt};
14
15/// A trait for creating new layout writers given a DType.
16pub trait LayoutStrategy: 'static + Send + Sync {
17    fn new_writer(&self, ctx: &ArrayContext, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>>;
18}
19
20/// Implement the [`LayoutStrategy`] trait for the [`FlatLayout`] for easy use.
21impl LayoutStrategy for FlatLayout {
22    fn new_writer(&self, ctx: &ArrayContext, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>> {
23        Ok(FlatLayoutWriter::new(ctx.clone(), dtype.clone(), Default::default()).boxed())
24    }
25}
26
27/// A layout strategy that preserves struct arrays and writes everything else as flat.
28pub struct StructStrategy;
29
30impl LayoutStrategy for StructStrategy {
31    fn new_writer(&self, ctx: &ArrayContext, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>> {
32        if let DType::Struct(..) = dtype {
33            StructLayoutWriter::try_new_with_strategy(ctx, dtype, StructStrategy).map(|w| w.boxed())
34        } else {
35            Ok(FlatLayoutWriter::new(ctx.clone(), dtype.clone(), Default::default()).boxed())
36        }
37    }
38}