Add support to cast from `UnionArray` by friendlymatthew · Pull Request #9544 · apache/arrow-rs

friendlymatthew · 2026-03-12T17:39:52Z

Which issue does this PR close?

Closes Add Union type casting support to arrow-cast #9225

Rationale for this change

This PR adds native support for casting Union arrays in the cast kernel. Previously, can_cast_types and cast_with_options had no handling for union types at all

friendlymatthew

self review

friendlymatthew · 2026-03-12T17:41:23Z

+        (_, Union(_, _)) => Err(ArrowError::CastError(format!(
+            "Casting from {from_type} to {to_type} not supported"
+        ))),


This PR does not support casting scalar types into a Union type

It's not a feature that we currently need and will probably need more complex work to get it correct. Though, I think this is a cool/valid feature. I can file an issue if liked

friendlymatthew · 2026-03-12T17:43:36Z

+    let type_ids = array.type_ids().clone();
+    let offsets = array.offsets().cloned();
+
+    let new_children: Vec<ArrayRef> = from_fields
+        .iter()
+        .map(|(from_id, _from_field)| {
+            let (_, to_field) = to_fields
+                .iter()
+                .find(|(to_id, _)| *to_id == from_id)
+                .ok_or_else(|| {
+                    ArrowError::CastError(format!(
+                        "Cannot cast union: type_id {from_id} not found in target union"
+                    ))
+                })?;
+            let child = array.child(from_id);
+            cast_with_options(child.as_ref(), to_field.data_type(), cast_options)
+        })
+        .collect::<Result<_, _>>()?;


look ups like this can go from n^2 to nlogn if we implement #8937, since sorting by type_id will give us logn searching

friendlymatthew · 2026-03-16T14:22:43Z

cc @Jefffrey @alamb curious to get your thoughts

This is a feature we use in production at Pydantic, and we actively maintain a vendored version that we're eager to upstream

alamb · 2026-03-16T20:11:00Z

Thanks -- I hope to go through arrow prs more carefully starting tomorrow

alamb

thank you @friendlymatthew

I like where this is headed.

I left some suggestions - let me know if it makes sense

alamb · 2026-03-17T20:45:33Z

+    array: &UnionArray,
+    from_fields: &UnionFields,
+    to_fields: &UnionFields,
+    _to_mode: UnionMode,


why is the _to_mode parameter ignored?

alamb · 2026-03-17T20:47:57Z

+
+    #[test]
+    fn test_cast_union_prefers_exact_type_match() {
+        // Union(Int32, Int64) → Int64: Int64 is an exact match, so only


Can you please document the expected union casting behavior in the documentation?

Specifically, somewhere here:

https://docs.rs/arrow/latest/arrow/compute/fn.cast_with_options.html#durations-and-intervals

That will make it easier to understand the intended semantics for union casting as well as verify that this code implements the semantics correctly

Jefffrey · 2026-03-18T14:22:52Z

+/// Since union extraction inherently introduces nulls for non-matching rows,
+/// the target type's inner fields are made nullable to avoid validation errors
+/// (e.g., casting to `List(non-nullable Utf8)` becomes `List(nullable Utf8)`).
+fn cast_union_to_type(


I do wonder on the semantics being introduced here; specifically about how we apparently allow casting if at least one field can be cast, but we null the other fields. It might not seem intuitive; I do wonder if we have some existing precedent in other systems we could follow 🤔

Another potential option is to introduce the behavior into union_extract: https://docs.rs/arrow/latest/arrow/compute/kernels/union_extract/fn.union_extract.html

For example, maybe union_extract_with_option that takes some sort of options struct that allows better control over the extraction

friendlymatthew

@Jefffrey @alamb

I restarted this PR to cut down on the scope.

This PR adds support for casting union arrays to scalar types only. The casting doesn't reimplement the extraction logic, but now builds on the existing union_extract kernel

The core addition is union_extract_by_type which selects the best-matching variant from a union for a given target type.

The variant selection uses a 3-pass heuristic: exact type match, same type family, then any castable variant

Curious to hear your thoughts!

friendlymatthew

self review

friendlymatthew · 2026-03-23T16:18:58Z

+// this is used during variant selection to prefer a "close" type over a distant cast
+// for example: when targeting Utf8View, a Utf8 variant is preferred over Int32 despite both being castable
+fn same_type_family(a: &DataType, b: &DataType) -> bool {
+    use DataType::*;
+    matches!(
+        (a, b),
+        (Utf8 | LargeUtf8 | Utf8View, Utf8 | LargeUtf8 | Utf8View)
+            | (
+                Binary | LargeBinary | BinaryView,
+                Binary | LargeBinary | BinaryView
+            )
+            | (Int8 | Int16 | Int32 | Int64, Int8 | Int16 | Int32 | Int64)
+            | (
+                UInt8 | UInt16 | UInt32 | UInt64,
+                UInt8 | UInt16 | UInt32 | UInt64
+            )
+            | (Float16 | Float32 | Float64, Float16 | Float32 | Float64)
+    )
+}


I wonder if this is something that would be useful as a DataType api...

Equivalence classes let you answer "are these the same logical data" without caring about the physical repr

I think the notion of "same logical type" gets a bit subtle -- for example LargeUtf8 and Utf8View might seem like the same logical type, but there are some values that LargeUTF8 can store that Utf8View can not (e.g. individual strings over 4GB)

friendlymatthew · 2026-03-23T16:20:18Z

+// variant selection heuristic — 3 passes with decreasing specificity:
+//
+// first pass: field type == target type
+// second pass: field and target are in the same equivalence class
+//              (e.g., Utf8 and Utf8View are both strings)
+// third pass: field can be cast to target
+//      note: this is the most permissive and may lose information
+//      also, the matching logic is greedy so it will pick the first 'castable' variant
+//
+// each pass picks the first matching variant by type_id order.


the main heuristic!

I think pass 2 is necessary because can_cast_types is permissive enough that without it, Union(Int32, Utf8) cast to a Utf8View would pick Int32 over Utf8...

Though we can fix this by defining the Utf8 variant first. But I wonder if we should leave it up to the user...

I think the user can have direct control if they want and can use union_extract and cast

friendlymatthew · 2026-04-01T14:28:25Z

@alamb @Jefffrey curious if you have any thoughts. Work is pretty eager to see this through!

alamb

I think this one is good to go from my perspective. Thank you @friendlymatthew

Some comments;

terminology

I found the use of the term "variant" with UnionArray confusing. I think The Arrow Spec refers to this concept as "child array" and I think the arrow-rs code refers to it astype_id

Maybe we can use the same terminology

alamb · 2026-04-01T21:32:36Z

+// this is used during variant selection to prefer a "close" type over a distant cast
+// for example: when targeting Utf8View, a Utf8 variant is preferred over Int32 despite both being castable
+fn same_type_family(a: &DataType, b: &DataType) -> bool {
+    use DataType::*;
+    matches!(
+        (a, b),
+        (Utf8 | LargeUtf8 | Utf8View, Utf8 | LargeUtf8 | Utf8View)
+            | (
+                Binary | LargeBinary | BinaryView,
+                Binary | LargeBinary | BinaryView
+            )
+            | (Int8 | Int16 | Int32 | Int64, Int8 | Int16 | Int32 | Int64)
+            | (
+                UInt8 | UInt16 | UInt32 | UInt64,
+                UInt8 | UInt16 | UInt32 | UInt64
+            )
+            | (Float16 | Float32 | Float64, Float16 | Float32 | Float64)
+    )
+}


I think the notion of "same logical type" gets a bit subtle -- for example LargeUtf8 and Utf8View might seem like the same logical type, but there are some values that LargeUTF8 can store that Utf8View can not (e.g. individual strings over 4GB)

alamb · 2026-04-01T21:41:45Z

+// variant selection heuristic — 3 passes with decreasing specificity:
+//
+// first pass: field type == target type
+// second pass: field and target are in the same equivalence class
+//              (e.g., Utf8 and Utf8View are both strings)
+// third pass: field can be cast to target
+//      note: this is the most permissive and may lose information
+//      also, the matching logic is greedy so it will pick the first 'castable' variant
+//
+// each pass picks the first matching variant by type_id order.


I think the user can have direct control if they want and can use union_extract and cast

alamb · 2026-04-01T21:47:54Z

+        )));
+    };
+
+    let extracted = union_extract(union_array, field.name())?;


