vortex_layout/

layout.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::any::Any;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::sync::Arc;

use arcref::ArcRef;
use itertools::Itertools;
use vortex_array::SerializeMetadata;
use vortex_dtype::DType;
use vortex_dtype::FieldName;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_err;
use vortex_session::VortexSession;

use crate::LayoutEncodingId;
use crate::LayoutEncodingRef;
use crate::LayoutReaderRef;
use crate::VTable;
use crate::display::DisplayLayoutTree;
use crate::segments::SegmentId;
use crate::segments::SegmentSource;

pub type LayoutId = ArcRef<str>;

pub type LayoutRef = Arc<dyn Layout>;

pub trait Layout: 'static + Send + Sync + Debug + private::Sealed {
    fn as_any(&self) -> &dyn Any;

    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync>;

    fn to_layout(&self) -> LayoutRef;

    /// Returns the [`crate::LayoutEncoding`] for this layout.
    fn encoding(&self) -> LayoutEncodingRef;

    /// The number of rows in this layout.
    fn row_count(&self) -> u64;

    /// The dtype of this layout when projected with the root scope.
    fn dtype(&self) -> &DType;

    /// The number of children in this layout.
    fn nchildren(&self) -> usize;

    /// Get the child at the given index.
    fn child(&self, idx: usize) -> VortexResult<LayoutRef>;

    /// Get the relative row offset of the child at the given index, returning `None` for
    /// any auxiliary children, e.g. dictionary values, zone maps, etc.
    fn child_type(&self, idx: usize) -> LayoutChildType;

    /// Get the metadata for this layout.
    fn metadata(&self) -> Vec<u8>;

    /// Get the segment IDs for this layout.
    fn segment_ids(&self) -> Vec<SegmentId>;

    #[cfg(gpu_unstable)]
    fn new_gpu_reader(
        &self,
        name: Arc<str>,
        segment_source: Arc<dyn SegmentSource>,
        ctx: Arc<cudarc::driver::CudaContext>,
    ) -> VortexResult<crate::gpu::GpuLayoutReaderRef>;

    fn new_reader(
        &self,
        name: Arc<str>,
        segment_source: Arc<dyn SegmentSource>,
        session: &VortexSession,
    ) -> VortexResult<LayoutReaderRef>;
}

pub trait IntoLayout {
    /// Converts this type into a [`LayoutRef`].
    fn into_layout(self) -> LayoutRef;
}

/// A type that allows us to identify how a layout child relates to its parent.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LayoutChildType {
    /// A layout child that retains the same schema and row offset position in the dataset.
    Transparent(Arc<str>),
    /// A layout child that provides auxiliary data, e.g. dictionary values, zone maps, etc.
    /// Contains a human-readable name of the child.
    Auxiliary(Arc<str>),
    /// A layout child that represents a row-based chunk of data.
    /// Contains the chunk index and relative row offset of the child.
    Chunk((usize, u64)),
    /// A layout child that represents a single field of data.
    /// Contains the field name of the child.
    Field(FieldName),
}

impl LayoutChildType {
    /// Returns the name of this child.
    pub fn name(&self) -> Arc<str> {
        match self {
            LayoutChildType::Chunk((idx, _offset)) => format!("[{idx}]").into(),
            LayoutChildType::Auxiliary(name) => name.clone(),
            LayoutChildType::Transparent(name) => name.clone(),
            LayoutChildType::Field(name) => name.clone().into(),
        }
    }

    /// Returns the relative row offset of this child.
    /// For auxiliary children, this is `None`.
    pub fn row_offset(&self) -> Option<u64> {
        match self {
            LayoutChildType::Chunk((_idx, offset)) => Some(*offset),
            LayoutChildType::Auxiliary(_) => None,
            LayoutChildType::Transparent(_) => Some(0),
            LayoutChildType::Field(_) => Some(0),
        }
    }
}

