vortex_layout/layouts/chunked/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4pub(crate) mod reader;
5pub mod writer;
6
7use std::sync::Arc;
8
9use vortex_array::EmptyMetadata;
10use vortex_array::dtype::DType;
11use vortex_error::VortexExpect;
12use vortex_error::VortexResult;
13use vortex_error::vortex_bail;
14use vortex_error::vortex_err;
15use vortex_session::VortexSession;
16use vortex_session::registry::CachedId;
17
18use crate::Layout;
19use crate::LayoutChildType;
20use crate::LayoutChildren;
21use crate::LayoutDeserializeArgs;
22use crate::LayoutId;
23use crate::LayoutParts;
24use crate::LayoutReaderContext;
25use crate::LayoutReaderRef;
26use crate::LayoutRef;
27use crate::VTable;
28use crate::children::OwnedLayoutChildren;
29use crate::layouts::chunked::reader::ChunkedReader;
30use crate::segments::SegmentSource;
31
32/// Chunked layout vtable.
33#[derive(Clone, Debug)]
34pub struct Chunked;
35
36/// Backwards-compatible name for the chunked layout plugin.
37pub use Chunked as ChunkedLayoutEncoding;
38
39/// Chunked-layout-specific data.
40#[derive(Clone, Debug)]
41pub struct ChunkedData {
42    chunk_offsets: Vec<u64>,
43}
44
45/// A layout partitioned into independently readable row chunks.
46pub type ChunkedLayout = Layout<Chunked>;
47
48impl VTable for Chunked {
49    type LayoutData = ChunkedData;
50    type Metadata = EmptyMetadata;
51
52    fn id(&self) -> LayoutId {
53        static ID: CachedId = CachedId::new("vortex.chunked");
54        *ID
55    }
56
57    fn metadata(_layout: &Layout<Self>) -> Self::Metadata {
58        EmptyMetadata
59    }
60
61    fn deserialize(
62        &self,
63        args: &LayoutDeserializeArgs<'_>,
64        _metadata: &EmptyMetadata,
65    ) -> VortexResult<Self::LayoutData> {
66        let chunk_offsets = chunk_offsets(args.children)?;
67        if chunk_offsets.last().copied() != Some(args.row_count) {
68            vortex_bail!("Chunked child row counts do not add up to parent row count");
69        }
70        Ok(ChunkedData { chunk_offsets })
71    }
72
73    fn child_dtype(layout: &Layout<Self>, _idx: usize) -> VortexResult<DType> {
74        Ok(layout.dtype().clone())
75    }
76
77    fn child_type(layout: &Layout<Self>, idx: usize) -> LayoutChildType {
78        LayoutChildType::Chunk((idx, layout.chunk_offsets[idx]))
79    }
80
81    fn new_reader(
82        layout: &Layout<Self>,
83        name: Arc<str>,
84        segment_source: Arc<dyn SegmentSource>,
85        session: &VortexSession,
86        ctx: &LayoutReaderContext,
87    ) -> VortexResult<LayoutReaderRef> {
88        Ok(Arc::new(ChunkedReader::new(
89            layout.clone(),
90            name,
91            segment_source,
92            session,
93            ctx.clone(),
94        )))
95    }
96}
97
98impl Layout<Chunked> {
99    /// Construct a chunked layout.
100    pub fn new(row_count: u64, dtype: DType, children: Arc<dyn LayoutChildren>) -> Self {
101        let offsets = chunk_offsets(children.as_ref()).vortex_expect("chunk row counts overflow");
102        assert_eq!(
103            offsets.last().copied(),
104            Some(row_count),
105            "Row count mismatch"
106        );
107        LayoutParts::new(
108            Chunked,
109            dtype,
110            row_count,
111            Vec::new(),
112            children,
113            ChunkedData {
114                chunk_offsets: offsets,
115            },
116        )
117        .into_typed()
118    }
119
120    /// Rebuild this layout with owned children.
121    pub fn with_children(&self, children: Vec<LayoutRef>) -> Self {
122        Self::new(
123            self.row_count(),
124            self.dtype().clone(),
125            OwnedLayoutChildren::layout_children(children),
126        )
127    }
128}
129
130fn chunk_offsets(children: &dyn LayoutChildren) -> VortexResult<Vec<u64>> {
131    let mut offsets = Vec::with_capacity(children.nchildren() + 1);
132    offsets.push(0u64);
133    for idx in 0..children.nchildren() {
134        offsets.push(
135            offsets[idx]
136                .checked_add(children.child_row_count(idx))
137                .ok_or_else(|| vortex_err!("Chunked child row counts overflow"))?,
138        );
139    }
140    Ok(offsets)
141}
vortex_layout/layouts/chunked/mod.rs

vortex_layout/layouts/chunked/
mod.rs