technically this will not work if there are repeated field names in the Union (as resolve_variant will return a string and union_extract will return the first field with that name)

You could potentially avoid this by returning the type_id / index instead

We can also handle it in a follow on PR

Filed: #9664

alamb · 2026-04-06T14:13:28Z

Thank you @friendlymatthew

github-actions bot added the arrow Changes to the arrow crate label Mar 12, 2026

friendlymatthew commented Mar 12, 2026

View reviewed changes

friendlymatthew force-pushed the friendlymatthew/cast-union branch 2 times, most recently from 9c19cb1 to 8dbe39f Compare March 12, 2026 18:24

alamb mentioned this pull request Mar 13, 2026

Add Union type casting support to arrow-cast #9225

Closed

alamb reviewed Mar 17, 2026

View reviewed changes

Jefffrey reviewed Mar 18, 2026

View reviewed changes

friendlymatthew force-pushed the friendlymatthew/cast-union branch 2 times, most recently from 096984a to 3056c06 Compare March 23, 2026 16:08

rework union casting to scalar

1cf2085

friendlymatthew force-pushed the friendlymatthew/cast-union branch from 3056c06 to 1cf2085 Compare March 23, 2026 16:13

friendlymatthew commented Mar 23, 2026

View reviewed changes

fix up a test behavior

b47a20a

friendlymatthew force-pushed the friendlymatthew/cast-union branch from 0dc9c7b to b47a20a Compare March 23, 2026 16:45

alamb changed the title ~~add native union type casting support~~ Add support to cast from UnionArray Apr 1, 2026

alamb approved these changes Apr 1, 2026

View reviewed changes

friendlymatthew added 3 commits April 6, 2026 09:19

use is_nested

1f2aad7

explain heuristic

cb9ef07

rename variant -> child array

ecdefa8

friendlymatthew force-pushed the friendlymatthew/cast-union branch from ad3291c to ecdefa8 Compare April 6, 2026 13:36

friendlymatthew mentioned this pull request Apr 6, 2026

Union cast is incorrect for duplicate field names #9664

Closed

alamb merged commit 9c6b42e into apache:main Apr 6, 2026
26 checks passed

Conversation

friendlymatthew commented Mar 12, 2026

Which issue does this PR close?

Rationale for this change

Uh oh!

friendlymatthew left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

friendlymatthew commented Mar 16, 2026

Uh oh!

alamb commented Mar 16, 2026

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

friendlymatthew left a comment

Choose a reason for hiding this comment

Uh oh!

friendlymatthew left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

friendlymatthew commented Apr 1, 2026

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

terminology

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

alamb commented Apr 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants