vortex_array/arrays/variant/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod vtable;
5
6pub(crate) mod compute;
7
8use vortex_error::VortexResult;
9use vortex_error::vortex_ensure;
10
11pub use self::vtable::Variant;
12pub use self::vtable::VariantArray;
13
14pub(crate) fn initialize(session: &vortex_session::VortexSession) {
15    vtable::initialize(session);
16}
17
18use crate::ArrayRef;
19use crate::array::Array;
20use crate::array::ArrayParts;
21use crate::array::EmptyArrayData;
22use crate::array_slots;
23use crate::dtype::DType;
24
25/// Slots for canonical variant storage.
26///
27/// A canonical variant array keeps the full variant value for every row in `core_storage` and may
28/// carry a row-aligned, storage-agnostic `shredded` typed tree for selected paths.
29///
30/// `core_storage` is a logical `DType::Variant` array, not a specific physical encoding: it may be
31/// chunked, constant, or otherwise encoded. Callers must use normal array operations instead of
32/// assuming a particular slot layout. The shredded child may have any dtype; its dtype is recorded
33/// during serialization and validated by normal child deserialization.
34#[array_slots(Variant)]
35pub struct VariantSlots {
36    /// The logical variant storage that preserves the full value for every row.
37    pub core_storage: ArrayRef,
38    /// The optional row-aligned typed shredded tree for selected variant paths.
39    /// This slot is `Some` only if the array was canonicalized and the shredded data
40    /// was pulled out of the underlying variant storage.
41    pub shredded: Option<ArrayRef>,
42}
43
44impl Array<Variant> {
45    /// Creates a new `VariantArray` with logical variant core storage and optional shredded storage.
46    ///
47    /// `core_storage` must have `DType::Variant`, but it may use any Variant-typed physical
48    /// encoding. See [`VariantSlots`] for the higher-level storage contract.
49    ///
50    /// `shredded`, when present, must be row-aligned with `core_storage` and stores typed values for
51    /// selected variant paths.
52    pub fn try_new(core_storage: ArrayRef, shredded: Option<ArrayRef>) -> VortexResult<Self> {
53        let dtype = core_storage.dtype().clone();
54        vortex_ensure!(
55            matches!(dtype, DType::Variant(_)),
56            "VariantArray core_storage dtype must be Variant, found {dtype}"
57        );
58        let len = core_storage.len();
59        let stats = core_storage.statistics().to_owned();
60        Ok(Array::try_from_parts(
61            ArrayParts::new(Variant, dtype, len, EmptyArrayData).with_slots(
62                VariantSlots {
63                    core_storage,
64                    shredded,
65                }
66                .into_slots(),
67            ),
68        )?
69        .with_stats_set(stats))
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use vortex_buffer::buffer;
76    use vortex_error::VortexResult;
77    use vortex_error::vortex_err;
78    use vortex_mask::Mask;
79
80    use crate::ArrayRef;
81    use crate::Canonical;
82    use crate::IntoArray;
83    use crate::VortexSessionExecute;
84    use crate::array_session;
85    use crate::arrays::BoolArray;
86    use crate::arrays::ChunkedArray;
87    use crate::arrays::ConstantArray;
88    use crate::arrays::PrimitiveArray;
89    use crate::arrays::StructArray;
90    use crate::arrays::VariantArray;
91    use crate::arrays::variant::VariantArraySlotsExt;
92    use crate::assert_arrays_eq;
93    use crate::builtins::ArrayBuiltins;
94    use crate::dtype::DType;
95    use crate::dtype::Nullability;
96    use crate::dtype::PType;
97    use crate::expr::root;
98    use crate::expr::variant_get;
99    use crate::scalar::Scalar;
100    use crate::scalar_fn::fns::variant_get::VariantPath;
101
102    fn core_storage(len: usize) -> ArrayRef {
103        ConstantArray::new(
104            Scalar::variant(Scalar::primitive(1i32, Nullability::NonNullable)),
105            len,
106        )
107        .into_array()
108    }
109
110    fn row_storage(values: impl IntoIterator<Item = i32>) -> VortexResult<ArrayRef> {
111        let chunks = values
112            .into_iter()
113            .map(|value| {
114                ConstantArray::new(
115                    Scalar::variant(Scalar::primitive(value, Nullability::NonNullable)),
116                    1,
117                )
118                .into_array()
119            })
120            .collect();
121
122        Ok(ChunkedArray::try_new(chunks, DType::Variant(Nullability::NonNullable))?.into_array())
123    }
124
125    fn variant_with_shredded(
126        core_values: impl IntoIterator<Item = i32>,
127        shredded_values: impl IntoIterator<Item = i32>,
128    ) -> VortexResult<VariantArray> {
129        VariantArray::try_new(
130            row_storage(core_values)?,
131            Some(PrimitiveArray::from_iter(shredded_values).into_array()),
132        )
133    }
134
135    fn execute_variant(array: ArrayRef) -> VortexResult<VariantArray> {
136        let mut ctx = array_session().create_execution_ctx();
137        let Canonical::Variant(variant) = array.execute::<Canonical>(&mut ctx)? else {
138            return Err(vortex_err!("expected canonical variant array"));
139        };
140        Ok(variant)
141    }
142
143    fn assert_variant_rows(
144        array: &VariantArray,
145        expected_core: &[Option<i32>],
146        expected_shredded: &[Option<i32>],
147    ) -> VortexResult<()> {
148        assert_variant_core_rows(array, expected_core)?;
149        assert_eq!(array.len(), expected_shredded.len());
150
151        let shredded = array
152            .shredded()
153            .ok_or_else(|| vortex_err!("expected shredded child"))?;
154        let mut ctx = array_session().create_execution_ctx();
155        let shredded = shredded.clone().execute::<PrimitiveArray>(&mut ctx)?;
156        let expected_shredded_array = if let Some(values) = expected_shredded
157            .iter()
158            .copied()
159            .collect::<Option<Vec<_>>>()
160        {
161            PrimitiveArray::from_iter(values)
162        } else {
163            PrimitiveArray::from_option_iter(expected_shredded.iter().copied())
164        };
165        assert_arrays_eq!(shredded, expected_shredded_array, &mut ctx);
166
167        Ok(())
168    }
169
170    fn assert_variant_core_rows(
171        array: &VariantArray,
172        expected_core: &[Option<i32>],
173    ) -> VortexResult<()> {
174        assert_eq!(array.len(), expected_core.len());
175
176        let mut ctx = array_session().create_execution_ctx();
177        for (idx, expected) in expected_core.iter().enumerate() {
178            let scalar = array.core_storage().execute_scalar(idx, &mut ctx)?;
179            let variant = scalar.as_variant();
180            match expected {
181                Some(expected) => {
182                    let value = variant
183                        .value()
184                        .ok_or_else(|| vortex_err!("expected non-null variant row"))?;
185                    assert_eq!(value.as_primitive().typed_value::<i32>(), Some(*expected));
186                }
187                None => assert!(variant.is_null()),
188            }
189        }
190
191        Ok(())
192    }
193
194    #[test]
195    fn try_new_exposes_core_storage_without_shredded() -> VortexResult<()> {
196        let core_storage = core_storage(2);
197
198        let variant = VariantArray::try_new(core_storage.clone(), None)?;
199
200        assert_eq!(variant.dtype(), core_storage.dtype());
201        assert_eq!(variant.len(), 2);
202        assert_eq!(variant.core_storage().dtype(), core_storage.dtype());
203        assert!(variant.shredded().is_none());
204
205        Ok(())
206    }
207
208    #[test]
209    fn try_new_exposes_core_storage_and_shredded() -> VortexResult<()> {
210        let core_storage = core_storage(3);
211        let shredded = buffer![10i32, 20, 30].into_array();
212
213        let variant = VariantArray::try_new(core_storage.clone(), Some(shredded.clone()))?;
214
215        assert_eq!(variant.dtype(), &DType::Variant(Nullability::NonNullable));
216        assert_eq!(variant.len(), 3);
217        assert_eq!(variant.core_storage().dtype(), core_storage.dtype());
218        assert_eq!(variant.core_storage().len(), core_storage.len());
219        assert_eq!(
220            variant.shredded().map(|child| child.dtype()),
221            Some(shredded.dtype())
222        );
223        assert_eq!(
224            variant.shredded().map(|child| child.len()),
225            Some(shredded.len())
226        );
227        assert_eq!(variant.as_ref().slot_name(0), "core_storage");
228        assert_eq!(variant.as_ref().slot_name(1), "shredded");
229
230        Ok(())
231    }
232
233    #[test]
234    fn try_new_rejects_non_variant_core_storage() {
235        let core_storage = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
236
237        assert!(VariantArray::try_new(core_storage, None).is_err());
238    }
239
240    #[test]
241    fn try_new_rejects_shredded_length_mismatch() {
242        let core_storage = core_storage(3);
243        let shredded = buffer![10i32, 20].into_array();
244
245        assert!(VariantArray::try_new(core_storage, Some(shredded)).is_err());
246    }
247
248    #[test]
249    fn scalar_at_merges_shredded_with_core_storage() -> VortexResult<()> {
250        let dtype = DType::Variant(Nullability::Nullable);
251        let core_chunks = [Some(1i32), None, Some(3)]
252            .into_iter()
253            .map(|value| {
254                let scalar = match value {
255                    Some(value) => {
256                        Scalar::variant(Scalar::primitive(value, Nullability::NonNullable))
257                            .cast(&dtype)?
258                    }
259                    None => Scalar::null(dtype.clone()),
260                };
261                Ok(ConstantArray::new(scalar, 1).into_array())
262            })
263            .collect::<VortexResult<Vec<_>>>()?;
264        let core_storage = ChunkedArray::try_new(core_chunks, dtype)?.into_array();
265        let shredded = PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]).into_array();
266        let variant = VariantArray::try_new(core_storage, Some(shredded))?;
267
268        let mut ctx = array_session().create_execution_ctx();
269        for (idx, expected) in [Some(10i32), None, Some(3)].into_iter().enumerate() {
270            let scalar = variant.execute_scalar(idx, &mut ctx)?;
271            let variant = scalar.as_variant();
272            match expected {
273                Some(expected) => {
274                    let value = variant
275                        .value()
276                        .ok_or_else(|| vortex_err!("expected non-null variant row"))?;
277                    assert_eq!(value.as_primitive().typed_value::<i32>(), Some(expected));
278                }
279                None => assert!(variant.is_null()),
280            }
281        }
282
283        Ok(())
284    }
285
286    #[test]
287    fn slice_preserves_core_storage_and_shredded_rows() -> VortexResult<()> {
288        let variant = variant_with_shredded(0..5, 10..15)?;
289
290        let sliced = execute_variant(variant.into_array().slice(1..4)?)?;
291
292        assert_variant_rows(
293            &sliced,
294            &[Some(1), Some(2), Some(3)],
295            &[Some(11), Some(12), Some(13)],
296        )
297    }
298
299    #[test]
300    fn filter_preserves_core_storage_and_shredded_rows() -> VortexResult<()> {
301        let variant = variant_with_shredded(0..5, 10..15)?;
302
303        let filtered = execute_variant(
304            variant
305                .into_array()
306                .filter(Mask::from_iter([true, false, true, false, true]))?,
307        )?;
308
309        assert_variant_rows(
310            &filtered,
311            &[Some(0), Some(2), Some(4)],
312            &[Some(10), Some(12), Some(14)],
313        )
314    }
315
316    #[test]
317    fn take_preserves_core_storage_and_shredded_rows() -> VortexResult<()> {
318        let variant = variant_with_shredded(0..5, 10..15)?;
319
320        let taken = execute_variant(
321            variant
322                .into_array()
323                .take(buffer![4u64, 1, 3].into_array())?,
324        )?;
325
326        assert_variant_rows(
327            &taken,
328            &[Some(4), Some(1), Some(3)],
329            &[Some(14), Some(11), Some(13)],
330        )
331    }
332
333    #[test]
334    fn mask_preserves_core_storage_and_shredded_rows() -> VortexResult<()> {
335        let variant = variant_with_shredded(0..5, 10..15)?;
336        let mask = BoolArray::from_iter([true, false, true, false, true]).into_array();
337
338        let masked = execute_variant(variant.into_array().mask(mask)?)?;
339
340        assert_variant_rows(
341            &masked,
342            &[Some(0), None, Some(2), None, Some(4)],
343            &[Some(10), None, Some(12), None, Some(14)],
344        )
345    }
346
347    #[test]
348    fn mask_preserves_chunked_core_storage_validity() -> VortexResult<()> {
349        let dtype = DType::Variant(Nullability::Nullable);
350        let core_chunks = [Some(1i32), None, Some(3), Some(4)]
351            .into_iter()
352            .map(|value| {
353                let scalar = match value {
354                    Some(value) => {
355                        Scalar::variant(Scalar::primitive(value, Nullability::NonNullable))
356                            .cast(&dtype)?
357                    }
358                    None => Scalar::null(dtype.clone()),
359                };
360                Ok(ConstantArray::new(scalar, 1).into_array())
361            })
362            .collect::<VortexResult<Vec<_>>>()?;
363        let core_storage = ChunkedArray::try_new(core_chunks, dtype)?.into_array();
364        let variant = VariantArray::try_new(core_storage, None)?;
365        let mask = BoolArray::from_iter([true, true, false, true]).into_array();
366
367        let masked = execute_variant(variant.into_array().mask(mask)?)?;
368
369        assert_variant_core_rows(&masked, &[Some(1), None, None, Some(4)])
370    }
371
372    #[test]
373    fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> {
374        let mut ctx = array_session().create_execution_ctx();
375        let core_storage = row_storage([1, 2, 3])?;
376        let shredded = StructArray::try_from_iter([(
377            "a",
378            PrimitiveArray::from_iter([10i32, 20, 30]).into_array(),
379        )])?;
380        let variant = VariantArray::try_new(core_storage, Some(shredded.into_array()))?;
381        let expr = variant_get(
382            root(),
383            VariantPath::field("a"),
384            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
385        );
386
387        let result = variant
388            .into_array()
389            .apply(&expr)?
390            .execute::<PrimitiveArray>(&mut array_session().create_execution_ctx())?;
391
392        assert_arrays_eq!(
393            result,
394            PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]),
395            &mut ctx
396        );
397        Ok(())
398    }
399
400    #[test]
401    fn variant_get_treats_value_and_typed_value_as_logical_field_names() -> VortexResult<()> {
402        let mut ctx = array_session().create_execution_ctx();
403        let core_storage = row_storage([1, 2, 3])?;
404        let shredded = StructArray::try_from_iter([
405            (
406                "value",
407                PrimitiveArray::from_iter([10i32, 20, 30]).into_array(),
408            ),
409            (
410                "typed_value",
411                PrimitiveArray::from_iter([40i32, 50, 60]).into_array(),
412            ),
413        ])?;
414        let variant = VariantArray::try_new(core_storage, Some(shredded.into_array()))?;
415
416        let value_expr = variant_get(
417            root(),
418            VariantPath::field("value"),
419            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
420        );
421        let value_result = variant
422            .clone()
423            .into_array()
424            .apply(&value_expr)?
425            .execute::<PrimitiveArray>(&mut array_session().create_execution_ctx())?;
426        assert_arrays_eq!(
427            value_result,
428            PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]),
429            &mut ctx
430        );
431
432        let typed_value_expr = variant_get(
433            root(),
434            VariantPath::field("typed_value"),
435            Some(DType::Primitive(PType::I32, Nullability::NonNullable)),
436        );
437        let typed_value_result = variant
438            .into_array()
439            .apply(&typed_value_expr)?
440            .execute::<PrimitiveArray>(&mut array_session().create_execution_ctx())?;
441        assert_arrays_eq!(
442            typed_value_result,
443            PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]),
444            &mut ctx
445        );
446        Ok(())
447    }
448}
vortex_array/arrays/variant/mod.rs

vortex_array/arrays/variant/
mod.rs
