vortex_layout/strategies/
mod.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
//! This is a collection of built-in layout strategies designed to be used in conjunction with one
//! another to develop an overall strategy.
//!
//! Each [`LayoutWriter`] is passed horizontal chunks of a Vortex array one-by-one, and is
//! eventually asked to return a [`LayoutData`]. The writers can buffer, re-chunk, flush, or
//! otherwise manipulate the chunks of data enabling experimentation with different strategies
//! all while remaining independent of the read code.

mod struct_of_chunks;

pub use struct_of_chunks::*;
use vortex_array::ArrayData;
use vortex_dtype::DType;
use vortex_error::VortexResult;

use crate::layouts::flat::writer::FlatLayoutWriter;
use crate::layouts::flat::FlatLayout;
use crate::segments::SegmentWriter;
use crate::LayoutData;

/// A strategy for writing chunks of an array into a layout.
/// FIXME(ngates): move this into writer.rs
pub trait LayoutWriter: Send {
    fn push_chunk(
        &mut self,
        segments: &mut dyn SegmentWriter,
        chunk: ArrayData,
    ) -> VortexResult<()>;

    fn finish(&mut self, segments: &mut dyn SegmentWriter) -> VortexResult<LayoutData>;
}

pub trait LayoutWriterExt: LayoutWriter {
    /// Box the layout writer.
    fn boxed(self) -> Box<dyn LayoutWriter>
    where
        Self: Sized + 'static,
    {
        Box::new(self)
    }

    /// Push a single chunk into the layout writer and return the finished [`LayoutData`].
    fn push_one(
        &mut self,
        segments: &mut dyn SegmentWriter,
        chunk: ArrayData,
    ) -> VortexResult<LayoutData> {
        self.push_chunk(segments, chunk)?;
        self.finish(segments)
    }

    /// Push all chunks of the iterator into the layout writer and return the finished
    /// [`LayoutData`].
    fn push_all<I: IntoIterator<Item = VortexResult<ArrayData>>>(
        &mut self,
        segments: &mut dyn SegmentWriter,
        iter: I,
    ) -> VortexResult<LayoutData> {
        for chunk in iter.into_iter() {
            self.push_chunk(segments, chunk?)?
        }
        self.finish(segments)
    }
}

impl<L: LayoutWriter> LayoutWriterExt for L {}

/// A trait for creating new layout writers given a DType.
pub trait LayoutStrategy: Send {
    fn new_writer(&self, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>>;
}

/// Implement the [`LayoutStrategy`] trait for the [`FlatLayout`] for easy use.
impl LayoutStrategy for FlatLayout {
    fn new_writer(&self, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>> {
        Ok(FlatLayoutWriter::new(dtype.clone()).boxed())
    }
}