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 serialized (present) children.
134    pub fn nchildren(&self) -> usize {
135        self.inner.children.nchildren()
136    }
137
138    /// Returns the number of logical child slots, including any that are absent.
139    pub fn nslots(&self) -> usize {
140        V::nslots(self)
141    }
142
143    /// Maps a logical `slot` to the index of its serialized child, or `None` if absent.
144    pub fn slot_to_child(&self, slot: usize) -> Option<usize> {
145        V::slot_to_child(self, slot)
146    }
147
148    /// Materialize the child in logical `slot`, or `None` if the slot is absent.
149    pub fn slot(&self, slot: usize) -> VortexResult<Option<LayoutRef>> {
150        match V::slot_to_child(self, slot) {
151            Some(idx) => self
152                .inner
153                .children
154                .child(idx, &V::child_dtype(self, slot)?)
155                .map(Some),
156            None => Ok(None),
157        }
158    }
159
160    /// Returns the relationship of the child in logical `slot` to this layout, or `None` if the
161    /// slot is absent.
162    pub fn slot_type(&self, slot: usize) -> Option<LayoutChildType> {
163        V::slot_to_child(self, slot).map(|_| V::child_type(self, slot))
164    }
165
166    /// Returns a child's serialized row count without materializing it.
167    pub fn child_row_count(&self, idx: usize) -> u64 {
168        self.inner.children.child_row_count(idx)
169    }
170
171    /// Erase this typed layout into a shared layout reference.
172    pub fn to_layout(&self) -> LayoutRef {
173        self.clone().into_layout()
174    }
175
176    /// Erase this typed layout into a shared layout reference.
177    pub fn into_layout(self) -> LayoutRef {
178        Arc::new(self)
179    }
180
181    /// Construct a reader for this layout.
182    pub fn new_reader(
183        &self,
184        name: Arc<str>,
185        segment_source: Arc<dyn SegmentSource>,
186        session: &VortexSession,
187        ctx: &LayoutReaderContext,
188    ) -> VortexResult<LayoutReaderRef> {
189        V::new_reader(self, name, segment_source, session, ctx)
190    }
191}
192
193impl<V: VTable> Clone for Layout<V> {
194    fn clone(&self) -> Self {
195        Self {
196            inner: Arc::clone(&self.inner),
197        }
198    }
199}
200
201impl<V: VTable> Debug for Layout<V> {
202    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
203        f.debug_struct("Layout")
204            .field("encoding_id", &self.vtable().id())
205            .field("dtype", &self.inner.dtype)
206            .field("row_count", &self.inner.row_count)
207            .field("segment_ids", &self.inner.segment_ids)
208            .field("data", &self.inner.data)
209            .finish()
210    }
211}
212
213impl<V: VTable> Deref for Layout<V> {
214    type Target = V::LayoutData;
215
216    fn deref(&self) -> &Self::Target {
217        self.data()
218    }
219}
220
221impl<V: VTable> From<Layout<V>> for LayoutRef {
222    fn from(value: Layout<V>) -> Self {
223        value.into_layout()
224    }
225}
226
227/// Erased layout behavior used by [`LayoutRef`].
228pub trait DynLayout: 'static + Send + Sync + Debug {
229    /// Returns this layout as [`Any`] for downcasting.
230    fn as_any(&self) -> &dyn Any;
231
232    /// Clone this layout as an erased reference.
233    fn dyn_to_layout(&self) -> LayoutRef;
234
235    /// Returns the layout ID.
236    fn dyn_encoding_id(&self) -> LayoutId;
237
238    /// Returns the row count.
239    fn dyn_row_count(&self) -> u64;
240
241    /// Returns the logical dtype.
242    fn dyn_dtype(&self) -> &DType;
243
244    /// Returns the number of serialized (present) children.
245    fn dyn_nchildren(&self) -> usize;
246
247    /// Returns the number of logical child slots, including any that are absent.
248    fn dyn_nslots(&self) -> usize;
249
250    /// Materializes the child in logical `slot`, or `None` if the slot is absent.
251    fn dyn_slot(&self, slot: usize) -> VortexResult<Option<LayoutRef>>;
252
253    /// Returns the relationship of the child in logical `slot`, or `None` if the slot is absent.
254    fn dyn_slot_type(&self, slot: usize) -> Option<LayoutChildType>;
255
256    /// Serializes layout-specific metadata.
257    fn dyn_metadata(&self) -> Vec<u8>;
258
259    /// Returns directly referenced segment IDs.
260    fn dyn_segment_ids(&self) -> Vec<SegmentId>;
261
262    /// Constructs a reader.
263    fn dyn_new_reader(
264        &self,
265        name: Arc<str>,
266        segment_source: Arc<dyn SegmentSource>,
267        session: &VortexSession,
268        ctx: &LayoutReaderContext,
269    ) -> VortexResult<LayoutReaderRef>;
270
271    /// Returns `true` if this layout is indivisible: its readers never register natural split
272    /// boundaries strictly inside their row range (see [`crate::VTable::is_indivisible`]).
273    fn dyn_is_indivisible(&self) -> bool {
274        false
275    }
276}
277
278impl<V: VTable> DynLayout for Layout<V> {
279    fn as_any(&self) -> &dyn Any {
280        self
281    }
282
283    fn dyn_to_layout(&self) -> LayoutRef {
284        Layout::to_layout(self)
285    }
286
287    fn dyn_encoding_id(&self) -> LayoutId {
288        self.vtable().id()
289    }
290
291    fn dyn_row_count(&self) -> u64 {
292        Layout::row_count(self)
293    }
294
295    fn dyn_dtype(&self) -> &DType {
296        Layout::dtype(self)
297    }
298
299    fn dyn_nchildren(&self) -> usize {
300        Layout::nchildren(self)
301    }
302
303    fn dyn_nslots(&self) -> usize {
304        Layout::nslots(self)
305    }
306
307    fn dyn_slot(&self, slot: usize) -> VortexResult<Option<LayoutRef>> {
308        Layout::slot(self, slot)
309    }
310
311    fn dyn_slot_type(&self, slot: usize) -> Option<LayoutChildType> {
312        Layout::slot_type(self, slot)
313    }
314
315    fn dyn_metadata(&self) -> Vec<u8> {
316        V::metadata(self).serialize()
317    }
318
319    fn dyn_segment_ids(&self) -> Vec<SegmentId> {
320        self.inner.segment_ids.clone()
321    }
322
323    fn dyn_new_reader(
324        &self,
325        name: Arc<str>,
326        segment_source: Arc<dyn SegmentSource>,
327        session: &VortexSession,
328        ctx: &LayoutReaderContext,
329    ) -> VortexResult<LayoutReaderRef> {
330        Layout::new_reader(self, name, segment_source, session, ctx)
331    }
332
333    fn dyn_is_indivisible(&self) -> bool {
334        self.vtable().is_indivisible()
335    }
336}
337
338/// Identifies how a layout child relates to its parent.
339#[derive(Debug, Clone, PartialEq, Eq)]
340pub enum LayoutChildType {
341    /// A child retaining the parent's schema and row offset.
342    Transparent(Arc<str>),
343    /// Auxiliary data, such as dictionary values or zone maps.
344    Auxiliary(Arc<str>),
345    /// A row-based chunk with its relative row offset.
346    Chunk((usize, u64)),
347    /// A single field of a struct.
348    Field(FieldName),
349}
350
351impl LayoutChildType {
352    /// Returns the child name.
353    pub fn name(&self) -> Arc<str> {
354        match self {
355            Self::Chunk((idx, _)) => format!("[{idx}]").into(),
356            Self::Auxiliary(name) | Self::Transparent(name) => Arc::clone(name),
357            Self::Field(name) => name.clone().into(),
358        }
359    }
360
361    /// Returns the relative row offset, or `None` for auxiliary children.
362    pub fn row_offset(&self) -> Option<u64> {
363        match self {
364            Self::Chunk((_, offset)) => Some(*offset),
365            Self::Auxiliary(_) => None,
366            Self::Transparent(_) | Self::Field(_) => Some(0),
367        }
368    }
369}
370
371impl dyn DynLayout + '_ {
372    /// Returns a cloned erased layout reference.
373    pub fn to_layout(&self) -> LayoutRef {
374        self.dyn_to_layout()
375    }
376
377    /// Returns the layout ID.
378    pub fn encoding_id(&self) -> LayoutId {
379        self.dyn_encoding_id()
380    }
381
382    /// Returns the logical dtype.
383    pub fn dtype(&self) -> &DType {
384        self.dyn_dtype()
385    }
386
387    /// Returns the number of rows.
388    pub fn row_count(&self) -> u64 {
389        self.dyn_row_count()
390    }
391
392    /// Returns the number of serialized (present) children.
393    pub fn nchildren(&self) -> usize {
394        self.dyn_nchildren()
395    }
396
397    /// Returns the number of logical child slots, including any that are absent.
398    pub fn nslots(&self) -> usize {
399        self.dyn_nslots()
400    }
401
402    /// Materializes the child in logical `slot`, or `None` if the slot is absent.
403    pub fn slot(&self, slot: usize) -> VortexResult<Option<LayoutRef>> {
404        self.dyn_slot(slot)
405    }
406
407    /// Returns the relationship of the child in logical `slot`, or `None` if the slot is absent.
408    pub fn slot_type(&self, slot: usize) -> Option<LayoutChildType> {
409        self.dyn_slot_type(slot)
410    }
411
412    /// Returns serialized layout-specific metadata.
413    pub fn metadata(&self) -> Vec<u8> {
414        self.dyn_metadata()
415    }
416
417    /// Returns directly referenced segment IDs.
418    pub fn segment_ids(&self) -> Vec<SegmentId> {
419        self.dyn_segment_ids()
420    }
421
422    /// Constructs a reader for this layout.
423    pub fn new_reader(
424        &self,
425        name: Arc<str>,
426        segment_source: Arc<dyn SegmentSource>,
427        session: &VortexSession,
428        ctx: &LayoutReaderContext,
429    ) -> VortexResult<LayoutReaderRef> {
430        self.dyn_new_reader(name, segment_source, session, ctx)
431    }
432
433    /// Returns all serialized (present) children, in slot order.
434    pub fn children(&self) -> VortexResult<Vec<LayoutRef>> {
435        (0..self.nslots())
436            .filter_map(|slot| self.slot(slot).transpose())
437            .try_collect()
438    }
439
440    /// Returns the types of all serialized (present) children, in slot order.
441    pub fn child_types(&self) -> impl Iterator<Item = LayoutChildType> + '_ {
442        (0..self.nslots()).filter_map(|slot| self.slot_type(slot))
443    }
444
445    /// Returns all child names.
446    pub fn child_names(&self) -> impl Iterator<Item = Arc<str>> + '_ {
447        self.child_types().map(|child| child.name())
448    }
449
450    /// Returns all child row offsets.
451    pub fn child_row_offsets(&self) -> impl Iterator<Item = Option<u64>> + '_ {
452        self.child_types().map(|child| child.row_offset())
453    }
454
455    /// Returns whether this layout uses vtable `V`.
456    pub fn is<V: VTable>(&self) -> bool {
457        self.as_opt::<V>().is_some()
458    }
459
460    /// Downcasts this layout to vtable `V`.
461    pub fn as_<V: VTable>(&self) -> &Layout<V> {
462        self.as_opt::<V>().vortex_expect("Failed to downcast")
463    }
464
465    /// Attempts to downcast this layout to vtable `V`.
466    pub fn as_opt<V: VTable>(&self) -> Option<&Layout<V>> {
467        self.as_any().downcast_ref()
468    }
469
470    /// Returns a depth-first pre-order traversal.
471    pub fn depth_first_traversal(&self) -> impl Iterator<Item = VortexResult<LayoutRef>> {
472        struct ChildrenIterator {
473            stack: Vec<LayoutRef>,
474        }
475
476        impl Iterator for ChildrenIterator {
477            type Item = VortexResult<LayoutRef>;
478
479            fn next(&mut self) -> Option<Self::Item> {
480                let next = self.stack.pop()?;
481                let Ok(children) = next.children() else {
482                    return Some(Ok(next));
483                };
484                self.stack.extend(children.into_iter().rev());
485                Some(Ok(next))
486            }
487        }
488
489        ChildrenIterator {
490            stack: vec![self.to_layout()],
491        }
492    }
493
494    /// Displays the layout as a tree.
495    pub fn display_tree(&self) -> DisplayLayoutTree {
496        DisplayLayoutTree::new(self.to_layout(), false)
497    }
498
499    /// Displays the layout as a tree with optional verbose metadata.
500    pub fn display_tree_verbose(&self, verbose: bool) -> DisplayLayoutTree {
501        DisplayLayoutTree::new(self.to_layout(), verbose)
502    }
503
504    /// Displays the tree after fetching segment sizes.
505    pub async fn display_tree_with_segments(
506        &self,
507        segment_source: Arc<dyn SegmentSource>,
508    ) -> VortexResult<DisplayLayoutTree> {
509        display_tree_with_segment_sizes(self.to_layout(), segment_source).await
510    }
511}
512
513impl Display for dyn DynLayout + '_ {
514    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
515        let segments = self.segment_ids();
516        if segments.is_empty() {
517            write!(
518                f,
519                "{}({}, rows={})",
520                self.encoding_id(),
521                self.dtype(),
522                self.row_count()
523            )
524        } else {
525            write!(
526                f,
527                "{}({}, rows={}, segments=[{}])",
528                self.encoding_id(),
529                self.dtype(),
530                self.row_count(),
531                segments.iter().map(|s| format!("{}", **s)).join(", ")
532            )
533        }
534    }
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540
541    #[test]
542    fn layout_child_type_names_and_offsets() {
543        let chunk = LayoutChildType::Chunk((5, 100));
544        assert_eq!(chunk.name().as_ref(), "[5]");
545        assert_eq!(chunk.row_offset(), Some(100));
546
547        let field = LayoutChildType::Field(FieldName::from("customer_id"));
548        assert_eq!(field.name().as_ref(), "customer_id");
549        assert_eq!(field.row_offset(), Some(0));
550
551        let auxiliary = LayoutChildType::Auxiliary("zone_map".into());
552        assert_eq!(auxiliary.row_offset(), None);
553        let transparent = LayoutChildType::Transparent("compressed".into());
554        assert_eq!(transparent.row_offset(), Some(0));
555    }
556}
vortex_layout/layout.rs

vortex_layout/
layout.rs
