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

vortex_layout/layouts/chunked/
mod.rs
