vortex_layout/
writer.rs

1use vortex_array::ArrayRef;
2use vortex_error::VortexResult;
3
4use crate::Layout;
5use crate::segments::SegmentWriter;
6
7/// A strategy for writing chunks of an array into a layout.
8// [layout writer]
9pub trait LayoutWriter: Send {
10    /// Push a chunk into the layout writer.
11    fn push_chunk(
12        &mut self,
13        segment_writer: &mut dyn SegmentWriter,
14        chunk: ArrayRef,
15    ) -> VortexResult<()>;
16
17    /// Flush any buffered chunks.
18    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()>;
19
20    /// Write any final data (e.g. stats) and return the finished [`Layout`].
21    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout>;
22}
23// [layout writer]
24
25pub trait LayoutWriterExt: LayoutWriter {
26    /// Box the layout writer.
27    fn boxed(self) -> Box<dyn LayoutWriter>
28    where
29        Self: Sized + 'static,
30    {
31        Box::new(self)
32    }
33
34    /// Push a single chunk into the layout writer and return the finished [`Layout`].
35    fn push_one(
36        &mut self,
37        segment_writer: &mut dyn SegmentWriter,
38        chunk: ArrayRef,
39    ) -> VortexResult<Layout> {
40        self.push_chunk(segment_writer, chunk)?;
41        self.flush(segment_writer)?;
42        self.finish(segment_writer)
43    }
44
45    /// Push all chunks of the iterator into the layout writer and return the finished
46    /// [`Layout`].
47    fn push_all<I: IntoIterator<Item = VortexResult<ArrayRef>>>(
48        &mut self,
49        segment_writer: &mut dyn SegmentWriter,
50        iter: I,
51    ) -> VortexResult<Layout> {
52        for chunk in iter.into_iter() {
53            self.push_chunk(segment_writer, chunk?)?
54        }
55        self.flush(segment_writer)?;
56        self.finish(segment_writer)
57    }
58}
59
60impl<L: LayoutWriter + ?Sized> LayoutWriterExt for L {}