1use serde::{Deserialize, Serialize};
2use vortex_array::serde::ArrayParts;
3use vortex_array::vtable::EncodingVTable;
4use vortex_array::{
5 Array, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl, Canonical,
6 DeserializeMetadata, EncodingId, SerdeMetadata,
7};
8use vortex_dtype::{DType, Nullability, PType};
9use vortex_error::{VortexExpect, VortexResult};
10
11use crate::compress::runend_encode;
12use crate::{RunEndArray, RunEndEncoding};
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct RunEndMetadata {
16 ends_ptype: PType,
17 num_runs: usize,
18 offset: usize,
19}
20
21impl EncodingVTable for RunEndEncoding {
22 fn id(&self) -> EncodingId {
23 EncodingId::new_ref("vortex.runend")
24 }
25
26 fn decode(
27 &self,
28 parts: &ArrayParts,
29 ctx: &ArrayContext,
30 dtype: DType,
31 len: usize,
32 ) -> VortexResult<ArrayRef> {
33 let metadata = SerdeMetadata::<RunEndMetadata>::deserialize(parts.metadata())?;
34
35 let ends_dtype = DType::Primitive(metadata.ends_ptype, Nullability::NonNullable);
36 let ends = parts.child(0).decode(ctx, ends_dtype, metadata.num_runs)?;
37
38 let values = parts.child(1).decode(ctx, dtype, metadata.num_runs)?;
39
40 Ok(RunEndArray::with_offset_and_length(ends, values, metadata.offset, len)?.into_array())
41 }
42
43 fn encode(
44 &self,
45 input: &Canonical,
46 _like: Option<&dyn Array>,
47 ) -> VortexResult<Option<ArrayRef>> {
48 let parray = input.clone().into_primitive()?;
49
50 let (ends, values) = runend_encode(&parray)?;
51
52 Ok(Some(
53 RunEndArray::try_new(ends.to_array(), values)?.to_array(),
54 ))
55 }
56}
57
58impl ArrayVisitorImpl<SerdeMetadata<RunEndMetadata>> for RunEndArray {
59 fn _visit_children(&self, visitor: &mut dyn ArrayChildVisitor) {
60 visitor.visit_child("ends", self.ends());
61 visitor.visit_child("values", self.values());
62 }
63
64 fn _metadata(&self) -> SerdeMetadata<RunEndMetadata> {
65 SerdeMetadata(RunEndMetadata {
66 ends_ptype: PType::try_from(self.ends().dtype()).vortex_expect("Must be a valid PType"),
67 num_runs: self.ends().len(),
68 offset: self.offset(),
69 })
70 }
71}
72
73#[cfg(test)]
74mod tests {
75 use vortex_array::SerdeMetadata;
76 use vortex_array::test_harness::check_metadata;
77 use vortex_dtype::PType;
78
79 use super::*;
80
81 #[cfg_attr(miri, ignore)]
82 #[test]
83 fn test_runend_metadata() {
84 check_metadata(
85 "runend.metadata",
86 SerdeMetadata(RunEndMetadata {
87 offset: usize::MAX,
88 ends_ptype: PType::U64,
89 num_runs: usize::MAX,
90 }),
91 );
92 }
93}