impl dyn Layout + '_ {
    /// The ID of the encoding for this layout.
    pub fn encoding_id(&self) -> LayoutEncodingId {
        self.encoding().id()
    }

    /// The children of this layout.
    pub fn children(&self) -> VortexResult<Vec<LayoutRef>> {
        (0..self.nchildren()).map(|i| self.child(i)).try_collect()
    }

    /// The child types of this layout.
    pub fn child_types(&self) -> impl Iterator<Item = LayoutChildType> {
        (0..self.nchildren()).map(|i| self.child_type(i))
    }

    /// The names of the children of this layout.
    pub fn child_names(&self) -> impl Iterator<Item = Arc<str>> {
        self.child_types().map(|child| child.name())
    }

    /// The row offsets of the children of this layout, where `None` indicates an auxiliary child.
    pub fn child_row_offsets(&self) -> impl Iterator<Item = Option<u64>> {
        self.child_types().map(|child| child.row_offset())
    }

    pub fn is<V: VTable>(&self) -> bool {
        self.as_opt::<V>().is_some()
    }

    /// Downcast a layout to a specific type.
    pub fn as_<V: VTable>(&self) -> &V::Layout {
        self.as_opt::<V>().vortex_expect("Failed to downcast")
    }

    /// Downcast a layout to a specific type.
    pub fn as_opt<V: VTable>(&self) -> Option<&V::Layout> {
        self.as_any()
            .downcast_ref::<LayoutAdapter<V>>()
            .map(|adapter| &adapter.0)
    }

    /// Downcast a layout to a specific type.
    pub fn into<V: VTable>(self: Arc<Self>) -> Arc<V::Layout> {
        let layout_adapter = self
            .as_any_arc()
            .downcast::<LayoutAdapter<V>>()
            .map_err(|_| vortex_err!("Invalid layout type"))
            .vortex_expect("Invalid layout type");

        // SAFETY: LayoutAdapter<V> is #[repr(transparent)] (see line 192) which guarantees
        // it has the same memory layout as V::Layout. The downcast above ensures we have
        // the correct type. This transmute is safe because both Arc types point to data
        // with identical layout and alignment.
        unsafe { std::mem::transmute::<Arc<LayoutAdapter<V>>, Arc<V::Layout>>(layout_adapter) }
    }

    /// Depth-first traversal of the layout and its children.
    pub fn depth_first_traversal(&self) -> impl Iterator<Item = VortexResult<LayoutRef>> {
        /// A depth-first pre-order iterator over a layout.
        struct ChildrenIterator {
            stack: Vec<LayoutRef>,
        }

        impl Iterator for ChildrenIterator {
            type Item = VortexResult<LayoutRef>;

            fn next(&mut self) -> Option<Self::Item> {
                let next = self.stack.pop()?;
                let Ok(children) = next.children() else {
                    return Some(Ok(next));
                };
                for child in children.into_iter().rev() {
                    self.stack.push(child);
                }
                Some(Ok(next))
            }
        }

        ChildrenIterator {
            stack: vec![self.to_layout()],
        }
    }

    /// Display the layout as a tree structure.
    pub fn display_tree(&self) -> DisplayLayoutTree {
        DisplayLayoutTree::new(self.to_layout(), false)
    }

    /// Display the layout as a tree structure with optional verbose metadata.
    pub fn display_tree_verbose(&self, verbose: bool) -> DisplayLayoutTree {
        DisplayLayoutTree::new(self.to_layout(), verbose)
    }
}

#[repr(transparent)]
pub struct LayoutAdapter<V: VTable>(V::Layout);

impl<V: VTable> Debug for LayoutAdapter<V> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        self.0.fmt(f)
    }
}

impl<V: VTable> Layout for LayoutAdapter<V> {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
        self
    }

    fn to_layout(&self) -> LayoutRef {
        Arc::new(LayoutAdapter::<V>(self.0.clone()))
    }

    fn encoding(&self) -> LayoutEncodingRef {
        V::encoding(&self.0)
    }

    fn row_count(&self) -> u64 {
        V::row_count(&self.0)
    }

    fn dtype(&self) -> &DType {
        V::dtype(&self.0)
    }

    fn nchildren(&self) -> usize {
        V::nchildren(&self.0)
    }

    fn child(&self, idx: usize) -> VortexResult<LayoutRef> {
        V::child(&self.0, idx)
    }

    fn child_type(&self, idx: usize) -> LayoutChildType {
        V::child_type(&self.0, idx)
    }

    fn metadata(&self) -> Vec<u8> {
        V::metadata(&self.0).serialize()
    }

    fn segment_ids(&self) -> Vec<SegmentId> {
        V::segment_ids(&self.0)
    }

    #[cfg(gpu_unstable)]
    fn new_gpu_reader(
        &self,
        name: Arc<str>,
        segment_source: Arc<dyn SegmentSource>,
        ctx: Arc<cudarc::driver::CudaContext>,
    ) -> VortexResult<crate::gpu::GpuLayoutReaderRef> {
        V::new_gpu_reader(&self.0, name, segment_source, ctx)
    }

    fn new_reader(
        &self,
        name: Arc<str>,
        segment_source: Arc<dyn SegmentSource>,
        session: &VortexSession,
    ) -> VortexResult<LayoutReaderRef> {
        V::new_reader(&self.0, name, segment_source, session)
    }
}

mod private {
    use super::*;

    pub trait Sealed {}

    impl<V: VTable> Sealed for LayoutAdapter<V> {}
}

#[cfg(test)]
mod tests {
    use rstest::rstest;

    use super::*;

    #[test]
    fn test_layout_child_type_name() {
        // Test Chunk variant
        let chunk = LayoutChildType::Chunk((5, 100));
        assert_eq!(chunk.name().as_ref(), "[5]");

        // Test Field variant
        let field = LayoutChildType::Field(FieldName::from("customer_id"));
        assert_eq!(field.name().as_ref(), "customer_id");

        // Test Auxiliary variant
        let aux = LayoutChildType::Auxiliary(Arc::from("zone_map"));
        assert_eq!(aux.name().as_ref(), "zone_map");

        // Test Transparent variant
        let transparent = LayoutChildType::Transparent(Arc::from("compressed"));
        assert_eq!(transparent.name().as_ref(), "compressed");
    }

