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 = try_placeholder_scalar(&dt).ok_or_else(|| {
+                                vortex_err!(
+                                    "cannot canonicalize null constant struct: field dtype {dt} \
+                                     has no placeholder 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.child());
 
     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),
+                _ => try_placeholder_scalar(&dtype).ok_or_else(|| {
+                    vortex_err!(
+                        "cannot canonicalize constant union: unselected variant {} ({dtype}) has \
+                         no placeholder 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),
@@ -234,6 +248,52 @@ fn constant_canonical_union(
     )
 }
 
+/// Builds an arbitrary valid scalar for array positions ignored by a parent validity or type ID.
+/// This does not give the dtype a semantic default value.
+fn try_placeholder_scalar(dtype: &DType) -> Option<Scalar> {
+    if let Some(default) = Scalar::try_default_value(dtype) {
+        return Some(default);
+    }
+
+    match dtype {
+        DType::FixedSizeList(element_dtype, size, nullability) => {
+            let element = try_placeholder_scalar(element_dtype)?;
+            Some(Scalar::fixed_size_list(
+                Arc::clone(element_dtype),
+                vec![element; *size as usize],
+                *nullability,
+            ))
+        }
+        DType::Struct(fields, _) => {
+            let children = fields
+                .fields()
+                .map(|field| try_placeholder_scalar(&field))
+                .collect::<Option<Vec<_>>>()?;
+            Some(Scalar::struct_(dtype.clone(), children))
+        }
+        DType::Union(variants, nullability) => {
+            variants
+                .variants()
+                .enumerate()
+                .find_map(|(index, child_dtype)| {
+                    let child = try_placeholder_scalar(&child_dtype)?;
+                    Scalar::union(
+                        variants.clone(),
+                        variants.child_index_to_tag(index),
+                        child,
+                        *nullability,
+                    )
+                    .ok()
+                })
+        }
+        DType::Extension(ext_dtype) => {
+            let storage = try_placeholder_scalar(ext_dtype.storage_dtype())?;
+            Some(Scalar::extension_ref(ext_dtype.clone(), storage))
+        }
+        _ => None,
+    }
+}
+
 fn constant_canonical_byte_view(
     scalar_bytes: Option<&[u8]>,
     dtype: &DType,

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;
@@ -302,12 +305,18 @@ fn constant_union_builds_nested_union_placeholders() -> VortexResult<()> {
         vec![DType::Primitive(PType::I64, Nullability::NonNullable)],
         vec![3],
     )?;
+    let nested_placeholder = Scalar::union(
+        nested_variants.clone(),
+        3,
+        0_i64.into(),
+        Nullability::NonNullable,
+    )?;
     let nested_dtype = DType::Union(nested_variants, Nullability::NonNullable);
     let outer_variants = UnionVariants::try_new(
         ["number", "nested"].into(),
         vec![
             DType::Primitive(PType::I32, Nullability::NonNullable),
-            nested_dtype.clone(),
+            nested_dtype,
         ],
         vec![5, 9],
     )?;
@@ -324,7 +333,7 @@ fn constant_union_builds_nested_union_placeholders() -> VortexResult<()> {
         .execute::<UnionArray>(&mut ctx)?;
     assert_eq!(
         array.child_by_name("nested")?.execute_scalar(0, &mut ctx)?,
-        Scalar::zero_value(&nested_dtype)
+        nested_placeholder
     );
 
     let outer_null = Scalar::null(DType::Union(outer_variants, Nullability::Nullable));
@@ -333,12 +342,95 @@ fn constant_union_builds_nested_union_placeholders() -> VortexResult<()> {
         .execute::<UnionArray>(&mut ctx)?;
     assert_eq!(
         array.child_by_name("nested")?.execute_scalar(0, &mut ctx)?,
-        Scalar::zero_value(&nested_dtype)
+        nested_placeholder
     );
 
     Ok(())
 }
 
+#[test]
+fn constant_union_builds_deeply_nested_union_placeholders() -> VortexResult<()> {
+    let nested_variants = UnionVariants::try_new(
+        ["value"].into(),
+        vec![DType::Primitive(PType::I64, Nullability::NonNullable)],
+        vec![3],
+    )?;
+    let nested_placeholder = Scalar::union(
+        nested_variants.clone(),
+        3,
+        0_i64.into(),
+        Nullability::NonNullable,
+    )?;
+    let nested_dtype = DType::Union(nested_variants, Nullability::NonNullable);
+    let wrapper_dtype = DType::struct_([("nested", nested_dtype)], 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(nested_placeholder)
+    );
+
+    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(),