vortex_layout/
children.rs

1use std::fmt::{Debug, Formatter};
2use std::sync::Arc;
3
4use flatbuffers::Follow;
5use itertools::Itertools;
6use vortex_dtype::DType;
7use vortex_error::{VortexResult, vortex_bail, vortex_err, vortex_panic};
8use vortex_flatbuffers::{FlatBuffer, layout as fbl};
9
10use crate::segments::SegmentId;
11use crate::{LayoutContext, LayoutRef};
12
13/// Abstract way of accessing the children of a layout.
14///
15/// This allows us to abstract over the lazy flatbuffer-based layouts, as well as the in-memory
16/// layout trees.
17pub trait LayoutChildren: 'static + Send + Sync {
18    fn to_arc(&self) -> Arc<dyn LayoutChildren>;
19
20    fn child(&self, idx: usize, dtype: &DType) -> VortexResult<LayoutRef>;
21
22    fn child_row_count(&self, idx: usize) -> u64;
23
24    fn nchildren(&self) -> usize;
25}
26
27impl Debug for dyn LayoutChildren {
28    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
29        f.debug_struct("LayoutChildren")
30            .field("nchildren", &self.nchildren())
31            .finish()
32    }
33}
34
35impl LayoutChildren for Arc<dyn LayoutChildren> {
36    fn to_arc(&self) -> Arc<dyn LayoutChildren> {
37        self.clone()
38    }
39
40    fn child(&self, idx: usize, dtype: &DType) -> VortexResult<LayoutRef> {
41        self.as_ref().child(idx, dtype)
42    }
43
44    fn child_row_count(&self, idx: usize) -> u64 {
45        self.as_ref().child_row_count(idx)
46    }
47
48    fn nchildren(&self) -> usize {
49        self.as_ref().nchildren()
50    }
51}
52
53/// An implementation of [`LayoutChildren`] for in-memory owned children.
54/// See also [`ViewLayoutChildren`] for lazily deserialized children from flatbuffers.
55#[derive(Clone)]
56pub(crate) struct OwnedLayoutChildren(Vec<LayoutRef>);
57
58impl OwnedLayoutChildren {
59    pub fn layout_children(children: Vec<LayoutRef>) -> Arc<dyn LayoutChildren> {
60        Arc::new(Self(children))
61    }
62}
63
64/// In-memory implementation of [`LayoutChildren`].
65impl LayoutChildren for OwnedLayoutChildren {
66    fn to_arc(&self) -> Arc<dyn LayoutChildren> {
67        Arc::new(self.clone())
68    }
69
70    fn child(&self, idx: usize, dtype: &DType) -> VortexResult<LayoutRef> {
71        if idx >= self.0.len() {
72            vortex_bail!("Child index out of bounds: {} of {}", idx, self.0.len());
73        }
74        let child = &self.0[idx];
75        if child.dtype() != dtype {
76            vortex_panic!("Child dtype mismatch: {} != {}", child.dtype(), dtype);
77        }
78        Ok(child.clone())
79    }
80
81    fn child_row_count(&self, idx: usize) -> u64 {
82        self.0[idx].row_count()
83    }
84
85    fn nchildren(&self) -> usize {
86        self.0.len()
87    }
88}
89
90#[derive(Clone)]
91pub(crate) struct ViewedLayoutChildren {
92    flatbuffer: FlatBuffer,
93    flatbuffer_loc: usize,
94    ctx: LayoutContext,
95}
96
97impl ViewedLayoutChildren {
98    /// Create a new [`ViewedLayoutChildren`] from the given parameters.
99    ///
100    /// # Safety
101    ///
102    /// Assumes the flatbuffer is validated and that the `flatbuffer_loc` is the correct offset
103    pub(super) unsafe fn new_unchecked(
104        flatbuffer: FlatBuffer,
105        flatbuffer_loc: usize,
106        ctx: LayoutContext,
107    ) -> Self {
108        Self {
109            flatbuffer,
110            flatbuffer_loc,
111            ctx,
112        }
113    }
114
115    /// Return the flatbuffer layout message.
116    fn flatbuffer(&self) -> fbl::Layout<'_> {
117        unsafe { fbl::Layout::follow(self.flatbuffer.as_ref(), self.flatbuffer_loc) }
118    }
119}
120
121impl LayoutChildren for ViewedLayoutChildren {
122    fn to_arc(&self) -> Arc<dyn LayoutChildren> {
123        Arc::new(self.clone())
124    }
125
126    fn child(&self, idx: usize, dtype: &DType) -> VortexResult<LayoutRef> {
127        if idx >= self.nchildren() {
128            vortex_bail!("Child index out of bounds: {} of {}", idx, self.nchildren());
129        }
130        let fb_child = self.flatbuffer().children().unwrap_or_default().get(idx);
131
132        let viewed_children = ViewedLayoutChildren {
133            flatbuffer: self.flatbuffer.clone(),
134            flatbuffer_loc: fb_child._tab.loc(),
135            ctx: self.ctx.clone(),
136        };
137        let encoding = self
138            .ctx
139            .lookup_encoding(fb_child.encoding())
140            .ok_or_else(|| vortex_err!("Encoding not found: {}", fb_child.encoding()))?;
141
142        encoding.build(
143            dtype,
144            fb_child.row_count(),
145            fb_child
146                .metadata()
147                .map(|m| m.bytes())
148                .unwrap_or_else(|| &[]),
149            fb_child
150                .segments()
151                .unwrap_or_default()
152                .iter()
153                .map(SegmentId::from)
154                .collect_vec(),
155            &viewed_children,
156        )
157    }
158
159    fn child_row_count(&self, idx: usize) -> u64 {
160        // Efficiently get the row count of the child at the given index, without a full
161        // deserialization.
162        self.flatbuffer()
163            .children()
164            .unwrap_or_default()
165            .get(idx)
166            .row_count()
167    }
168
169    fn nchildren(&self) -> usize {
170        self.flatbuffer().children().unwrap_or_default().len()
171    }
172}
vortex_layout/children.rs

vortex_layout/
children.rs
