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, remainder) = offsets.as_chunks::<N>();
407    chunks.iter().enumerate().all(|(i, chunk)| {
408        //checks if values within chunk are sequential
409        chunk
410            .iter()
411            .copied()
412            .enumerate()
413            .fold(true, |acc, (i, offset)| {
414                acc & (offset == chunk[0] + i as i32)
415            })
416            && offsets[0] + (i * N) as i32 == chunk[0] //checks if chunk is sequential relative to the first offset
417    }) && remainder
418        .iter()
419        .copied()
420        .enumerate()
421        .fold(true, |acc, (i, offset)| {
422            acc & (offset == remainder[0] + i as i32)
423        }) //if the remainder is sequential relative to the first offset is checked at the start of the function
424}
425
426#[cfg(test)]
427mod tests {
428    use super::{
429        BoolValue, eq_scalar_inner, is_sequential_generic, union_extract, union_extract_by_id,
430    };
431    use arrow_array::{Array, Int32Array, NullArray, StringArray, UnionArray, new_null_array};
432    use arrow_buffer::{BooleanBuffer, ScalarBuffer};
433    use arrow_schema::{ArrowError, DataType, Field, UnionFields, UnionMode};
434    use std::sync::Arc;
435
436    #[test]
437    #[cfg_attr(miri, ignore)] // Takes too long
438    fn test_eq_scalar() {
439        //multiple all equal chunks, so it's loop and sum logic it's tested
440        //multiple chunks after, so it's loop logic it's tested
441        const ARRAY_LEN: usize = 64 * 4;
442
443        //so out of 64 boundaries chunks can be generated and checked for
444        const EQ_SCALAR_CHUNK_SIZE: usize = 3;
445
446        fn eq_scalar(type_ids: &[i8], target: i8) -> BoolValue {
447            eq_scalar_inner(EQ_SCALAR_CHUNK_SIZE, type_ids, target)
448        }
449
450        fn cross_check(left: &[i8], right: i8) -> BooleanBuffer {
451            BooleanBuffer::collect_bool(left.len(), |i| left[i] == right)
452        }
453
454        assert_eq!(eq_scalar(&[], 1), BoolValue::Scalar(true));
455
456        assert_eq!(eq_scalar(&[1], 1), BoolValue::Scalar(true));
457        assert_eq!(eq_scalar(&[2], 1), BoolValue::Scalar(false));
458
459        let mut values = [1; ARRAY_LEN];
460
461        assert_eq!(eq_scalar(&values, 1), BoolValue::Scalar(true));
462        assert_eq!(eq_scalar(&values, 2), BoolValue::Scalar(false));
463
464        //every subslice should return the same value
465        for i in 1..ARRAY_LEN {
466            assert_eq!(eq_scalar(&values[..i], 1), BoolValue::Scalar(true));
467            assert_eq!(eq_scalar(&values[..i], 2), BoolValue::Scalar(false));
468        }
469
470        // test that a single change anywhere is checked for
471        for i in 0..ARRAY_LEN {
472            values[i] = 2;
473
474            assert_eq!(
475                eq_scalar(&values, 1),
476                BoolValue::Buffer(cross_check(&values, 1))
477            );
478            assert_eq!(
479                eq_scalar(&values, 2),
480                BoolValue::Buffer(cross_check(&values, 2))
481            );
482
483            values[i] = 1;
484        }
485    }
486
487    #[test]
488    fn test_is_sequential() {
489        /*
490        the smallest value that satisfies:
491        >1 so the fold logic of a exact chunk executes
492        >2 so a >1 non-exact remainder can exist, and it's fold logic executes
493         */
494        const CHUNK_SIZE: usize = 3;
495        //we test arrays of size up to 8 = 2 * CHUNK_SIZE + 2:
496        //multiple(2) exact chunks, so the AND logic between them executes
497        //a >1(2) remainder, so:
498        //    the AND logic between all exact chunks and the remainder executes
499        //    the remainder fold logic executes
500
501        fn is_sequential(v: &[i32]) -> bool {
502            is_sequential_generic::<CHUNK_SIZE>(v)
503        }
504
505        assert!(is_sequential(&[])); //empty
506        assert!(is_sequential(&[1])); //single
507
508        assert!(is_sequential(&[1, 2]));
509        assert!(is_sequential(&[1, 2, 3]));
510        assert!(is_sequential(&[1, 2, 3, 4]));
511        assert!(is_sequential(&[1, 2, 3, 4, 5]));
512        assert!(is_sequential(&[1, 2, 3, 4, 5, 6]));
513        assert!(is_sequential(&[1, 2, 3, 4, 5, 6, 7]));
514        assert!(is_sequential(&[1, 2, 3, 4, 5, 6, 7, 8]));
515
516        assert!(!is_sequential(&[8, 7]));
517        assert!(!is_sequential(&[8, 7, 6]));
518        assert!(!is_sequential(&[8, 7, 6, 5]));
519        assert!(!is_sequential(&[8, 7, 6, 5, 4]));
520        assert!(!is_sequential(&[8, 7, 6, 5, 4, 3]));
521        assert!(!is_sequential(&[8, 7, 6, 5, 4, 3, 2]));
522        assert!(!is_sequential(&[8, 7, 6, 5, 4, 3, 2, 1]));
523
524        assert!(!is_sequential(&[0, 2]));
525        assert!(!is_sequential(&[1, 0]));
526
527        assert!(!is_sequential(&[0, 2, 3]));
528        assert!(!is_sequential(&[1, 0, 3]));
529        assert!(!is_sequential(&[1, 2, 0]));
530
531        assert!(!is_sequential(&[0, 2, 3, 4]));
532        assert!(!is_sequential(&[1, 0, 3, 4]));
533        assert!(!is_sequential(&[1, 2, 0, 4]));
534        assert!(!is_sequential(&[1, 2, 3, 0]));
535
536        assert!(!is_sequential(&[0, 2, 3, 4, 5]));
537        assert!(!is_sequential(&[1, 0, 3, 4, 5]));
538        assert!(!is_sequential(&[1, 2, 0, 4, 5]));
539        assert!(!is_sequential(&[1, 2, 3, 0, 5]));
540        assert!(!is_sequential(&[1, 2, 3, 4, 0]));
541
542        assert!(!is_sequential(&[0, 2, 3, 4, 5, 6]));
543        assert!(!is_sequential(&[1, 0, 3, 4, 5, 6]));
544        assert!(!is_sequential(&[1, 2, 0, 4, 5, 6]));
545        assert!(!is_sequential(&[1, 2, 3, 0, 5, 6]));
546        assert!(!is_sequential(&[1, 2, 3, 4, 0, 6]));
547        assert!(!is_sequential(&[1, 2, 3, 4, 5, 0]));
548
549        assert!(!is_sequential(&[0, 2, 3, 4, 5, 6, 7]));
550        assert!(!is_sequential(&[1, 0, 3, 4, 5, 6, 7]));
551        assert!(!is_sequential(&[1, 2, 0, 4, 5, 6, 7]));
552        assert!(!is_sequential(&[1, 2, 3, 0, 5, 6, 7]));
553        assert!(!is_sequential(&[1, 2, 3, 4, 0, 6, 7]));
554        assert!(!is_sequential(&[1, 2, 3, 4, 5, 0, 7]));
555        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 0]));
556
557        assert!(!is_sequential(&[0, 2, 3, 4, 5, 6, 7, 8]));
558        assert!(!is_sequential(&[1, 0, 3, 4, 5, 6, 7, 8]));
559        assert!(!is_sequential(&[1, 2, 0, 4, 5, 6, 7, 8]));
560        assert!(!is_sequential(&[1, 2, 3, 0, 5, 6, 7, 8]));
561        assert!(!is_sequential(&[1, 2, 3, 4, 0, 6, 7, 8]));
562        assert!(!is_sequential(&[1, 2, 3, 4, 5, 0, 7, 8]));
563        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 0, 8]));
564        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 7, 0]));
565
566        // checks increments at the chunk boundary
567        assert!(!is_sequential(&[1, 2, 3, 5]));
568        assert!(!is_sequential(&[1, 2, 3, 5, 6]));
569        assert!(!is_sequential(&[1, 2, 3, 5, 6, 7]));
570        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 8]));
571        assert!(!is_sequential(&[1, 2, 3, 4, 5, 6, 8, 9]));
572    }
573
574    fn str1() -> UnionFields {
575        UnionFields::try_new(vec![1], vec![Field::new("str", DataType::Utf8, true)]).unwrap()
576    }
577
578    fn str1_int3() -> UnionFields {
579        UnionFields::try_new(
580            vec![1, 3],
581            vec![
582                Field::new("str", DataType::Utf8, true),
583                Field::new("int", DataType::Int32, true),
584            ],
585        )
586        .unwrap()
587    }
588
589    #[test]
590    fn sparse_1_1_single_field() {
591        let union = UnionArray::try_new(
592            //single field
593            str1(),
594            ScalarBuffer::from(vec![1, 1]), // non empty, every type id must match
595            None,                           //sparse
596            vec![
597                Arc::new(StringArray::from(vec!["a", "b"])), // not null
598            ],
599        )
600        .unwrap();
601
602        let expected = StringArray::from(vec!["a", "b"]);
603        let extracted = union_extract(&union, "str").unwrap();
604
605        assert_eq!(extracted.into_data(), expected.into_data());
606    }
607
608    #[test]
609    fn sparse_1_2_empty() {
610        let union = UnionArray::try_new(
611            // multiple fields
612            str1_int3(),
613            ScalarBuffer::from(vec![]), //empty union
614            None,                       // sparse
615            vec![
616                Arc::new(StringArray::new_null(0)),
617                Arc::new(Int32Array::new_null(0)),
618            ],
619        )
620        .unwrap();
621
622        let expected = StringArray::new_null(0);
623        let extracted = union_extract(&union, "str").unwrap(); //target type is not Null
624
625        assert_eq!(extracted.into_data(), expected.into_data());
626    }
627
628    #[test]
629    fn sparse_1_3a_null_target() {
630        let union = UnionArray::try_new(
631            // multiple fields
632            UnionFields::try_new(
633                vec![1, 3],
634                vec![
635                    Field::new("str", DataType::Utf8, true),
636                    Field::new("null", DataType::Null, true), // target type is Null
637                ],
638            )
639            .unwrap(),
640            ScalarBuffer::from(vec![1]), //not empty
641            None,                        // sparse
642            vec![
643                Arc::new(StringArray::new_null(1)),
644                Arc::new(NullArray::new(1)), // null data type
645            ],
646        )
647        .unwrap();
648
649        let expected = NullArray::new(1);
650        let extracted = union_extract(&union, "null").unwrap();
651
652        assert_eq!(extracted.into_data(), expected.into_data());
653    }
654
655    #[test]
656    fn sparse_1_3b_null_target() {
657        let union = UnionArray::try_new(
658            // multiple fields
659            str1_int3(),
660            ScalarBuffer::from(vec![1]), //not empty
661            None,                        // sparse
662            vec![
663                Arc::new(StringArray::new_null(1)), //all null
664                Arc::new(Int32Array::new_null(1)),
665            ],
666        )
667        .unwrap();
668
669        let expected = StringArray::new_null(1);
670        let extracted = union_extract(&union, "str").unwrap(); //target type is not Null
671
672        assert_eq!(extracted.into_data(), expected.into_data());
673    }
674
675    #[test]
676    fn sparse_2_all_types_match() {
677        let union = UnionArray::try_new(
678            //multiple fields
679            str1_int3(),
680            ScalarBuffer::from(vec![3, 3]), // all types match
681            None,                           //sparse
682            vec![
683                Arc::new(StringArray::new_null(2)),
684                Arc::new(Int32Array::from(vec![1, 4])), // not null
685            ],
686        )
687        .unwrap();
688
689        let expected = Int32Array::from(vec![1, 4]);
690        let extracted = union_extract(&union, "int").unwrap();
691
692        assert_eq!(extracted.into_data(), expected.into_data());
693    }
694
695    #[test]
696    fn sparse_3_1_none_match_target_can_contain_null_mask() {
697        let union = UnionArray::try_new(
698            //multiple fields
699            str1_int3(),
700            ScalarBuffer::from(vec![1, 1, 1, 1]), // none match
701            None,                                 // sparse
702            vec![
703                Arc::new(StringArray::new_null(4)),
704                Arc::new(Int32Array::from(vec![None, Some(4), None, Some(8)])), // target is not null
705            ],
706        )
707        .unwrap();
708
709        let expected = Int32Array::new_null(4);
710        let extracted = union_extract(&union, "int").unwrap();
711
712        assert_eq!(extracted.into_data(), expected.into_data());
713    }
714
715    fn str1_union3(union3_datatype: DataType) -> UnionFields {
716        UnionFields::try_new(
717            vec![1, 3],
718            vec![
719                Field::new("str", DataType::Utf8, true),
720                Field::new("union", union3_datatype, true),
721            ],
722        )
723        .unwrap()
724    }
725
726    #[test]
727    fn sparse_3_2_none_match_cant_contain_null_mask_union_target() {
728        let target_fields = str1();
729        let target_type = DataType::Union(target_fields.clone(), UnionMode::Sparse);
730
731        let union = UnionArray::try_new(
732            //multiple fields
733            str1_union3(target_type.clone()),
734            ScalarBuffer::from(vec![1, 1]), // none match
735            None,                           //sparse
736            vec![
737                Arc::new(StringArray::new_null(2)),
738                //target is not null
739                Arc::new(
740                    UnionArray::try_new(
741                        target_fields.clone(),
742                        ScalarBuffer::from(vec![1, 1]),
743                        None,
744                        vec![Arc::new(StringArray::from(vec!["a", "b"]))],
745                    )
746                    .unwrap(),
747                ),
748            ],
749        )
750        .unwrap();
751
752        let expected = new_null_array(&target_type, 2);
753        let extracted = union_extract(&union, "union").unwrap();
754
755        assert_eq!(extracted.into_data(), expected.into_data());
756    }
757
758    #[test]
759    fn sparse_4_1_1_target_with_nulls() {
760        let union = UnionArray::try_new(
761            //multiple fields
762            str1_int3(),
763            ScalarBuffer::from(vec![3, 3, 1, 1]), // multiple selected types
764            None,                                 // sparse
765            vec![
766                Arc::new(StringArray::new_null(4)),
767                Arc::new(Int32Array::from(vec![None, Some(4), None, Some(8)])), // target with nulls
768            ],
769        )
770        .unwrap();
771
772        let expected = Int32Array::from(vec![None, Some(4), None, None]);
773        let extracted = union_extract(&union, "int").unwrap();
774
775        assert_eq!(extracted.into_data(), expected.into_data());
776    }
777
778    #[test]
779    fn sparse_4_1_2_target_without_nulls() {
780        let union = UnionArray::try_new(
781            //multiple fields
782            str1_int3(),
783            ScalarBuffer::from(vec![1, 3, 3]), // multiple selected types
784            None,                              // sparse
785            vec![
786                Arc::new(StringArray::new_null(3)),
787                Arc::new(Int32Array::from(vec![2, 4, 8])), // target without nulls
788            ],
789        )
790        .unwrap();
791
792        let expected = Int32Array::from(vec![None, Some(4), Some(8)]);
793        let extracted = union_extract(&union, "int").unwrap();
794
795        assert_eq!(extracted.into_data(), expected.into_data());
796    }
797
798    #[test]
799    fn sparse_4_2_some_match_target_cant_contain_null_mask() {
800        let target_fields = str1();
801        let target_type = DataType::Union(target_fields.clone(), UnionMode::Sparse);
802
803        let union = UnionArray::try_new(
804            //multiple fields
805            str1_union3(target_type),
806            ScalarBuffer::from(vec![3, 1]), // some types match, but not all
807            None,                           //sparse
808            vec![
809                Arc::new(StringArray::new_null(2)),
810                Arc::new(
811                    UnionArray::try_new(
812                        target_fields.clone(),
813                        ScalarBuffer::from(vec![1, 1]),
814                        None,
815                        vec![Arc::new(StringArray::from(vec!["a", "b"]))],
816                    )
817                    .unwrap(),
818                ),
819            ],
820        )
821        .unwrap();
822
823        let expected = UnionArray::try_new(
824            target_fields,
825            ScalarBuffer::from(vec![1, 1]),
826            None,
827            vec![Arc::new(StringArray::from(vec![Some("a"), None]))],
828        )
829        .unwrap();
830        let extracted = union_extract(&union, "union").unwrap();
831
832        assert_eq!(extracted.into_data(), expected.into_data());
833    }
834
835    #[test]
836    fn dense_1_1_both_empty() {
837        let union = UnionArray::try_new(
838            str1_int3(),
839            ScalarBuffer::from(vec![]),       //empty union
840            Some(ScalarBuffer::from(vec![])), // dense
841            vec![
842                Arc::new(StringArray::new_null(0)), //empty target
843                Arc::new(Int32Array::new_null(0)),
844            ],
845        )
846        .unwrap();
847
848        let expected = StringArray::new_null(0);
849        let extracted = union_extract(&union, "str").unwrap();
850
851        assert_eq!(extracted.into_data(), expected.into_data());
852    }
853
854    #[test]
855    fn dense_1_2_empty_union_target_non_empty() {
856        let union = UnionArray::try_new(
857            str1_int3(),
858            ScalarBuffer::from(vec![]),       //empty union
859            Some(ScalarBuffer::from(vec![])), // dense
860            vec![
861                Arc::new(StringArray::new_null(1)), //non empty target
862                Arc::new(Int32Array::new_null(0)),
863            ],
864        )
865        .unwrap();
866
867        let expected = StringArray::new_null(0);
868        let extracted = union_extract(&union, "str").unwrap();
869
870        assert_eq!(extracted.into_data(), expected.into_data());
871    }
872
873    #[test]
874    fn dense_2_non_empty_union_target_empty() {
875        let union = UnionArray::try_new(
876            str1_int3(),
877            ScalarBuffer::from(vec![3, 3]),       //non empty union
878            Some(ScalarBuffer::from(vec![0, 1])), // dense
879            vec![
880                Arc::new(StringArray::new_null(0)), //empty target
881                Arc::new(Int32Array::new_null(2)),
882            ],
883        )
884        .unwrap();
885
886        let expected = StringArray::new_null(2);
887        let extracted = union_extract(&union, "str").unwrap();
888
889        assert_eq!(extracted.into_data(), expected.into_data());
890    }
891
892    #[test]
893    fn dense_3_1_null_target_smaller_len() {
894        let union = UnionArray::try_new(
895            str1_int3(),
896            ScalarBuffer::from(vec![3, 3]),       //non empty union
897            Some(ScalarBuffer::from(vec![0, 0])), //dense
898            vec![
899                Arc::new(StringArray::new_null(1)), //smaller target
900                Arc::new(Int32Array::new_null(2)),
901            ],
902        )
903        .unwrap();
904
905        let expected = StringArray::new_null(2);
906        let extracted = union_extract(&union, "str").unwrap();
907
908        assert_eq!(extracted.into_data(), expected.into_data());
909    }
910
911    #[test]
912    fn dense_3_2_null_target_equal_len() {
913        let union = UnionArray::try_new(
914            str1_int3(),
915            ScalarBuffer::from(vec![3, 3]),       //non empty union
916            Some(ScalarBuffer::from(vec![0, 0])), //dense
917            vec![
918                Arc::new(StringArray::new_null(2)), //equal len
919                Arc::new(Int32Array::new_null(2)),
920            ],
921        )
922        .unwrap();
923
924        let expected = StringArray::new_null(2);
925        let extracted = union_extract(&union, "str").unwrap();
926
927        assert_eq!(extracted.into_data(), expected.into_data());
928    }
929
930    #[test]
931    fn dense_3_3_null_target_bigger_len() {
932        let union = UnionArray::try_new(
933            str1_int3(),
934            ScalarBuffer::from(vec![3, 3]),       //non empty union
935            Some(ScalarBuffer::from(vec![0, 0])), //dense
936            vec![
937                Arc::new(StringArray::new_null(3)), //bigger len
938                Arc::new(Int32Array::new_null(3)),
939            ],
940        )
941        .unwrap();
942
943        let expected = StringArray::new_null(2);
944        let extracted = union_extract(&union, "str").unwrap();
945
946        assert_eq!(extracted.into_data(), expected.into_data());
947    }
948
949    #[test]
950    fn dense_4_1a_single_type_sequential_offsets_equal_len() {
951        let union = UnionArray::try_new(
952            // single field
953            str1(),
954            ScalarBuffer::from(vec![1, 1]),       //non empty union
955            Some(ScalarBuffer::from(vec![0, 1])), //sequential
956            vec![
957                Arc::new(StringArray::from(vec!["a1", "b2"])), //equal len, non null
958            ],
959        )
960        .unwrap();
961
962        let expected = StringArray::from(vec!["a1", "b2"]);
963        let extracted = union_extract(&union, "str").unwrap();
964
965        assert_eq!(extracted.into_data(), expected.into_data());
966    }
967
968    #[test]
969    fn dense_4_2a_single_type_sequential_offsets_bigger() {
970        let union = UnionArray::try_new(
971            // single field
972            str1(),
973            ScalarBuffer::from(vec![1, 1]),       //non empty union
974            Some(ScalarBuffer::from(vec![0, 1])), //sequential
975            vec![
976                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), //equal len, non null
977            ],
978        )
979        .unwrap();
980
981        let expected = StringArray::from(vec!["a1", "b2"]);
982        let extracted = union_extract(&union, "str").unwrap();
983
984        assert_eq!(extracted.into_data(), expected.into_data());
985    }
986
987    #[test]
988    fn dense_4_3a_single_type_non_sequential() {
989        let union = UnionArray::try_new(
990            // single field
991            str1(),
992            ScalarBuffer::from(vec![1, 1]),       //non empty union
993            Some(ScalarBuffer::from(vec![0, 2])), //non sequential
994            vec![
995                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), //equal len, non null
996            ],
997        )
998        .unwrap();
999
1000        let expected = StringArray::from(vec!["a1", "c3"]);
1001        let extracted = union_extract(&union, "str").unwrap();
1002
1003        assert_eq!(extracted.into_data(), expected.into_data());
1004    }
1005
1006    #[test]
1007    fn dense_4_1b_empty_siblings_sequential_equal_len() {
1008        let union = UnionArray::try_new(
1009            // multiple fields
1010            str1_int3(),
1011            ScalarBuffer::from(vec![1, 1]),       //non empty union
1012            Some(ScalarBuffer::from(vec![0, 1])), //sequential
1013            vec![
1014                Arc::new(StringArray::from(vec!["a", "b"])), //equal len, non null
1015                Arc::new(Int32Array::new_null(0)),           //empty sibling
1016            ],
1017        )
1018        .unwrap();
1019
1020        let expected = StringArray::from(vec!["a", "b"]);
1021        let extracted = union_extract(&union, "str").unwrap();
1022
1023        assert_eq!(extracted.into_data(), expected.into_data());
1024    }
1025
1026    #[test]
1027    fn dense_4_2b_empty_siblings_sequential_bigger_len() {
1028        let union = UnionArray::try_new(
1029            // multiple fields
1030            str1_int3(),
1031            ScalarBuffer::from(vec![1, 1]),       //non empty union
1032            Some(ScalarBuffer::from(vec![0, 1])), //sequential
1033            vec![
1034                Arc::new(StringArray::from(vec!["a", "b", "c"])), //bigger len, non null
1035                Arc::new(Int32Array::new_null(0)),                //empty sibling
1036            ],
1037        )
1038        .unwrap();
1039
1040        let expected = StringArray::from(vec!["a", "b"]);
1041        let extracted = union_extract(&union, "str").unwrap();
1042
1043        assert_eq!(extracted.into_data(), expected.into_data());
1044    }
1045
1046    #[test]
1047    fn dense_4_3b_empty_sibling_non_sequential() {
1048        let union = UnionArray::try_new(
1049            // multiple fields
1050            str1_int3(),
1051            ScalarBuffer::from(vec![1, 1]),       //non empty union
1052            Some(ScalarBuffer::from(vec![0, 2])), //non sequential
1053            vec![
1054                Arc::new(StringArray::from(vec!["a", "b", "c"])), //non null
1055                Arc::new(Int32Array::new_null(0)),                //empty sibling
1056            ],
1057        )
1058        .unwrap();
1059
1060        let expected = StringArray::from(vec!["a", "c"]);
1061        let extracted = union_extract(&union, "str").unwrap();
1062
1063        assert_eq!(extracted.into_data(), expected.into_data());
1064    }
1065
1066    #[test]
1067    fn dense_4_1c_all_types_match_sequential_equal_len() {
1068        let union = UnionArray::try_new(
1069            // multiple fields
1070            str1_int3(),
1071            ScalarBuffer::from(vec![1, 1]),       //all types match
1072            Some(ScalarBuffer::from(vec![0, 1])), //sequential
1073            vec![
1074                Arc::new(StringArray::from(vec!["a1", "b2"])), //equal len
1075                Arc::new(Int32Array::new_null(2)),             //non empty sibling
1076            ],
1077        )
1078        .unwrap();
1079
1080        let expected = StringArray::from(vec!["a1", "b2"]);
1081        let extracted = union_extract(&union, "str").unwrap();
1082
1083        assert_eq!(extracted.into_data(), expected.into_data());
1084    }
1085
1086    #[test]
1087    fn dense_4_2c_all_types_match_sequential_bigger_len() {
1088        let union = UnionArray::try_new(
1089            // multiple fields
1090            str1_int3(),
1091            ScalarBuffer::from(vec![1, 1]),       //all types match
1092            Some(ScalarBuffer::from(vec![0, 1])), //sequential
1093            vec![
1094                Arc::new(StringArray::from(vec!["a1", "b2", "b3"])), //bigger len
1095                Arc::new(Int32Array::new_null(2)),                   //non empty sibling
1096            ],
1097        )
1098        .unwrap();
1099
1100        let expected = StringArray::from(vec!["a1", "b2"]);
1101        let extracted = union_extract(&union, "str").unwrap();
1102
1103        assert_eq!(extracted.into_data(), expected.into_data());
1104    }
1105
1106    #[test]
1107    fn dense_4_3c_all_types_match_non_sequential() {
1108        let union = UnionArray::try_new(
1109            // multiple fields
1110            str1_int3(),
1111            ScalarBuffer::from(vec![1, 1]),       //all types match
1112            Some(ScalarBuffer::from(vec![0, 2])), //non sequential
1113            vec![
1114                Arc::new(StringArray::from(vec!["a1", "b2", "b3"])),
1115                Arc::new(Int32Array::new_null(2)), //non empty sibling
1116            ],
1117        )
1118        .unwrap();
1119
1120        let expected = StringArray::from(vec!["a1", "b3"]);
1121        let extracted = union_extract(&union, "str").unwrap();
1122
1123        assert_eq!(extracted.into_data(), expected.into_data());
1124    }
1125
1126    #[test]
1127    fn dense_5_1a_none_match_less_len() {
1128        let union = UnionArray::try_new(
1129            // multiple fields
1130            str1_int3(),
1131            ScalarBuffer::from(vec![3, 3, 3, 3, 3]), //none matches
1132            Some(ScalarBuffer::from(vec![0, 0, 0, 1, 1])), // dense
1133            vec![
1134                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), // less len
1135                Arc::new(Int32Array::from(vec![1, 2])),
1136            ],
1137        )
1138        .unwrap();
1139
1140        let expected = StringArray::new_null(5);
1141        let extracted = union_extract(&union, "str").unwrap();
1142
1143        assert_eq!(extracted.into_data(), expected.into_data());
1144    }
1145
1146    #[test]
1147    fn dense_5_1b_cant_contain_null_mask() {
1148        let target_fields = str1();
1149        let target_type = DataType::Union(target_fields.clone(), UnionMode::Sparse);
1150
1151        let union = UnionArray::try_new(
1152            // multiple fields
1153            str1_union3(target_type.clone()),
1154            ScalarBuffer::from(vec![1, 1, 1, 1, 1]), //none matches
1155            Some(ScalarBuffer::from(vec![0, 0, 0, 1, 1])), // dense
1156            vec![
1157                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), // less len
1158                Arc::new(
1159                    UnionArray::try_new(
1160                        target_fields.clone(),
1161                        ScalarBuffer::from(vec![1]),
1162                        None,
1163                        vec![Arc::new(StringArray::from(vec!["a"]))],
1164                    )
1165                    .unwrap(),
1166                ), // non empty
1167            ],
1168        )
1169        .unwrap();
1170
1171        let expected = new_null_array(&target_type, 5);
1172        let extracted = union_extract(&union, "union").unwrap();
1173
1174        assert_eq!(extracted.into_data(), expected.into_data());
1175    }
1176
1177    #[test]
1178    fn dense_5_2_none_match_equal_len() {
1179        let union = UnionArray::try_new(
1180            // multiple fields
1181            str1_int3(),
1182            ScalarBuffer::from(vec![3, 3, 3, 3, 3]), //none matches
1183            Some(ScalarBuffer::from(vec![0, 0, 0, 1, 1])), // dense
1184            vec![
1185                Arc::new(StringArray::from(vec!["a1", "b2", "c3", "d4", "e5"])), // equal len
1186                Arc::new(Int32Array::from(vec![1, 2])),
1187            ],
1188        )
1189        .unwrap();
1190
1191        let expected = StringArray::new_null(5);
1192        let extracted = union_extract(&union, "str").unwrap();
1193
1194        assert_eq!(extracted.into_data(), expected.into_data());
1195    }
1196
1197    #[test]
1198    fn dense_5_3_none_match_greater_len() {
1199        let union = UnionArray::try_new(
1200            // multiple fields
1201            str1_int3(),
1202            ScalarBuffer::from(vec![3, 3, 3, 3, 3]), //none matches
1203            Some(ScalarBuffer::from(vec![0, 0, 0, 1, 1])), // dense
1204            vec![
1205                Arc::new(StringArray::from(vec!["a1", "b2", "c3", "d4", "e5", "f6"])), // greater len
1206                Arc::new(Int32Array::from(vec![1, 2])),                                //non null
1207            ],
1208        )
1209        .unwrap();
1210
1211        let expected = StringArray::new_null(5);
1212        let extracted = union_extract(&union, "str").unwrap();
1213
1214        assert_eq!(extracted.into_data(), expected.into_data());
1215    }
1216
1217    #[test]
1218    fn dense_6_some_matches() {
1219        let union = UnionArray::try_new(
1220            // multiple fields
1221            str1_int3(),
1222            ScalarBuffer::from(vec![3, 3, 1, 1, 1]), //some matches
1223            Some(ScalarBuffer::from(vec![0, 1, 0, 1, 2])), // dense
1224            vec![
1225                Arc::new(StringArray::from(vec!["a1", "b2", "c3"])), // non null
1226                Arc::new(Int32Array::from(vec![1, 2])),
1227            ],
1228        )
1229        .unwrap();
1230
1231        let expected = Int32Array::from(vec![Some(1), Some(2), None, None, None]);
1232        let extracted = union_extract(&union, "int").unwrap();
1233
1234        assert_eq!(extracted.into_data(), expected.into_data());
1235    }
1236
1237    #[test]
1238    fn empty_sparse_union() {
1239        let union = UnionArray::try_new(
1240            UnionFields::empty(),
1241            ScalarBuffer::from(vec![]),
1242            None,
1243            vec![],
1244        )
1245        .unwrap();
1246
1247        assert_eq!(
1248            union_extract(&union, "a").unwrap_err().to_string(),
1249            ArrowError::InvalidArgumentError("field a not found on union".into()).to_string()
1250        );
1251    }
1252
1253    #[test]
1254    fn empty_dense_union() {
1255        let union = UnionArray::try_new(
1256            UnionFields::empty(),
1257            ScalarBuffer::from(vec![]),
1258            Some(ScalarBuffer::from(vec![])),
1259            vec![],
1260        )
1261        .unwrap();
1262
1263        assert_eq!(
1264            union_extract(&union, "a").unwrap_err().to_string(),
1265            ArrowError::InvalidArgumentError("field a not found on union".into()).to_string()
1266        );
1267    }
1268
1269    #[test]
1270    fn extract_by_id_sparse_duplicate_names() {
1271        // Two fields with the same name "val" but different type_ids and types
1272        let fields = UnionFields::try_new(
1273            [0, 1],
1274            [
1275                Field::new("val", DataType::Int32, true),
1276                Field::new("val", DataType::Utf8, true),
1277            ],
1278        )
1279        .unwrap();
1280
1281        let union = UnionArray::try_new(
1282            fields,
1283            vec![0_i8, 1, 0, 1].into(),
1284            None,
1285            vec![
1286                Arc::new(Int32Array::from(vec![Some(42), None, Some(99), None])) as _,
1287                Arc::new(StringArray::from(vec![
1288                    None,
1289                    Some("hello"),
1290                    None,
1291                    Some("world"),
1292                ])),
1293            ],
1294        )
1295        .unwrap();
1296
1297        // union_extract by name always returns type_id 0 (first match)
1298        let by_name = union_extract(&union, "val").unwrap();
1299        assert_eq!(
1300            *by_name,
1301            Int32Array::from(vec![Some(42), None, Some(99), None])
1302        );
1303
1304        // union_extract_by_id can select type_id 1 (the Utf8 child)
1305        let by_id = union_extract_by_id(&union, 1).unwrap();
1306        assert_eq!(
1307            *by_id,
1308            StringArray::from(vec![None, Some("hello"), None, Some("world")])
1309        );
1310    }
1311
1312    #[test]
1313    fn extract_by_id_dense_duplicate_names() {
1314        let fields = UnionFields::try_new(
1315            [0, 1],
1316            [
1317                Field::new("val", DataType::Int32, true),
1318                Field::new("val", DataType::Utf8, true),
1319            ],
1320        )
1321        .unwrap();
1322
1323        let union = UnionArray::try_new(
1324            fields,
1325            vec![0_i8, 1, 0].into(),
1326            Some(vec![0_i32, 0, 1].into()),
1327            vec![
1328                Arc::new(Int32Array::from(vec![Some(42), Some(99)])) as _,
1329                Arc::new(StringArray::from(vec![Some("hello")])),
1330            ],
1331        )
1332        .unwrap();
1333
1334        // by type_id 0 → Int32 child
1335        let by_id_0 = union_extract_by_id(&union, 0).unwrap();
1336        assert_eq!(*by_id_0, Int32Array::from(vec![Some(42), None, Some(99)]));
1337
1338        // by type_id 1 → Utf8 child
1339        let by_id_1 = union_extract_by_id(&union, 1).unwrap();
1340        assert_eq!(*by_id_1, StringArray::from(vec![None, Some("hello"), None]));
1341    }
1342
1343    #[test]
1344    fn extract_by_id_not_found() {
1345        let fields = UnionFields::try_new(
1346            [0, 1],
1347            [
1348                Field::new("a", DataType::Int32, true),
1349                Field::new("b", DataType::Utf8, true),
1350            ],
1351        )
1352        .unwrap();
1353
1354        let union = UnionArray::try_new(
1355            fields,
1356            vec![0_i8, 1].into(),
1357            None,
1358            vec![
1359                Arc::new(Int32Array::from(vec![Some(1), None])) as _,
1360                Arc::new(StringArray::from(vec![None, Some("x")])),
1361            ],
1362        )
1363        .unwrap();
1364
1365        assert_eq!(
1366            union_extract_by_id(&union, 5).unwrap_err().to_string(),
1367            ArrowError::InvalidArgumentError("type_id 5 not found on union".into()).to_string()
1368        );
1369    }
1370}
arrow_select/union_extract.rs

arrow_select/
union_extract.rs
