Add Union scalar support

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

Author: connortsui20

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

@@ -18,3 +18,46 @@ impl OperationsVTable<Constant> for Constant {
         Ok(array.scalar.clone())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_error::VortexResult;
+
+    use crate::IntoArray;
+    use crate::VortexSessionExecute;
+    use crate::arrays::ConstantArray;
+    use crate::dtype::DType;
+    use crate::dtype::Nullability;
+    use crate::dtype::PType;
+    use crate::dtype::UnionVariants;
+    use crate::scalar::Scalar;
+
+    #[test]
+    fn scalar_at_preserves_union_scalar() -> VortexResult<()> {
+        let variants = UnionVariants::try_new(
+            ["int", "string"].into(),
+            vec![
+                DType::Primitive(PType::I32, Nullability::Nullable),
+                DType::Utf8(Nullability::NonNullable),
+            ],
+            vec![5, 9],
+        )?;
+
+        let scalar = Scalar::union(
+            variants.clone(),
+            5,
+            Scalar::primitive(42_i32, Nullability::Nullable),
+        )?;
+        let array = ConstantArray::new(scalar.clone(), 3).into_array();
+        let mut ctx = crate::array_session().create_execution_ctx();
+
+        assert_eq!(array.execute_scalar(1, &mut ctx)?, scalar);
+
+        let null = Scalar::null(DType::Union(variants));
+        let array = ConstantArray::new(null.clone(), 3).into_array();
+
+        assert_eq!(array.execute_scalar(1, &mut ctx)?, null);
+
+        Ok(())
+    }
+}

vortex-array/src/scalar/arbitrary.rs

@@ -95,7 +95,21 @@ pub fn random_scalar(u: &mut Unstructured, dtype: &DType) -> Result<Scalar> {
             )),
         )
         .vortex_expect("unable to construct random `Scalar`_"),