    #[test]
    fn test_layout_child_type_row_offset() {
        // Chunk should return the offset
        let chunk = LayoutChildType::Chunk((0, 42));
        assert_eq!(chunk.row_offset(), Some(42));

        // Field should return 0
        let field = LayoutChildType::Field(FieldName::from("field1"));
        assert_eq!(field.row_offset(), Some(0));

        // Auxiliary should return None
        let aux = LayoutChildType::Auxiliary(Arc::from("metadata"));
        assert_eq!(aux.row_offset(), None);

        // Transparent should return 0
        let transparent = LayoutChildType::Transparent(Arc::from("wrapper"));
        assert_eq!(transparent.row_offset(), Some(0));
    }

    #[test]
    fn test_layout_child_type_equality() {
        // Test Chunk equality
        let chunk1 = LayoutChildType::Chunk((1, 100));
        let chunk2 = LayoutChildType::Chunk((1, 100));
        let chunk3 = LayoutChildType::Chunk((2, 100));
        let chunk4 = LayoutChildType::Chunk((1, 200));

        assert_eq!(chunk1, chunk2);
        assert_ne!(chunk1, chunk3);
        assert_ne!(chunk1, chunk4);

        // Test Field equality
        let field1 = LayoutChildType::Field(FieldName::from("name"));
        let field2 = LayoutChildType::Field(FieldName::from("name"));
        let field3 = LayoutChildType::Field(FieldName::from("age"));

        assert_eq!(field1, field2);
        assert_ne!(field1, field3);

        // Test Auxiliary equality
        let aux1 = LayoutChildType::Auxiliary(Arc::from("stats"));
        let aux2 = LayoutChildType::Auxiliary(Arc::from("stats"));
        let aux3 = LayoutChildType::Auxiliary(Arc::from("index"));

        assert_eq!(aux1, aux2);
        assert_ne!(aux1, aux3);

        // Test Transparent equality
        let trans1 = LayoutChildType::Transparent(Arc::from("enc"));
        let trans2 = LayoutChildType::Transparent(Arc::from("enc"));
        let trans3 = LayoutChildType::Transparent(Arc::from("dec"));

        assert_eq!(trans1, trans2);
        assert_ne!(trans1, trans3);

        // Test cross-variant inequality
        assert_ne!(chunk1, field1);
        assert_ne!(field1, aux1);
        assert_ne!(aux1, trans1);
    }

    #[rstest]
    #[case(LayoutChildType::Chunk((0, 0)), "[0]", Some(0))]
    #[case(LayoutChildType::Chunk((999, 1000000)), "[999]", Some(1000000))]
    #[case(LayoutChildType::Field(FieldName::from("")), "", Some(0))]
    #[case(
        LayoutChildType::Field(FieldName::from("very_long_field_name_that_is_quite_lengthy")),
        "very_long_field_name_that_is_quite_lengthy",
        Some(0)
    )]
    #[case(LayoutChildType::Auxiliary(Arc::from("aux")), "aux", None)]
    #[case(LayoutChildType::Transparent(Arc::from("t")), "t", Some(0))]
    fn test_layout_child_type_parameterized(
        #[case] child_type: LayoutChildType,
        #[case] expected_name: &str,
        #[case] expected_offset: Option<u64>,
    ) {
        assert_eq!(child_type.name().as_ref(), expected_name);
        assert_eq!(child_type.row_offset(), expected_offset);
    }

    #[test]
    fn test_chunk_with_different_indices_and_offsets() {
        let chunks = [
            LayoutChildType::Chunk((0, 0)),
            LayoutChildType::Chunk((1, 100)),
            LayoutChildType::Chunk((2, 200)),
            LayoutChildType::Chunk((100, 10000)),
        ];

        for chunk in chunks.iter() {
            let name = chunk.name();
            assert!(name.starts_with('['));
            assert!(name.ends_with(']'));

            if let LayoutChildType::Chunk((idx, offset)) = chunk {
                assert_eq!(name.as_ref(), format!("[{}]", idx));
                assert_eq!(chunk.row_offset(), Some(*offset));
            }
        }
    }

    #[test]
    fn test_field_names_with_special_characters() {
        let special_fields: Vec<Arc<str>> = vec![
            Arc::from("field-with-dashes"),
            Arc::from("field_with_underscores"),
            Arc::from("field.with.dots"),
            Arc::from("field::with::colons"),
            Arc::from("field/with/slashes"),
            Arc::from("field@with#symbols"),
        ];

        for field_name in special_fields {
            let field = LayoutChildType::Field(field_name.clone().into());
            assert_eq!(field.name(), field_name);
            assert_eq!(field.row_offset(), Some(0));
        }
    }
}