vortex_layout/layouts/flat/
writer.rs

1use vortex_array::serde::SerializeOptions;
2use vortex_array::stats::{STATS_TO_WRITE, Stat};
3use vortex_array::{Array, ArrayContext, ArrayRef};
4use vortex_dtype::DType;
5use vortex_error::{VortexResult, vortex_bail, vortex_err};
6
7use crate::layouts::flat::FlatLayout;
8use crate::segments::SegmentWriter;
9use crate::writer::LayoutWriter;
10use crate::{Layout, LayoutStrategy, LayoutVTableRef, LayoutWriterExt};
11
12#[derive(Clone)]
13pub struct FlatLayoutStrategy {
14    /// Stats to preserve when writing arrays
15    pub array_stats: Vec<Stat>,
16    /// Whether to include padding for memory-mapped reads.
17    pub include_padding: bool,
18}
19
20impl Default for FlatLayoutStrategy {
21    fn default() -> Self {
22        Self {
23            array_stats: STATS_TO_WRITE.to_vec(),
24            include_padding: true,
25        }
26    }
27}
28
29impl LayoutStrategy for FlatLayoutStrategy {
30    fn new_writer(&self, ctx: &ArrayContext, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>> {
31        Ok(FlatLayoutWriter::new(ctx.clone(), dtype.clone(), self.clone()).boxed())
32    }
33}
34
35/// Writer for a [`FlatLayout`].
36pub struct FlatLayoutWriter {
37    ctx: ArrayContext,
38    dtype: DType,
39    options: FlatLayoutStrategy,
40    layout: Option<Layout>,
41}
42
43impl FlatLayoutWriter {
44    pub fn new(ctx: ArrayContext, dtype: DType, options: FlatLayoutStrategy) -> Self {
45        Self {
46            ctx,
47            dtype,
48            options,
49            layout: None,
50        }
51    }
52}
53
54fn update_stats(array: &dyn Array, stats: &[Stat]) -> VortexResult<()> {
55    // TODO(ngates): consider whether we want to do this
56    // array.statistics().compute_all(stats)?;
57    array.statistics().retain(stats);
58    for child in array.children() {
59        update_stats(&child, stats)?
60    }
61    Ok(())
62}
63
64impl LayoutWriter for FlatLayoutWriter {
65    fn push_chunk(
66        &mut self,
67        segment_writer: &mut dyn SegmentWriter,
68        chunk: ArrayRef,
69    ) -> VortexResult<()> {
70        if self.layout.is_some() {
71            vortex_bail!("FlatLayoutStrategy::push_batch called after finish");
72        }
73        let row_count = chunk.len() as u64;
74        update_stats(&chunk, &self.options.array_stats)?;
75
76        let buffers = chunk.serialize(
77            &self.ctx,
78            &SerializeOptions {
79                offset: 0,
80                include_padding: self.options.include_padding,
81            },
82        );
83        let segment_id = segment_writer.put(&buffers);
84
85        self.layout = Some(Layout::new_owned(
86            "flat".into(),
87            LayoutVTableRef::new_ref(&FlatLayout),
88            self.dtype.clone(),
89            row_count,
90            vec![segment_id],
91            vec![],
92            None,
93        ));
94        Ok(())
95    }
96
97    fn flush(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
98        Ok(())
99    }
100
101    fn finish(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
102        self.layout
103            .take()
104            .ok_or_else(|| vortex_err!("FlatLayoutStrategy::finish called without push_batch"))
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use std::sync::Arc;
111
112    use futures::executor::block_on;
113    use vortex_array::arrays::PrimitiveArray;
114    use vortex_array::stats::{Precision, Stat};
115    use vortex_array::validity::Validity;
116    use vortex_array::{Array, ArrayContext};
117    use vortex_buffer::buffer;
118    use vortex_expr::ident;
119    use vortex_mask::Mask;
120
121    use crate::ExprEvaluator;
122    use crate::layouts::flat::writer::FlatLayoutWriter;
123    use crate::segments::{SegmentSource, TestSegments};
124    use crate::writer::LayoutWriterExt;
125
126    // Currently, flat layouts do not force compute stats during write, they only retain
127    // pre-computed stats.
128    #[should_panic]
129    #[test]
130    fn flat_stats() {
131        block_on(async {
132            let ctx = ArrayContext::empty();
133            let mut segments = TestSegments::default();
134            let array = PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid);
135            let layout =
136                FlatLayoutWriter::new(ctx.clone(), array.dtype().clone(), Default::default())
137                    .push_one(&mut segments, array.into_array())
138                    .unwrap();
139            let segments: Arc<dyn SegmentSource> = Arc::new(segments);
140
141            let result = layout
142                .reader(&segments, &ctx)
143                .unwrap()
144                .projection_evaluation(&(0..layout.row_count()), &ident())
145                .unwrap()
146                .invoke(Mask::new_true(layout.row_count().try_into().unwrap()))
147                .await
148                .unwrap();
149
150            assert_eq!(
151                result.statistics().get_as::<bool>(Stat::IsSorted),
152                Some(Precision::Exact(true))
153            );
154        })
155    }
156}
vortex_layout/layouts/flat/writer.rs

vortex_layout/layouts/flat/
writer.rs
