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::writer::FlatLayoutWriter;
11use crate::layouts::struct_::writer::StructLayoutWriter;
12use crate::writer::{LayoutWriter, LayoutWriterExt};
13
14/// A trait for creating new layout writers given a DType.
15pub trait LayoutStrategy: 'static + Send + Sync {
16    fn new_writer(&self, ctx: &ArrayContext, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>>;
17}
18
19/// A layout strategy that preserves struct arrays and writes everything else as flat.
20pub struct StructStrategy;
21
22impl LayoutStrategy for StructStrategy {
23    fn new_writer(&self, ctx: &ArrayContext, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>> {
24        if let DType::Struct(..) = dtype {
25            StructLayoutWriter::try_new_with_strategy(ctx, dtype, StructStrategy).map(|w| w.boxed())
26        } else {
27            Ok(FlatLayoutWriter::new(ctx.clone(), dtype.clone(), Default::default()).boxed())
28        }
29    }
30}