vortex_runend/
serde.rs

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