Add UnionArray compute functions

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs

@@ -9,6 +9,7 @@ mod list_view;
 mod null;
 mod primitive;
 mod struct_;
+mod union;
 mod varbinview;
 
 use std::mem::size_of;
@@ -21,6 +22,7 @@ use list_view::list_view_uncompressed_size_in_bytes;
 use null::null_uncompressed_size_in_bytes;
 use primitive::primitive_uncompressed_size_in_bytes;
 use struct_::struct_uncompressed_size_in_bytes;
+use union::union_uncompressed_size_in_bytes;
 use varbinview::varbinview_uncompressed_size_in_bytes;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -199,9 +201,7 @@ pub(crate) fn canonical_uncompressed_size_in_bytes(
         Canonical::List(array) => list_view_uncompressed_size_in_bytes(array, ctx),
         Canonical::FixedSizeList(array) => fixed_size_list_uncompressed_size_in_bytes(array, ctx),
         Canonical::Struct(array) => struct_uncompressed_size_in_bytes(array, ctx),
-        Canonical::Union(_) => {
-            todo!("TODO(connor)[Union]: implement UncompressedSizeInBytes for Union arrays")
-        }
+        Canonical::Union(array) => union_uncompressed_size_in_bytes(array, ctx),
         Canonical::Extension(array) => extension_uncompressed_size_in_bytes(array, ctx),
         Canonical::Variant(_) => {
             vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays")
@@ -236,7 +236,12 @@ pub(crate) fn constant_uncompressed_size_in_bytes(
             let canonical = array.array().clone().execute::<Canonical>(ctx)?;
             return canonical_uncompressed_size_in_bytes(&canonical, ctx);
         }
-        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
+        DType::Union(..) => {
+            todo!(
+                "TODO(connor)[Union]: support constant Union size accounting after constant Union \
+                 canonicalization defines inactive sparse-child placeholders"
+            )
+        }
         DType::Variant(_) => {
             vortex_bail!("UncompressedSizeInBytes is not supported for Variant arrays")
         }
@@ -342,6 +347,7 @@ mod tests {
     use crate::arrays::NullArray;
     use crate::arrays::PrimitiveArray;
     use crate::arrays::StructArray;
+    use crate::arrays::UnionArray;
     use crate::arrays::VarBinViewArray;
     use crate::arrays::VariantArray;
     use crate::builders::builder_with_capacity;
@@ -350,6 +356,7 @@ mod tests {
     use crate::dtype::FieldNames;
     use crate::dtype::Nullability;
     use crate::dtype::PType;
+    use crate::dtype::UnionVariants;
     use crate::expr::stats::Precision;
     use crate::expr::stats::Stat;
     use crate::expr::stats::StatsProvider;
@@ -539,6 +546,26 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn union_sums_type_ids_and_sparse_children() -> VortexResult<()> {
+        let type_ids = PrimitiveArray::from_iter([5_u8, 9, 5]).into_array();
+        let numbers = PrimitiveArray::from_iter([10_i32, 0, 30]).into_array();
+        let flags = BoolArray::from_iter([false, true, false]).into_array();
+        let expected = aggregate(&type_ids)? + aggregate(&numbers)? + aggregate(&flags)?;
+        let variants = UnionVariants::try_new(
+            ["number", "flag"].into(),
+            vec![
+                DType::Primitive(PType::I32, Nullability::NonNullable),
+                DType::Bool(Nullability::NonNullable),
+            ],
+            vec![5, 9],
+        )?;
+        let array = UnionArray::try_new(type_ids, variants, vec![numbers, flags])?.into_array();
+
+        assert_eq!(aggregate(&array)?, expected);
+        Ok(())
+    }
+
     #[test]
     fn extension_matches_materialized_size() -> VortexResult<()> {
         let storage = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array();

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/union.rs

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use super::uncompressed_size_in_bytes_u64;
+use crate::ExecutionCtx;
+use crate::arrays::UnionArray;
+use crate::arrays::union::UnionArrayExt;
+
+pub(super) fn union_uncompressed_size_in_bytes(
+    array: &UnionArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<u64> {
+    let mut size = uncompressed_size_in_bytes_u64(array.type_ids(), ctx)?;
+
+    for child in array.iter_children() {
+        size = size
+            .checked_add(uncompressed_size_in_bytes_u64(child, ctx)?)
+            .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?;
+    }
+
+    Ok(size)
+}

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

@@ -251,6 +251,12 @@ impl VTable for Chunked {
 
     fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
         match array.dtype() {
+            DType::Union(..) => {
+                todo!(
+                    "TODO(connor)[Union]: canonicalize chunked Union arrays by packing type IDs and \
+                     every sparse child along identical chunk boundaries"
+                )
+            }
             // Struct, List, FixedSizeList, and Variant need child swizzling that the builder path
             // cannot express.
             DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Variant(..) => {

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

@@ -164,7 +164,13 @@ pub(crate) fn constant_canonicalize(
                 StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity)
             })
         }
-        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
+        DType::Union(..) => {
+            todo!(
+                "TODO(connor)[Union]: canonicalize constant Union arrays in a focused follow-up \
+                 after defining placeholder values for every inactive sparse child, including \
+                 nested Struct and Union variants"
+            )
+        }
         DType::Variant(_) => Canonical::Variant(VariantArray::try_new(
             array.array().clone().into_array(),
             None,

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

@@ -54,7 +54,10 @@ pub(crate) fn take_canonical(
         }
         Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)),
         Canonical::Union(_) => {
-            todo!("TODO(connor)[Union]: implement dictionary execution for Union arrays")
+            todo!(
+                "TODO(connor)[Union]: implement dictionary execution after Union take supports \
+                 nullable indices and outer null propagation"
+            )
         }
         Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)),
         Canonical::Variant(a) => {

vortex-array/src/arrays/filter/execute/mod.rs

@@ -28,17 +28,21 @@ use crate::arrays::variant::VariantArrayExt;
 use crate::scalar::Scalar;
 use crate::validity::Validity;
 
+pub(crate) mod byte_compress;
+
+mod slice;
+mod take;
+
 mod bitbuffer;
-mod bool;
 mod buffer;
-pub(crate) mod byte_compress;
+
+mod bool;
 mod decimal;
 mod fixed_size_list;
 mod listview;
 mod primitive;
-mod slice;
 mod struct_;
-pub mod take;
+mod union;
 mod varbinview;
 
 /// A helper function that lazily filters a [`Validity`] with selection mask values.
@@ -95,9 +99,7 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc<MaskValues>) -> Ca
             Canonical::FixedSizeList(fixed_size_list::filter_fixed_size_list(&a, mask))
         }
         Canonical::Struct(a) => Canonical::Struct(struct_::filter_struct(&a, mask)),
-        Canonical::Union(_) => {
-            todo!("TODO(connor)[Union]: implement filter for Union arrays")
-        }
+        Canonical::Union(a) => Canonical::Union(union::filter_union(&a, mask)),
         Canonical::Extension(a) => {
             let filtered_storage = a
                 .storage_array()

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

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::sync::Arc;
+
+use vortex_error::VortexExpect;
+use vortex_mask::Mask;
+use vortex_mask::MaskValues;
+
+use crate::ArrayRef;
+use crate::arrays::UnionArray;
+use crate::arrays::union::UnionArrayExt;
+
+pub fn filter_union(array: &UnionArray, mask: &Arc<MaskValues>) -> UnionArray {
+    let filter_mask = Mask::Values(Arc::clone(mask));
+
+    let type_ids = array
+        .type_ids()
+        .filter(filter_mask.clone())
+        .vortex_expect("UnionArray type IDs are guaranteed to support filter");
+
+    let children: Vec<ArrayRef> = array
+        .iter_children()
+        .map(|child| {
+            child
+                .filter(filter_mask.clone())
+                .vortex_expect("UnionArray children are guaranteed to support filter")
+        })
+        .collect();
+
+    UnionArray::try_new(type_ids, array.variants().clone(), children)
+        .vortex_expect("filtered UnionArray children have consistent dtypes and lengths")
+}

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

@@ -17,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;
@@ -50,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(_) => {
-            todo!("TODO(connor)[Union]: implement masking for Union arrays")
-        }
+        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)?),
     })
@@ -69,7 +69,7 @@ fn mask_validity_primitive(
 ) -> VortexResult<PrimitiveArray> {
     let ptype = array.ptype();
     let new_validity = Validity::and(array.validity()?, validity)?;
-    // SAFETY: validity has same length as values
+    // SAFETY: We're only changing validity, not the data structure.
     Ok(unsafe {
         PrimitiveArray::new_unchecked_from_handle(
             array.buffer_handle().clone(),
@@ -81,7 +81,7 @@ fn mask_validity_primitive(
 
 fn mask_validity_decimal(array: DecimalArray, validity: Validity) -> VortexResult<DecimalArray> {
     let new_validity = Validity::and(array.validity()?, validity)?;
-    // SAFETY: We're only changing validity, not the data structure
+    // SAFETY: We're only changing validity, not the data structure.
     Ok(unsafe {
         DecimalArray::new_unchecked_handle(
             array.buffer_handle().clone(),
@@ -99,7 +99,7 @@ fn mask_validity_varbinview(
 ) -> VortexResult<VarBinViewArray> {
     let dtype = array.dtype().as_nullable();
     let new_validity = Validity::and(array.validity()?, validity)?;
-    // SAFETY: We're only changing validity, not the data structure
+    // SAFETY: We're only changing validity, not the data structure.
     Ok(unsafe {
         VarBinViewArray::new_handle_unchecked(
             array.views_handle().clone(),
@@ -112,7 +112,7 @@ fn mask_validity_varbinview(
 
 fn mask_validity_listview(array: ListViewArray, validity: Validity) -> VortexResult<ListViewArray> {
     let new_validity = Validity::and(array.validity()?, validity)?;
-    // SAFETY: We're only changing validity, not the data structure
+    // SAFETY: We're only changing validity, not the data structure.
     let is_zctl = array.is_zero_copy_to_list();
     Ok(unsafe {
         ListViewArray::new_unchecked(
@@ -132,7 +132,7 @@ fn mask_validity_fixed_size_list(
     let len = array.len();
     let list_size = array.list_size();
     let new_validity = Validity::and(array.validity()?, validity)?;
-    // SAFETY: We're only changing validity, not the data structure
+    // SAFETY: We're only changing validity, not the data structure.
     Ok(unsafe {
         FixedSizeListArray::new_unchecked(array.elements().clone(), list_size, new_validity, len)
     })
@@ -143,16 +143,28 @@ fn mask_validity_struct(array: StructArray, validity: Validity) -> VortexResult<
     let new_validity = Validity::and(array.validity()?, validity)?;
     let fields = array.unmasked_fields();
     let struct_fields = array.struct_fields();
-    // SAFETY: We're only changing validity, not the data structure
+    // SAFETY: We're only changing validity, not the data structure.
     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: We're only changing validity, not the data structure.
+    Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, children) })
+}
+
 fn mask_validity_extension(
     array: ExtensionArray,
     validity: Validity,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ExtensionArray> {
-    // For extension arrays, we need to mask the underlying storage
+    // For extension arrays, we need to mask the underlying storage.
     let storage = array.storage_array().clone().execute::<Canonical>(ctx)?;
     let masked_storage = mask_validity_canonical(storage, validity, ctx)?;
     let masked_storage = masked_storage.into_array();

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

@@ -20,6 +20,7 @@ use crate::arrays::PrimitiveArray;
 use crate::arrays::Union;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
+use crate::dtype::PType;
 use crate::dtype::UnionVariants;
 
 /// The row-aligned array of type IDs selecting a union child.
@@ -144,10 +145,7 @@ impl Array<Union> {
         children: impl Into<Arc<[ArrayRef]>>,
     ) -> VortexResult<Self> {
         vortex_ensure!(
-            matches!(
-                type_ids.dtype(),
-                DType::Primitive(crate::dtype::PType::U8, _)
-            ),
+            matches!(type_ids.dtype(), DType::Primitive(PType::U8, _)),
             "UnionArray type_ids must be u8, got {}",
             type_ids.dtype()
         );

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

@@ -0,0 +1,24 @@
+// 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>> {
+        UnionArray::try_new(
+            array.type_ids().clone().mask(mask.clone())?,
+            array.variants().clone(),
+            array.children(),
+        )
+        .map(|a| Some(a.into_array()))
+    }
+}

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

@@ -0,0 +1,7 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+pub(crate) mod rules;
+
+mod mask;
+mod slice;

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

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::arrays::Union;
+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(&MaskReduceAdaptor(Union)),
+    ParentRuleSet::lift(&SliceReduceAdaptor(Union)),
+]);
+
+// TODO(connor)[Union]: Register TakeReduce only once nullable indices can introduce outer Union
+// nulls while preserving the sparse children's dtypes and row alignment.