vortex_array/arrays/struct_/vtable/
operator.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_error::VortexResult;
7use vortex_vector::Vector;
8use vortex_vector::struct_::StructVector;
9
10use crate::ArrayRef;
11use crate::arrays::{StructArray, StructVTable};
12use crate::execution::{BatchKernelRef, BindCtx, kernel};
13use crate::vtable::{OperatorVTable, ValidityHelper};
14
15impl OperatorVTable<StructVTable> for StructVTable {
16    fn bind(
17        array: &StructArray,
18        selection: Option<&ArrayRef>,
19        ctx: &mut dyn BindCtx,
20    ) -> VortexResult<BatchKernelRef> {
21        // Bind all child field arrays with the selection.
22        let field_kernels: Vec<_> = array
23            .fields()
24            .iter()
25            .map(|field| ctx.bind(field, selection))
26            .collect::<VortexResult<_>>()?;
27        let validity = ctx.bind_validity(array.validity(), array.len(), selection)?;
28
29        Ok(kernel(move || {
30            // Execute all child field kernels.
31            let fields: Vec<Vector> = field_kernels
32                .into_iter()
33                .map(|k| k.execute())
34                .collect::<VortexResult<_>>()?;
35            let validity_mask = validity.execute()?;
36
37            Ok(StructVector::try_new(Arc::new(fields.into_boxed_slice()), validity_mask)?.into())
38        }))
39    }
40}
41
42#[cfg(test)]
43mod tests {
44    use vortex_dtype::{FieldNames, PTypeDowncast};
45    use vortex_mask::Mask;
46    use vortex_vector::VectorOps;
47
48    use crate::IntoArray;
49    use crate::arrays::{BoolArray, PrimitiveArray, StructArray};
50    use crate::validity::Validity;
51
52    #[test]
53    fn test_struct_operator_basic() {
54        // Create a struct array with two fields: integers and booleans.
55        let int_field = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]);
56        let bool_field = BoolArray::from_iter([true, false, true, false, true]);
57
58        let struct_array = StructArray::try_new(
59            FieldNames::from(["ints", "bools"]),
60            vec![int_field.into_array(), bool_field.into_array()],
61            5,
62            Validity::AllValid,
63        )
64        .unwrap();
65
66        // Execute without selection.
67        let result = struct_array.execute().unwrap();
68        assert_eq!(result.len(), 5);
69
70        // Verify the struct vector fields.
71        let struct_vector = result.as_struct();
72        let fields = struct_vector.fields();
73        assert_eq!(fields.len(), 2);
74
75        // Verify the integer field values match the original.
76        let int_vector = fields[0].as_primitive().clone().into_i32();
77        assert_eq!(int_vector.elements().as_slice(), &[1, 2, 3, 4, 5]);
78
79        // Verify the boolean field values match the original.
80        let bool_vector = fields[1].as_bool();
81        let bool_values: Vec<bool> = (0..5).map(|i| bool_vector.bits().value(i)).collect();
82        assert_eq!(bool_values, vec![true, false, true, false, true]);
83    }
84
85    #[test]
86    fn test_struct_operator_with_mask() {
87        // Create a struct array with two fields.
88        let int_field = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]);
89        let bool_field = BoolArray::from_iter([true, false, true, false, true, false]);
90
91        let struct_array = StructArray::try_new(
92            FieldNames::from(["numbers", "flags"]),
93            vec![int_field.into_array(), bool_field.into_array()],
94            6,
95            Validity::AllValid,
96        )
97        .unwrap();
98
99        // Create a selection mask that selects indices 0, 2, 4 (alternating pattern).
100        let selection = Mask::from_iter([true, false, true, false, true, false]);
101
102        // Execute with selection mask.
103        let result = struct_array.execute_with_selection(&selection).unwrap();
104
105        // Verify the result has the filtered length.
106        assert_eq!(result.len(), 3);
107
108        // Verify the struct vector fields.
109        let struct_vector = result.as_struct();
110        let fields = struct_vector.fields();
111        assert_eq!(fields.len(), 2);
112
113        // Verify the integer field has the correct filtered values (indices 0, 2, 4).
114        let int_vector = fields[0].as_primitive().clone().into_i32();
115        assert_eq!(int_vector.elements().as_slice(), &[10, 30, 50]);
116
117        // Verify the boolean field has the correct filtered values (indices 0, 2, 4).
118        let bool_vector = fields[1].as_bool();
119        let bool_values: Vec<bool> = (0..3).map(|i| bool_vector.bits().value(i)).collect();
120        assert_eq!(bool_values, vec![true, true, true]);
121    }
122
123    #[test]
124    fn test_struct_operator_null_handling() {
125        // Create fields with nulls.
126        let int_field = PrimitiveArray::from_option_iter([
127            Some(100i32),
128            None,
129            Some(200),
130            Some(300),
131            None,
132            Some(400),
133        ]);
134
135        // Create bool field with its own validity.
136        let bool_array = BoolArray::from_iter([true, false, true, false, true, false]);
137        let bool_validity = Validity::from_iter([true, true, false, true, true, false]);
138        let bool_field = BoolArray::from_bit_buffer(bool_array.bit_buffer().clone(), bool_validity);
139
140        // Create struct with its own validity mask (rows 1 and 4 are null).
141        let struct_validity = Validity::from_iter([true, false, true, true, false, true]);
142
143        let struct_array = StructArray::try_new(
144            FieldNames::from(["values", "flags"]),
145            vec![int_field.into_array(), bool_field.into_array()],
146            6,
147            struct_validity,
148        )
149        .unwrap();
150
151        // Create a selection mask that selects indices 0, 1, 2, 4, 5.
152        let selection = Mask::from_iter([true, true, true, false, true, true]);
153
154        // Execute with selection mask.
155        let result = struct_array.execute_with_selection(&selection).unwrap();
156
157        assert_eq!(result.len(), 5);
158
159        // Verify the struct vector fields.
160        let struct_vector = result.as_struct();
161        let fields = struct_vector.fields();
162        assert_eq!(fields.len(), 2);
163
164        // Verify integer field has the correct filtered values with nulls.
165        // Selected indices: 0, 1, 2, 4, 5 from [Some(100), None, Some(200), Some(300), None, Some(400)].
166        let int_vector = fields[0].as_primitive().clone().into_i32();
167        let int_values: Vec<Option<i32>> = (0..5).map(|i| int_vector.get(i).copied()).collect();
168        assert_eq!(
169            int_values,
170            vec![Some(100), None, Some(200), None, Some(400)]
171        );
172
173        // Verify boolean field values.
174        // Selected indices: 0, 1, 2, 4, 5 from [T, F, T, F, T, F].
175        let bool_vector = fields[1].as_bool();
176        let bool_values: Vec<bool> = (0..5).map(|i| bool_vector.bits().value(i)).collect();
177        assert_eq!(bool_values, vec![true, false, true, true, false]);
178
179        // Verify the struct-level validity is correctly propagated.
180        // Original struct validity: [T, F, T, T, F, T]
181        // Selected indices: 0, 1, 2, 4, 5 -> validity: [T, F, T, F, T].
182        let validity_mask = struct_vector.validity();
183        let struct_validity_values: Vec<bool> = (0..5).map(|i| validity_mask.value(i)).collect();
184        assert_eq!(struct_validity_values, vec![true, false, true, false, true]);
185    }
186}
vortex_array/arrays/struct_/vtable/operator.rs

vortex_array/arrays/struct_/vtable/
operator.rs
