vortex_array/arrays/list/
serde.rs1use vortex_dtype::{DType, Nullability, PType};
2use vortex_error::{VortexExpect, VortexResult, vortex_bail};
3
4use super::{ListArray, ListEncoding};
5use crate::serde::ArrayParts;
6use crate::validity::Validity;
7use crate::vtable::EncodingVTable;
8use crate::{
9 Array, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl, DeserializeMetadata,
10 EncodingId, RkyvMetadata,
11};
12
13impl EncodingVTable for ListEncoding {
14 fn id(&self) -> EncodingId {
15 EncodingId::new_ref("vortex.list")
16 }
17
18 fn decode(
19 &self,
20 parts: &ArrayParts,
21 ctx: &ArrayContext,
22 dtype: DType,
23 len: usize,
24 ) -> VortexResult<ArrayRef> {
25 let metadata = RkyvMetadata::<ListMetadata>::deserialize(parts.metadata())?;
26
27 let validity = if parts.nchildren() == 2 {
28 Validity::from(dtype.nullability())
29 } else if parts.nchildren() == 3 {
30 let validity = parts.child(2).decode(ctx, Validity::DTYPE, len)?;
31 Validity::Array(validity)
32 } else {
33 vortex_bail!("Expected 2 or 3 children, got {}", parts.nchildren());
34 };
35
36 let DType::List(element_dtype, _) = &dtype else {
37 vortex_bail!("Expected List dtype, got {:?}", dtype);
38 };
39 let elements =
40 parts
41 .child(0)
42 .decode(ctx, element_dtype.as_ref().clone(), metadata.elements_len)?;
43
44 let offsets = parts.child(1).decode(
45 ctx,
46 DType::Primitive(metadata.offset_ptype, Nullability::NonNullable),
47 len + 1,
48 )?;
49
50 Ok(ListArray::try_new(elements, offsets, validity)?.into_array())
51 }
52}
53
54#[derive(Debug, Clone, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
55pub struct ListMetadata {
56 elements_len: usize,
57 offset_ptype: PType,
58}
59
60impl ArrayVisitorImpl<RkyvMetadata<ListMetadata>> for ListArray {
61 fn _visit_children(&self, visitor: &mut dyn ArrayChildVisitor) {
62 visitor.visit_child("elements", self.elements());
63 visitor.visit_child("offsets", self.offsets());
64 visitor.visit_validity(self.validity(), self.len());
65 }
66
67 fn _metadata(&self) -> RkyvMetadata<ListMetadata> {
68 RkyvMetadata(ListMetadata {
69 elements_len: self.elements().len(),
70 offset_ptype: PType::try_from(self.offsets().dtype())
71 .vortex_expect("Must be a valid PType"),
72 })
73 }
74}