Fix UnionArray canonicalization edge cases

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

Author: connortsui20

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

@@ -8,6 +8,7 @@ use vortex_buffer::Buffer;
 use vortex_buffer::buffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_err;
 
 use crate::Canonical;
 use crate::ExecutionCtx;
@@ -149,16 +150,22 @@ pub(crate) fn constant_canonicalize(
                     .collect(),
                 None => {
                     assert!(matches!(validity, Validity::AllInvalid));
-                    // The struct is entirely null, so fields just need placeholder values with the
-                    // correct dtype. We use `default_value` which returns a zero for non-nullable
-                    // dtypes and null for nullable dtypes, preserving each field's nullability.
                     struct_dtype
                         .fields()
                         .map(|dt| {
-                            let scalar = Scalar::default_value(&dt);
-                            ConstantArray::new(scalar, array.len()).into_array()
+                            if array.is_empty() {
+                                return Ok(Canonical::empty(&dt).into_array());
+                            }
+
+                            let scalar = Scalar::try_default_value(&dt).ok_or_else(|| {
+                                vortex_err!(
+                                    "cannot canonicalize null constant struct: field dtype {dt} \
+                                     has no default value"
+                                )
+                            })?;
+                            Ok(ConstantArray::new(scalar, array.len()).into_array())
                         })
-                        .collect()
+                        .collect::<VortexResult<Vec<_>>>()?
                 }
             };
             // SAFETY: Fields are constructed from the same struct scalar, all have same
