vortex_fastlanes/rle/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use prost::Message;
5use vortex_array::ArrayRef;
6use vortex_array::ProstMetadata;
7use vortex_array::buffer::BufferHandle;
8use vortex_array::serde::ArrayChildren;
9use vortex_array::vtable;
10use vortex_array::vtable::ArrayId;
11use vortex_array::vtable::ArrayVTable;
12use vortex_array::vtable::ArrayVTableExt;
13use vortex_array::vtable::NotSupported;
14use vortex_array::vtable::VTable;
15use vortex_array::vtable::ValidityVTableFromChildSliceHelper;
16use vortex_dtype::DType;
17use vortex_dtype::Nullability;
18use vortex_dtype::PType;
19use vortex_error::VortexResult;
20use vortex_error::vortex_ensure;
21
22use crate::RLEArray;
23
24mod array;
25mod canonical;
26mod encode;
27mod operations;
28mod validity;
29mod visitor;
30
31vtable!(RLE);
32
33#[derive(Clone, prost::Message)]
34pub struct RLEMetadata {
35    #[prost(uint64, tag = "1")]
36    pub values_len: u64,
37    #[prost(uint64, tag = "2")]
38    pub indices_len: u64,
39    #[prost(enumeration = "PType", tag = "3")]
40    pub indices_ptype: i32,
41    #[prost(uint64, tag = "4")]
42    pub values_idx_offsets_len: u64,
43    #[prost(enumeration = "PType", tag = "5")]
44    pub values_idx_offsets_ptype: i32,
45    #[prost(uint64, tag = "6", default = "0")]
46    pub offset: u64,
47}
48
49impl VTable for RLEVTable {
50    type Array = RLEArray;
51
52    type Metadata = ProstMetadata<RLEMetadata>;
53
54    type ArrayVTable = Self;
55    type CanonicalVTable = Self;
56    type OperationsVTable = Self;
57    type ValidityVTable = ValidityVTableFromChildSliceHelper;
58    type VisitorVTable = Self;
59    type ComputeVTable = NotSupported;
60    type EncodeVTable = Self;
61
62    fn id(&self) -> ArrayId {
63        ArrayId::new_ref("fastlanes.rle")
64    }
65
66    fn encoding(_array: &Self::Array) -> ArrayVTable {
67        RLEVTable.as_vtable()
68    }
69
70    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
71        // RLEArray children order (from visit_children):
72        // 1. values
73        // 2. indices
74        // 3. values_idx_offsets
75
76        vortex_ensure!(
77            children.len() == 3,
78            "Expected 3 children for RLE encoding, got {}",
79            children.len()
80        );
81
82        array.values = children[0].clone();
83        array.indices = children[1].clone();
84        array.values_idx_offsets = children[2].clone();
85
86        Ok(())
87    }
88
89    fn metadata(array: &RLEArray) -> VortexResult<Self::Metadata> {
90        Ok(ProstMetadata(RLEMetadata {
91            values_len: array.values().len() as u64,
92            indices_len: array.indices().len() as u64,
93            indices_ptype: PType::try_from(array.indices().dtype())? as i32,
94            values_idx_offsets_len: array.values_idx_offsets().len() as u64,
95            values_idx_offsets_ptype: PType::try_from(array.values_idx_offsets().dtype())? as i32,
96            offset: array.offset() as u64,
97        }))
98    }
99
100    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
101        Ok(Some(metadata.0.encode_to_vec()))
102    }
103
104    fn deserialize(buffer: &[u8]) -> VortexResult<Self::Metadata> {
105        Ok(ProstMetadata(RLEMetadata::decode(buffer)?))
106    }
107
108    fn build(
109        &self,
110        dtype: &DType,
111        len: usize,
112        metadata: &Self::Metadata,
113        _buffers: &[BufferHandle],
114        children: &dyn ArrayChildren,
115    ) -> VortexResult<RLEArray> {
116        let metadata = &metadata.0;
117        let values = children.get(
118            0,
119            &DType::Primitive(dtype.as_ptype(), Nullability::NonNullable),
120            usize::try_from(metadata.values_len)?,
121        )?;
122
123        let indices = children.get(
124            1,
125            &DType::Primitive(metadata.indices_ptype(), dtype.nullability()),
126            usize::try_from(metadata.indices_len)?,
127        )?;
128
129        let values_idx_offsets = children.get(
130            2,
131            &DType::Primitive(
132                metadata.values_idx_offsets_ptype(),
133                Nullability::NonNullable,
134            ),
135            usize::try_from(metadata.values_idx_offsets_len)?,
136        )?;
137
138        RLEArray::try_new(
139            values,
140            indices,
141            values_idx_offsets,
142            metadata.offset as usize,
143            len,
144        )
145    }
146}
147
148#[derive(Debug)]
149pub struct RLEVTable;
150
151#[cfg(test)]
152mod tests {
153    use vortex_array::test_harness::check_metadata;
154
155    use super::ProstMetadata;
156    use super::RLEMetadata;
157
158    #[cfg_attr(miri, ignore)]
159    #[test]
160    fn test_rle_metadata() {
161        check_metadata(
162            "rle.metadata",
163            ProstMetadata(RLEMetadata {
164                values_len: u64::MAX,
165                indices_len: u64::MAX,
166                indices_ptype: i32::MAX,
167                values_idx_offsets_len: u64::MAX,
168                values_idx_offsets_ptype: i32::MAX,
169                offset: u64::MAX,
170            }),
171        );
172    }
173}
vortex_fastlanes/rle/vtable/mod.rs

vortex_fastlanes/rle/vtable/
mod.rs
