1use 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}