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 FlatLayoutOptions {
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 FlatLayoutOptions {
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 FlatLayoutOptions {
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: FlatLayoutOptions,
40    layout: Option<Layout>,
41}
42
43impl FlatLayoutWriter {
44    pub fn new(ctx: ArrayContext, dtype: DType, options: FlatLayoutOptions) -> 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
120    use crate::RowMask;
121    use crate::layouts::flat::writer::FlatLayoutWriter;
122    use crate::segments::test::TestSegments;
123    use crate::writer::LayoutWriterExt;
124
125    // Currently, flat layouts do not force compute stats during write, they only retain
126    // pre-computed stats.
127    #[should_panic]
128    #[test]
129    fn flat_stats() {
130        block_on(async {
131            let ctx = ArrayContext::empty();
132            let mut segments = TestSegments::default();
133            let array = PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid);
134            let layout =
135                FlatLayoutWriter::new(ctx.clone(), array.dtype().clone(), Default::default())
136                    .push_one(&mut segments, array.into_array())
137                    .unwrap();
138
139            let result = layout
140                .reader(Arc::new(segments), ctx)
141                .unwrap()
142                .evaluate_expr(RowMask::new_valid_between(0, layout.row_count()), ident())
143                .await
144                .unwrap();
145
146            assert_eq!(
147                result.statistics().get_as::<bool>(Stat::IsSorted),
148                Some(Precision::Exact(true))
149            );
150        })
151    }
152}
vortex_layout/layouts/flat/writer.rs

vortex_layout/layouts/flat/
writer.rs
