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::ArrayContext;
10use vortex_array::DeserializeMetadata;
11use vortex_array::EmptyMetadata;
12use vortex_dtype::DType;
13use vortex_error::VortexResult;
14use vortex_session::VortexSession;
15
16use crate::LayoutChildType;
17use crate::LayoutEncodingRef;
18use crate::LayoutId;
19use crate::LayoutReaderRef;
20use crate::LayoutRef;
21use crate::VTable;
22use crate::children::LayoutChildren;
23use crate::layouts::chunked::reader::ChunkedReader;
24use crate::segments::SegmentId;
25use crate::segments::SegmentSource;
26use crate::vtable;
27
28vtable!(Chunked);
29
30impl VTable for ChunkedVTable {
31    type Layout = ChunkedLayout;
32    type Encoding = ChunkedLayoutEncoding;
33    type Metadata = EmptyMetadata;
34
35    fn id(_encoding: &Self::Encoding) -> LayoutId {
36        LayoutId::new_ref("vortex.chunked")
37    }
38
39    fn encoding(_layout: &Self::Layout) -> LayoutEncodingRef {
40        LayoutEncodingRef::new_ref(ChunkedLayoutEncoding.as_ref())
41    }
42
43    fn row_count(layout: &Self::Layout) -> u64 {
44        layout.row_count
45    }
46
47    fn dtype(layout: &Self::Layout) -> &DType {
48        &layout.dtype
49    }
50
51    fn metadata(_layout: &Self::Layout) -> Self::Metadata {
52        EmptyMetadata
53    }
54
55    fn segment_ids(_layout: &Self::Layout) -> Vec<SegmentId> {
56        vec![]
57    }
58
59    fn nchildren(layout: &Self::Layout) -> usize {
60        layout.children.nchildren()
61    }
62
63    fn child(layout: &Self::Layout, idx: usize) -> VortexResult<LayoutRef> {
64        layout.children.child(idx, Self::dtype(layout))
65    }
66
67    fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType {
68        LayoutChildType::Chunk((idx, layout.chunk_offsets[idx]))
69    }
70
71    fn new_reader(
72        layout: &Self::Layout,
73        name: Arc<str>,
74        segment_source: Arc<dyn SegmentSource>,
75        session: &VortexSession,
76    ) -> VortexResult<LayoutReaderRef> {
77        Ok(Arc::new(ChunkedReader::new(
78            layout.clone(),
79            name,
80            segment_source,
81            session,
82        )))
83    }
84
85    #[cfg(gpu_unstable)]
86    fn new_gpu_reader(
87        layout: &Self::Layout,
88        name: Arc<str>,
89        segment_source: Arc<dyn SegmentSource>,
90        ctx: Arc<cudarc::driver::CudaContext>,
91    ) -> VortexResult<crate::gpu::GpuLayoutReaderRef> {
92        Ok(Arc::new(
93            crate::gpu::layouts::chunked::GpuChunkedLayoutReader::new(
94                layout.clone(),
95                name,
96                segment_source,
97                ctx,
98            ),
99        ))
100    }
101
102    fn build(
103        _encoding: &Self::Encoding,
104        dtype: &DType,
105        row_count: u64,
106        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
107        _segment_ids: Vec<SegmentId>,
108        children: &dyn LayoutChildren,
109        _ctx: ArrayContext,
110    ) -> VortexResult<Self::Layout> {
111        Ok(ChunkedLayout::new(
112            row_count,
113            dtype.clone(),
114            children.to_arc(),
115        ))
116    }
117}
118
119#[derive(Debug)]
120pub struct ChunkedLayoutEncoding;
121
122#[derive(Clone, Debug)]
123pub struct ChunkedLayout {
124    row_count: u64,
125    dtype: DType,
126    children: Arc<dyn LayoutChildren>,
127    chunk_offsets: Vec<u64>,
128}
129
130impl ChunkedLayout {
131    pub fn new(row_count: u64, dtype: DType, children: Arc<dyn LayoutChildren>) -> Self {
132        let mut chunk_offsets = Vec::with_capacity(children.nchildren() + 1);
133
134        chunk_offsets.push(0);
135        for i in 0..children.nchildren() {
136            chunk_offsets.push(chunk_offsets[i] + children.child_row_count(i));
137        }
138        assert_eq!(
139            chunk_offsets[children.nchildren()],
140            row_count,
141            "Row count mismatch"
142        );
143        Self {
144            row_count,
145            dtype,
146            children,
147            chunk_offsets,
148        }
149    }
150
151    pub fn children(&self) -> &Arc<dyn LayoutChildren> {
152        &self.children
153    }
154}
vortex_layout/layouts/chunked/mod.rs

vortex_layout/layouts/chunked/
mod.rs
