feat: implement take for canonical sparse Union arrays

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

@@ -242,16 +242,11 @@ pub(crate) fn constant_uncompressed_size_in_bytes(
         | DType::Map(..)
         | DType::FixedSizeList(..)
         | DType::Struct(..)
+        | DType::Union(..)
         | DType::Extension(_) => {
             let canonical = array.array().clone().execute::<Canonical>(ctx)?;
             return canonical_uncompressed_size_in_bytes(&canonical, ctx);
         }
-        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")
         }

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

@@ -24,6 +24,7 @@ use crate::arrays::MapArray;
 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::arrays::varbinview::BinaryView;
@@ -175,13 +176,7 @@ pub(crate) fn constant_canonicalize(
                 StructArray::new_unchecked(fields, struct_dtype.clone(), array.len(), validity)
             })
         }
-        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::Union(..) => Canonical::Union(UnionArray::constant(scalar, array.len())?),
         DType::Variant(_) => Canonical::Variant(VariantArray::try_new(
             array.array().clone().into_array(),
             None,

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

@@ -28,6 +28,8 @@ use crate::arrays::Primitive;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::Struct;
 use crate::arrays::StructArray;
+use crate::arrays::Union;
+use crate::arrays::UnionArray;
 use crate::arrays::VarBinView;
 use crate::arrays::VarBinViewArray;
 use crate::arrays::VariantArray;
@@ -56,12 +58,7 @@ pub(crate) fn take_canonical(
             Canonical::FixedSizeList(take_fixed_size_list(a, codes, ctx))
         }
         CanonicalView::Struct(a) => Canonical::Struct(take_struct(a, codes)),
-        CanonicalView::Union(_) => {
-            todo!(
-                "TODO(connor)[Union]: implement dictionary execution after Union take supports \
-                 nullable indices and outer null propagation"
-            )
-        }
+        CanonicalView::Union(a) => Canonical::Union(take_union(a, codes)),
         CanonicalView::Extension(a) => Canonical::Extension(take_extension(a, codes, ctx)),
         CanonicalView::Variant(a) => {
             let indices = codes.array().clone();
@@ -178,6 +175,15 @@ fn take_struct(array: ArrayView<'_, Struct>, codes: ArrayView<'_, Primitive>) ->
         .into_owned()
 }
 
+fn take_union(array: ArrayView<'_, Union>, codes: ArrayView<'_, Primitive>) -> UnionArray {
+    let codes_ref = codes.array();
+    <Union as TakeReduce>::take(array, codes_ref)
+        .vortex_expect("take union array")
+        .vortex_expect("take union should not return None")
+        .as_::<Union>()
+        .into_owned()
+}
+
 fn take_extension(
     array: ArrayView<'_, Extension>,
     codes: ArrayView<'_, Primitive>,

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

@@ -14,12 +14,15 @@ use crate::array::ArrayParts;
 use crate::array::EmptyArrayData;
 use crate::array::TypedArrayRef;
 use crate::array_slots;
+use crate::arrays::ConstantArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::Union;
+use crate::arrays::union::union_type_ids_dtype;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::dtype::PType;
 use crate::dtype::UnionVariants;
+use crate::scalar::Scalar;
 
 /// Slot layout of a canonical sparse union array.
 #[array_slots(Union)]
@@ -174,6 +177,48 @@ impl Array<Union> {
         }
     }
 
+    /// Construct a `len`-row union in which every row holds `scalar`.
+    ///
+    /// Unselected children are filled with their variant's default value, a null for a nullable
+    /// variant and a zero for a non-nullable one. An outer null `scalar` selects no child at all.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if `scalar` does not have a union dtype.
+    pub fn constant(scalar: &Scalar, len: usize) -> VortexResult<Self> {
+        let union = scalar
+            .as_union_opt()
+            .ok_or_else(|| vortex_err!("Expected a union scalar, got {}", scalar.dtype()))?;
+        let variants = union.variants().clone();
+        let nullability = union.nullability();
+
+        let type_ids = match union.type_id() {
+            Some(type_id) => Scalar::primitive(type_id, nullability),
+            None => Scalar::null(union_type_ids_dtype(nullability)),
+        };
+
+        let selected = union.child_index().zip(union.child());
+
+        let children = variants
+            .variants()
+            .enumerate()
+            .map(|(index, dtype)| {
+                let value = match &selected {
+                    Some((selected_index, child)) if *selected_index == index => child.clone(),
+                    _ => Scalar::default_value(&dtype),
+                };
+
+                ConstantArray::new(value, len).into_array()
+            })
+            .collect::<Vec<_>>();
+
+        Self::try_new(
+            ConstantArray::new(type_ids, len).into_array(),
+            variants,
+            children,
+        )
+    }
+
     /// 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();

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

@@ -5,3 +5,4 @@ pub(crate) mod rules;
 
 mod mask;
 mod slice;
+mod take;

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

@@ -2,14 +2,13 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use crate::arrays::Union;
+use crate::arrays::dict::TakeReduceAdaptor;
 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)),
