Implement nullable UnionArray semantics

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-array/src/arrays/chunked/vtable/canonical.rs

@@ -25,8 +25,8 @@ use crate::arrays::chunked::ChunkedArrayExt;
 use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
 use crate::arrays::listview::ListViewArrayExt;
 use crate::arrays::listview::ListViewRebuildMode;
-use crate::arrays::union::TYPE_IDS_DTYPE;
 use crate::arrays::union::UnionArrayExt;
+use crate::arrays::union::type_ids_dtype;
 use crate::arrays::variant::VariantArrayExt;
 use crate::builders::builder_with_capacity_in;
 use crate::builtins::ArrayBuiltins;
@@ -105,7 +105,7 @@ fn pack_union_chunks(chunks: Vec<ArrayRef>, ctx: &mut ExecutionCtx) -> VortexRes
             .iter()
             .map(|chunk| chunk.type_ids().clone())
             .collect(),
-        TYPE_IDS_DTYPE,
+        type_ids_dtype(union_chunks[0].dtype().nullability()),
     )?
     .into_array();
     let children = variants

vortex-array/src/arrays/constant/vtable/canonical.rs

@@ -8,7 +8,6 @@ use vortex_buffer::Buffer;
 use vortex_buffer::buffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
 
 use crate::Canonical;
 use crate::ExecutionCtx;
@@ -32,6 +31,7 @@ use crate::builders::builder_with_capacity;
 use crate::dtype::DType;
 use crate::dtype::DecimalType;
 use crate::dtype::Nullability;
+use crate::dtype::PType;
 use crate::dtype::UnionVariants;
 use crate::match_each_decimal_value;
 use crate::match_each_decimal_value_type;
@@ -167,9 +167,12 @@ pub(crate) fn constant_canonicalize(
                 StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity)
             })
         }
