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
79
80
81
82
83
84
85
86

use flatbuffers::{FlatBufferBuilder, WIPOffset};
use vortex_buffer::Buffer;
use vortex_flatbuffers::{footer as fb, FlatBufferRoot, WriteFlatBuffer};

use crate::byte_range::ByteRange;
use crate::{LayoutId, CHUNKED_LAYOUT_ID, COLUMNAR_LAYOUT_ID, FLAT_LAYOUT_ID};

#[derive(Debug, Clone)]
pub struct LayoutSpec {
    id: LayoutId,
    buffers: Option<Vec<ByteRange>>,
    children: Option<Vec<LayoutSpec>>,
    row_count: u64,
    metadata: Option<Buffer>,
}

impl LayoutSpec {
    pub fn flat(buffer: ByteRange, row_count: u64) -> Self {
        Self {
            id: FLAT_LAYOUT_ID,
            buffers: Some(vec![buffer]),
            children: None,
            row_count,
            metadata: None,
        }
    }

    /// Create a chunked layout with children.
    ///
    /// has_metadata indicates whether first child is a layout containing metadata about other children.
    pub fn chunked(children: Vec<LayoutSpec>, row_count: u64, metadata: Option<Buffer>) -> Self {
        Self {
            id: CHUNKED_LAYOUT_ID,
            buffers: None,
            children: Some(children),
            row_count,
            metadata,
        }
    }

    pub fn column(children: Vec<LayoutSpec>, row_count: u64) -> Self {
        Self {
            id: COLUMNAR_LAYOUT_ID,
            buffers: None,
            children: Some(children),
            row_count,
            metadata: None,
        }
    }
}

impl FlatBufferRoot for LayoutSpec {}

impl WriteFlatBuffer for LayoutSpec {
    type Target<'a> = fb::Layout<'a>;

    fn write_flatbuffer<'fb>(
        &self,
        fbb: &mut FlatBufferBuilder<'fb>,
    ) -> WIPOffset<Self::Target<'fb>> {
        let buffer_offsets = self.buffers.as_ref().map(|buf| {
            buf.iter()
                .map(|b| fb::Buffer::new(b.begin, b.end))
                .collect::<Vec<_>>()
        });
        let buffers = buffer_offsets.map(|bufs| fbb.create_vector(&bufs));
        let metadata = self.metadata.as_ref().map(|b| fbb.create_vector(b));
        let child_offsets = self.children.as_ref().map(|children| {
            children
                .iter()
                .map(|layout| layout.write_flatbuffer(fbb))
                .collect::<Vec<_>>()
        });
        let children = child_offsets.map(|c| fbb.create_vector(&c));
        fb::Layout::create(
            fbb,
            &fb::LayoutArgs {
                encoding: self.id.0,
                buffers,
                children,
                row_count: self.row_count,
                metadata,
            },
        )
    }
}