+    ParentRuleSet::lift(&TakeReduceAdaptor(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.

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

@@ -0,0 +1,52 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use itertools::Itertools;
+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::dict::TakeReduce;
+use crate::arrays::union::UnionArrayExt;
+use crate::arrays::union::UnionArraySlotsExt;
+use crate::builtins::ArrayBuiltins;
+use crate::scalar::Scalar;
+
+/// Gathers the type IDs and every sparse child at `indices`.
+///
+/// Sparse children are row-aligned with the union, so a gather must visit all of them. Take costs
+/// `O(variants * indices)`, which only the dense encoding fixes.
+///
+/// The type IDs carry the union's validity, so gathering them with the original `indices` turns a
+/// null index into an outer union null. The children are gathered with the nulls filled in, which
+/// keeps their declared variant dtypes.
+impl TakeReduce for Union {
+    fn take(array: ArrayView<'_, Union>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
+        // An empty union has no row to point at, so the indices must be all null.
+        if array.is_empty() {
+            return UnionArray::constant(&Scalar::null(array.dtype().as_nullable()), indices.len())
+                .map(UnionArray::into_array)
+                .map(Some);
+        }
+
+        let type_ids = array.type_ids().take(indices.clone())?;
+
+        // This stays a lazy node, so the fill runs once per child. `TakeReduce` has no
+        // `ExecutionCtx` to materialize it with, and the cost scales with the indices, not the
+        // data behind them.
+        let fill_scalar = Scalar::zero_value(&indices.dtype().as_nonnullable());
+        let child_indices = indices.clone().fill_null(fill_scalar)?;
+
+        let children: Vec<ArrayRef> = array
+            .iter_children()
+            .map(|child| child.take(child_indices.clone()))
+            .try_collect()?;
+
+        UnionArray::try_new(type_ids, array.variants().clone(), children)
+            .map(UnionArray::into_array)
+            .map(Some)
+    }
+}

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

@@ -10,6 +10,9 @@
 //!
 //! Type ID values are not validated during construction. Accessing a non-null row whose type ID
 //! is not declared by the union variants will panic.
+//!
+//! Mask rewrites only the type IDs. Slice and take rewrite every child to keep them row-aligned,
+//! so each costs `O(variants)` child operations.
 
 use crate::dtype::DType;
 use crate::dtype::Nullability;

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

@@ -5,14 +5,18 @@ use vortex_buffer::ByteBufferMut;
 use vortex_buffer::buffer;
 use vortex_error::VortexExpect;
 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::IntoArray;
 use crate::VortexSessionExecute;
+use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes;
 use crate::array_session;
 use crate::arrays::BoolArray;
+use crate::arrays::ConstantArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::Union;
 use crate::arrays::UnionArray;
@@ -27,6 +31,8 @@ use crate::serde::SerializeOptions;
 use crate::serde::SerializedArray;
 use crate::validity::Validity;
 
+mod take;
+
 fn variants() -> VortexResult<UnionVariants> {
     UnionVariants::try_new(
         ["number", "flag"].into(),
@@ -234,6 +240,68 @@ fn slice_and_filter_preserve_sparse_alignment() -> VortexResult<()> {
     Ok(())
 }
 
+/// A constant union canonicalizes into a sparse union whose selected child repeats the scalar's
+/// value and whose unselected children hold placeholders.
+#[test]
+fn constant_union_canonicalizes_to_sparse_union() -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let i32_variant = DType::Primitive(PType::I32, Nullability::NonNullable);
+    let i64_variant = DType::Primitive(PType::I64, Nullability::Nullable);
+
+    // Each case pairs a constant scalar with the dtype of a variant it does not select.
+    let cases = [
+        (
+            Scalar::union(variants()?, 9, true.into(), Nullability::NonNullable)?,
+            i32_variant.clone(),
+        ),
+        (
+            Scalar::union(nullable_variants()?, 5, 7i32.into(), Nullability::Nullable)?,
+            i64_variant,
+        ),
+        (
+            Scalar::null(DType::Union(variants()?, Nullability::Nullable)),
+            i32_variant,
+        ),
+    ];
+
+    for (scalar, unselected_dtype) in cases {
+        let canonical = ConstantArray::new(scalar.clone(), 3)
+            .into_array()
+            .execute::<Canonical>(&mut ctx)?
+            .into_union();
+
+        assert_eq!(canonical.dtype(), scalar.dtype());
+
+        let unselected = canonical
+            .iter_children()
+            .find(|child| child.dtype() == &unselected_dtype)
+            .ok_or_else(|| vortex_err!("No child with dtype {unselected_dtype}"))?;
+        assert_eq!(
+            unselected.execute_scalar(0, &mut ctx)?,
+            Scalar::default_value(&unselected_dtype)
+        );
+
+        for index in 0..canonical.len() {
+            assert_eq!(canonical.execute_scalar(index, &mut ctx)?, scalar);
+        }
+    }
+
+    Ok(())
+}
+
+#[test]
+fn constant_union_reports_uncompressed_size() -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+    let scalar = Scalar::union(variants()?, 5, 10i32.into(), Nullability::NonNullable)?;
+
+    // Four `u8` type IDs, four `i32` rows, and four `bool` placeholder bits rounded up to a byte.
+    let size = uncompressed_size_in_bytes(&ConstantArray::new(scalar, 4).into_array(), &mut ctx)?;
+
+    assert_eq!(size, 4 + 4 * 4 + 1);
+
+    Ok(())
+}
+
 #[test]
 fn serde_roundtrip() -> VortexResult<()> {
     let session = array_session();