vortex_sparse/
variants.rs

1use vortex_array::variants::{
2    BinaryArrayTrait, BoolArrayTrait, DecimalArrayTrait, ExtensionArrayTrait, ListArrayTrait,
3    NullArrayTrait, PrimitiveArrayTrait, StructArrayTrait, Utf8ArrayTrait,
4};
5use vortex_array::{Array, ArrayVariants, ArrayVariantsImpl};
6use vortex_dtype::FieldName;
7use vortex_error::{VortexExpect, VortexResult, vortex_err};
8use vortex_scalar::StructScalar;
9
10use crate::{ArrayRef, SparseArray};
11
12/// Sparse arrays support all DTypes
13impl ArrayVariantsImpl for SparseArray {
14    fn _as_null_typed(&self) -> Option<&dyn NullArrayTrait> {
15        Some(self)
16    }
17
18    fn _as_bool_typed(&self) -> Option<&dyn BoolArrayTrait> {
19        Some(self)
20    }
21
22    fn _as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait> {
23        Some(self)
24    }
25
26    fn _as_decimal_typed(&self) -> Option<&dyn DecimalArrayTrait> {
27        Some(self)
28    }
29
30    fn _as_utf8_typed(&self) -> Option<&dyn Utf8ArrayTrait> {
31        Some(self)
32    }
33
34    fn _as_binary_typed(&self) -> Option<&dyn BinaryArrayTrait> {
35        Some(self)
36    }
37
38    fn _as_struct_typed(&self) -> Option<&dyn StructArrayTrait> {
39        Some(self)
40    }
41
42    fn _as_list_typed(&self) -> Option<&dyn ListArrayTrait> {
43        Some(self)
44    }
45
46    fn _as_extension_typed(&self) -> Option<&dyn ExtensionArrayTrait> {
47        Some(self)
48    }
49}
50
51impl NullArrayTrait for SparseArray {}
52
53impl BoolArrayTrait for SparseArray {}
54
55impl PrimitiveArrayTrait for SparseArray {}
56
57impl Utf8ArrayTrait for SparseArray {}
58
59impl BinaryArrayTrait for SparseArray {}
60
61impl DecimalArrayTrait for SparseArray {}
62
63impl StructArrayTrait for SparseArray {
64    fn maybe_null_field_by_idx(&self, idx: usize) -> VortexResult<ArrayRef> {
65        let new_patches = self.patches().clone().map_values(|values| {
66            values
67                .as_struct_typed()
68                .vortex_expect("Expected struct array")
69                .maybe_null_field_by_idx(idx)
70        })?;
71        let scalar = StructScalar::try_from(self.fill_scalar())?.field_by_idx(idx)?;
72
73        Ok(SparseArray::try_new_from_patches(new_patches, scalar)?.into_array())
74    }
75
76    fn project(&self, projection: &[FieldName]) -> VortexResult<ArrayRef> {
77        let new_patches = self.patches().clone().map_values(|values| {
78            values
79                .as_struct_typed()
80                .ok_or_else(|| vortex_err!("Chunk was not a StructArray"))?
81                .project(projection)
82        })?;
83        let scalar = StructScalar::try_from(self.fill_scalar())?.project(projection)?;
84
85        Ok(SparseArray::try_new_from_patches(new_patches, scalar)?.into_array())
86    }
87}
88
89impl ListArrayTrait for SparseArray {}
90
91impl ExtensionArrayTrait for SparseArray {
92    fn storage_data(&self) -> ArrayRef {
93        SparseArray::try_new_from_patches(
94            self.patches()
95                .clone()
96                .map_values(|values| {
97                    Ok(values
98                        .as_extension_typed()
99                        .vortex_expect("Expected extension array")
100                        .storage_data())
101                })
102                .vortex_expect("as_extension_array preserves the length"),
103            self.fill_scalar().clone(),
104        )
105        .vortex_expect("Failed to create new sparse array")
106        .into_array()
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use vortex_array::arrays::BoolArray;
113    use vortex_array::compute::invert;
114    use vortex_array::{Array, IntoArray, ToCanonical};
115    use vortex_buffer::buffer;
116    use vortex_scalar::Scalar;
117
118    use crate::SparseArray;
119
120    #[test]
121    fn invert_bools_non_null_fill() {
122        let sparse_bools = SparseArray::try_new(
123            buffer![0u64].into_array(),
124            BoolArray::from_iter([false]).into_array(),
125            2,
126            Scalar::from(true),
127        )
128        .unwrap();
129        let inverted = invert(&sparse_bools).unwrap();
130        assert_eq!(
131            inverted
132                .to_bool()
133                .unwrap()
134                .boolean_buffer()
135                .iter()
136                .collect::<Vec<_>>(),
137            vec![true, false]
138        );
139    }
140}
vortex_sparse/variants.rs

vortex_sparse/
variants.rs
