vortex_layout/
layout.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::fmt::Debug;
6use std::fmt::Display;
7use std::fmt::Formatter;
8use std::ops::Deref;
9use std::sync::Arc;
10
11use itertools::Itertools;
12use vortex_array::SerializeMetadata;
13use vortex_array::dtype::DType;
14use vortex_array::dtype::FieldName;
15use vortex_error::VortexExpect;
16use vortex_error::VortexResult;
17use vortex_session::VortexSession;
18use vortex_session::registry::Id;
19
20use crate::LayoutReaderContext;
21use crate::LayoutReaderRef;
22use crate::children::LayoutChildren;
23use crate::display::DisplayLayoutTree;
24use crate::display::display_tree_with_segment_sizes;
25use crate::segments::SegmentId;
26use crate::segments::SegmentSource;
27use crate::vtable::LayoutRef;
28use crate::vtable::VTable;
29
30/// A unique identifier for a layout encoding.
31pub type LayoutId = Id;
32
33/// Pieces used to construct a typed layout.
34pub struct LayoutParts<V: VTable> {
35    vtable: V,
36    dtype: DType,
37    row_count: u64,
38    segment_ids: Vec<SegmentId>,
39    children: Arc<dyn LayoutChildren>,
40    data: V::LayoutData,
41}
42
43impl<V: VTable> LayoutParts<V> {
44    /// Create layout parts from common fields and vtable-specific data.
45    pub fn new(
46        vtable: V,
47        dtype: DType,
48        row_count: u64,
49        segment_ids: Vec<SegmentId>,
50        children: Arc<dyn LayoutChildren>,
51        data: V::LayoutData,
52    ) -> Self {
53        Self {
54            vtable,
55            dtype,
56            row_count,
57            segment_ids,
58            children,
59            data,
60        }
61    }
62
63    /// Convert these parts into a typed layout.
64    pub fn into_typed(self) -> Layout<V> {
65        Layout::from_parts(self)
66    }
67
68    /// Erase these parts into a layout reference.
69    pub fn into_layout(self) -> LayoutRef {
70        self.into_typed().into_layout()
71    }
72}
73
74/// A typed layout node.
75pub struct Layout<V: VTable> {
76    inner: Arc<LayoutInner<V>>,
77}
78
79struct LayoutInner<V: VTable> {
80    vtable: V,
81    dtype: DType,
82    row_count: u64,
83    segment_ids: Vec<SegmentId>,
84    children: Arc<dyn LayoutChildren>,
85    data: V::LayoutData,
86}
87
88impl<V: VTable> Layout<V> {
89    /// Construct a layout from explicit parts.
90    pub fn from_parts(parts: LayoutParts<V>) -> Self {
91        Self {
92            inner: Arc::new(LayoutInner {
93                vtable: parts.vtable,
94                dtype: parts.dtype,
95                row_count: parts.row_count,
96                segment_ids: parts.segment_ids,
97                children: parts.children,
98                data: parts.data,
99            }),
100        }
101    }
102
103    /// Returns the vtable.
104    pub fn vtable(&self) -> &V {
105        &self.inner.vtable
106    }
107
108    /// Returns layout-specific data.
109    pub fn data(&self) -> &V::LayoutData {
110        &self.inner.data
111    }
112
113    /// Returns the logical dtype.
114    pub fn dtype(&self) -> &DType {
115        &self.inner.dtype
116    }
117
118    /// Returns the number of rows.
119    pub fn row_count(&self) -> u64 {
120        self.inner.row_count
121    }
122
123    /// Returns directly referenced segment IDs.
124    pub fn segment_ids(&self) -> &[SegmentId] {
125        &self.inner.segment_ids
126    }
127
128    /// Returns the child adapter.
129    pub fn children(&self) -> &Arc<dyn LayoutChildren> {
130        &self.inner.children
131    }
132
133    /// Returns the number of children.
134    pub fn nchildren(&self) -> usize {
135        self.inner.children.nchildren()
136    }
137
138    /// Materialize child `idx` with its expected dtype.
139    pub fn child(&self, idx: usize) -> VortexResult<LayoutRef> {
140        self.inner.children.child(idx, &V::child_dtype(self, idx)?)
141    }
142
143    /// Returns a child's serialized row count without materializing it.
144    pub fn child_row_count(&self, idx: usize) -> u64 {
145        self.inner.children.child_row_count(idx)
146    }
147
148    /// Returns a child's relationship to this layout.
149    pub fn child_type(&self, idx: usize) -> LayoutChildType {
150        V::child_type(self, idx)
151    }
152
153    /// Erase this typed layout into a shared layout reference.
154    pub fn to_layout(&self) -> LayoutRef {
155        self.clone().into_layout()
156    }
157
158    /// Erase this typed layout into a shared layout reference.
159    pub fn into_layout(self) -> LayoutRef {
160        Arc::new(self)
161    }
162
163    /// Construct a reader for this layout.
164    pub fn new_reader(
165        &self,
166        name: Arc<str>,
167        segment_source: Arc<dyn SegmentSource>,
168        session: &VortexSession,
169        ctx: &LayoutReaderContext,
170    ) -> VortexResult<LayoutReaderRef> {
171        V::new_reader(self, name, segment_source, session, ctx)
172    }
173}
174
175impl<V: VTable> Clone for Layout<V> {
176    fn clone(&self) -> Self {
177        Self {
178            inner: Arc::clone(&self.inner),
179        }
180    }
181}
182
183impl<V: VTable> Debug for Layout<V> {
184    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
185        f.debug_struct("Layout")
186            .field("encoding_id", &self.vtable().id())
187            .field("dtype", &self.inner.dtype)
188            .field("row_count", &self.inner.row_count)
189            .field("segment_ids", &self.inner.segment_ids)
190            .field("data", &self.inner.data)
191            .finish()
192    }
193}
194
195impl<V: VTable> Deref for Layout<V> {
196    type Target = V::LayoutData;
197
198    fn deref(&self) -> &Self::Target {
199        self.data()
200    }
201}
202
203impl<V: VTable> From<Layout<V>> for LayoutRef {
204    fn from(value: Layout<V>) -> Self {
205        value.into_layout()
206    }
207}
208
209/// Erased layout behavior used by [`LayoutRef`].
210pub trait DynLayout: 'static + Send + Sync + Debug {
211    /// Returns this layout as [`Any`] for downcasting.
212    fn as_any(&self) -> &dyn Any;
213
214    /// Clone this layout as an erased reference.
215    fn dyn_to_layout(&self) -> LayoutRef;
216
217    /// Returns the layout ID.
218    fn dyn_encoding_id(&self) -> LayoutId;
219
220    /// Returns the row count.
221    fn dyn_row_count(&self) -> u64;
222
223    /// Returns the logical dtype.
224    fn dyn_dtype(&self) -> &DType;
225
226    /// Returns the number of children.
227    fn dyn_nchildren(&self) -> usize;
228
229    /// Returns child `idx`.
230    fn dyn_child(&self, idx: usize) -> VortexResult<LayoutRef>;
231
232    /// Returns the relationship of child `idx`.
233    fn dyn_child_type(&self, idx: usize) -> LayoutChildType;
234
235    /// Serializes layout-specific metadata.
236    fn dyn_metadata(&self) -> Vec<u8>;
237
238    /// Returns directly referenced segment IDs.
239    fn dyn_segment_ids(&self) -> Vec<SegmentId>;
240
241    /// Constructs a reader.
242    fn dyn_new_reader(
243        &self,
244        name: Arc<str>,
245        segment_source: Arc<dyn SegmentSource>,
246        session: &VortexSession,
247        ctx: &LayoutReaderContext,
248    ) -> VortexResult<LayoutReaderRef>;
249}
250
251impl<V: VTable> DynLayout for Layout<V> {
252    fn as_any(&self) -> &dyn Any {
253        self
254    }
255
256    fn dyn_to_layout(&self) -> LayoutRef {
257        Layout::to_layout(self)
258    }
259
260    fn dyn_encoding_id(&self) -> LayoutId {
261        self.vtable().id()
262    }
263
264    fn dyn_row_count(&self) -> u64 {
265        Layout::row_count(self)
266    }
267
268    fn dyn_dtype(&self) -> &DType {
269        Layout::dtype(self)
270    }
271
272    fn dyn_nchildren(&self) -> usize {
273        Layout::nchildren(self)
274    }
275
276    fn dyn_child(&self, idx: usize) -> VortexResult<LayoutRef> {
277        Layout::child(self, idx)
278    }
279
280    fn dyn_child_type(&self, idx: usize) -> LayoutChildType {
281        Layout::child_type(self, idx)
282    }
283
284    fn dyn_metadata(&self) -> Vec<u8> {
285        V::metadata(self).serialize()
286    }
287
288    fn dyn_segment_ids(&self) -> Vec<SegmentId> {
289        self.inner.segment_ids.clone()
290    }
291
292    fn dyn_new_reader(
293        &self,
294        name: Arc<str>,
295        segment_source: Arc<dyn SegmentSource>,
296        session: &VortexSession,
297        ctx: &LayoutReaderContext,
298    ) -> VortexResult<LayoutReaderRef> {
299        Layout::new_reader(self, name, segment_source, session, ctx)
300    }
301}
302
303/// Identifies how a layout child relates to its parent.
304#[derive(Debug, Clone, PartialEq, Eq)]
305pub enum LayoutChildType {
306    /// A child retaining the parent's schema and row offset.
307    Transparent(Arc<str>),
308    /// Auxiliary data, such as dictionary values or zone maps.
309    Auxiliary(Arc<str>),
310    /// A row-based chunk with its relative row offset.
311    Chunk((usize, u64)),
312    /// A single field of a struct.
313    Field(FieldName),
314}
315
316impl LayoutChildType {
317    /// Returns the child name.
318    pub fn name(&self) -> Arc<str> {
319        match self {
320            Self::Chunk((idx, _)) => format!("[{idx}]").into(),
321            Self::Auxiliary(name) | Self::Transparent(name) => Arc::clone(name),
322            Self::Field(name) => name.clone().into(),
323        }
324    }
325
326    /// Returns the relative row offset, or `None` for auxiliary children.
327    pub fn row_offset(&self) -> Option<u64> {
328        match self {
329            Self::Chunk((_, offset)) => Some(*offset),
330            Self::Auxiliary(_) => None,
331            Self::Transparent(_) | Self::Field(_) => Some(0),
332        }
333    }
334}
335
336impl dyn DynLayout + '_ {
337    /// Returns a cloned erased layout reference.
338    pub fn to_layout(&self) -> LayoutRef {
339        self.dyn_to_layout()
340    }
341
342    /// Returns the layout ID.
343    pub fn encoding_id(&self) -> LayoutId {
344        self.dyn_encoding_id()
345    }
346
347    /// Returns the logical dtype.
348    pub fn dtype(&self) -> &DType {
349        self.dyn_dtype()
350    }
351
352    /// Returns the number of rows.
353    pub fn row_count(&self) -> u64 {
354        self.dyn_row_count()
355    }
356
357    /// Returns the number of children.
358    pub fn nchildren(&self) -> usize {
359        self.dyn_nchildren()
360    }
361
362    /// Returns child `idx`.
363    pub fn child(&self, idx: usize) -> VortexResult<LayoutRef> {
364        self.dyn_child(idx)
365    }
366
367    /// Returns the relationship of child `idx`.
368    pub fn child_type(&self, idx: usize) -> LayoutChildType {
369        self.dyn_child_type(idx)
370    }
371
372    /// Returns serialized layout-specific metadata.
373    pub fn metadata(&self) -> Vec<u8> {
374        self.dyn_metadata()
375    }
376
377    /// Returns directly referenced segment IDs.
378    pub fn segment_ids(&self) -> Vec<SegmentId> {
379        self.dyn_segment_ids()
380    }
381
382    /// Constructs a reader for this layout.
383    pub fn new_reader(
384        &self,
385        name: Arc<str>,
386        segment_source: Arc<dyn SegmentSource>,
387        session: &VortexSession,
388        ctx: &LayoutReaderContext,
389    ) -> VortexResult<LayoutReaderRef> {
390        self.dyn_new_reader(name, segment_source, session, ctx)
391    }
392
393    /// Returns all children.
394    pub fn children(&self) -> VortexResult<Vec<LayoutRef>> {
395        (0..self.nchildren())
396            .map(|idx| self.child(idx))
397            .try_collect()
398    }
399
400    /// Returns all child types.
401    pub fn child_types(&self) -> impl Iterator<Item = LayoutChildType> + '_ {
402        (0..self.nchildren()).map(|idx| self.child_type(idx))
403    }
404
405    /// Returns all child names.
406    pub fn child_names(&self) -> impl Iterator<Item = Arc<str>> + '_ {
407        self.child_types().map(|child| child.name())
408    }
409
410    /// Returns all child row offsets.
411    pub fn child_row_offsets(&self) -> impl Iterator<Item = Option<u64>> + '_ {
412        self.child_types().map(|child| child.row_offset())
413    }
414
415    /// Returns whether this layout uses vtable `V`.
416    pub fn is<V: VTable>(&self) -> bool {
417        self.as_opt::<V>().is_some()
418    }
419
420    /// Downcasts this layout to vtable `V`.
421    pub fn as_<V: VTable>(&self) -> &Layout<V> {
422        self.as_opt::<V>().vortex_expect("Failed to downcast")
423    }
424
425    /// Attempts to downcast this layout to vtable `V`.
426    pub fn as_opt<V: VTable>(&self) -> Option<&Layout<V>> {
427        self.as_any().downcast_ref()
428    }
429
430    /// Returns a depth-first pre-order traversal.
431    pub fn depth_first_traversal(&self) -> impl Iterator<Item = VortexResult<LayoutRef>> {
432        struct ChildrenIterator {
433            stack: Vec<LayoutRef>,
434        }
435
436        impl Iterator for ChildrenIterator {
437            type Item = VortexResult<LayoutRef>;
438
439            fn next(&mut self) -> Option<Self::Item> {
440                let next = self.stack.pop()?;
441                let Ok(children) = next.children() else {
442                    return Some(Ok(next));
443                };
444                self.stack.extend(children.into_iter().rev());
445                Some(Ok(next))
446            }
447        }
448
449        ChildrenIterator {
450            stack: vec![self.to_layout()],
451        }
452    }
453
454    /// Displays the layout as a tree.
455    pub fn display_tree(&self) -> DisplayLayoutTree {
456        DisplayLayoutTree::new(self.to_layout(), false)
457    }
458
459    /// Displays the layout as a tree with optional verbose metadata.
460    pub fn display_tree_verbose(&self, verbose: bool) -> DisplayLayoutTree {
461        DisplayLayoutTree::new(self.to_layout(), verbose)
462    }
463
464    /// Displays the tree after fetching segment sizes.
465    pub async fn display_tree_with_segments(
466        &self,
467        segment_source: Arc<dyn SegmentSource>,
468    ) -> VortexResult<DisplayLayoutTree> {
469        display_tree_with_segment_sizes(self.to_layout(), segment_source).await
470    }
471}
472
473impl Display for dyn DynLayout + '_ {
474    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
475        let segments = self.segment_ids();
476        if segments.is_empty() {
477            write!(
478                f,
479                "{}({}, rows={})",
480                self.encoding_id(),
481                self.dtype(),
482                self.row_count()
483            )
484        } else {
485            write!(
486                f,
487                "{}({}, rows={}, segments=[{}])",
488                self.encoding_id(),
489                self.dtype(),
490                self.row_count(),
491                segments.iter().map(|s| format!("{}", **s)).join(", ")
492            )
493        }
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500
501    #[test]
502    fn layout_child_type_names_and_offsets() {
503        let chunk = LayoutChildType::Chunk((5, 100));
504        assert_eq!(chunk.name().as_ref(), "[5]");
505        assert_eq!(chunk.row_offset(), Some(100));
506
507        let field = LayoutChildType::Field(FieldName::from("customer_id"));
508        assert_eq!(field.name().as_ref(), "customer_id");
509        assert_eq!(field.row_offset(), Some(0));
510
511        let auxiliary = LayoutChildType::Auxiliary("zone_map".into());
512        assert_eq!(auxiliary.row_offset(), None);
513        let transparent = LayoutChildType::Transparent("compressed".into());
514        assert_eq!(transparent.row_offset(), Some(0));
515    }
516}
vortex_layout/layout.rs

vortex_layout/
layout.rs
