Skip to main content

arrow_select/
union_extract.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//! Defines union_extract kernel for [UnionArray]
19
20use crate::take::take;
21use arrow_array::{
22    Array, ArrayRef, BooleanArray, Int32Array, Scalar, UnionArray, make_array, new_empty_array,
23    new_null_array,
24};
25use arrow_buffer::{BooleanBuffer, MutableBuffer, NullBuffer, ScalarBuffer, bit_util};
26use arrow_data::layout;
27use arrow_schema::{ArrowError, DataType, UnionFields};
28use std::cmp::Ordering;
29use std::sync::Arc;
30
31/// Returns the value of the target field when selected, or NULL otherwise.
32/// ```text
33/// ┌─────────────────┐                                   ┌─────────────────┐
34/// │       A=1       │                                   │        1        │
35/// ├─────────────────┤                                   ├─────────────────┤
36/// │      A=NULL     │                                   │       NULL      │
37/// ├─────────────────┤    union_extract(values, 'A')     ├─────────────────┤
38/// │      B='t'      │  ────────────────────────────▶    │       NULL      │
39/// ├─────────────────┤                                   ├─────────────────┤
40/// │       A=3       │                                   │        3        │
41/// ├─────────────────┤                                   ├─────────────────┤
42/// │      B=NULL     │                                   │       NULL      │
43/// └─────────────────┘                                   └─────────────────┘
44///    union array                                              result
45/// ```
46/// # Errors
47///
48/// Returns error if target field is not found
49///
50/// # Examples
51/// ```
52/// # use std::sync::Arc;
53/// # use arrow_schema::{DataType, Field, UnionFields};
54/// # use arrow_array::{UnionArray, StringArray, Int32Array};
55/// # use arrow_select::union_extract::union_extract;
56/// let fields = UnionFields::try_new(
57///     [1, 3],
58///     [
59///         Field::new("A", DataType::Int32, true),
60///         Field::new("B", DataType::Utf8, true)
61///     ]
62/// ).unwrap();
63///
64/// let union = UnionArray::try_new(
65///     fields,
66///     vec![1, 1, 3, 1, 3].into(),
67///     None,
68///     vec![
69///         Arc::new(Int32Array::from(vec![Some(1), None, None, Some(3), Some(0)])),
70///         Arc::new(StringArray::from(vec![None, None, Some("t"), Some("."), None]))
71///     ]
72/// ).unwrap();
73///
74/// // Extract field A
75/// let extracted = union_extract(&union, "A").unwrap();
76///
77/// assert_eq!(*extracted, Int32Array::from(vec![Some(1), None, None, Some(3), None]));
78/// ```
79pub fn union_extract(union_array: &UnionArray, target: &str) -> Result<ArrayRef, ArrowError> {
80    let DataType::Union(fields, _) = union_array.data_type() else {
81        unreachable!()
82    };
83
84    let (target_type_id, _) = fields
85        .iter()
86        .find(|field| field.1.name() == target)
87        .ok_or_else(|| {
88            ArrowError::InvalidArgumentError(format!("field {target} not found on union"))
89        })?;
90
91    union_extract_impl(union_array, fields, target_type_id)
92}
93
94/// Like [`union_extract`], but selects the child by `type_id` rather than by
95/// field name.
96///
97/// This avoids ambiguity when the union contains duplicate field names.
98///
99/// # Errors
100///
101/// Returns error if `target_type_id` does not correspond to a field in the union.
102pub fn union_extract_by_id(
103    union_array: &UnionArray,
104    target_type_id: i8,
105) -> Result<ArrayRef, ArrowError> {
106    let DataType::Union(fields, _) = union_array.data_type() else {
107        unreachable!()
108    };
109
110    if fields.iter().all(|(id, _)| id != target_type_id) {
111        return Err(ArrowError::InvalidArgumentError(format!(
112            "type_id {target_type_id} not found on union"
113        )));
114    }
115
116    union_extract_impl(union_array, fields, target_type_id)
117}
118
119fn union_extract_impl(
120    union_array: &UnionArray,
121    fields: &UnionFields,
122    target_type_id: i8,
123) -> Result<ArrayRef, ArrowError> {
124    match union_array.offsets() {
125        Some(_) => extract_dense(union_array, fields, target_type_id),
126        None => extract_sparse(union_array, fields, target_type_id),
127    }
128}
129
130fn extract_sparse(
131    union_array: &UnionArray,
132    fields: &UnionFields,
133    target_type_id: i8,
134) -> Result<ArrayRef, ArrowError> {
135    let target = union_array.child(target_type_id);
136
137    if fields.len() == 1 // case 1.1: if there is a single field, all type ids are the same, and since union doesn't have a null mask, the result array is exactly the same as it only child
138        || union_array.is_empty() // case 1.2: sparse union length and childrens length must match, if the union is empty, so is any children
139        || target.null_count() == target.len() || target.data_type().is_null()
140    // case 1.3: if all values of the target children are null, regardless of selected type ids, the result will also be completely null
141    {
142        Ok(Arc::clone(target))
143    } else {
144        match eq_scalar(union_array.type_ids(), target_type_id) {
145            // case 2: all type ids equals our target, and since unions doesn't have a null mask, the result array is exactly the same as our target
146            BoolValue::Scalar(true) => Ok(Arc::clone(target)),
147            // case 3: none type_id matches our target, the result is a null array
148            BoolValue::Scalar(false) => {
149                if layout(target.data_type()).can_contain_null_mask {
150                    // case 3.1: target array can contain a null mask
151                    //SAFETY: The only change to the array data is the addition of a null mask, and if the target data type can contain a null mask was just checked above
152                    let data = unsafe {
153                        target
154                            .into_data()
155                            .into_builder()
156                            .nulls(Some(NullBuffer::new_null(target.len())))
157                            .build_unchecked()
158                    };
159
160                    Ok(make_array(data))
161                } else {
162                    // case 3.2: target can't contain a null mask
163                    Ok(new_null_array(target.data_type(), target.len()))
164                }
165            }
166            // case 4: some but not all type_id matches our target
167            BoolValue::Buffer(selected) => {
168                if layout(target.data_type()).can_contain_null_mask {
169                    // case 4.1: target array can contain a null mask
170                    let nulls = match target.nulls().filter(|n| n.null_count() > 0) {
171                        // case 4.1.1: our target child has nulls and types other than our target are selected, union the masks
172                        // the case where n.null_count() == n.len() is cheaply handled at case 1.3
173                        Some(nulls) => &selected & nulls.inner(),
174                        // case 4.1.2: target child has no nulls, but types other than our target are selected, use the selected mask as a null mask
175                        None => selected,
176                    };
177
178                    //SAFETY: The only change to the array data is the addition of a null mask, and if the target data type can contain a null mask was just checked above
179                    let data = unsafe {
180                        assert_eq!(nulls.len(), target.len());
181
182                        target
183                            .into_data()
184                            .into_builder()
185                            .nulls(Some(nulls.into()))
186                            .build_unchecked()
187                    };
188
189                    Ok(make_array(data))
190                } else {
191                    // case 4.2: target can't contain a null mask, zip the values that match with a null value
192                    Ok(crate::zip::zip(
193                        &BooleanArray::new(selected, None),
194                        target,
195                        &Scalar::new(new_null_array(target.data_type(), 1)),
196                    )?)
197                }
198            }
199        }
200    }
201}
202
203fn extract_dense(
204    union_array: &UnionArray,
205    fields: &UnionFields,
206    target_type_id: i8,
207) -> Result<ArrayRef, ArrowError> {
208    let target = union_array.child(target_type_id);
209    let offsets = union_array.offsets().unwrap();
210
211    if union_array.is_empty() {
212        // case 1: the union is empty
213        if target.is_empty() {
214            // case 1.1: the target is also empty, do a cheap Arc::clone instead of allocating a new empty array
215            Ok(Arc::clone(target))
216        } else {
217            // case 1.2: the target is not empty, allocate a new empty array
218            Ok(new_empty_array(target.data_type()))
219        }
220    } else if target.is_empty() {
221        // case 2: the union is not empty but the target is, which implies that none type_id points to it. The result is a null array
222        Ok(new_null_array(target.data_type(), union_array.len()))
223    } else if target.null_count() == target.len() || target.data_type().is_null() {
224        // case 3: since all values on our target are null, regardless of selected type ids and offsets, the result is a null array
225        match target.len().cmp(&union_array.len()) {
226            // case 3.1: since the target is smaller than the union, allocate a new correctly sized null array
227            Ordering::Less => Ok(new_null_array(target.data_type(), union_array.len())),
228            // case 3.2: target equals the union len, return it directly
229            Ordering::Equal => Ok(Arc::clone(target)),
230            // case 3.3: target len is bigger than the union len, slice it
231            Ordering::Greater => Ok(target.slice(0, union_array.len())),
232        }
233    } else if fields.len() == 1 // case A: since there's a single field, our target, every type id must matches our target
234        || fields
235            .iter()
236            .filter(|(field_type_id, _)| *field_type_id != target_type_id)
237            .all(|(sibling_type_id, _)| union_array.child(sibling_type_id).is_empty())
238    // case B: since siblings are empty, every type id must matches our target
239    {
240        // case 4: every type id matches our target
241        Ok(extract_dense_all_selected(union_array, target, offsets)?)
242    } else {
243        match eq_scalar(union_array.type_ids(), target_type_id) {
244            // case 4C: all type ids matches our target.
245            // Non empty sibling without any selected value may happen after slicing the parent union,
246            // since only type_ids and offsets are sliced, not the children
247            BoolValue::Scalar(true) => {
248                Ok(extract_dense_all_selected(union_array, target, offsets)?)
249            }
250            BoolValue::Scalar(false) => {
251                // case 5: none type_id matches our target, so the result array will be completely null
252                // Non empty target without any selected value may happen after slicing the parent union,
253                // since only type_ids and offsets are sliced, not the children
254                match (target.len().cmp(&union_array.len()), layout(target.data_type()).can_contain_null_mask) {
255                    (Ordering::Less, _) // case 5.1A: our target is smaller than the parent union, allocate a new correctly sized null array
256                    | (_, false) => { // case 5.1B: target array can't contain a null mask
257                        Ok(new_null_array(target.data_type(), union_array.len()))
258                    }
259                    // case 5.2: target and parent union lengths are equal, and the target can contain a null mask, let's set it to a all-null null-buffer
260                    (Ordering::Equal, true) => {
261                        //SAFETY: The only change to the array data is the addition of a null mask, and if the target data type can contain a null mask was just checked above
262                        let data = unsafe {
263                            target
264                                .into_data()
265                                .into_builder()
266                                .nulls(Some(NullBuffer::new_null(union_array.len())))
267                                .build_unchecked()
268                        };
269
270                        Ok(make_array(data))
271                    }
272                    // case 5.3: target is bigger than it's parent union and can contain a null mask, let's slice it, and set it's nulls to a all-null null-buffer
273                    (Ordering::Greater, true) => {
274                        //SAFETY: The only change to the array data is the addition of a null mask, and if the target data type can contain a null mask was just checked above
275                        let data = unsafe {
276                            target
277                                .into_data()
278                                .slice(0, union_array.len())
279                                .into_builder()
280                                .nulls(Some(NullBuffer::new_null(union_array.len())))
281                                .build_unchecked()
282                        };
283
284                        Ok(make_array(data))
285                    }
286                }
287            }
288            BoolValue::Buffer(selected) => {
289                //case 6: some type_ids matches our target, but not all. For selected values, take the value pointed by the offset. For unselected, use a valid null
290                Ok(take(
291                    target,
292                    &Int32Array::try_new(offsets.clone(), Some(selected.into()))?,
293                    None,
294                )?)
295            }
296        }
297    }
298}
299
300fn extract_dense_all_selected(
301    union_array: &UnionArray,
302    target: &Arc<dyn Array>,
303    offsets: &ScalarBuffer<i32>,
304) -> Result<ArrayRef, ArrowError> {
305    let sequential =
306        target.len() - offsets[0] as usize >= union_array.len() && is_sequential(offsets);
307
308    if sequential && target.len() == union_array.len() {
309        // case 1: all offsets are sequential and both lengths match, return the array directly
310        Ok(Arc::clone(target))
311    } else if sequential && target.len() > union_array.len() {
312        // case 2: All offsets are sequential, but our target is bigger than our union, slice it, starting at the first offset
313        Ok(target.slice(offsets[0] as usize, union_array.len()))
314    } else {
315        // case 3: Since offsets are not sequential, take them from the child to a new sequential and correctly sized array
316        let indices = Int32Array::try_new(offsets.clone(), None)?;
317
318        Ok(take(target, &indices, None)?)
319    }
320}
321
322const EQ_SCALAR_CHUNK_SIZE: usize = 512;
323
324/// The result of checking which type_ids matches the target type_id
325#[derive(Debug, PartialEq)]
326enum BoolValue {
327    /// If true, all type_ids matches the target type_id
328    /// If false, none type_ids matches the target type_id
329    Scalar(bool),
330    /// A mask representing which type_ids matches the target type_id
331    Buffer(BooleanBuffer),
332}
333
334fn eq_scalar(type_ids: &[i8], target: i8) -> BoolValue {
335    eq_scalar_inner(EQ_SCALAR_CHUNK_SIZE, type_ids, target)
336}
337
338fn count_first_run(chunk_size: usize, type_ids: &[i8], mut f: impl FnMut(i8) -> bool) -> usize {
339    type_ids
340        .chunks(chunk_size)
341        .take_while(|chunk| chunk.iter().copied().fold(true, |b, v| b & f(v)))
342        .map(|chunk| chunk.len())
343        .sum()
344}
345
346// This is like MutableBuffer::collect_bool(type_ids.len(), |i| type_ids[i] == target) with fast paths for all true or all false values.
347fn eq_scalar_inner(chunk_size: usize, type_ids: &[i8], target: i8) -> BoolValue {
348    let true_bits = count_first_run(chunk_size, type_ids, |v| v == target);
349
350    let (set_bits, val) = if true_bits == type_ids.len() {
351        return BoolValue::Scalar(true);
352    } else if true_bits == 0 {
353        let false_bits = count_first_run(chunk_size, type_ids, |v| v != target);
354
355        if false_bits == type_ids.len() {
356            return BoolValue::Scalar(false);
357        }
358        (false_bits, false)
359    } else {
360        (true_bits, true)
361    };
362
363    // restrict to chunk boundaries
364    let set_bits = set_bits - set_bits % 64;
365
366    let mut buffer =
367        MutableBuffer::new(bit_util::ceil(type_ids.len(), 8)).with_bitset(set_bits / 8, val);
368
369    buffer.extend(type_ids[set_bits..].chunks(64).map(|chunk| {
370        chunk
371            .iter()
372            .copied()
373            .enumerate()
374            .fold(0, |packed, (bit_idx, v)| {
375                packed | (((v == target) as u64) << bit_idx)
376            })
377    }));
378
379    BoolValue::Buffer(BooleanBuffer::new(buffer.into(), 0, type_ids.len()))
380}
381
382const IS_SEQUENTIAL_CHUNK_SIZE: usize = 64;
383
384fn is_sequential(offsets: &[i32]) -> bool {
385    is_sequential_generic::<IS_SEQUENTIAL_CHUNK_SIZE>(offsets)
386}
387
388fn is_sequential_generic<const N: usize>(offsets: &[i32]) -> bool {
389    if offsets.is_empty() {
390        return true;
391    }
392
393    // fast check this common combination:
394    // 1: sequential nulls are represented as a single null value on the values array, pointed by the same offset multiple times
395    // 2: valid values offsets increase one by one.
396    // example for an union with a single field A with type_id 0:
397    // union    = A=7 A=NULL A=NULL A=5 A=9
398    // a values = 7 NULL 5 9
399    // offsets  = 0 1 1 2 3
400    // type_ids = 0 0 0 0 0
401    // this also checks if the last chunk/remainder is sequential relative to the first offset
402    if offsets[0] + offsets.len() as i32 - 1 != offsets[offsets.len() - 1] {
403        return false;
404    }
405
406    let chunks = offsets.chunks_exact(N);
407
408    let remainder = chunks.remainder();
409
410    chunks.enumerate().all(|(i, chunk)| {
411        let chunk_array = <&[i32; N]>::try_from(chunk).unwrap();
412
413        //checks if values within chunk are sequential
414        chunk_array
415            .iter()
416            .copied()
417            .enumerate()
418            .fold(true, |acc, (i, offset)| {
419                acc & (offset == chunk_array[0] + i as i32)
420            })
421            && offsets[0] + (i * N) as i32 == chunk_array[0] //checks if chunk is sequential relative to the first offset
422    }) && remainder
423        .iter()
424        .copied()
425        .enumerate()
426        .fold(true, |acc, (i, offset)| {
427            acc & (offset == remainder[0] + i as i32)
428        }) //if the remainder is sequential relative to the first offset is checked at the start of the function
429}
430
431#[cfg(test)]
432mod tests {
433    use super::{
434        BoolValue, eq_scalar_inner, is_sequential_generic, union_extract, union_extract_by_id,
435    };
436    use arrow_array::{Array, Int32Array, NullArray, StringArray, UnionArray, new_null_array};
437    use arrow_buffer::{BooleanBuffer, ScalarBuffer};
438    use arrow_schema::{ArrowError, DataType, Field, UnionFields, UnionMode};
439    use std::sync::Arc;
440
441    #[test]
442    #[cfg_attr(miri, ignore)] // Takes too long
443    fn test_eq_scalar() {
444        //multiple all equal chunks, so it's loop and sum logic it's tested
445        //multiple chunks after, so it's loop logic it's tested
446        const ARRAY_LEN: usize = 64 * 4;
447
448        //so out of 64 boundaries chunks can be generated and checked for
449        const EQ_SCALAR_CHUNK_SIZE: usize = 3;
450
451        fn eq_scalar(type_ids: &[i8], target: i8) -> BoolValue {
452            eq_scalar_inner(EQ_SCALAR_CHUNK_SIZE, type_ids, target)
453        }
454
455        fn cross_check(left: &[i8], right: i8) -> BooleanBuffer {
456            BooleanBuffer::collect_bool(left.len(), |i| left[i] == right)
457        }
458
459        assert_eq!(eq_scalar(&[], 1), BoolValue::Scalar(true));
460
461        assert_eq!(eq_scalar(&[1], 1), BoolValue::Scalar(true));
462        assert_eq!(eq_scalar(&[2], 1), BoolValue::Scalar(false));
463
464        let mut values = [1; ARRAY_LEN];
465
466        assert_eq!(eq_scalar(&values, 1), BoolValue::Scalar(true));
467        assert_eq!(eq_scalar(&values, 2), BoolValue::Scalar(false));
468
469        //every subslice should return the same value
470        for i in 1..ARRAY_LEN {
471            assert_eq!(eq_scalar(&values[..i], 1), BoolValue::Scalar(true));
472            assert_eq!(eq_scalar(&values[..i], 2), BoolValue::Scalar(false));
473        }
474
475        // test that a single change anywhere is checked for
476        for i in 0..ARRAY_LEN {
477            values[i] = 2;
478
479            assert_eq!(
480                eq_scalar(&values, 1),
481                BoolValue::Buffer(cross_check(&values, 1))
482            );
483            assert_eq!(
484                eq_scalar(&values, 2),
485                BoolValue::Buffer(cross_check(&values, 2))
486            );
487
488            values[i] = 1;
489        }
490    }
491
492    #[test]
493    fn test_is_sequential() {
494        /*
495        the smallest value that satisfies:
496        >1 so the fold logic of a exact chunk executes
497        >2 so a >1 non-exact remainder can exist, and it's fold logic executes
498         */
499        const CHUNK_SIZE: usize = 3;
500        //we test arrays of size up to 8 = 2 * CHUNK_SIZE + 2:
501        //multiple(2) exact chunks, so the AND logic between them executes
502        //a >1(2) remainder, so:
503        //    the AND logic between all exact chunks and the remainder executes
504        //    the remainder fold logic executes
505
506        fn is_sequential(v: &[i32]) -> bool {
507            is_sequential_generic::<CHUNK_SIZE>(v)
508        }
509
510        assert!(is_sequential(&[])); //empty
511        assert!(is_sequential(&[1])); //single
512
513        assert!(is_sequential(&[1, 2]));
514        assert!(is_sequential(&[1, 2, 3]));
515        assert!(is_sequential(&[1, 2, 3, 4]));
516        assert!(is_sequential(&[1, 2, 3, 4, 5]));
517        assert!(is_sequential(&[1, 2, 3, 4, 5, 6]));
518        assert!(is_sequential(&[1, 2, 3, 4, 5, 6, 7]));
519        assert!(is_sequential(&[1, 2, 3, 4, 5, 6, 7, 8]));
520
521        assert!(!is_sequential(&[8, 7]));
522        assert!(!is_sequential(&[8, 7, 6]));
523        assert!(!is_sequential(&[8, 7, 6, 5]));
524        assert!(!is_sequential(&[8, 7, 6, 5, 4]));
525        assert!(!is_sequential(&[8, 7, 6, 5, 4, 3]));
526        assert!(!is_sequential(&[8, 7, 6, 5, 4, 3, 2]));
527        assert!(!is_sequential(&[8, 7, 6, 5, 4, 3, 2, 1]));
528
529        assert!(!is_sequential(&[0, 2]));
530        assert!(!is_sequential(&[1, 0]));
531
532        assert!(!is_sequential(&[0, 2, 3]));
533        assert!(!is_sequential(&[1, 0, 3]));
534        assert!(!is_sequential(&[1, 2, 0]));
535
536        assert!(!is_sequential(&[0, 2, 3, 4]));
537        assert!(!is_sequential(&[1, 0, 3, 4]));
538        assert!(!is_sequential(&[1, 2, 0, 4]));
539        assert!(!is_sequential(&[1, 2, 3, 0]));
540
541        assert!(!is_sequential(&[0, 2, 3, 4, 5]));
542        assert!(!is_sequential(&[1, 0, 3, 4, 5]));
543        assert!(!is_sequential(&[1, 2, 0, 4, 5]));
544        assert!(!is_sequential(&[1, 2, 3, 0, 5]));
545        assert!(!is_sequential(&[1, 2, 3, 4, 0]));
546
547        assert!(!is_sequential(&[0, 2, 3, 4, 5, 6]));
548        assert!(!is_sequential(&[1, 0, 3, 4, 5, 6]));
549        assert!(!is_sequential(&[1, 2, 0, 4, 5, 6]));
550        assert!(!is_sequential(&[1, 2, 3, 0, 5, 6]));
551        assert!(!is_sequential(&[1, 2, 3, 4, 0, 6]));
552        assert!(!is_sequential(&[1, 2, 3, 4, 5, 0]));
553
554        assert!(!is_sequential(&[0, 2, 3, 4, 5, 6, 7]));
555        assert!(!is_sequential(&[1, 0, 3, 4, 5, 6, 7]));
556        assert!(!is_sequential(&[1, 2, 0, 4, 5, 6, 7]));
557        assert!(!is_sequential(&[1, 2, 3, 0, 5, 6, 7]));
558        assert!(!is_sequential(&[1, 2, 3, 4, 0, 6, 7]));
559        assert!(!is_sequential(&[1, 2, 3, 4, 5, 0, 7]));
560        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 0]));
561
562        assert!(!is_sequential(&[0, 2, 3, 4, 5, 6, 7, 8]));
563        assert!(!is_sequential(&[1, 0, 3, 4, 5, 6, 7, 8]));
564        assert!(!is_sequential(&[1, 2, 0, 4, 5, 6, 7, 8]));
565        assert!(!is_sequential(&[1, 2, 3, 0, 5, 6, 7, 8]));
566        assert!(!is_sequential(&[1, 2, 3, 4, 0, 6, 7, 8]));
567        assert!(!is_sequential(&[1, 2, 3, 4, 5, 0, 7, 8]));
568        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 0, 8]));
569        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 7, 0]));
570
571        // checks increments at the chunk boundary
572        assert!(!is_sequential(&[1, 2, 3, 5]));
573        assert!(!is_sequential(&[1, 2, 3, 5, 6]));
574        assert!(!is_sequential(&[1, 2, 3, 5, 6, 7]));
575        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 8]));
576        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 8, 9]));
577    }
578
579    fn str1() -> UnionFields {
580        UnionFields::try_new(vec![1], vec![Field::new("str", DataType::Utf8, true)]).unwrap()
581    }
582
583    fn str1_int3() -> UnionFields {
584        UnionFields::try_new(
585            vec![1, 3],
586            vec![
587                Field::new("str", DataType::Utf8, true),
588                Field::new("int", DataType::Int32, true),
589            ],
590        )
591        .unwrap()
592    }
593
594    #[test]
595    fn sparse_1_1_single_field() {
596        let union = UnionArray::try_new(
597            //single field
598            str1(),
599            ScalarBuffer::from(vec![1, 1]), // non empty, every type id must match
600            None,                           //sparse
601            vec![
602                Arc::new(StringArray::from(vec!["a", "b"])), // not null
603            ],
604        )
605        .unwrap();
606
607        let expected = StringArray::from(vec!["a", "b"]);
608        let extracted = union_extract(&union, "str").unwrap();
609
610        assert_eq!(extracted.into_data(), expected.into_data());
611    }
612
613    #[test]
614    fn sparse_1_2_empty() {
615        let union = UnionArray::try_new(
616            // multiple fields
617            str1_int3(),
618            ScalarBuffer::from(vec![]), //empty union
619            None,                       // sparse
620            vec![
621                Arc::new(StringArray::new_null(0)),
622                Arc::new(Int32Array::new_null(0)),
623            ],
624        )
625        .unwrap();
626
627        let expected = StringArray::new_null(0);
628        let extracted = union_extract(&union, "str").unwrap(); //target type is not Null
629
630        assert_eq!(extracted.into_data(), expected.into_data());
631    }
632
633    #[test]
634    fn sparse_1_3a_null_target() {
635        let union = UnionArray::try_new(
636            // multiple fields
637            UnionFields::try_new(
638                vec![1, 3],
639                vec![
640                    Field::new("str", DataType::Utf8, true),
641                    Field::new("null", DataType::Null, true), // target type is Null
642                ],
643            )
644            .unwrap(),
645            ScalarBuffer::from(vec![1]), //not empty
646            None,                        // sparse
647            vec![
648                Arc::new(StringArray::new_null(1)),
649                Arc::new(NullArray::new(1)), // null data type
650            ],
651        )
652        .unwrap();
653
654        let expected = NullArray::new(1);
655        let extracted = union_extract(&union, "null").unwrap();
656
657        assert_eq!(extracted.into_data(), expected.into_data());
658    }
659
660    #[test]
661    fn sparse_1_3b_null_target() {
662        let union = UnionArray::try_new(
663            // multiple fields
664            str1_int3(),
665            ScalarBuffer::from(vec![1]), //not empty
666            None,                        // sparse
667            vec![
668                Arc::new(StringArray::new_null(1)), //all null
669                Arc::new(Int32Array::new_null(1)),
670            ],
671        )
672        .unwrap();
673
674        let expected = StringArray::new_null(1);
675        let extracted = union_extract(&union, "str").unwrap(); //target type is not Null
676
677        assert_eq!(extracted.into_data(), expected.into_data());
678    }
679
680    #[test]
681    fn sparse_2_all_types_match() {
682        let union = UnionArray::try_new(
683            //multiple fields
684            str1_int3(),
685            ScalarBuffer::from(vec![3, 3]), // all types match
686            None,                           //sparse
687            vec![
688                Arc::new(StringArray::new_null(2)),
689                Arc::new(Int32Array::from(vec![1, 4])), // not null
690            ],
691        )
692        .unwrap();
693
694        let expected = Int32Array::from(vec![1, 4]);
695        let extracted = union_extract(&union, "int").unwrap();
696
697        assert_eq!(extracted.into_data(), expected.into_data());
698    }
699
700    #[test]
701    fn sparse_3_1_none_match_target_can_contain_null_mask() {
702        let union = UnionArray::try_new(
703            //multiple fields
704            str1_int3(),
705            ScalarBuffer::from(vec![1, 1, 1, 1]), // none match
706            None,                                 // sparse
707            vec![
708                Arc::new(StringArray::new_null(4)),
709                Arc::new(Int32Array::from(vec![None, Some(4), None, Some(8)])), // target is not null
710            ],
711        )
712        .unwrap();
713
714        let expected = Int32Array::new_null(4);
715        let extracted = union_extract(&union, "int").unwrap();
716
717        assert_eq!(extracted.into_data(), expected.into_data());
718    }
719
720    fn str1_union3(union3_datatype: DataType) -> UnionFields {
721        UnionFields::try_new(
722            vec![1, 3],
723            vec![
724                Field::new("str", DataType::Utf8, true),
725                Field::new("union", union3_datatype, true),
726            ],
727        )
728        .unwrap()
729    }
730
731    #[test]
732    fn sparse_3_2_none_match_cant_contain_null_mask_union_target() {
733        let target_fields = str1();
734        let target_type = DataType::Union(target_fields.clone(), UnionMode::Sparse);
735
736        let union = UnionArray::try_new(
737            //multiple fields
738            str1_union3(target_type.clone()),
739            ScalarBuffer::from(vec![1, 1]), // none match
740            None,                           //sparse
741            vec![
742                Arc::new(StringArray::new_null(2)),
743                //target is not null
744                Arc::new(
745                    UnionArray::try_new(
746                        target_fields.clone(),
747                        ScalarBuffer::from(vec![1, 1]),
748                        None,
749                        vec![Arc::new(StringArray::from(vec!["a", "b"]))],
750                    )
751                    .unwrap(),
752                ),
753            ],
754        )
755        .unwrap();
756
757        let expected = new_null_array(&target_type, 2);
758        let extracted = union_extract(&union, "union").unwrap();
759
760        assert_eq!(extracted.into_data(), expected.into_data());
761    }
762
763    #[test]
764    fn sparse_4_1_1_target_with_nulls() {
765        let union = UnionArray::try_new(
766            //multiple fields
767            str1_int3(),
768            ScalarBuffer::from(vec![3, 3, 1, 1]), // multiple selected types
769            None,                                 // sparse
770            vec![
771                Arc::new(StringArray::new_null(4)),
772                Arc::new(Int32Array::from(vec![None, Some(4), None, Some(8)])), // target with nulls
773            ],
774        )
775        .unwrap();
776
777        let expected = Int32Array::from(vec![None, Some(4), None, None]);
778        let extracted = union_extract(&union, "int").unwrap();
779
780        assert_eq!(extracted.into_data(), expected.into_data());
781    }
782
783    #[test]
784    fn sparse_4_1_2_target_without_nulls() {
785        let union = UnionArray::try_new(
786            //multiple fields
787            str1_int3(),
788            ScalarBuffer::from(vec![1, 3, 3]), // multiple selected types
789            None,                              // sparse
790            vec![
791                Arc::new(StringArray::new_null(3)),
792                Arc::new(Int32Array::from(vec![2, 4, 8])), // target without nulls
793            ],
794        )
795        .unwrap();
796
797        let expected = Int32Array::from(vec![None, Some(4), Some(8)]);
798        let extracted = union_extract(&union, "int").unwrap();
799
800        assert_eq!(extracted.into_data(), expected.into_data());
801    }
802
803    #[test]
804    fn sparse_4_2_some_match_target_cant_contain_null_mask() {
805        let target_fields = str1();
806        let target_type = DataType::Union(target_fields.clone(), UnionMode::Sparse);
807
808        let union = UnionArray::try_new(
809            //multiple fields
810            str1_union3(target_type),
811            ScalarBuffer::from(vec![3, 1]), // some types match, but not all
812            None,                           //sparse
813            vec![
814                Arc::new(StringArray::new_null(2)),
815                Arc::new(
816                    UnionArray::try_new(
817                        target_fields.clone(),
818                        ScalarBuffer::from(vec![1, 1]),
819                        None,
820                        vec![Arc::new(StringArray::from(vec!["a", "b"]))],
821                    )
822                    .unwrap(),
823                ),
824            ],
825        )
826        .unwrap();
827
828        let expected = UnionArray::try_new(
829            target_fields,
830            ScalarBuffer::from(vec![1, 1]),
831            None,
832            vec![Arc::new(StringArray::from(vec![Some("a"), None]))],
833        )
834        .unwrap();
835        let extracted = union_extract(&union, "union").unwrap();
836
837        assert_eq!(extracted.into_data(), expected.into_data());
838    }
839
840    #[test]
841    fn dense_1_1_both_empty() {
842        let union = UnionArray::try_new(
843            str1_int3(),
844            ScalarBuffer::from(vec![]),       //empty union
845            Some(ScalarBuffer::from(vec![])), // dense
846            vec![
847                Arc::new(StringArray::new_null(0)), //empty target
848                Arc::new(Int32Array::new_null(0)),
849            ],
850        )
851        .unwrap();
852
853        let expected = StringArray::new_null(0);
854        let extracted = union_extract(&union, "str").unwrap();
855
856        assert_eq!(extracted.into_data(), expected.into_data());
857    }
858
859    #[test]
860    fn dense_1_2_empty_union_target_non_empty() {
861        let union = UnionArray::try_new(
862            str1_int3(),
863            ScalarBuffer::from(vec![]),       //empty union
864            Some(ScalarBuffer::from(vec![])), // dense
865            vec![
866                Arc::new(StringArray::new_null(1)), //non empty target
867                Arc::new(Int32Array::new_null(0)),
868            ],
869        )
870        .unwrap();
871
872        let expected = StringArray::new_null(0);
873        let extracted = union_extract(&union, "str").unwrap();
874
875        assert_eq!(extracted.into_data(), expected.into_data());
876    }
877
878    #[test]
879    fn dense_2_non_empty_union_target_empty() {
880        let union = UnionArray::try_new(
881            str1_int3(),
882            ScalarBuffer::from(vec![3, 3]),       //non empty union
883            Some(ScalarBuffer::from(vec![0, 1])), // dense
884            vec![
885                Arc::new(StringArray::new_null(0)), //empty target
886                Arc::new(Int32Array::new_null(2)),
887            ],
888        )
889        .unwrap();
890
891        let expected = StringArray::new_null(2);
892        let extracted = union_extract(&union, "str").unwrap();
893
894        assert_eq!(extracted.into_data(), expected.into_data());
895    }
896
897    #[test]
898    fn dense_3_1_null_target_smaller_len() {
899        let union = UnionArray::try_new(
900            str1_int3(),
901            ScalarBuffer::from(vec![3, 3]),       //non empty union
902            Some(ScalarBuffer::from(vec![0, 0])), //dense
903            vec![
904                Arc::new(StringArray::new_null(1)), //smaller target
905                Arc::new(Int32Array::new_null(2)),
906            ],
907        )
908        .unwrap();
909
910        let expected = StringArray::new_null(2);
911        let extracted = union_extract(&union, "str").unwrap();
912
913        assert_eq!(extracted.into_data(), expected.into_data());
914    }
915
916    #[test]
917    fn dense_3_2_null_target_equal_len() {
918        let union = UnionArray::try_new(
919            str1_int3(),
920            ScalarBuffer::from(vec![3, 3]),       //non empty union
921            Some(ScalarBuffer::from(vec![0, 0])), //dense
922            vec![
923                Arc::new(StringArray::new_null(2)), //equal len
924                Arc::new(Int32Array::new_null(2)),
925            ],
926        )
927        .unwrap();
928
929        let expected = StringArray::new_null(2);
930        let extracted = union_extract(&union, "str").unwrap();
931
932        assert_eq!(extracted.into_data(), expected.into_data());
933    }
934
935    #[test]
936    fn dense_3_3_null_target_bigger_len() {
937        let union = UnionArray::try_new(
938            str1_int3(),
939            ScalarBuffer::from(vec![3, 3]),       //non empty union
940            Some(ScalarBuffer::from(vec![0, 0])), //dense
941            vec![
942                Arc::new(StringArray::new_null(3)), //bigger len
943                Arc::new(Int32Array::new_null(3)),
944            ],
945        )
946        .unwrap();
947
948        let expected = StringArray::new_null(2);
949        let extracted = union_extract(&union, "str").unwrap();
950
951        assert_eq!(extracted.into_data(), expected.into_data());
952    }
953
954    #[test]
955    fn dense_4_1a_single_type_sequential_offsets_equal_len() {
956        let union = UnionArray::try_new(
957            // single field
958            str1(),
959            ScalarBuffer::from(vec![1, 1]),       //non empty union
960            Some(ScalarBuffer::from(vec![0, 1])), //sequential
961            vec![
962                Arc::new(StringArray::from(vec!["a1", "b2"])), //equal len, non null
963            ],
964        )
965        .unwrap();
966
967        let expected = StringArray::from(vec!["a1", "b2"]);
968        let extracted = union_extract(&union, "str").unwrap();
969
970        assert_eq!(extracted.into_data(), expected.into_data());
971    }
972
973    #[test]
974    fn dense_4_2a_single_type_sequential_offsets_bigger() {
975        let union = UnionArray::try_new(
976            // single field
977            str1(),
978            ScalarBuffer::from(vec![1, 1]),       //non empty union
979            Some(ScalarBuffer::from(vec![0, 1])), //sequential
980            vec![
981                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), //equal len, non null
982            ],
983        )
984        .unwrap();
985
986        let expected = StringArray::from(vec!["a1", "b2"]);
987        let extracted = union_extract(&union, "str").unwrap();
988
989        assert_eq!(extracted.into_data(), expected.into_data());
990    }
991
992    #[test]
993    fn dense_4_3a_single_type_non_sequential() {
994        let union = UnionArray::try_new(
995            // single field
996            str1(),
997            ScalarBuffer::from(vec![1, 1]),       //non empty union
998            Some(ScalarBuffer::from(vec![0, 2])), //non sequential
999            vec![
1000                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), //equal len, non null
1001            ],
1002        )
1003        .unwrap();
1004
1005        let expected = StringArray::from(vec!["a1", "c3"]);
1006        let extracted = union_extract(&union, "str").unwrap();
1007
1008        assert_eq!(extracted.into_data(), expected.into_data());
1009    }
1010
1011    #[test]
1012    fn dense_4_1b_empty_siblings_sequential_equal_len() {
1013        let union = UnionArray::try_new(
1014            // multiple fields
1015            str1_int3(),
1016            ScalarBuffer::from(vec![1, 1]),       //non empty union
1017            Some(ScalarBuffer::from(vec![0, 1])), //sequential
1018            vec![
1019                Arc::new(StringArray::from(vec!["a", "b"])), //equal len, non null
1020                Arc::new(Int32Array::new_null(0)),           //empty sibling
1021            ],
1022        )
1023        .unwrap();
1024
1025        let expected = StringArray::from(vec!["a", "b"]);
1026        let extracted = union_extract(&union, "str").unwrap();
1027
1028        assert_eq!(extracted.into_data(), expected.into_data());
1029    }
1030
1031    #[test]
1032    fn dense_4_2b_empty_siblings_sequential_bigger_len() {
1033        let union = UnionArray::try_new(
1034            // multiple fields
1035            str1_int3(),
1036            ScalarBuffer::from(vec![1, 1]),       //non empty union
1037            Some(ScalarBuffer::from(vec![0, 1])), //sequential
1038            vec![
1039                Arc::new(StringArray::from(vec!["a", "b", "c"])), //bigger len, non null
1040                Arc::new(Int32Array::new_null(0)),                //empty sibling
1041            ],
1042        )
1043        .unwrap();
1044
1045        let expected = StringArray::from(vec!["a", "b"]);
1046        let extracted = union_extract(&union, "str").unwrap();
1047
1048        assert_eq!(extracted.into_data(), expected.into_data());
1049    }
1050
1051    #[test]
1052    fn dense_4_3b_empty_sibling_non_sequential() {
1053        let union = UnionArray::try_new(
1054            // multiple fields
1055            str1_int3(),
1056            ScalarBuffer::from(vec![1, 1]),       //non empty union
1057            Some(ScalarBuffer::from(vec![0, 2])), //non sequential
1058            vec![
1059                Arc::new(StringArray::from(vec!["a", "b", "c"])), //non null
1060                Arc::new(Int32Array::new_null(0)),                //empty sibling
1061            ],
1062        )
1063        .unwrap();
1064
1065        let expected = StringArray::from(vec!["a", "c"]);
1066        let extracted = union_extract(&union, "str").unwrap();
1067
1068        assert_eq!(extracted.into_data(), expected.into_data());
1069    }
1070
1071    #[test]
1072    fn dense_4_1c_all_types_match_sequential_equal_len() {
1073        let union = UnionArray::try_new(
1074            // multiple fields
1075            str1_int3(),
1076            ScalarBuffer::from(vec![1, 1]),       //all types match
1077            Some(ScalarBuffer::from(vec![0, 1])), //sequential
1078            vec![
1079                Arc::new(StringArray::from(vec!["a1", "b2"])), //equal len
1080                Arc::new(Int32Array::new_null(2)),             //non empty sibling
1081            ],
1082        )
1083        .unwrap();
1084
1085        let expected = StringArray::from(vec!["a1", "b2"]);
1086        let extracted = union_extract(&union, "str").unwrap();
1087
1088        assert_eq!(extracted.into_data(), expected.into_data());
1089    }
1090
1091    #[test]
1092    fn dense_4_2c_all_types_match_sequential_bigger_len() {
1093        let union = UnionArray::try_new(
1094            // multiple fields
1095            str1_int3(),
1096            ScalarBuffer::from(vec![1, 1]),       //all types match
1097            Some(ScalarBuffer::from(vec![0, 1])), //sequential
1098            vec![
1099                Arc::new(StringArray::from(vec!["a1", "b2", "b3"])), //bigger len
1100                Arc::new(Int32Array::new_null(2)),                   //non empty sibling
1101            ],
1102        )
1103        .unwrap();
1104
1105        let expected = StringArray::from(vec!["a1", "b2"]);
1106        let extracted = union_extract(&union, "str").unwrap();
1107
1108        assert_eq!(extracted.into_data(), expected.into_data());
1109    }
1110
1111    #[test]
1112    fn dense_4_3c_all_types_match_non_sequential() {
1113        let union = UnionArray::try_new(
1114            // multiple fields
1115            str1_int3(),
1116            ScalarBuffer::from(vec![1, 1]),       //all types match
1117            Some(ScalarBuffer::from(vec![0, 2])), //non sequential
1118            vec![
1119                Arc::new(StringArray::from(vec!["a1", "b2", "b3"])),
1120                Arc::new(Int32Array::new_null(2)), //non empty sibling
1121            ],
1122        )
1123        .unwrap();
1124
1125        let expected = StringArray::from(vec!["a1", "b3"]);
1126        let extracted = union_extract(&union, "str").unwrap();
1127
1128        assert_eq!(extracted.into_data(), expected.into_data());
1129    }
1130
1131    #[test]
1132    fn dense_5_1a_none_match_less_len() {
1133        let union = UnionArray::try_new(
1134            // multiple fields
1135            str1_int3(),
1136            ScalarBuffer::from(vec![3, 3, 3, 3, 3]), //none matches
1137            Some(ScalarBuffer::from(vec![0, 0, 0, 1, 1])), // dense
1138            vec![
1139                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), // less len
1140                Arc::new(Int32Array::from(vec![1, 2])),
1141            ],
1142        )
1143        .unwrap();
1144
1145        let expected = StringArray::new_null(5);
1146        let extracted = union_extract(&union, "str").unwrap();
1147
1148        assert_eq!(extracted.into_data(), expected.into_data());
1149    }
1150
1151    #[test]
1152    fn dense_5_1b_cant_contain_null_mask() {
1153        let target_fields = str1();
1154        let target_type = DataType::Union(target_fields.clone(), UnionMode::Sparse);
1155
1156        let union = UnionArray::try_new(
1157            // multiple fields
1158            str1_union3(target_type.clone()),
1159            ScalarBuffer::from(vec![1, 1, 1, 1, 1]), //none matches
1160            Some(ScalarBuffer::from(vec![0, 0, 0, 1, 1])), // dense
1161            vec![
1162                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), // less len
1163                Arc::new(
1164                    UnionArray::try_new(
1165                        target_fields.clone(),
1166                        ScalarBuffer::from(vec![1]),
1167                        None,
1168                        vec![Arc::new(StringArray::from(vec!["a"]))],
1169                    )
1170                    .unwrap(),
1171                ), // non empty
1172            ],
1173        )
1174        .unwrap();
1175
1176        let expected = new_null_array(&target_type, 5);
1177        let extracted = union_extract(&union, "union").unwrap();
1178
1179        assert_eq!(extracted.into_data(), expected.into_data());
1180    }
1181
1182    #[test]
1183    fn dense_5_2_none_match_equal_len() {
1184        let union = UnionArray::try_new(
1185            // multiple fields
1186            str1_int3(),
1187            ScalarBuffer::from(vec![3, 3, 3, 3, 3]), //none matches
1188            Some(ScalarBuffer::from(vec![0, 0, 0, 1, 1])), // dense
1189            vec![
1190                Arc::new(StringArray::from(vec!["a1", "b2", "c3", "d4", "e5"])), // equal len
1191                Arc::new(Int32Array::from(vec![1, 2])),
1192            ],
1193        )
1194        .unwrap();
1195
1196        let expected = StringArray::new_null(5);
1197        let extracted = union_extract(&union, "str").unwrap();
1198
1199        assert_eq!(extracted.into_data(), expected.into_data());
1200    }
1201
1202    #[test]
1203    fn dense_5_3_none_match_greater_len() {
1204        let union = UnionArray::try_new(
1205            // multiple fields
1206            str1_int3(),
1207            ScalarBuffer::from(vec![3, 3, 3, 3, 3]), //none matches
1208            Some(ScalarBuffer::from(vec![0, 0, 0, 1, 1])), // dense
1209            vec![
1210                Arc::new(StringArray::from(vec!["a1", "b2", "c3", "d4", "e5", "f6"])), // greater len
1211                Arc::new(Int32Array::from(vec![1, 2])),                                //non null
1212            ],
1213        )
1214        .unwrap();
1215
1216        let expected = StringArray::new_null(5);
1217        let extracted = union_extract(&union, "str").unwrap();
1218
1219        assert_eq!(extracted.into_data(), expected.into_data());
1220    }
1221
1222    #[test]
1223    fn dense_6_some_matches() {
1224        let union = UnionArray::try_new(
1225            // multiple fields
1226            str1_int3(),
1227            ScalarBuffer::from(vec![3, 3, 1, 1, 1]), //some matches
1228            Some(ScalarBuffer::from(vec![0, 1, 0, 1, 2])), // dense
1229            vec![
1230                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), // non null
1231                Arc::new(Int32Array::from(vec![1, 2])),
1232            ],
1233        )
1234        .unwrap();
1235
1236        let expected = Int32Array::from(vec![Some(1), Some(2), None, None, None]);
1237        let extracted = union_extract(&union, "int").unwrap();
1238
1239        assert_eq!(extracted.into_data(), expected.into_data());
1240    }
1241
1242    #[test]
1243    fn empty_sparse_union() {
1244        let union = UnionArray::try_new(
1245            UnionFields::empty(),
1246            ScalarBuffer::from(vec![]),
1247            None,
1248            vec![],
1249        )
1250        .unwrap();
1251
1252        assert_eq!(
1253            union_extract(&union, "a").unwrap_err().to_string(),
1254            ArrowError::InvalidArgumentError("field a not found on union".into()).to_string()
1255        );
1256    }
1257
1258    #[test]
1259    fn empty_dense_union() {
1260        let union = UnionArray::try_new(
1261            UnionFields::empty(),
1262            ScalarBuffer::from(vec![]),
1263            Some(ScalarBuffer::from(vec![])),
1264            vec![],
1265        )
1266        .unwrap();
1267
1268        assert_eq!(
1269            union_extract(&union, "a").unwrap_err().to_string(),
1270            ArrowError::InvalidArgumentError("field a not found on union".into()).to_string()
1271        );
1272    }
1273
1274    #[test]
1275    fn extract_by_id_sparse_duplicate_names() {
1276        // Two fields with the same name "val" but different type_ids and types
1277        let fields = UnionFields::try_new(
1278            [0, 1],
1279            [
1280                Field::new("val", DataType::Int32, true),
1281                Field::new("val", DataType::Utf8, true),
1282            ],
1283        )
1284        .unwrap();
1285
1286        let union = UnionArray::try_new(
1287            fields,
1288            vec![0_i8, 1, 0, 1].into(),
1289            None,
1290            vec![
1291                Arc::new(Int32Array::from(vec![Some(42), None, Some(99), None])) as _,
1292                Arc::new(StringArray::from(vec![
1293                    None,
1294                    Some("hello"),
1295                    None,
1296                    Some("world"),
1297                ])),
1298            ],
1299        )
1300        .unwrap();
1301
1302        // union_extract by name always returns type_id 0 (first match)
1303        let by_name = union_extract(&union, "val").unwrap();
1304        assert_eq!(
1305            *by_name,
1306            Int32Array::from(vec![Some(42), None, Some(99), None])
1307        );
1308
1309        // union_extract_by_id can select type_id 1 (the Utf8 child)
1310        let by_id = union_extract_by_id(&union, 1).unwrap();
1311        assert_eq!(
1312            *by_id,
1313            StringArray::from(vec![None, Some("hello"), None, Some("world")])
1314        );
1315    }
1316
1317    #[test]
1318    fn extract_by_id_dense_duplicate_names() {
1319        let fields = UnionFields::try_new(
1320            [0, 1],
1321            [
1322                Field::new("val", DataType::Int32, true),
1323                Field::new("val", DataType::Utf8, true),
1324            ],
1325        )
1326        .unwrap();
1327
1328        let union = UnionArray::try_new(
1329            fields,
1330            vec![0_i8, 1, 0].into(),
1331            Some(vec![0_i32, 0, 1].into()),
1332            vec![
1333                Arc::new(Int32Array::from(vec![Some(42), Some(99)])) as _,
1334                Arc::new(StringArray::from(vec![Some("hello")])),
1335            ],
1336        )
1337        .unwrap();
1338
1339        // by type_id 0 → Int32 child
1340        let by_id_0 = union_extract_by_id(&union, 0).unwrap();
1341        assert_eq!(*by_id_0, Int32Array::from(vec![Some(42), None, Some(99)]));
1342
1343        // by type_id 1 → Utf8 child
1344        let by_id_1 = union_extract_by_id(&union, 1).unwrap();
1345        assert_eq!(*by_id_1, StringArray::from(vec![None, Some("hello"), None]));
1346    }
1347
1348    #[test]
1349    fn extract_by_id_not_found() {
1350        let fields = UnionFields::try_new(
1351            [0, 1],
1352            [
1353                Field::new("a", DataType::Int32, true),
1354                Field::new("b", DataType::Utf8, true),
1355            ],
1356        )
1357        .unwrap();
1358
1359        let union = UnionArray::try_new(
1360            fields,
1361            vec![0_i8, 1].into(),
1362            None,
1363            vec![
1364                Arc::new(Int32Array::from(vec![Some(1), None])) as _,
1365                Arc::new(StringArray::from(vec![None, Some("x")])),
1366            ],
1367        )
1368        .unwrap();
1369
1370        assert_eq!(
1371            union_extract_by_id(&union, 5).unwrap_err().to_string(),
1372            ArrowError::InvalidArgumentError("type_id 5 not found on union".into()).to_string()
1373        );
1374    }
1375}