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

vortex_fastlanes/rle/vtable/
mod.rs
