vortex_array/arrays/struct_/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_dtype::DType;
6use vortex_error::{VortexExpect, VortexResult, vortex_bail};
7
8use crate::arrays::{StructArray, StructVTable};
9use crate::compute::{CastKernel, CastKernelAdapter, cast};
10use crate::vtable::ValidityHelper;
11use crate::{ArrayRef, IntoArray, register_kernel};
12
13impl CastKernel for StructVTable {
14    fn cast(&self, array: &StructArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
15        let Some(target_sdtype) = dtype.as_struct_opt() else {
16            return Ok(None);
17        };
18
19        let source_sdtype = array
20            .dtype()
21            .as_struct_opt()
22            .vortex_expect("struct array must have struct dtype");
23
24        if target_sdtype.names() != source_sdtype.names() {
25            vortex_bail!("cannot cast {} to {}", array.dtype(), dtype);
26        }
27
28        let validity = array
29            .validity()
30            .clone()
31            .cast_nullability(dtype.nullability())?;
32
33        StructArray::try_new(
34            target_sdtype.names().clone(),
35            array
36                .fields()
37                .iter()
38                .zip_eq(target_sdtype.fields())
39                .map(|(field, dtype)| cast(field, &dtype))
40                .try_collect()?,
41            array.len(),
42            validity,
43        )
44        .map(|a| Some(a.into_array()))
45    }
46}
47
48register_kernel!(CastKernelAdapter(StructVTable).lift());
49
50#[cfg(test)]
51mod tests {
52    use rstest::rstest;
53    use vortex_buffer::buffer;
54    use vortex_dtype::{DType, FieldNames, Nullability};
55
56    use crate::IntoArray;
57    use crate::arrays::{StructArray, VarBinArray};
58    use crate::compute::conformance::cast::test_cast_conformance;
59
60    #[rstest]
61    #[case(create_test_struct(false))]
62    #[case(create_test_struct(true))]
63    #[case(create_nested_struct())]
64    #[case(create_simple_struct())]
65    fn test_cast_struct_conformance(#[case] array: StructArray) {
66        test_cast_conformance(array.as_ref());
67    }
68
69    fn create_test_struct(nullable: bool) -> StructArray {
70        let names: FieldNames = vec!["a".into(), "b".into()].into();
71
72        let a = buffer![1i32, 2, 3].into_array();
73        let b = VarBinArray::from_iter(
74            vec![Some("x"), None, Some("z")],
75            DType::Utf8(Nullability::Nullable),
76        )
77        .into_array();
78
79        StructArray::try_new(
80            names,
81            vec![a, b],
82            3,
83            if nullable {
84                crate::validity::Validity::AllValid
85            } else {
86                crate::validity::Validity::NonNullable
87            },
88        )
89        .unwrap()
90    }
91
92    fn create_nested_struct() -> StructArray {
93        // Create inner struct
94        let inner_names: FieldNames = vec!["x".into(), "y".into()].into();
95
96        let x = buffer![1.0f32, 2.0, 3.0].into_array();
97        let y = buffer![4.0f32, 5.0, 6.0].into_array();
98        let inner_struct = StructArray::try_new(
99            inner_names,
100            vec![x, y],
101            3,
102            crate::validity::Validity::NonNullable,
103        )
104        .unwrap()
105        .into_array();
106
107        // Create outer struct with inner struct as a field
108        let outer_names: FieldNames = vec!["id".into(), "point".into()].into();
109        // Outer struct would have fields: id (I64) and point (inner struct)
110
111        let ids = buffer![100i64, 200, 300].into_array();
112
113        StructArray::try_new(
114            outer_names,
115            vec![ids, inner_struct],
116            3,
117            crate::validity::Validity::NonNullable,
118        )
119        .unwrap()
120    }
121
122    fn create_simple_struct() -> StructArray {
123        let names: FieldNames = vec!["value".into()].into();
124        // Simple struct with a single U8 field
125
126        let values = buffer![42u8].into_array();
127
128        StructArray::try_new(
129            names,
130            vec![values],
131            1,
132            crate::validity::Validity::NonNullable,
133        )
134        .unwrap()
135    }
136}
vortex_array/arrays/struct_/compute/cast.rs

vortex_array/arrays/struct_/compute/
cast.rs
