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}