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::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 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}