vortex_array/arrays/struct_/compute/
zip.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::BitAnd;
5use std::ops::BitOr;
6use std::ops::Not;
7
8use vortex_error::VortexResult;
9use vortex_mask::Mask;
10
11use crate::ArrayRef;
12use crate::ExecutionCtx;
13use crate::IntoArray;
14use crate::array::ArrayView;
15use crate::arrays::Struct;
16use crate::arrays::StructArray;
17use crate::arrays::struct_::StructArrayExt;
18use crate::builtins::ArrayBuiltins;
19use crate::scalar_fn::fns::zip::ZipKernel;
20use crate::validity::Validity;
21
22impl ZipKernel for Struct {
23    fn zip(
24        if_true: ArrayView<'_, Struct>,
25        if_false: &ArrayRef,
26        mask: &ArrayRef,
27        ctx: &mut ExecutionCtx,
28    ) -> VortexResult<Option<ArrayRef>> {
29        let Some(if_false) = if_false.as_opt::<Struct>() else {
30            return Ok(None);
31        };
32        assert_eq!(
33            if_true.names(),
34            if_false.names(),
35            "input arrays to zip must have the same field names",
36        );
37
38        let fields = if_true
39            .iter_unmasked_fields()
40            .zip(if_false.iter_unmasked_fields())
41            .map(|(t, f)| mask.zip(t.clone(), f.clone()))
42            .collect::<VortexResult<Vec<_>>>()?;
43
44        let v1 = if_true.validity()?;
45        let v2 = if_false.validity()?;
46        let validity = match (&v1, &v2) {
47            (Validity::NonNullable, Validity::NonNullable) => Validity::NonNullable,
48            (Validity::AllValid, Validity::AllValid) => Validity::AllValid,
49            (Validity::AllInvalid, Validity::AllInvalid) => Validity::AllInvalid,
50
51            (v1, v2) => {
52                let mask_mask = mask.clone().fill_null(false)?.execute::<Mask>(ctx)?;
53                let v1m = v1.execute_mask(if_true.len(), ctx)?;
54                let v2m = v2.execute_mask(if_false.len(), ctx)?;
55
56                let combined = v1m.bitand(&mask_mask).bitor(&v2m.bitand(&mask_mask.not()));
57                Validity::from_mask(
58                    combined,
59                    if_true.dtype().nullability() | if_false.dtype().nullability(),
60                )
61            }
62        };
63
64        Ok(Some(
65            StructArray::try_new(if_true.names().clone(), fields, if_true.len(), validity)?
66                .into_array(),
67        ))
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use vortex_mask::Mask;
74
75    use crate::IntoArray;
76    use crate::VortexSessionExecute;
77    use crate::array_session;
78    use crate::arrays::PrimitiveArray;
79    use crate::arrays::StructArray;
80    use crate::builtins::ArrayBuiltins;
81    use crate::dtype::FieldNames;
82    use crate::validity::Validity;
83
84    #[test]
85    fn test_validity_zip_both_validity_array() {
86        let ctx = array_session().create_execution_ctx();
87        // Both structs have Validity::Array
88        let if_true = StructArray::new(
89            FieldNames::from_iter(["field"]),
90            vec![PrimitiveArray::from_iter([1, 2, 3, 4]).into_array()],
91            4,
92            Validity::from_iter([true, false, true, false]),
93        )
94        .into_array();
95
96        let if_false = StructArray::new(
97            FieldNames::from_iter(["field"]),
98            vec![PrimitiveArray::from_iter([10, 20, 30, 40]).into_array()],
99            4,
100            Validity::from_iter([false, true, false, true]),
101        )
102        .into_array();
103
104        let mask = Mask::from_iter([false, false, true, false]);
105
106        let result = mask.into_array().zip(if_true, if_false).unwrap();
107
108        insta::assert_snapshot!(result.display_table(&ctx), @r"
109        ┌───────┐
110        │ field │
111        ├───────┤
112        │ null  │
113        ├───────┤
114        │ 20i32 │
115        ├───────┤
116        │ 3i32  │
117        ├───────┤
118        │ 40i32 │
119        └───────┘
120        ");
121    }
122
123    #[test]
124    fn test_validity_zip_allvalid_and_array() {
125        let ctx = array_session().create_execution_ctx();
126        let if_true = StructArray::new(
127            FieldNames::from_iter(["a"]),
128            vec![PrimitiveArray::from_iter([1, 2, 3, 4]).into_array()],
129            4,
130            Validity::AllValid,
131        )
132        .into_array();
133
134        let if_false = StructArray::new(
135            FieldNames::from_iter(["a"]),
136            vec![PrimitiveArray::from_iter([10, 20, 30, 40]).into_array()],
137            4,
138            Validity::from_iter([false, false, true, true]),
139        )
140        .into_array();
141
142        let mask = Mask::from_iter([true, false, false, false]);
143
144        let result = mask.into_array().zip(if_true, if_false).unwrap();
145
146        insta::assert_snapshot!(result.display_table(&ctx), @r"
147        ┌───────┐
148        │   a   │
149        ├───────┤
150        │ 1i32  │
151        ├───────┤
152        │ null  │
153        ├───────┤
154        │ 30i32 │
155        ├───────┤
156        │ 40i32 │
157        └───────┘
158        ");
159    }
160}
vortex_array/arrays/struct_/compute/zip.rs

vortex_array/arrays/struct_/compute/
zip.rs
