vortex_runend/
serde.rs

1use vortex_array::serde::ArrayParts;
2use vortex_array::vtable::EncodingVTable;
3use vortex_array::{
4    Array, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl, Canonical,
5    DeserializeMetadata, EncodingId, ProstMetadata,
6};
7use vortex_dtype::{DType, Nullability, PType};
8use vortex_error::{VortexExpect, VortexResult};
9
10use crate::compress::runend_encode;
11use crate::{RunEndArray, RunEndEncoding};
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 EncodingVTable for RunEndEncoding {
24    fn id(&self) -> EncodingId {
25        EncodingId::new_ref("vortex.runend")
26    }
27
28    fn decode(
29        &self,
30        parts: &ArrayParts,
31        ctx: &ArrayContext,
32        dtype: DType,
33        len: usize,
34    ) -> VortexResult<ArrayRef> {
35        let metadata = ProstMetadata::<RunEndMetadata>::deserialize(parts.metadata())?;
36
37        let ends_dtype = DType::Primitive(metadata.ends_ptype(), Nullability::NonNullable);
38        let runs = usize::try_from(metadata.num_runs).vortex_expect("Must be a valid usize");
39        let ends = parts.child(0).decode(ctx, ends_dtype, runs)?;
40
41        let values = parts.child(1).decode(ctx, dtype, runs)?;
42
43        Ok(RunEndArray::with_offset_and_length(
44            ends,
45            values,
46            usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize"),
47            len,
48        )?
49        .into_array())
50    }
51
52    fn encode(
53        &self,
54        input: &Canonical,
55        _like: Option<&dyn Array>,
56    ) -> VortexResult<Option<ArrayRef>> {
57        let parray = input.clone().into_primitive()?;
58
59        let (ends, values) = runend_encode(&parray)?;
60
61        Ok(Some(
62            RunEndArray::try_new(ends.to_array(), values)?.to_array(),
63        ))
64    }
65}
66
67impl ArrayVisitorImpl<ProstMetadata<RunEndMetadata>> for RunEndArray {
68    fn _visit_children(&self, visitor: &mut dyn ArrayChildVisitor) {
69        visitor.visit_child("ends", self.ends());
70        visitor.visit_child("values", self.values());
71    }
72
73    fn _metadata(&self) -> ProstMetadata<RunEndMetadata> {
74        ProstMetadata(RunEndMetadata {
75            ends_ptype: PType::try_from(self.ends().dtype()).vortex_expect("Must be a valid PType")
76                as i32,
77            num_runs: self.ends().len() as u64,
78            offset: self.offset() as u64,
79        })
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use vortex_array::ProstMetadata;
86    use vortex_array::test_harness::check_metadata;
87    use vortex_dtype::PType;
88
89    use super::*;
90
91    #[cfg_attr(miri, ignore)]
92    #[test]
93    fn test_runend_metadata() {
94        check_metadata(
95            "runend.metadata",
96            ProstMetadata(RunEndMetadata {
97                ends_ptype: PType::U64 as i32,
98                num_runs: u64::MAX,
99                offset: u64::MAX,
100            }),
101        );
102    }
103}
vortex_runend/serde.rs

vortex_runend/
serde.rs
