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::try_new_offset_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        // SAFETY: runend_decode implementation must return valid RunEndArray
70        //  components.
71        unsafe {
72            Ok(Some(RunEndArray::new_unchecked(
73                ends.to_array(),
74                values,
75                0,
76                parray.len(),
77            )))
78        }
79    }
80}
81
82impl VisitorVTable<RunEndVTable> for RunEndVTable {
83    fn visit_buffers(_array: &RunEndArray, _visitor: &mut dyn ArrayBufferVisitor) {}
84
85    fn visit_children(array: &RunEndArray, visitor: &mut dyn ArrayChildVisitor) {
86        visitor.visit_child("ends", array.ends());
87        visitor.visit_child("values", array.values());
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use vortex_array::ProstMetadata;
94    use vortex_array::test_harness::check_metadata;
95    use vortex_dtype::PType;
96
97    use super::*;
98
99    #[cfg_attr(miri, ignore)]
100    #[test]
101    fn test_runend_metadata() {
102        check_metadata(
103            "runend.metadata",
104            ProstMetadata(RunEndMetadata {
105                ends_ptype: PType::U64 as i32,
106                num_runs: u64::MAX,
107                offset: u64::MAX,
108            }),
109        );
110    }
111}