vortex_layout/layouts/chunked/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod reader;
5pub mod writer;
6
7use std::sync::Arc;
8
9use vortex_array::DeserializeMetadata;
10use vortex_array::EmptyMetadata;
11use vortex_array::dtype::DType;
12use vortex_error::VortexResult;
13use vortex_session::VortexSession;
14use vortex_session::registry::CachedId;
15
16use crate::LayoutBuildContext;
17use crate::LayoutChildType;
18use crate::LayoutEncodingRef;
19use crate::LayoutId;
20use crate::LayoutReaderContext;
21use crate::LayoutReaderRef;
22use crate::LayoutRef;
23use crate::VTable;
24use crate::children::LayoutChildren;
25use crate::children::OwnedLayoutChildren;
26use crate::layouts::chunked::reader::ChunkedReader;
27use crate::segments::SegmentId;
28use crate::segments::SegmentSource;
29use crate::vtable;
30
31vtable!(Chunked);
32
33impl VTable for Chunked {
34    type Layout = ChunkedLayout;
35    type Encoding = ChunkedLayoutEncoding;
36    type Metadata = EmptyMetadata;
37
38    fn id(_encoding: &Self::Encoding) -> LayoutId {
39        static ID: CachedId = CachedId::new("vortex.chunked");
40        *ID
41    }
42
43    fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef {
44        LayoutEncodingRef::new_ref(ChunkedLayoutEncoding.as_ref())
45    }
46
47    fn row_count(layout: &Self::Layout) -> u64 {
48        layout.row_count
49    }
50
51    fn dtype(layout: &Self::Layout) -> &DType {
52        &layout.dtype
53    }
54
55    fn metadata(_layout: &Self::Layout) -> Self::Metadata {
56        EmptyMetadata
57    }
58
59    fn segment_ids(_layout: &Self::Layout) -> Vec<SegmentId> {
60        vec![]
61    }
62
63    fn nchildren(layout: &Self::Layout) -> usize {
64        layout.children.nchildren()
65    }
66
67    fn child(layout: &Self::Layout, idx: usize) -> VortexResult<LayoutRef> {
68        layout.children.child(idx, Self::dtype(layout))
69    }
70
71    fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType {
72        LayoutChildType::Chunk((idx, layout.chunk_offsets[idx]))
73    }
74
75    fn new_reader(
76        layout: &Self::Layout,
77        name: Arc<str>,
78        segment_source: Arc<dyn SegmentSource>,
79        session: &VortexSession,
80        ctx: &LayoutReaderContext,
81    ) -> VortexResult<LayoutReaderRef> {
82        Ok(Arc::new(ChunkedReader::new(
83            layout.clone(),
84            name,
85            segment_source,
86            session,
87            ctx.clone(),
88        )))
89    }
90
91    fn build(
92        _encoding: &Self::Encoding,
93        dtype: &DType,
94        row_count: u64,
95        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
96        _segment_ids: Vec<SegmentId>,
97        children: &dyn LayoutChildren,
98        _build_ctx: &LayoutBuildContext<'_>,
99    ) -> VortexResult<Self::Layout> {
100        Ok(ChunkedLayout::new(
101            row_count,
102            dtype.clone(),
103            children.to_arc(),
104        ))
105    }
106
107    fn with_children(layout: &mut Self::Layout, children: Vec<LayoutRef>) -> VortexResult<()> {
108        let new_children = OwnedLayoutChildren::layout_children(children);
109
110        // Recalculate chunk offsets based on new children
111        let mut chunk_offsets = vec![0; new_children.nchildren() + 1];
112        for i in 0..new_children.nchildren() {
113            chunk_offsets[i + 1] = chunk_offsets[i] + new_children.child_row_count(i);
114        }
115
116        layout.children = new_children;
117        layout.chunk_offsets = chunk_offsets;
118        Ok(())
119    }
120}
121
122#[derive(Debug)]
123pub struct ChunkedLayoutEncoding;
124
125/// Partitions a column into row-based chunks so that each chunk can be read independently.
126///
127/// Used to break large columns into smaller pieces for parallel I/O and to limit memory
128/// usage when scanning.
129#[derive(Clone, Debug)]
130pub struct ChunkedLayout {
131    row_count: u64,
132    dtype: DType,
133    children: Arc<dyn LayoutChildren>,
134    chunk_offsets: Vec<u64>,
135}
136
137impl ChunkedLayout {
138    pub fn new(row_count: u64, dtype: DType, children: Arc<dyn LayoutChildren>) -> Self {
139        let mut chunk_offsets = vec![0; children.nchildren() + 1];
140        for i in 0..children.nchildren() {
141            chunk_offsets[i + 1] = chunk_offsets[i] + children.child_row_count(i);
142        }
143
144        assert_eq!(
145            chunk_offsets[children.nchildren()],
146            row_count,
147            "Row count mismatch"
148        );
149        Self {
150            row_count,
151            dtype,
152            children,
153            chunk_offsets,
154        }
155    }
156
157    pub fn children(&self) -> &Arc<dyn LayoutChildren> {
158        &self.children
159    }
160}
vortex_layout/layouts/chunked/mod.rs

vortex_layout/layouts/chunked/
mod.rs
