arrow_cast/cast/
union.rs

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

arrow_cast/cast/
union.rs
