arrow_cast/cast/
union.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Cast support for union arrays.
6
7use crate::cast::can_cast_types;
8use crate::cast_with_options;
9use arrow_array::{Array, ArrayRef, UnionArray};
10use arrow_schema::{ArrowError, DataType, FieldRef, UnionFields};
11use arrow_select::union_extract::union_extract_by_id;
12
13use super::CastOptions;
14
15// this is used during child array selection to prefer a "close" type over a distant cast
16// for example: when targeting Utf8View, a Utf8 child is preferred over Int32 despite both being castable
17fn same_type_family(a: &DataType, b: &DataType) -> bool {
18    use DataType::*;
19    matches!(
20        (a, b),
21        (Utf8 | LargeUtf8 | Utf8View, Utf8 | LargeUtf8 | Utf8View)
22            | (
23                Binary | LargeBinary | BinaryView,
24                Binary | LargeBinary | BinaryView
25            )
26            | (Int8 | Int16 | Int32 | Int64, Int8 | Int16 | Int32 | Int64)
27            | (
28                UInt8 | UInt16 | UInt32 | UInt64,
29                UInt8 | UInt16 | UInt32 | UInt64
30            )
31            | (Float16 | Float32 | Float64, Float16 | Float32 | Float64)
32    )
33}
34
35/// Selects the best-matching child array from a [`UnionArray`] for a given target type
36///
37/// The goal is to find the source field whose type is closest to the target,
38/// so that the subsequent cast is as lossless as possible. The heuristic uses
39/// three passes with decreasing specificity:
40///
41/// 1. **Exact match**: field type equals the target type.
42/// 2. **Same type family**: field and target belong to the same logical family
43///    (e.g. `Utf8` and `Utf8View` are both strings). This avoids a greedy
44///    cross-family cast in pass 3 (e.g. picking `Int32` over `Utf8` when the
45///    target is `Utf8View`, since `can_cast_types(Int32, Utf8View)` is true)
46/// 3. **Castable**:`can_cast_types` reports the field can be cast to the target
47///    Nested target types are skipped here because union extraction introduces
48///    nulls, which can conflict with non-nullable inner fields
49///
50/// Each pass greedily picks the first matching field by type_id order
51pub(crate) fn resolve_child_array<'a>(
52    fields: &'a UnionFields,
53    target_type: &DataType,
54) -> Option<(i8, &'a FieldRef)> {
55    fields
56        .iter()
57        .find(|(_, f)| f.data_type() == target_type)
58        .or_else(|| {
59            fields
60                .iter()
61                .find(|(_, f)| same_type_family(f.data_type(), target_type))
62        })
63        .or_else(|| {
64            // skip nested types in pass 3 because union extraction introduces nulls,
65            // and casting nullable arrays to nested types like List/Struct/Map can fail
66            // when inner fields are non-nullable.
67            if target_type.is_nested() {
68                return None;
69            }
70            fields
71                .iter()
72                .find(|(_, f)| can_cast_types(f.data_type(), target_type))
73        })
74}
75
76/// Extracts the best-matching child array from a [`UnionArray`] for a given target type,
77/// and casts it to that type.
78///
79/// Rows where a different child array is active become NULL.
80///
81/// # Errors
82///
83/// Returns [`ArrowError::InvalidArgumentError`] if no child type can be selected for
84/// the target type, or conversion from the selected child type is unsupported.
85///
86/// [`ArrowError::CastError`] indicates that the conversion is supported, but a
87/// value in the selected child array cannot be converted, for example because
88/// of an invalid string format or numeric overflow. When [`CastOptions::safe`]
89/// is `true`, failed values are replaced with nulls where supported.
90///
91/// # Example
92///
93/// ```
94/// # use std::sync::Arc;
95/// # use arrow_schema::{DataType, Field, UnionFields};
96/// # use arrow_array::{UnionArray, StringArray, Int32Array, Array};
97/// # use arrow_cast::cast::union_extract_by_type;
98/// # use arrow_cast::CastOptions;
99/// let fields = UnionFields::try_new(
100///     [0, 1],
101///     [
102///         Field::new("int", DataType::Int32, true),
103///         Field::new("str", DataType::Utf8, true),
104///     ],
105/// ).unwrap();
106///
107/// let union = UnionArray::try_new(
108///     fields,
109///     vec![0, 1, 0].into(),
110///     None,
111///     vec![
112///         Arc::new(Int32Array::from(vec![Some(42), None, Some(99)])),
113///         Arc::new(StringArray::from(vec![None, Some("hello"), None])),
114///     ],
115/// )
116/// .unwrap();
117///
118/// // extract the Utf8 child array and cast to Utf8View
119/// let result = union_extract_by_type(&union, &DataType::Utf8View, &CastOptions::default()).unwrap();
120/// assert_eq!(result.data_type(), &DataType::Utf8View);
121/// assert!(result.is_null(0));   // Int32 row -> NULL
122/// assert!(!result.is_null(1));  // Utf8 row -> "hello"
123/// assert!(result.is_null(2));   // Int32 row -> NULL
124/// ```
125pub fn union_extract_by_type(
126    union_array: &UnionArray,
127    target_type: &DataType,
128    cast_options: &CastOptions,
129) -> Result<ArrayRef, ArrowError> {
130    let DataType::Union(fields, _) = union_array.data_type() else {
131        unreachable!("union_extract_by_type called on non-union array")
132    };
133
134    let Some((type_id, _)) = resolve_child_array(fields, target_type) else {
135        return Err(ArrowError::InvalidArgumentError(format!(
136            "cannot cast Union with fields {} to {}",
137            fields
138                .iter()
139                .map(|(_, f)| f.data_type().to_string())
140                .collect::<Vec<_>>()
141                .join(", "),
142            target_type
143        )));
144    };
145
146    let extracted = union_extract_by_id(union_array, type_id)?;
147
148    if extracted.data_type() == target_type {
149        return Ok(extracted);
150    }
151
152    cast_with_options(&extracted, target_type, cast_options)
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158    use crate::cast;
159    use arrow_array::*;
160    use arrow_schema::{Field, UnionFields, UnionMode};
161    use std::sync::Arc;
162
163    fn int_str_fields() -> UnionFields {
164        UnionFields::try_new(
165            [0, 1],
166            [
167                Field::new("int", DataType::Int32, true),
168                Field::new("str", DataType::Utf8, true),
169            ],
170        )
171        .unwrap()
172    }
173
174    fn int_str_union_type(mode: UnionMode) -> DataType {
175        DataType::Union(int_str_fields(), mode)
176    }
177
178    // pass 1: exact type match.
179    // Union(Int32, Utf8) targeting Utf8. The Utf8 child matches exactly.
180    // Int32 rows become NULL. tested for both sparse and dense.
181    #[test]
182    fn test_exact_type_match() {
183        let target = DataType::Utf8;
184
185        // sparse
186        assert!(can_cast_types(
187            &int_str_union_type(UnionMode::Sparse),
188            &target
189        ));
190
191        let sparse = UnionArray::try_new(
192            int_str_fields(),
193            vec![1_i8, 0, 1].into(),
194            None,
195            vec![
196                Arc::new(Int32Array::from(vec![None, Some(42), None])) as ArrayRef,
197                Arc::new(StringArray::from(vec![Some("hello"), None, Some("world")])),
198            ],
199        )
200        .unwrap();
201
202        let result = cast::cast(&sparse, &target).unwrap();
203        assert_eq!(result.data_type(), &target);
204        let arr = result.as_any().downcast_ref::<StringArray>().unwrap();
205        assert_eq!(arr.value(0), "hello");
206        assert!(arr.is_null(1));
207        assert_eq!(arr.value(2), "world");
208
209        // dense
210        assert!(can_cast_types(
211            &int_str_union_type(UnionMode::Dense),
212            &target
213        ));
214
215        let dense = UnionArray::try_new(
216            int_str_fields(),
217            vec![1_i8, 0, 1].into(),
218            Some(vec![0_i32, 0, 1].into()),
219            vec![
220                Arc::new(Int32Array::from(vec![Some(42)])) as ArrayRef,
221                Arc::new(StringArray::from(vec![Some("hello"), Some("world")])),
222            ],
223        )
224        .unwrap();
225
226        let result = cast::cast(&dense, &target).unwrap();
227        let arr = result.as_any().downcast_ref::<StringArray>().unwrap();
228        assert_eq!(arr.value(0), "hello");
229        assert!(arr.is_null(1));
230        assert_eq!(arr.value(2), "world");
231    }
232
233    // pass 2: same type family match.
234    // Union(Int32, Utf8) targeting Utf8View. No exact match, but Utf8 and Utf8View
235    // are in the same family. picks the Utf8 child array and casts to Utf8View.
236    // this is the bug that motivated this work: without pass 2, pass 3 would
237    // greedily pick Int32 (since can_cast_types(Int32, Utf8View) is true).
238    #[test]
239    fn test_same_family_utf8_to_utf8view() {
240        let target = DataType::Utf8View;
241
242        // sparse
243        assert!(can_cast_types(
244            &int_str_union_type(UnionMode::Sparse),
245            &target
246        ));
247
248        let sparse = UnionArray::try_new(
249            int_str_fields(),
250            vec![1_i8, 0, 1, 1].into(),
251            None,
252            vec![
253                Arc::new(Int32Array::from(vec![None, Some(42), None, None])) as ArrayRef,
254                Arc::new(StringArray::from(vec![
255                    Some("agent_alpha"),
256                    None,
257                    Some("agent_beta"),
258                    None,
259                ])),
260            ],
261        )
262        .unwrap();
263
264        let result = cast::cast(&sparse, &target).unwrap();
265        assert_eq!(result.data_type(), &target);
266        let arr = result.as_any().downcast_ref::<StringViewArray>().unwrap();
267        assert_eq!(arr.value(0), "agent_alpha");
268        assert!(arr.is_null(1));
269        assert_eq!(arr.value(2), "agent_beta");
270        assert!(arr.is_null(3));
271
272        // dense
273        assert!(can_cast_types(
274            &int_str_union_type(UnionMode::Dense),
275            &target
276        ));
277
278        let dense = UnionArray::try_new(
279            int_str_fields(),
280            vec![1_i8, 0, 1].into(),
281            Some(vec![0_i32, 0, 1].into()),
282            vec![
283                Arc::new(Int32Array::from(vec![Some(42)])) as ArrayRef,
284                Arc::new(StringArray::from(vec![Some("alpha"), Some("beta")])),
285            ],
286        )
287        .unwrap();
288
289        let result = cast::cast(&dense, &target).unwrap();
290        let arr = result.as_any().downcast_ref::<StringViewArray>().unwrap();
291        assert_eq!(arr.value(0), "alpha");
292        assert!(arr.is_null(1));
293        assert_eq!(arr.value(2), "beta");
294    }
295
296    // pass 3: one-directional cast across type families.
297    // Union(Int32, Utf8) targeting Boolean — no exact match, no family match.
298    // pass 3 picks Int32 (first child array where can_cast_types is true) and
299    // casts to Boolean (0 → false, nonzero → true). Utf8 rows become NULL.
300    #[test]
301    fn test_one_directional_cast() {
302        let target = DataType::Boolean;
303
304        // sparse
305        assert!(can_cast_types(
306            &int_str_union_type(UnionMode::Sparse),
307            &target
308        ));
309
310        let sparse = UnionArray::try_new(
311            int_str_fields(),
312            vec![0_i8, 1, 0].into(),
313            None,
314            vec![
315                Arc::new(Int32Array::from(vec![Some(42), None, Some(0)])) as ArrayRef,
316                Arc::new(StringArray::from(vec![None, Some("hello"), None])),
317            ],
318        )
319        .unwrap();
320
321        let result = cast::cast(&sparse, &target).unwrap();
322        assert_eq!(result.data_type(), &target);
323        let arr = result.as_any().downcast_ref::<BooleanArray>().unwrap();
324        assert!(arr.value(0));
325        assert!(arr.is_null(1));
326        assert!(!arr.value(2));
327
328        // dense
329        assert!(can_cast_types(
330            &int_str_union_type(UnionMode::Dense),
331            &target
332        ));
333
334        let dense = UnionArray::try_new(
335            int_str_fields(),
336            vec![0_i8, 1, 0].into(),
337            Some(vec![0_i32, 0, 1].into()),
338            vec![
339                Arc::new(Int32Array::from(vec![Some(42), Some(0)])) as ArrayRef,
340                Arc::new(StringArray::from(vec![Some("hello")])),
341            ],
342        )
343        .unwrap();
344
345        let result = cast::cast(&dense, &target).unwrap();
346        let arr = result.as_any().downcast_ref::<BooleanArray>().unwrap();
347        assert!(arr.value(0));
348        assert!(arr.is_null(1));
349        assert!(!arr.value(2));
350    }
351
352    // duplicate field names: ensure we resolve by type_id, not field name.
353    // Union has two children both named "val" — Int32 (type_id 0) and Utf8 (type_id 1).
354    // Casting to Utf8 should select the Utf8 child (type_id 1), not the Int32 child (type_id 0).
355    #[test]
356    fn test_duplicate_field_names() {
357        let fields = UnionFields::try_new(
358            [0, 1],
359            [
360                Field::new("val", DataType::Int32, true),
361                Field::new("val", DataType::Utf8, true),
362            ],
363        )
364        .unwrap();
365
366        let target = DataType::Utf8;
367
368        let sparse = UnionArray::try_new(
369            fields.clone(),
370            vec![0_i8, 1, 0, 1].into(),
371            None,
372            vec![
373                Arc::new(Int32Array::from(vec![Some(42), None, Some(99), None])) as ArrayRef,
374                Arc::new(StringArray::from(vec![
375                    None,
376                    Some("hello"),
377                    None,
378                    Some("world"),
379                ])),
380            ],
381        )
382        .unwrap();
383
384        let result = cast::cast(&sparse, &target).unwrap();
385        let arr = result.as_any().downcast_ref::<StringArray>().unwrap();
386        assert!(arr.is_null(0));
387        assert_eq!(arr.value(1), "hello");
388        assert!(arr.is_null(2));
389        assert_eq!(arr.value(3), "world");
390
391        let dense = UnionArray::try_new(
392            fields,
393            vec![0_i8, 1, 1].into(),
394            Some(vec![0_i32, 0, 1].into()),
395            vec![
396                Arc::new(Int32Array::from(vec![Some(42)])) as ArrayRef,
397                Arc::new(StringArray::from(vec![Some("hello"), Some("world")])),
398            ],
399        )
400        .unwrap();
401
402        let result = cast::cast(&dense, &target).unwrap();
403        let arr = result.as_any().downcast_ref::<StringArray>().unwrap();
404        assert!(arr.is_null(0));
405        assert_eq!(arr.value(1), "hello");
406        assert_eq!(arr.value(2), "world");
407    }
408
409    // no matching child array, all three passes fail.
410    // Union(Int32, Utf8) targeting Struct({x: Int32}). neither Int32 nor Utf8
411    // can be cast to a Struct, so both can_cast_types and cast return errors.
412    #[test]
413    fn test_no_match_errors() {
414        let target = DataType::Struct(vec![Field::new("x", DataType::Int32, true)].into());
415
416        assert!(!can_cast_types(
417            &int_str_union_type(UnionMode::Sparse),
418            &target
419        ));
420
421        let union = UnionArray::try_new(
422            int_str_fields(),
423            vec![0_i8, 1].into(),
424            None,
425            vec![
426                Arc::new(Int32Array::from(vec![Some(42), None])) as ArrayRef,
427                Arc::new(StringArray::from(vec![None, Some("hello")])),
428            ],
429        )
430        .unwrap();
431
432        assert!(cast::cast(&union, &target).is_err());
433    }
434
435    // priority: exact match (pass 1) wins over family match (pass 2).
436    // Union(Utf8, Utf8View) targeting Utf8View. Both child arrays are in the string
437    // family, but Utf8View is an exact match. pass 1 should pick it, not Utf8.
438    #[test]
439    fn test_exact_match_preferred_over_family() {
440        let fields = UnionFields::try_new(
441            [0, 1],
442            [
443                Field::new("a", DataType::Utf8, true),
444                Field::new("b", DataType::Utf8View, true),
445            ],
446        )
447        .unwrap();
448        let target = DataType::Utf8View;
449
450        assert!(can_cast_types(
451            &DataType::Union(fields.clone(), UnionMode::Sparse),
452            &target,
453        ));
454
455        // [Utf8("from_a"), Utf8View("from_b"), Utf8("also_a")]
456        let union = UnionArray::try_new(
457            fields,
458            vec![0_i8, 1, 0].into(),
459            None,
460            vec![
461                Arc::new(StringArray::from(vec![
462                    Some("from_a"),
463                    None,
464                    Some("also_a"),
465                ])) as ArrayRef,
466                Arc::new(StringViewArray::from(vec![None, Some("from_b"), None])),
467            ],
468        )
469        .unwrap();
470
471        let result = cast::cast(&union, &target).unwrap();
472        assert_eq!(result.data_type(), &target);
473        let arr = result.as_any().downcast_ref::<StringViewArray>().unwrap();
474
475        // pass 1 picks child "b" (Utf8View), so child "a" rows become NULL
476        assert!(arr.is_null(0));
477        assert_eq!(arr.value(1), "from_b");
478        assert!(arr.is_null(2));
479    }
480
481    // null values within the selected child array stay null.
482    // this is distinct from "wrong child array -> NULL": here the correct child array
483    // is active but its value is null.
484    #[test]
485    fn test_null_in_selected_child_array() {
486        let target = DataType::Utf8;
487
488        assert!(can_cast_types(
489            &int_str_union_type(UnionMode::Sparse),
490            &target
491        ));
492
493        // ["hello", NULL(str), "world"]
494        // all rows are the Utf8 child array, but row 1 has a null value
495        let union = UnionArray::try_new(
496            int_str_fields(),
497            vec![1_i8, 1, 1].into(),
498            None,
499            vec![
500                Arc::new(Int32Array::from(vec![None, None, None])) as ArrayRef,
501                Arc::new(StringArray::from(vec![Some("hello"), None, Some("world")])),
502            ],
503        )
504        .unwrap();
505
506        let result = cast::cast(&union, &target).unwrap();
507        let arr = result.as_any().downcast_ref::<StringArray>().unwrap();
508        assert_eq!(arr.value(0), "hello");
509        assert!(arr.is_null(1));
510        assert_eq!(arr.value(2), "world");
511    }
512
513    // empty union array returns a zero-length result of the target type.
514    #[test]
515    fn test_empty_union() {
516        let target = DataType::Utf8View;
517
518        assert!(can_cast_types(
519            &int_str_union_type(UnionMode::Sparse),
520            &target
521        ));
522
523        let union = UnionArray::try_new(
524            int_str_fields(),
525            Vec::<i8>::new().into(),
526            None,
527            vec![
528                Arc::new(Int32Array::from(Vec::<Option<i32>>::new())) as ArrayRef,
529                Arc::new(StringArray::from(Vec::<Option<&str>>::new())),
530            ],
531        )
532        .unwrap();
533
534        let result = cast::cast(&union, &target).unwrap();
535        assert_eq!(result.data_type(), &target);
536        assert_eq!(result.len(), 0);
537    }
538}
arrow_cast/cast/union.rs

arrow_cast/cast/
union.rs
