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