-        DType::Union(variants) => {
-            Canonical::Union(constant_canonical_union(scalar, variants, array.len())?)
-        }
+        DType::Union(variants, nullability) => Canonical::Union(constant_canonical_union(
+            scalar,
+            variants,
+            *nullability,
+            array.len(),
+        )?),
         DType::Variant(_) => Canonical::Variant(VariantArray::try_new(
             array.array().clone().into_array(),
             None,
@@ -196,41 +199,31 @@ pub(crate) fn constant_canonicalize(
 fn constant_canonical_union(
     scalar: &Scalar,
     variants: &UnionVariants,
+    nullability: Nullability,
     len: usize,
 ) -> VortexResult<UnionArray> {
-    if scalar.is_null() {
-        vortex_bail!(
-            "Canonicalizing a null union scalar is not supported until null semantics are defined"
-        )
-    }
-    if variants.variants().any(|variant| variant.is_nullable()) {
-        vortex_bail!("Canonical UnionArray children must be non-nullable")
-    }
-
     let union = scalar.as_union();
-    let type_id = union
-        .type_id()
-        .vortex_expect("non-null union scalar must have a type ID");
-    let selected_child = union
-        .child_index()
-        .vortex_expect("validated union scalar must select a child");
-    let selected_value = union
-        .value()
-        .vortex_expect("non-null union scalar must have a value");
+    let selected = union.child_index().zip(union.value());
 
     let children = variants
         .variants()
         .enumerate()
         .map(|(index, dtype)| {
-            let value = if index == selected_child {
-                selected_value.clone()
-            } else {
-                Scalar::zero_value(&dtype)
+            let value = match selected {
+                Some((selected_child, selected_value)) if index == selected_child => {
+                    selected_value.clone()
+                }
+                _ => Scalar::default_value(&dtype),
             };
             ConstantArray::new(value, len).into_array()
         })
         .collect::<Vec<_>>();
 
+    let type_id = match union.type_id() {
+        Some(type_id) => Scalar::primitive(type_id, nullability),
+        None => Scalar::null(DType::Primitive(PType::U8, Nullability::Nullable)),
+    };
+
     UnionArray::try_new(
         ConstantArray::new(type_id, len).into_array(),
         variants.clone(),

vortex-array/src/arrays/masked/execute.rs

@@ -6,7 +6,6 @@
 use std::sync::Arc;
 
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
 
 use crate::Canonical;
 use crate::IntoArray;
@@ -18,13 +17,15 @@ use crate::arrays::ListViewArray;
 use crate::arrays::MaskedArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::StructArray;
+use crate::arrays::UnionArray;
 use crate::arrays::VarBinViewArray;
 use crate::arrays::VariantArray;
 use crate::arrays::bool::BoolArrayExt;
 use crate::arrays::extension::ExtensionArrayExt;
 use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
 use crate::arrays::listview::ListViewArrayExt;
 use crate::arrays::struct_::StructArrayExt;
+use crate::arrays::union::UnionArrayExt;
 use crate::arrays::variant::VariantArrayExt;
 use crate::builtins::ArrayBuiltins;
 use crate::executor::ExecutionCtx;
@@ -51,9 +52,7 @@ pub fn mask_validity_canonical(
             Canonical::FixedSizeList(mask_validity_fixed_size_list(a, validity)?)
         }
         Canonical::Struct(a) => Canonical::Struct(mask_validity_struct(a, validity)?),
-        Canonical::Union(_) => {
-            vortex_bail!("Masking UnionArray is not supported until null semantics are defined")
-        }
+        Canonical::Union(a) => Canonical::Union(mask_validity_union(a, validity)?),
         Canonical::Extension(a) => Canonical::Extension(mask_validity_extension(a, validity, ctx)?),
         Canonical::Variant(a) => Canonical::Variant(mask_validity_variant(a, validity, ctx)?),
     })
@@ -146,6 +145,19 @@ fn mask_validity_struct(array: StructArray, validity: Validity) -> VortexResult<
     Ok(unsafe { StructArray::new_unchecked(fields, struct_fields.clone(), len, new_validity) })
 }
 
+fn mask_validity_union(array: UnionArray, validity: Validity) -> VortexResult<UnionArray> {
+    let type_ids = array
+        .type_ids()
+        .clone()
+        .mask(validity.to_array(array.len()))?;
+    let variants = array.variants().clone();
+    let children = array.children();
+
+    // SAFETY: Masking type IDs changes only outer validity. Values at newly-null positions are no
+    // longer selected, and every row-aligned component retains its length and dtype.
+    Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, children) })
+}
+
 fn mask_validity_extension(
     array: ExtensionArray,
     validity: Validity,

vortex-array/src/arrays/union/array.rs

@@ -19,8 +19,8 @@ use crate::array::EmptyArrayData;
 use crate::array::TypedArrayRef;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::Union;
-use crate::arrays::union::TYPE_IDS_DTYPE;
 use crate::dtype::DType;
+use crate::dtype::Nullability;
 use crate::dtype::UnionVariants;
 use crate::legacy_session;
 
@@ -35,11 +35,18 @@ pub(super) fn make_union_parts(
     children: &[ArrayRef],
 ) -> ArrayParts<Union> {
     let len = type_ids.len();
+    let nullability = type_ids.dtype().nullability();
     let slots: ArraySlots = once(Some(type_ids))
         .chain(children.iter().cloned().map(Some))
         .collect();
 
-    ArrayParts::new(Union, DType::Union(variants), len, EmptyArrayData).with_slots(slots)
+    ArrayParts::new(
+        Union,
+        DType::Union(variants, nullability),
+        len,
+        EmptyArrayData,
+    )
+    .with_slots(slots)
 }
 
 /// Concrete parts of a [`UnionArray`](super::UnionArray).
@@ -57,12 +64,12 @@ pub trait UnionArrayExt: TypedArrayRef<Union> {
     /// The union's variant schema.
     fn variants(&self) -> &UnionVariants {
         match self.as_ref().dtype() {
-            DType::Union(variants) => variants,
+            DType::Union(variants, _) => variants,
             _ => unreachable!("UnionArrayExt requires a union dtype"),
         }
     }
 
-    /// The row-aligned non-nullable `i8` type IDs.
+    /// The row-aligned `u8` type IDs whose nulls represent outer union nulls.
     fn type_ids(&self) -> &ArrayRef {
         self.as_ref().slots()[TYPE_IDS_SLOT]
             .as_ref()
@@ -87,7 +94,7 @@ pub trait UnionArrayExt: TypedArrayRef<Union> {
     }
 
     /// Return a sparse child selected by a data-level type ID.
-    fn child_by_type_id(&self, type_id: i8) -> Option<&ArrayRef> {
+    fn child_by_type_id(&self, type_id: u8) -> Option<&ArrayRef> {
         self.child(self.variants().tag_to_child_index(type_id)?)
     }
 
@@ -118,13 +125,22 @@ fn validate_type_id_values(array: &Array<Union>) -> VortexResult<()> {
         return Ok(());
     }
 
+    let mut ctx = legacy_session().create_execution_ctx();
     let type_ids = array
         .type_ids()
         .clone()
-        .execute::<PrimitiveArray>(&mut legacy_session().create_execution_ctx())?;
-    for type_id in type_ids.as_slice::<i8>() {
+        .execute::<PrimitiveArray>(&mut ctx)?;
+    let values = type_ids.as_slice::<u8>();
+    let validity = type_ids.validity()?.execute_mask(array.len(), &mut ctx)?;
+    let valid_indices: Box<dyn Iterator<Item = usize> + '_> = match &validity {
+        vortex_mask::Mask::AllTrue(_) => Box::new(0..values.len()),
+        vortex_mask::Mask::AllFalse(_) => Box::new(std::iter::empty()),
+        vortex_mask::Mask::Values(mask) => Box::new(mask.indices().iter().copied()),
+    };
+    for index in valid_indices {
+        let type_id = values[index];
         vortex_ensure!(
-            array.variants().tag_to_child_index(*type_id).is_some(),
+            array.variants().tag_to_child_index(type_id).is_some(),
             "UnionArray type ID {type_id} is not present in {:?}",
             array.variants().type_ids()
         );
@@ -150,8 +166,6 @@ impl Array<Union> {
 
     /// Try to construct a canonical sparse union array.
     ///
-    /// Until Union nullability semantics are settled, every child must be non-nullable.
-    ///
     /// # Errors
     ///
     /// Returns an error if the type IDs and sparse children do not match the variant schema or do
@@ -162,8 +176,11 @@ impl Array<Union> {
         children: impl Into<Arc<[ArrayRef]>>,
     ) -> VortexResult<Self> {
         vortex_ensure!(
-            type_ids.dtype() == &TYPE_IDS_DTYPE,
-            "UnionArray type_ids must be non-nullable i8, got {}",
+            matches!(
+                type_ids.dtype(),
+                DType::Primitive(crate::dtype::PType::U8, _)
+            ),
+            "UnionArray type_ids must be u8, got {}",
             type_ids.dtype()
         );
 
@@ -178,9 +195,9 @@ impl Array<Union> {
     ///
     /// # Safety
     ///
-    /// The caller must ensure the type IDs are non-nullable `i8` values declared by `variants`,
-    /// every child has the corresponding variant dtype, and all arrays have the same length. The
-    /// children must currently be non-nullable.
+    /// The caller must ensure every non-null type ID is a `u8` value declared by `variants`, every
+    /// child has the corresponding variant dtype, and all arrays have the same length. Null type
+    /// IDs represent outer union nulls.
     pub unsafe fn new_unchecked(
         type_ids: ArrayRef,
         variants: UnionVariants,
@@ -202,9 +219,9 @@ impl Array<Union> {
         }
     }
 
-    /// Create an empty array for a non-nullable union dtype.
-    pub(crate) fn empty(variants: UnionVariants) -> Self {
-        let type_ids = PrimitiveArray::from_iter(Vec::<i8>::new()).into_array();
+    /// Create an empty array for a union dtype.
+    pub(crate) fn empty(variants: UnionVariants, nullability: Nullability) -> Self {
+        let type_ids = PrimitiveArray::empty::<u8>(nullability).into_array();
         let children: Vec<_> = variants
             .variants()
             .map(|dtype| crate::Canonical::empty(&dtype).into_array())

vortex-array/src/arrays/union/compute/mask.rs

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::ArrayRef;
+use crate::IntoArray;
+use crate::array::ArrayView;
+use crate::arrays::Union;
+use crate::arrays::UnionArray;
+use crate::arrays::union::UnionArrayExt;
+use crate::builtins::ArrayBuiltins;
+use crate::scalar_fn::fns::mask::MaskReduce;
+
+impl MaskReduce for Union {
+    fn mask(array: ArrayView<'_, Union>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
+        let type_ids = array.type_ids().clone().mask(mask.clone())?;
+        let children = array.children();
+
+        // SAFETY: Masking type IDs changes only outer validity. Values at newly-null positions are
+        // no longer selected, and every row-aligned component retains its length and dtype.
+        Ok(Some(
+            unsafe { UnionArray::new_unchecked(type_ids, array.variants().clone(), children) }
+                .into_array(),
+        ))
+    }
+}

vortex-array/src/arrays/union/compute/mod.rs

@@ -4,6 +4,8 @@
 mod filter;
 pub(crate) use filter::filter_union;
 
+mod mask;
+
 pub(crate) mod rules;
 
 mod slice;

vortex-array/src/arrays/union/compute/rules.rs

@@ -6,9 +6,11 @@ use crate::arrays::dict::TakeReduceAdaptor;
 use crate::arrays::filter::FilterReduceAdaptor;
 use crate::arrays::slice::SliceReduceAdaptor;
 use crate::optimizer::rules::ParentRuleSet;
+use crate::scalar_fn::fns::mask::MaskReduceAdaptor;
 
 pub(crate) const PARENT_RULES: ParentRuleSet<Union> = ParentRuleSet::new(&[
     ParentRuleSet::lift(&SliceReduceAdaptor(Union)),
     ParentRuleSet::lift(&FilterReduceAdaptor(Union)),
     ParentRuleSet::lift(&TakeReduceAdaptor(Union)),
+    ParentRuleSet::lift(&MaskReduceAdaptor(Union)),
 ]);

vortex-array/src/arrays/union/mod.rs

@@ -3,12 +3,10 @@
 
 //! Canonical sparse union arrays.
 //!
-//! A [`UnionArray`] stores one non-nullable `i8` type ID per row followed by one row-aligned child
-//! for each variant. The type ID selects which child's value is active for a row; values in all
-//! other children at that row are placeholders.
-//!
-//! Union nullability semantics are still being designed, so this encoding currently requires every
-//! variant child to be non-nullable.
+//! A [`UnionArray`] stores one `u8` type ID per row followed by one row-aligned child for each
+//! variant. The type ID selects which child's value is active for a row; values in all other
+//! children at that row are placeholders. Outer union nulls are stored as nulls in the type IDs
+//! child, independently of nulls in the selected variant child.
 
 use crate::dtype::DType;
 use crate::dtype::Nullability;
@@ -24,7 +22,9 @@ pub(crate) mod compute;
 mod vtable;
 pub use vtable::Union;
 
-pub(crate) const TYPE_IDS_DTYPE: DType = DType::Primitive(PType::I8, Nullability::NonNullable);
+pub(crate) fn type_ids_dtype(nullability: Nullability) -> DType {
+    DType::Primitive(PType::U8, nullability)
+}
 
 #[cfg(test)]
 mod tests;