vortex_runend/
serde.rs

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