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

vortex_array/arrays/variant/
mod.rs
