zarrs/array/element/
optional.rs

1use crate::array::{ArrayBytes, DataType};
2
3use super::{Element, ElementError, ElementOwned};
4
5use ElementError::IncompatibleElementType as IET;
6
7impl<T> Element for Option<T>
8where
9    T: Element + Default,
10{
11    fn validate_data_type(data_type: &DataType) -> Result<(), ElementError> {
12        let opt = data_type.as_optional().ok_or(IET)?;
13        T::validate_data_type(opt.data_type())
14    }
15
16    fn to_array_bytes<'a>(
17        data_type: &DataType,
18        elements: &'a [Self],
19    ) -> Result<ArrayBytes<'a>, ElementError> {
20        Self::validate_data_type(data_type)?;
21
22        let opt = data_type.as_optional().ok_or(IET)?;
23
24        let num_elements = elements.len();
25
26        // Create validity mask - one byte per element
27        let mut mask = Vec::with_capacity(num_elements);
28
29        // Create dense data - all elements, using default/zero for None values
30        // We need to use a placeholder value for None elements
31        let default_value = T::default();
32        let mut dense_elements = Vec::with_capacity(num_elements);
33
34        for element in elements {
35            if let Some(value) = element {
36                mask.push(1u8);
37                dense_elements.push(value.clone());
38            } else {
39                mask.push(0u8);
40                dense_elements.push(default_value.clone());
41            }
42        }
43
44        // Convert all elements (dense) to ArrayBytes
45        let data = T::into_array_bytes(opt.data_type(), dense_elements)?.into_owned();
46
47        // Create optional ArrayBytes by adding mask to the data
48        Ok(data.with_optional_mask(mask))
49    }
50
51    fn into_array_bytes(
52        data_type: &DataType,
53        elements: Vec<Self>,
54    ) -> Result<ArrayBytes<'static>, ElementError> {
55        Ok(Self::to_array_bytes(data_type, &elements)?.into_owned())
56    }
57}
58
59impl<T> ElementOwned for Option<T>
60where
61    T: ElementOwned + Clone + Default,
62{
63    fn from_array_bytes(
64        data_type: &DataType,
65        bytes: ArrayBytes<'_>,
66    ) -> Result<Vec<Self>, ElementError> {
67        Self::validate_data_type(data_type)?;
68
69        let opt = data_type.as_optional().ok_or(IET)?;
70
71        // Extract mask and dense data from optional ArrayBytes
72        let optional_bytes = bytes.into_optional()?;
73        let (data, mask) = optional_bytes.into_parts();
74
75        // Convert the dense inner data to a Vec<T>
76        let dense_values = T::from_array_bytes(opt.data_type(), *data)?;
77
78        // Build the result vector using mask to determine Some vs None
79        let mut elements = Vec::with_capacity(mask.len());
80        for (i, &mask_byte) in mask.iter().enumerate() {
81            if mask_byte == 0 {
82                // None value
83                elements.push(None);
84            } else {
85                // Some value - take from dense data
86                if i >= dense_values.len() {
87                    return Err(ElementError::Other(format!(
88                        "Not enough dense values for mask at index {i}"
89                    )));
90                }
91                elements.push(Some(dense_values[i].clone()));
92            }
93        }
94
95        Ok(elements)
96    }
97}
zarrs/array/element/optional.rs

zarrs/array/element/
optional.rs
