vortex_sparse/
variants.rs

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

vortex_sparse/
variants.rs