@@ -202,25 +209,32 @@ fn constant_canonical_union(
     nullability: Nullability,
     len: usize,
 ) -> VortexResult<UnionArray> {
+    if len == 0 {
+        return Ok(UnionArray::empty(variants.clone(), nullability));
+    }
+
     let union = scalar.as_union();
     let selected = union.child_index().zip(union.value());
 
     let children = variants
         .variants()
         .enumerate()
         .map(|(index, dtype)| {
-            let value = match selected {
-                Some((selected_child, selected_value)) if index == selected_child => {
+            let value = match &selected {
+                Some((selected_child, selected_value)) if index == *selected_child => {
                     selected_value.clone()
                 }
-                _ if matches!(&dtype, DType::Union(_, Nullability::NonNullable)) => {
-                    Scalar::zero_value(&dtype)
-                }
-                _ => Scalar::default_value(&dtype),
+                _ => Scalar::try_default_value(&dtype).ok_or_else(|| {
+                    vortex_err!(
+                        "cannot canonicalize constant union: unselected variant {} ({dtype}) has \
+                         no default value",
+                        variants.names()[index]
+                    )
+                })?,
             };
-            ConstantArray::new(value, len).into_array()
+            Ok(ConstantArray::new(value, len).into_array())
         })
-        .collect::<Vec<_>>();
+        .collect::<VortexResult<Vec<_>>>()?;
 
     let type_id = match union.type_id() {
         Some(type_id) => Scalar::primitive(type_id, nullability),

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

@@ -4,10 +4,13 @@
 use vortex_buffer::ByteBufferMut;
 use vortex_buffer::buffer;
 use vortex_error::VortexResult;
+use vortex_error::vortex_err;
 use vortex_mask::Mask;
 use vortex_session::registry::ReadContext;
 
 use crate::ArrayContext;
+use crate::Canonical;
+use crate::CanonicalValidity;
 use crate::IntoArray;
 use crate::VortexSessionExecute;
 use crate::array_session;
@@ -339,6 +342,121 @@ fn constant_union_builds_nested_union_placeholders() -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+fn constant_union_builds_deeply_nested_union_placeholders() -> VortexResult<()> {
+    let nested_dtype = DType::Union(
+        UnionVariants::try_new(
+            ["value"].into(),
+            vec![DType::Primitive(PType::I64, Nullability::NonNullable)],
+            vec![3],
+        )?,
+        Nullability::NonNullable,
+    );
+    let wrapper_dtype =
+        DType::struct_([("nested", nested_dtype.clone())], Nullability::NonNullable);
+    let outer_variants = UnionVariants::try_new(
+        ["number", "wrapper"].into(),
+        vec![
+            DType::Primitive(PType::I32, Nullability::NonNullable),
+            wrapper_dtype,
+        ],
+        vec![5, 9],
+    )?;
+    let selected_number =
+        Scalar::union(outer_variants, 5, 42_i32.into(), Nullability::NonNullable)?;
+    let mut ctx = array_session().create_execution_ctx();
+    let array = ConstantArray::new(selected_number, 2)
+        .into_array()
+        .execute::<UnionArray>(&mut ctx)?;
+    let wrapper = array
+        .child_by_name("wrapper")?
+        .execute_scalar(0, &mut ctx)?;
+
+    assert_eq!(
+        wrapper.as_struct().field("nested"),
+        Some(Scalar::default_value(&nested_dtype))
+    );
+
+    Ok(())
+}
+
+#[test]
+fn constant_union_rejects_uninhabited_placeholders_without_panicking() -> VortexResult<()> {
+    let uninhabited = DType::Union(
+        UnionVariants::new(Default::default(), vec![])?,
+        Nullability::NonNullable,
+    );
+    let outer_variants = UnionVariants::try_new(
+        ["number", "uninhabited"].into(),
+        vec![
+            DType::Primitive(PType::I32, Nullability::NonNullable),
+            uninhabited,
+        ],
+        vec![5, 9],
+    )?;
+    let selected_number =
+        Scalar::union(outer_variants, 5, 42_i32.into(), Nullability::NonNullable)?;
+    let mut ctx = array_session().create_execution_ctx();
+
+    assert!(
+        ConstantArray::new(selected_number.clone(), 1)
+            .into_array()
+            .execute::<UnionArray>(&mut ctx)
+            .is_err()
+    );
+    assert_eq!(
+        ConstantArray::new(selected_number, 0)
+            .into_array()
+            .execute::<UnionArray>(&mut ctx)?
+            .len(),
+        0
+    );
+
+    Ok(())
+}
+
+#[test]
+fn empty_union_supports_variant_children() -> VortexResult<()> {
+    let variants = UnionVariants::try_new(
+        ["dynamic"].into(),
+        vec![DType::Variant(Nullability::NonNullable)],
+        vec![5],
+    )?;
+    let array = Canonical::empty(&DType::Union(variants, Nullability::NonNullable)).into_union();
+
+    assert_eq!(array.len(), 0);
+    assert_eq!(
+        array.child(0).map(|child| child.dtype()),
+        Some(&DType::Variant(Nullability::NonNullable))
+    );
+
+    Ok(())
+}
+
+#[test]
+fn canonical_validity_canonicalizes_union_type_ids() -> VortexResult<()> {
+    let array = UnionArray::try_new(
+        ConstantArray::new(Scalar::primitive(5_u8, Nullability::Nullable), 2).into_array(),
+        variants()?,
+        vec![
+            PrimitiveArray::from_iter([10_i32, 20]).into_array(),
+            BoolArray::from_iter([false, true]).into_array(),
+        ],
+    )?;
+    let mut ctx = array_session().create_execution_ctx();
+    let Canonical::Union(array) = array.into_array().execute::<CanonicalValidity>(&mut ctx)?.0
+    else {
+        return Err(vortex_err!(
+            "UnionArray must remain canonical Union storage"
+        ));
+    };
+
+    assert!(array.type_ids().is::<crate::arrays::Primitive>());
+    assert_eq!(array.dtype().nullability(), Nullability::Nullable);
+
+    Ok(())
+}
+
 #[test]
 fn chunked_union_packs_components() -> VortexResult<()> {
     let first = union_array()?.into_array().slice(0..1)?;

vortex-array/src/canonical.rs

@@ -21,6 +21,7 @@ use crate::array::ArrayView;
 use crate::array::child_to_validity;
 use crate::arrays::Bool;
 use crate::arrays::BoolArray;
+use crate::arrays::ChunkedArray;
 use crate::arrays::Decimal;
 use crate::arrays::DecimalArray;
 use crate::arrays::Extension;
@@ -237,9 +238,15 @@ impl Canonical {
             DType::Union(variants, nullability) => {
                 Canonical::Union(UnionArray::empty(variants.clone(), *nullability))
             }
-            DType::Variant(_) => {
-                vortex_panic!(InvalidArgument: "Canonical empty is not supported for Variant")
-            }
+            DType::Variant(_) => Canonical::Variant(
+                VariantArray::try_new(
+                    ChunkedArray::try_new(vec![], dtype.clone())
+                        .vortex_expect("empty Variant core storage must be valid")
+                        .into_array(),
+                    None,
+                )
+                .vortex_expect("empty VariantArray must be valid"),
+            ),
             DType::Extension(ext_dtype) => Canonical::Extension(ExtensionArray::new(
                 ext_dtype.clone(),
                 Canonical::empty(ext_dtype.storage_dtype()).into_array(),
@@ -678,7 +685,18 @@ impl Executable for CanonicalValidity {
                     StructArray::new_unchecked(fields, struct_fields, len, validity.execute(ctx)?)
                 })))
             }
-            union @ Canonical::Union(_) => Ok(CanonicalValidity(union)),
+            Canonical::Union(union) => {
+                let UnionDataParts {
+                    variants,
+                    type_ids,
+                    children,
+                } = union.into_data_parts();
+                let type_ids = type_ids.execute::<CanonicalValidity>(ctx)?.0.into_array();
+
+                Ok(CanonicalValidity(Canonical::Union(unsafe {
+                    UnionArray::new_unchecked(type_ids, variants, children)
+                })))
+            }
             Canonical::Extension(ext) => Ok(CanonicalValidity(Canonical::Extension(
                 ExtensionArray::new(
                     ext.ext_dtype().clone(),