-        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
+        DType::Union(variants) => {
+            let child_index = u.choose_index(variants.len())?;
+
+            let child_dtype = variants
+                .variant_by_index(child_index)
+                .vortex_expect("chosen union child index must be valid");
+            let child = random_scalar(u, &child_dtype)?;
+
+            Scalar::union(
+                variants.clone(),
+                variants.child_index_to_tag(child_index),
+                child,
+            )
+            .vortex_expect("generated union scalar must be valid")
+        }
         DType::Variant(_) => todo!(),
         DType::Extension(..) => {
             unreachable!("Can't yet generate arbitrary scalars for ext dtype")

vortex-array/src/scalar/cast.rs

@@ -24,6 +24,16 @@ impl Scalar {
             return Ok(self.clone());
         }
 
+        // Union nullability is part of its variant dtypes, so the generic nullability-only cast
+        // below is not valid for unions. Keep all non-identity union casts unsupported until their
+        // semantics are defined.
+        if self.dtype().is_union() || target_dtype.is_union() {
+            vortex_bail!(
+                "non-identity union scalar cast from {} to {target_dtype} is not supported",
+                self.dtype()
+            );
+        }
+
         // Check for solely nullability casting.
         if self.dtype().eq_ignore_nullability(target_dtype) {
             // Cast from non-nullable to nullable or vice versa.
@@ -58,13 +68,16 @@ impl Scalar {
             DType::Binary(_) => self.as_binary().cast(target_dtype),
             DType::List(..) | DType::FixedSizeList(..) => self.as_list().cast(target_dtype),
             DType::Struct(..) => self.as_struct().cast(target_dtype),
-            DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
+            DType::Union(..) => unreachable!("union casts are handled before scalar dispatch"),
             DType::Variant(_) => vortex_bail!("Variant scalars can't be cast to {target_dtype}"),
             DType::Extension(..) => self.as_extension().cast(target_dtype),
         }
     }
 
     /// Cast the scalar into a nullable version of its current type.
+    ///
+    /// For a union this makes every variant nullable, since a union has no top-level nullability of
+    /// its own.
     pub fn into_nullable(self) -> Scalar {
         let (dtype, value) = self.into_parts();
         Self::try_new(dtype.as_nullable(), value)

vortex-array/src/scalar/constructor.rs

@@ -8,20 +8,25 @@ use std::sync::Arc;
 use vortex_buffer::BufferString;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure_eq;
+use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
 
 use crate::dtype::DType;
 use crate::dtype::DecimalDType;
 use crate::dtype::NativePType;
 use crate::dtype::Nullability;
 use crate::dtype::PType;
+use crate::dtype::UnionVariants;
 use crate::dtype::extension::ExtDType;
 use crate::dtype::extension::ExtDTypeRef;
 use crate::dtype::extension::ExtVTable;
 use crate::scalar::DecimalValue;
 use crate::scalar::PValue;
 use crate::scalar::Scalar;
 use crate::scalar::ScalarValue;
+use crate::scalar::UnionValue;
 
 // TODO(connor): Really, we want `try_` constructors that return errors instead of just panic.
 impl Scalar {
@@ -190,6 +195,40 @@ impl Scalar {
             .vortex_expect("unable to construct an extension `Scalar`")
     }
 
+    /// Creates a union scalar from a type ID and child scalar.
+    ///
+    /// The selected child scalar is stored in full, so an inner null child remains distinct from a
+    /// null at the outer union level.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the type ID is not present in `variants` or if the child scalar's dtype
+    /// does not exactly match the selected variant dtype.
+    pub fn union(variants: UnionVariants, type_id: u8, child: Scalar) -> VortexResult<Self> {
+        let child_index = variants.tag_to_child_index(type_id).ok_or_else(|| {
+            vortex_err!(
+                "union type ID {type_id} is not present in {:?}",
+                variants.type_ids()
+            )
+        })?;
+
+        let expected_dtype = variants
+            .variant_by_index(child_index)
+            .vortex_expect("type ID resolved to a valid child index");
+
+        vortex_ensure_eq!(
+            child.dtype(),
+            &expected_dtype,
+            "union type ID {type_id} selects child dtype {expected_dtype}, got {}",
+            child.dtype()
+        );
+
+        Self::try_new(
+            DType::Union(variants),
+            Some(ScalarValue::Union(UnionValue::new(type_id, child))),
+        )
+    }
+
     /// Creates a new variant scalar from a row-specific nested scalar.
     ///
     /// Use [`Scalar::null(DType::Variant(Nullability::Nullable))`][Scalar::null] for a top-level

vortex-array/src/scalar/display.rs

@@ -20,7 +20,7 @@ impl Display for Scalar {
             DType::Binary(_) => write!(f, "{}", self.as_binary()),
             DType::List(..) | DType::FixedSizeList(..) => write!(f, "{}", self.as_list()),
             DType::Struct(..) => write!(f, "{}", self.as_struct()),
-            DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
+            DType::Union(..) => write!(f, "{}", self.as_union()),
             DType::Variant(_) => write!(f, "{}", self.as_variant()),
             DType::Extension(_) => write!(f, "{}", self.as_extension()),
         }
@@ -30,13 +30,15 @@ impl Display for Scalar {
 #[cfg(test)]
 mod tests {
     use vortex_buffer::ByteBuffer;
+    use vortex_error::VortexResult;
 
     use crate::dtype::DType;
     use crate::dtype::FieldName;
     use crate::dtype::Nullability::NonNullable;
     use crate::dtype::Nullability::Nullable;
     use crate::dtype::PType;
     use crate::dtype::StructFields;
+    use crate::dtype::UnionVariants;
     use crate::extension::datetime::Date;
     use crate::extension::datetime::Time;
     use crate::extension::datetime::TimeUnit;
@@ -79,6 +81,30 @@ mod tests {
         );
     }
 
+    #[test]
+    fn display_union() -> VortexResult<()> {
+        let variants = UnionVariants::new(
+            ["int", "string"].into(),
+            vec![
+                DType::Primitive(PType::I32, Nullable),
+                DType::Utf8(NonNullable),
+            ],
+        )?;
+
+        let scalar = Scalar::union(variants.clone(), 0, Scalar::primitive(42_i32, Nullable))?;
+
+        assert_eq!(format!("{scalar}"), "int(42i32)");
+        let inner_null = Scalar::union(
+            variants.clone(),
+            0,
+            Scalar::null(DType::Primitive(PType::I32, Nullable)),
+        )?;
+        assert_eq!(format!("{inner_null}"), "int(null)");
+        assert_eq!(format!("{}", Scalar::null(DType::Union(variants))), "null");
+
+        Ok(())
+    }
+
     #[test]
     fn display_utf8() {
         assert_eq!(

vortex-array/src/scalar/downcast.rs

@@ -8,6 +8,7 @@ use vortex_buffer::ByteBuffer;
 use vortex_error::VortexExpect;
 use vortex_error::vortex_panic;
 
+use crate::dtype::DType;
 use crate::scalar::BinaryScalar;
 use crate::scalar::BoolScalar;
 use crate::scalar::DecimalScalar;
@@ -19,6 +20,8 @@ use crate::scalar::PrimitiveScalar;
 use crate::scalar::Scalar;
 use crate::scalar::ScalarValue;
 use crate::scalar::StructScalar;
+use crate::scalar::UnionScalar;
+use crate::scalar::UnionValue;
 use crate::scalar::Utf8Scalar;
 use crate::scalar::VariantScalar;
 
@@ -156,6 +159,26 @@ impl Scalar {
         Some(ExtScalar::new_unchecked(self.dtype(), self.value()))
     }
 
+    /// Returns a view of the scalar as a union scalar.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the scalar does not have a [`Union`](crate::dtype::DType::Union) type.
+    pub fn as_union(&self) -> UnionScalar<'_> {
+        self.as_union_opt()
+            .vortex_expect("Failed to convert scalar to union")
+    }
+
+    /// Returns a view of the scalar as a union scalar if it has a union type.
+    pub fn as_union_opt(&self) -> Option<UnionScalar<'_>> {
+        if !matches!(self.dtype(), DType::Union(_)) {
+            return None;
+        }
+
+        // Scalar construction has already validated the value against this union dtype.
+        Some(UnionScalar::new_unchecked(self.dtype(), self.value()))
+    }
+
     /// Returns a view of the scalar as a variant scalar.
     ///
     /// # Panics
@@ -273,6 +296,22 @@ impl ScalarValue {
         }
     }
 
+    /// Returns the union value, panicking if the value is not a [`Union`](ScalarValue::Union).
+    pub fn as_union(&self) -> &UnionValue {
+        match self {
+            ScalarValue::Union(value) => value,
+            _ => vortex_panic!("ScalarValue is not a Union"),
+        }
+    }
+
+    /// Returns the union value, panicking if the value is not a [`Union`](ScalarValue::Union).
+    pub fn into_union(self) -> UnionValue {
+        match self {
+            ScalarValue::Union(value) => value,
+            _ => vortex_panic!("ScalarValue is not a Union"),
+        }
+    }
+
     /// Returns the row-specific scalar wrapped by a variant, panicking if the value is not a
     /// variant.
     pub fn as_variant(&self) -> &Scalar {

vortex-array/src/scalar/mod.rs

@@ -4,8 +4,8 @@
 //! Scalar values and types for the Vortex system.
 //!
 //! This crate provides scalar types and values that can be used to represent individual data
-//! elements in the Vortex array system. [`Scalar`]s are composed of a logical data type ([`DType`])
-//! and an optional (encoding nullability) value ([`ScalarValue`]).
+//! elements in the Vortex array system. A [`Scalar`] pairs a logical data type ([`DType`]) with an
+//! optional non-null value ([`ScalarValue`]); [`None`] represents a null scalar.
 //!
 //! Note that the implementations of `Scalar` are split into several different modules.
 //!