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::{IntoLayout, LayoutRef, LayoutStrategy, 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<LayoutRef>,
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        assert_eq!(
71            chunk.dtype(),
72            &self.dtype,
73            "Can't push chunks of the wrong dtype into a LayoutWriter. Pushed {} but expected {}.",
74            chunk.dtype(),
75            self.dtype
76        );
77
78        if self.layout.is_some() {
79            vortex_bail!("FlatLayoutStrategy::push_batch called after finish");
80        }
81        let row_count = chunk.len() as u64;
82        update_stats(&chunk, &self.options.array_stats)?;
83
84        let buffers = chunk.serialize(
85            &self.ctx,
86            &SerializeOptions {
87                offset: 0,
88                include_padding: self.options.include_padding,
89            },
90        )?;
91        let segment_id = segment_writer.put(&buffers);
92
93        self.layout =
94            Some(FlatLayout::new(row_count, self.dtype.clone(), segment_id).into_layout());
95
96        Ok(())
97    }
98
99    fn flush(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
100        Ok(())
101    }
102
103    fn finish(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<LayoutRef> {
104        self.layout
105            .take()
106            .ok_or_else(|| vortex_err!("FlatLayoutStrategy::finish called without push_batch"))
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use std::sync::Arc;
113
114    use futures::executor::block_on;
115    use vortex_array::arrays::PrimitiveArray;
116    use vortex_array::stats::{Precision, Stat};
117    use vortex_array::validity::Validity;
118    use vortex_array::{Array, ArrayContext};
119    use vortex_buffer::buffer;
120    use vortex_expr::ident;
121    use vortex_mask::Mask;
122
123    use crate::layouts::flat::writer::FlatLayoutWriter;
124    use crate::segments::{SegmentSource, TestSegments};
125    use crate::writer::LayoutWriterExt;
126
127    // Currently, flat layouts do not force compute stats during write, they only retain
128    // pre-computed stats.
129    #[should_panic]
130    #[test]
131    fn flat_stats() {
132        block_on(async {
133            let ctx = ArrayContext::empty();
134            let mut segments = TestSegments::default();
135            let array = PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid);
136            let layout =
137                FlatLayoutWriter::new(ctx.clone(), array.dtype().clone(), Default::default())
138                    .push_one(&mut segments, array.to_array())
139                    .unwrap();
140            let segments: Arc<dyn SegmentSource> = Arc::new(segments);
141
142            let result = layout
143                .new_reader(&"".into(), &segments, &ctx)
144                .unwrap()
145                .projection_evaluation(&(0..layout.row_count()), &ident())
146                .unwrap()
147                .invoke(Mask::new_true(layout.row_count().try_into().unwrap()))
148                .await
149                .unwrap();
150
151            assert_eq!(
152                result.statistics().get_as::<bool>(Stat::IsSorted),
153                Some(Precision::Exact(true))
154            );
155        })
156    }
157}
vortex_layout/layouts/flat/writer.rs

vortex_layout/layouts/flat/
writer.rs
