vortex_layout/
children.rs

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

vortex_layout/
children.rs
