vortex_runend/
serde.rs

1use vortex_array::serde::ArrayChildren;
2use vortex_array::vtable::{EncodeVTable, SerdeVTable, VisitorVTable};
3use vortex_array::{
4    Array, ArrayBufferVisitor, ArrayChildVisitor, Canonical, DeserializeMetadata, ProstMetadata,
5};
6use vortex_buffer::ByteBuffer;
7use vortex_dtype::{DType, Nullability, PType};
8use vortex_error::{VortexExpect, VortexResult};
9
10use crate::compress::runend_encode;
11use crate::{RunEndArray, RunEndEncoding, RunEndVTable};
12
13#[derive(Clone, prost::Message)]
14pub struct RunEndMetadata {
15    #[prost(enumeration = "PType", tag = "1")]
16    ends_ptype: i32,
17    #[prost(uint64, tag = "2")]
18    num_runs: u64,
19    #[prost(uint64, tag = "3")]
20    offset: u64,
21}
22
23impl SerdeVTable<RunEndVTable> for RunEndVTable {
24    type Metadata = ProstMetadata<RunEndMetadata>;
25
26    fn metadata(array: &RunEndArray) -> VortexResult<Option<Self::Metadata>> {
27        Ok(Some(ProstMetadata(RunEndMetadata {
28            ends_ptype: PType::try_from(array.ends().dtype()).vortex_expect("Must be a valid PType")
29                as i32,
30            num_runs: array.ends().len() as u64,
31            offset: array.offset() as u64,
32        })))
33    }
34
35    fn build(
36        _encoding: &RunEndEncoding,
37        dtype: &DType,
38        len: usize,
39        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
40        _buffers: &[ByteBuffer],
41        children: &dyn ArrayChildren,
42    ) -> VortexResult<RunEndArray> {
43        let ends_dtype = DType::Primitive(metadata.ends_ptype(), Nullability::NonNullable);
44        let runs = usize::try_from(metadata.num_runs).vortex_expect("Must be a valid usize");
45        let ends = children.get(0, &ends_dtype, runs)?;
46
47        let values = children.get(1, dtype, runs)?;
48
49        RunEndArray::with_offset_and_length(
50            ends,
51            values,
52            usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize"),
53            len,
54        )
55    }
56}
57
58impl EncodeVTable<RunEndVTable> for RunEndVTable {
59    fn encode(
60        _encoding: &RunEndEncoding,
61        canonical: &Canonical,
62        _like: Option<&RunEndArray>,
63    ) -> VortexResult<Option<RunEndArray>> {
64        let parray = canonical.clone().into_primitive()?;
65        let (ends, values) = runend_encode(&parray)?;
66        Ok(Some(RunEndArray::try_new(ends.to_array(), values)?))
67    }
68}
69
70impl VisitorVTable<RunEndVTable> for RunEndVTable {
71    fn visit_buffers(_array: &RunEndArray, _visitor: &mut dyn ArrayBufferVisitor) {}
72
73    fn visit_children(array: &RunEndArray, visitor: &mut dyn ArrayChildVisitor) {
74        visitor.visit_child("ends", array.ends());
75        visitor.visit_child("values", array.values());
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use vortex_array::ProstMetadata;
82    use vortex_array::test_harness::check_metadata;
83    use vortex_dtype::PType;
84
85    use super::*;
86
87    #[cfg_attr(miri, ignore)]
88    #[test]
89    fn test_runend_metadata() {
90        check_metadata(
91            "runend.metadata",
92            ProstMetadata(RunEndMetadata {
93                ends_ptype: PType::U64 as i32,
94                num_runs: u64::MAX,
95                offset: u64::MAX,
96            }),
97        );
98    }
99}