Skip to main content

arrow_select/
interleave.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Interleave elements from multiple arrays
6
7use crate::concat::concat;
8use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values};
9use arrow_array::ListLikeArray;
10use arrow_array::builder::{BooleanBufferBuilder, PrimitiveBuilder};
11use arrow_array::cast::AsArray;
12use arrow_array::types::*;
13use arrow_array::*;
14use arrow_buffer::bit_mask::set_bits;
15use arrow_buffer::bit_util;
16use arrow_buffer::{
17    ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, OffsetBuffer,
18};
19use arrow_data::ByteView;
20use arrow_data::transform::MutableArrayData;
21use arrow_schema::{ArrowError, DataType, FieldRef, Fields};
22use std::collections::BTreeMap;
23use std::sync::Arc;
24
25macro_rules! primitive_helper {
26    ($t:ty, $values:ident, $indices:ident, $data_type:ident) => {
27        interleave_primitive::<$t>($values, $indices, $data_type)
28    };
29}
30
31macro_rules! dict_helper {
32    ($t:ty, $values:expr, $indices:expr) => {
33        interleave_dictionaries::<$t>($values, $indices)
34    };
35}
36
37///
38/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values.
39///
40/// Each element in `indices` is a pair of `usize` with the first identifying the index
41/// of the [`Array`] in `values`, and the second the index of the value within that [`Array`]
42///
43/// ```text
44/// ┌─────────────────┐      ┌─────────┐                                  ┌─────────────────┐
45/// │        A        │      │ (0, 0)  │        interleave(               │        A        │
46/// ├─────────────────┤      ├─────────┤          [values0, values1],     ├─────────────────┤
47/// │        D        │      │ (1, 0)  │          indices                 │        B        │
48/// └─────────────────┘      ├─────────┤        )                         ├─────────────────┤
49///   values array 0         │ (1, 1)  │      ─────────────────────────▶  │        C        │
50///                          ├─────────┤                                  ├─────────────────┤
51///                          │ (0, 1)  │                                  │        D        │
52///                          └─────────┘                                  └─────────────────┘
53/// ┌─────────────────┐       indices
54/// │        B        │        array
55/// ├─────────────────┤                                                    result
56/// │        C        │
57/// ├─────────────────┤
58/// │        E        │
59/// └─────────────────┘
60///   values array 1
61/// ```
62///
63/// For selecting values by index from a single array see [`crate::take`]
64pub fn interleave(
65    values: &[&dyn Array],
66    indices: &[(usize, usize)],
67) -> Result<ArrayRef, ArrowError> {
68    if values.is_empty() {
69        return Err(ArrowError::InvalidArgumentError(
70            "interleave requires input of at least one array".to_string(),
71        ));
72    }
73    let data_type = values[0].data_type();
74
75    for array in values.iter().skip(1) {
76        if array.data_type() != data_type {
77            return Err(ArrowError::InvalidArgumentError(format!(
78                "It is not possible to interleave arrays of different data types ({} and {})",
79                data_type,
80                array.data_type()
81            )));
82        }
83    }
84
85    if indices.is_empty() {
86        return Ok(new_empty_array(data_type));
87    }
88
89    downcast_primitive! {
90        data_type => (primitive_helper, values, indices, data_type),
91        DataType::Boolean => interleave_boolean(values, indices),
92        DataType::Utf8 => interleave_bytes::<Utf8Type>(values, indices),
93        DataType::LargeUtf8 => interleave_bytes::<LargeUtf8Type>(values, indices),
94        DataType::Binary => interleave_bytes::<BinaryType>(values, indices),
95        DataType::LargeBinary => interleave_bytes::<LargeBinaryType>(values, indices),
96        DataType::BinaryView => interleave_views::<BinaryViewType>(values, indices),
97        DataType::Utf8View => interleave_views::<StringViewType>(values, indices),
98        DataType::Dictionary(k, _) => downcast_integer! {
99            k.as_ref() => (dict_helper, values, indices),
100            _ => unreachable!("illegal dictionary key type {k}")
101        },
102        DataType::Struct(fields) => interleave_struct(fields, values, indices),
103        DataType::List(field) => interleave_list::<i32>(values, indices, field),
104        DataType::LargeList(field) => interleave_list::<i64>(values, indices, field),
105        DataType::FixedSizeList(field, size) => interleave_fixed_size_list(values, indices, field, *size),
106        DataType::Map(field, ordered) => interleave_map(values, indices, field, *ordered),
107        DataType::RunEndEncoded(r, _) => match r.data_type() {
108            DataType::Int16 => interleave_run_end::<Int16Type>(values, indices),
109            DataType::Int32 => interleave_run_end::<Int32Type>(values, indices),
110            DataType::Int64 => interleave_run_end::<Int64Type>(values, indices),
111            t => unreachable!("illegal run-end type {t}"),
112        },
113        DataType::ListView(field) => interleave_list_view::<i32>(values, indices, field),
114        DataType::LargeListView(field) => interleave_list_view::<i64>(values, indices, field),
115        _ => interleave_fallback(values, indices)
116    }
117}
118
119/// Common functionality for interleaving arrays
120///
121/// T is the concrete Array type
122struct Interleave<'a, T> {
123    /// The input arrays downcast to T
124    arrays: Vec<&'a T>,
125    /// The null buffer of the interleaved output
126    nulls: Option<NullBuffer>,
127}
128
129impl<'a, T: Array + 'static> Interleave<'a, T> {
130    fn new(values: &[&'a dyn Array], indices: &'a [(usize, usize)]) -> Self {
131        let mut has_nulls = false;
132        let arrays: Vec<&T> = values
133            .iter()
134            .map(|x| {
135                has_nulls = has_nulls || x.null_count() != 0;
136                x.as_any().downcast_ref().unwrap()
137            })
138            .collect();
139
140        let nulls = match has_nulls {
141            true => {
142                let nulls = BooleanBuffer::collect_bool(indices.len(), |i| {
143                    let (a, b) = indices[i];
144                    arrays[a].is_valid(b)
145                });
146                Some(nulls.into())
147            }
148            false => None,
149        };
150
151        Self { arrays, nulls }
152    }
153}
154
155fn interleave_primitive<T: ArrowPrimitiveType>(
156    values: &[&dyn Array],
157    indices: &[(usize, usize)],
158    data_type: &DataType,
159) -> Result<ArrayRef, ArrowError> {
160    let interleaved = Interleave::<'_, PrimitiveArray<T>>::new(values, indices);
161    let arrays = &interleaved.arrays;
162    let len = indices.len();
163
164    let mut output = Vec::with_capacity(len);
165    let dst: *mut T::Native = output.as_mut_ptr();
166    let mut base = 0;
167
168    // Process 8 elements at a time to issue multiple independent loads
169    // and increase memory-level parallelism for random access patterns.
170    let (chunks, remainder) = indices.as_chunks::<8>();
171    for chunk in chunks {
172        let v0 = arrays[chunk[0].0].value(chunk[0].1);
173        let v1 = arrays[chunk[1].0].value(chunk[1].1);
174        let v2 = arrays[chunk[2].0].value(chunk[2].1);
175        let v3 = arrays[chunk[3].0].value(chunk[3].1);
176        let v4 = arrays[chunk[4].0].value(chunk[4].1);
177        let v5 = arrays[chunk[5].0].value(chunk[5].1);
178        let v6 = arrays[chunk[6].0].value(chunk[6].1);
179        let v7 = arrays[chunk[7].0].value(chunk[7].1);
180
181        // SAFETY: base+7 < len == output capacity
182        debug_assert!(base + 7 < len);
183        unsafe {
184            dst.add(base).write(v0);
185            dst.add(base + 1).write(v1);
186            dst.add(base + 2).write(v2);
187            dst.add(base + 3).write(v3);
188            dst.add(base + 4).write(v4);
189            dst.add(base + 5).write(v5);
190            dst.add(base + 6).write(v6);
191            dst.add(base + 7).write(v7);
192        }
193        base += 8;
194    }
195
196    for idx in remainder {
197        // SAFETY: base < len == output capacity
198        debug_assert!(base < len);
199        unsafe { dst.add(base).write(arrays[idx.0].value(idx.1)) };
200        base += 1;
201    }
202
203    // SAFETY: all `len` elements have been initialized
204    debug_assert_eq!(base, len);
205    unsafe { output.set_len(len) };
206
207    let array = PrimitiveArray::<T>::try_new(output.into(), interleaved.nulls)?;
208    Ok(Arc::new(array.with_data_type(data_type.clone())))
209}
210
211fn interleave_boolean(
212    values: &[&dyn Array],
213    indices: &[(usize, usize)],
214) -> Result<ArrayRef, ArrowError> {
215    let interleaved = Interleave::<'_, BooleanArray>::new(values, indices);
216    let arrays = &interleaved.arrays;
217
218    let mut values_builder = BooleanBufferBuilder::new(indices.len());
219
220    for &(array_idx, row_idx) in indices {
221        values_builder.append(arrays[array_idx].value(row_idx));
222    }
223
224    let array = BooleanArray::new(values_builder.build(), interleaved.nulls);
225    Ok(Arc::new(array))
226}
227
228fn interleave_bytes<T: ByteArrayType>(
229    values: &[&dyn Array],
230    indices: &[(usize, usize)],
231) -> Result<ArrayRef, ArrowError> {
232    let interleaved = Interleave::<'_, GenericByteArray<T>>::new(values, indices);
233
234    let mut capacity = 0;
235    let mut offsets = Vec::with_capacity(indices.len() + 1);
236    offsets.push(T::Offset::from_usize(0).unwrap());
237    for (a, b) in indices {
238        let o = interleaved.arrays[*a].value_offsets();
239        let element_len = o[*b + 1].as_usize() - o[*b].as_usize();
240        capacity += element_len;
241        offsets.push(
242            T::Offset::from_usize(capacity)
243                .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
244        );
245    }
246
247    let mut values = Vec::with_capacity(capacity);
248    for (a, b) in indices {
249        values.extend_from_slice(interleaved.arrays[*a].value(*b).as_ref());
250    }
251
252    // Safety: safe by construction
253    let array = unsafe {
254        let offsets = OffsetBuffer::new_unchecked(offsets.into());
255        GenericByteArray::<T>::new_unchecked(offsets, values.into(), interleaved.nulls)
256    };
257    Ok(Arc::new(array))
258}
259
260fn interleave_dictionaries<K: ArrowDictionaryKeyType>(
261    arrays: &[&dyn Array],
262    indices: &[(usize, usize)],
263) -> Result<ArrayRef, ArrowError> {
264    let dictionaries: Vec<_> = arrays.iter().map(|x| x.as_dictionary::<K>()).collect();
265    let (should_merge, has_overflow) =
266        should_merge_dictionary_values::<K>(&dictionaries, indices.len());
267    if !should_merge {
268        return if has_overflow {
269            interleave_fallback(arrays, indices)
270        } else {
271            interleave_fallback_dictionary::<K>(&dictionaries, indices)
272        };
273    }
274
275    let masks_and_lookups: Result<Vec<_>, ArrowError> = dictionaries
276        .iter()
277        .enumerate()
278        .map(|(a_idx, dictionary)| {
279            let buffer = MutableBuffer::try_from_len_zeroed(bit_util::ceil(dictionary.len(), 8))
280                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
281            let mut key_mask = BooleanBufferBuilder::new_from_buffer(buffer, dictionary.len());
282
283            let mut lookups = 0;
284            for (_, key_idx) in indices.iter().filter(|(a, _)| *a == a_idx) {
285                key_mask.set_bit(*key_idx, true);
286                lookups += 1;
287            }
288            Ok((key_mask.finish(), lookups))
289        })
290        .collect();
291    let (masks, key_lookups): (Vec<_>, Vec<_>) = masks_and_lookups?.into_iter().unzip();
292
293    let merged = merge_dictionary_values(&dictionaries, Some(&masks), Some(&key_lookups))?;
294
295    // Recompute keys
296    let mut keys = PrimitiveBuilder::<K>::with_capacity(indices.len());
297    for (a, b) in indices {
298        let old_keys: &PrimitiveArray<K> = dictionaries[*a].keys();
299        match old_keys.is_valid(*b) {
300            true => {
301                let old_key = old_keys.values()[*b];
302                keys.append_value(
303                    merged.key_mappings[*a]
304                        .get(old_key.as_usize())
305                        .expect("referenced dictionary key should be mapped"),
306                )
307            }
308            false => keys.append_null(),
309        }
310    }
311    let array = unsafe { DictionaryArray::new_unchecked(keys.finish(), merged.values) };
312    Ok(Arc::new(array))
313}
314
315fn interleave_views<T: ByteViewType>(
316    values: &[&dyn Array],
317    indices: &[(usize, usize)],
318) -> Result<ArrayRef, ArrowError> {
319    let interleaved = Interleave::<'_, GenericByteViewArray<T>>::new(values, indices);
320    let mut buffers = Vec::new();
321
322    // Contains the offsets of start buffer in `buffer_to_new_index`
323    let mut offsets = Vec::with_capacity(interleaved.arrays.len() + 1);
324    offsets.push(0);
325    let mut total_buffers = 0;
326    for a in &interleaved.arrays {
327        total_buffers += a.data_buffers().len();
328        offsets.push(total_buffers);
329    }
330
331    // Marks a buffer in `buffer_to_new_index` that has not yet been referenced.
332    //
333    // The view's buffer index is a signed 32-bit integer in the Arrow specification,
334    // so `u32::MAX` can never be a valid buffer index
335    const UNASSIGNED: u32 = u32::MAX;
336
337    // contains the mapping from old buffer index to new buffer index
338    let mut buffer_to_new_index = vec![UNASSIGNED; total_buffers];
339
340    // Contains the index in `buffers` of each buffer already emitted, keyed by identity
341    let mut seen_buffers = BTreeMap::new();
342
343    let views: Vec<u128> = indices
344        .iter()
345        .map(|(array_idx, value_idx)| {
346            let array = interleaved.arrays[*array_idx];
347            let view = array.views().get(*value_idx).unwrap();
348            let view_len = *view as u32;
349            if view_len <= 12 {
350                return *view;
351            }
352            // value is big enough to be in a variadic buffer
353            let view = ByteView::from(*view);
354            let buffer_to_new_idx = offsets[*array_idx] + view.buffer_index as usize;
355            let new_buffer_idx = &mut buffer_to_new_index[buffer_to_new_idx];
356            if *new_buffer_idx == UNASSIGNED {
357                *new_buffer_idx = push_buffer(
358                    &mut seen_buffers,
359                    &mut buffers,
360                    &array.data_buffers()[view.buffer_index as usize],
361                );
362            }
363            view.with_buffer_index(*new_buffer_idx).as_u128()
364        })
365        .collect();
366
367    let array = unsafe {
368        GenericByteViewArray::<T>::new_unchecked(views.into(), buffers.into(), interleaved.nulls)
369    };
370    Ok(Arc::new(array))
371}
372
373/// Returns the index of `buffer` in `buffers`, adding it if not already present.
374///
375/// Multiple input arrays, e.g. slices of the same array, may reference the same
376/// buffer, so this checks by identity to ensure it is only emitted once
377#[inline(never)]
378fn push_buffer(
379    seen: &mut BTreeMap<(*const u8, usize), u32>,
380    buffers: &mut Vec<Buffer>,
381    buffer: &Buffer,
382) -> u32 {
383    *seen
384        .entry((buffer.as_ptr(), buffer.len()))
385        .or_insert_with(|| {
386            buffers.push(buffer.clone());
387            (buffers.len() - 1) as u32
388        })
389}
390
391fn interleave_struct(
392    fields: &Fields,
393    values: &[&dyn Array],
394    indices: &[(usize, usize)],
395) -> Result<ArrayRef, ArrowError> {
396    let interleaved = Interleave::<'_, StructArray>::new(values, indices);
397
398    if fields.is_empty() {
399        let array = StructArray::try_new_with_length(
400            fields.clone(),
401            vec![],
402            interleaved.nulls,
403            indices.len(),
404        )?;
405        return Ok(Arc::new(array));
406    }
407
408    let struct_fields_array: Result<Vec<_>, _> = (0..fields.len())
409        .map(|i| {
410            let field_values: Vec<&dyn Array> = interleaved
411                .arrays
412                .iter()
413                .map(|x| x.column(i).as_ref())
414                .collect();
415            interleave(&field_values, indices)
416        })
417        .collect();
418
419    let struct_array =
420        StructArray::try_new(fields.clone(), struct_fields_array?, interleaved.nulls)?;
421    Ok(Arc::new(struct_array))
422}
423
424fn interleave_list_like_primitive_child<L: ListLikeArray, T: ArrowPrimitiveType>(
425    interleaved: &Interleave<'_, L>,
426    indices: &[(usize, usize)],
427    capacity: usize,
428    data_type: &DataType,
429) -> Result<ArrayRef, ArrowError> {
430    let child_arrays: Vec<&PrimitiveArray<T>> = interleaved
431        .arrays
432        .iter()
433        .map(|list| list.values().as_primitive::<T>())
434        .collect();
435
436    let has_child_nulls = child_arrays.iter().any(|a| a.null_count() > 0);
437
438    // Build values buffer by copying contiguous slices
439    let mut values: Vec<T::Native> = Vec::with_capacity(capacity);
440    for &(array, row) in indices {
441        let range = interleaved.arrays[array].element_range(row);
442        if !range.is_empty() {
443            values.extend_from_slice(&child_arrays[array].values()[range]);
444        }
445    }
446
447    // Build null buffer. Pre-allocate with 0x00 (all null), then:
448    // - Sources with nulls: set_bits copies the source validity bits into the destination range.
449    // - Sources without nulls: set the bit range to all 1s directly.
450    let nulls = if has_child_nulls {
451        let null_byte_len = bit_util::ceil(capacity, 8);
452        let mut output_null_buf = MutableBuffer::try_from_len_zeroed(null_byte_len)
453            .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
454
455        let mut offset_write = 0;
456        let mut output_null_count = 0usize;
457        for &(array, row) in indices {
458            let range = interleaved.arrays[array].element_range(row);
459            let len = range.len();
460            if len > 0 {
461                match child_arrays[array].nulls() {
462                    Some(null_buffer) => {
463                        output_null_count += set_bits(
464                            output_null_buf.as_slice_mut(),
465                            null_buffer.validity(),
466                            offset_write,
467                            null_buffer.offset() + range.start,
468                            len,
469                        );
470                    }
471                    None => {
472                        // For a non-nullable source, set the bit range to all 1s directly.
473                        let buf = output_null_buf.as_slice_mut();
474                        (offset_write..offset_write + len).for_each(|i| bit_util::set_bit(buf, i));
475                    }
476                }
477            }
478            offset_write += len;
479        }
480
481        if output_null_count > 0 {
482            let bool_buf = BooleanBuffer::new(output_null_buf.into(), 0, capacity);
483            // SAFETY: null_count is accumulated from set_bits which correctly counts unset bits
484            Some(unsafe { NullBuffer::new_unchecked(bool_buf, output_null_count) })
485        } else {
486            None
487        }
488    } else {
489        None
490    };
491
492    Ok(Arc::new(
493        PrimitiveArray::<T>::new(values.into(), nulls).with_data_type(data_type.clone()),
494    ))
495}
496
497/// Interleave child values for non-primitive child types, shared by List and FixedSizeList.
498fn interleave_list_like_child<L: ListLikeArray>(
499    interleaved: &Interleave<'_, L>,
500    indices: &[(usize, usize)],
501    capacity: usize,
502) -> Result<ArrayRef, ArrowError> {
503    let mut child_indices = Vec::with_capacity(capacity);
504    for &(array, row) in indices {
505        let range = interleaved.arrays[array].element_range(row);
506        child_indices.extend(range.map(|i| (array, i)));
507    }
508
509    let child_arrays: Vec<&dyn Array> = interleaved
510        .arrays
511        .iter()
512        .map(|list| list.values().as_ref())
513        .collect();
514    interleave(&child_arrays, &child_indices)
515}
516
517fn interleave_list<O: OffsetSizeTrait>(
518    values: &[&dyn Array],
519    indices: &[(usize, usize)],
520    field: &FieldRef,
521) -> Result<ArrayRef, ArrowError> {
522    let interleaved = Interleave::<'_, GenericListArray<O>>::new(values, indices);
523
524    // Step 1: compute output offsets and total child capacity
525    let mut capacity = 0usize;
526    let mut offsets = Vec::with_capacity(indices.len() + 1);
527    offsets.push(O::from_usize(0).unwrap());
528    for (array, row) in indices {
529        let o = interleaved.arrays[*array].value_offsets();
530        let element_len = o[*row + 1].as_usize() - o[*row].as_usize();
531        capacity += element_len;
532        offsets.push(
533            O::from_usize(capacity).ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
534        );
535    }
536
537    // Step 2: build child values.
538    macro_rules! list_primitive_helper {
539        ($t:ty) => {
540            interleave_list_like_primitive_child::<GenericListArray<O>, $t>(
541                &interleaved,
542                indices,
543                capacity,
544                field.data_type(),
545            )?
546        };
547    }
548
549    let child_values = downcast_primitive! {
550        // For primitive child types, directly copy typed value slices and null bit
551        // ranges, avoiding both the intermediate child_indices Vec allocation and
552        // MutableArrayData's function pointer indirection.
553        field.data_type() => (list_primitive_helper),
554        _ => {
555            interleave_list_like_child(&interleaved, indices, capacity)?
556        }
557    };
558
559    let offsets = OffsetBuffer::new(offsets.into());
560    let list_array =
561        GenericListArray::<O>::new(field.clone(), offsets, child_values, interleaved.nulls);
562
563    Ok(Arc::new(list_array))
564}
565
566fn interleave_fixed_size_list(
567    values: &[&dyn Array],
568    indices: &[(usize, usize)],
569    field: &FieldRef,
570    size: i32,
571) -> Result<ArrayRef, ArrowError> {
572    let interleaved = Interleave::<'_, FixedSizeListArray>::new(values, indices);
573    let capacity = indices
574        .len()
575        .checked_mul(size as usize)
576        .ok_or_else(|| ArrowError::MemoryError("interleave output size overflow".into()))?;
577
578    macro_rules! fsl_primitive_helper {
579        ($t:ty) => {
580            interleave_list_like_primitive_child::<FixedSizeListArray, $t>(
581                &interleaved,
582                indices,
583                capacity,
584                field.data_type(),
585            )?
586        };
587    }
588
589    let interleaved_values = downcast_primitive! {
590        field.data_type() => (fsl_primitive_helper),
591        _ => {
592            interleave_list_like_child(&interleaved, indices, capacity)?
593        }
594    };
595
596    let array = FixedSizeListArray::try_new_with_length(
597        field.clone(),
598        size,
599        interleaved_values,
600        interleaved.nulls,
601        indices.len(),
602    )?;
603    Ok(Arc::new(array))
604}
605
606fn interleave_map(
607    values: &[&dyn Array],
608    indices: &[(usize, usize)],
609    field: &FieldRef,
610    ordered: bool,
611) -> Result<ArrayRef, ArrowError> {
612    let interleaved = Interleave::<'_, MapArray>::new(values, indices);
613
614    let mut capacity = 0usize;
615    let mut offsets = Vec::with_capacity(indices.len() + 1);
616    offsets.push(0i32);
617    for &(array, row) in indices {
618        let o = interleaved.arrays[array].value_offsets();
619        let element_len = (o[row + 1] - o[row]) as usize;
620        capacity += element_len;
621        offsets
622            .push(i32::try_from(capacity).map_err(|_| ArrowError::OffsetOverflowError(capacity))?);
623    }
624
625    let mut child_indices = Vec::with_capacity(capacity);
626    for &(array, row) in indices {
627        let o = interleaved.arrays[array].value_offsets();
628        let start = o[row] as usize;
629        let end = o[row + 1] as usize;
630        child_indices.extend((start..end).map(|i| (array, i)));
631    }
632
633    let entries_arrays: Vec<&dyn Array> = interleaved
634        .arrays
635        .iter()
636        .map(|m| m.entries() as &dyn Array)
637        .collect();
638    let interleaved_entries = interleave(&entries_arrays, &child_indices)?;
639
640    let offsets = OffsetBuffer::new(offsets.into());
641    let entries = interleaved_entries.as_struct().clone();
642    let array = MapArray::new(field.clone(), offsets, entries, interleaved.nulls, ordered);
643    Ok(Arc::new(array))
644}
645
646/// Specialized [`interleave`] for [`RunArray`].
647fn interleave_run_end<R: RunEndIndexType>(
648    values: &[&dyn Array],
649    indices: &[(usize, usize)],
650) -> Result<ArrayRef, ArrowError> {
651    if indices.is_empty() {
652        return Ok(new_empty_array(values[0].data_type()));
653    }
654
655    let n = indices.len();
656    R::Native::from_usize(n).ok_or_else(|| {
657        ArrowError::ComputeError(format!(
658            "interleave_run_end: output length {n} does not fit run-end type"
659        ))
660    })?;
661
662    let runs: Vec<&RunArray<R>> = values.iter().map(|a| a.as_run::<R>()).collect();
663    let value_arrays: Vec<&dyn Array> = runs.iter().map(|r| r.values().as_ref()).collect();
664
665    // Resolve each (array, logical_row) to (array, physical_row), so we can
666    // lookup physical indices by batch.
667    let mut phys_pairs: Vec<(usize, usize)> = vec![(0, 0); n];
668    let mut grouped: Vec<(Vec<R::Native>, Vec<usize>)> =
669        (0..runs.len()).map(|_| (Vec::new(), Vec::new())).collect();
670    for (out_pos, &(arr, row)) in indices.iter().enumerate() {
671        let row = R::Native::from_usize(row).ok_or_else(|| {
672            ArrowError::InvalidArgumentError(format!(
673                "interleave_run_end: row index {row} not representable as run-end type {}",
674                R::DATA_TYPE
675            ))
676        })?;
677        grouped[arr].0.push(row);
678        grouped[arr].1.push(out_pos);
679    }
680    for (arr_idx, (logical_rows, out_positions)) in grouped.into_iter().enumerate() {
681        let phys = runs[arr_idx].get_physical_indices(&logical_rows)?;
682        for (p, out_pos) in phys.iter().zip(out_positions.iter()) {
683            phys_pairs[*out_pos] = (arr_idx, *p);
684        }
685    }
686
687    // Coalesce by physical-pair equality only: emit a new run when the
688    // (array_idx, physical_idx) pair changes between adjacent output rows.
689    // TODO: We could perform an equality check across sources to extend the
690    // output run, but we can't call make_comparator from this crate.
691    let mut run_ends_buf: Vec<R::Native> = Vec::with_capacity(n);
692    let mut dedup_pairs: Vec<(usize, usize)> = Vec::with_capacity(n);
693    dedup_pairs.push(phys_pairs[0]);
694    for i in 1..n {
695        if phys_pairs[i] != phys_pairs[i - 1] {
696            run_ends_buf.push(R::Native::from_usize(i).unwrap());
697            dedup_pairs.push(phys_pairs[i]);
698        }
699    }
700    run_ends_buf.push(R::Native::from_usize(n).unwrap());
701
702    let taken_values = interleave(&value_arrays, &dedup_pairs)?;
703    let run_ends = PrimitiveArray::<R>::from_iter_values(run_ends_buf);
704
705    Ok(Arc::new(RunArray::<R>::try_new(
706        &run_ends,
707        taken_values.as_ref(),
708    )?))
709}
710
711fn interleave_list_view<O: OffsetSizeTrait>(
712    values: &[&dyn Array],
713    indices: &[(usize, usize)],
714    field: &FieldRef,
715) -> Result<ArrayRef, ArrowError> {
716    let interleaved = Interleave::<'_, GenericListViewArray<O>>::new(values, indices);
717
718    // Pick whichever strategy produces fewer child elements:
719    // - Per-row copy: total = sum of selected sizes. Better for sparse selections.
720    // - Concat + offset adjustment: total = sum of source backing array lengths.
721    //   Better when rows share backing elements via overlapping offset/size ranges.
722    let concat_cost: usize = interleaved.arrays.iter().map(|lv| lv.values().len()).sum();
723    let per_row_cost: usize = indices
724        .iter()
725        .map(|&(a, r)| interleaved.arrays[a].sizes()[r].as_usize())
726        .sum();
727
728    if per_row_cost <= concat_cost {
729        interleave_list_view_copy::<O>(&interleaved, indices, field)
730    } else {
731        interleave_list_view_concat::<O>(&interleaved, indices, field)
732    }
733}
734
735/// Per-row copy: copies each selected row's child elements into a new flat array.
736fn interleave_list_view_copy<O: OffsetSizeTrait>(
737    interleaved: &Interleave<'_, GenericListViewArray<O>>,
738    indices: &[(usize, usize)],
739    field: &FieldRef,
740) -> Result<ArrayRef, ArrowError> {
741    let mut capacity = 0usize;
742    let mut offsets = Vec::with_capacity(indices.len());
743    let mut sizes = Vec::with_capacity(indices.len());
744    for &(array_idx, row_idx) in indices {
745        let list = interleaved.arrays[array_idx];
746        let size = list.sizes()[row_idx].as_usize();
747        offsets.push(
748            O::from_usize(capacity).ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
749        );
750        sizes.push(O::from_usize(size).ok_or_else(|| ArrowError::OffsetOverflowError(size))?);
751        capacity += size;
752    }
753
754    let child_data: Vec<_> = interleaved
755        .arrays
756        .iter()
757        .map(|list| list.values().to_data())
758        .collect();
759    let child_data_refs: Vec<_> = child_data.iter().collect();
760    let mut mutable_child = MutableArrayData::try_new(child_data_refs, false, capacity)?;
761    for &(array_idx, row_idx) in indices {
762        let list = interleaved.arrays[array_idx];
763        let start = list.offsets()[row_idx].as_usize();
764        let size = list.sizes()[row_idx].as_usize();
765        if size > 0 {
766            mutable_child.try_extend(array_idx, start, start + size)?;
767        }
768    }
769
770    Ok(Arc::new(GenericListViewArray::<O>::new(
771        field.clone(),
772        offsets.into(),
773        sizes.into(),
774        make_array(mutable_child.freeze()),
775        interleaved.nulls.clone(),
776    )))
777}
778
779/// Concat backing arrays: concatenates all source value arrays and adjusts offsets.
780/// Preserves within-source element sharing.
781fn interleave_list_view_concat<O: OffsetSizeTrait>(
782    interleaved: &Interleave<'_, GenericListViewArray<O>>,
783    indices: &[(usize, usize)],
784    field: &FieldRef,
785) -> Result<ArrayRef, ArrowError> {
786    let child_arrays: Vec<&dyn Array> = interleaved
787        .arrays
788        .iter()
789        .map(|lv| lv.values().as_ref())
790        .collect();
791    let mut base_offsets = Vec::with_capacity(interleaved.arrays.len());
792    let mut running = 0usize;
793    for lv in &interleaved.arrays {
794        base_offsets.push(running);
795        running += lv.values().len();
796    }
797    let combined_values = concat(&child_arrays)?;
798
799    let mut new_offsets = Vec::with_capacity(indices.len());
800    let mut new_sizes = Vec::with_capacity(indices.len());
801    for &(array_idx, row_idx) in indices {
802        let lv = interleaved.arrays[array_idx];
803        let adjusted = lv.offsets()[row_idx].as_usize() + base_offsets[array_idx];
804        new_offsets.push(
805            O::from_usize(adjusted).ok_or_else(|| ArrowError::OffsetOverflowError(adjusted))?,
806        );
807        new_sizes.push(lv.sizes()[row_idx]);
808    }
809
810    Ok(Arc::new(GenericListViewArray::<O>::new(
811        field.clone(),
812        new_offsets.into(),
813        new_sizes.into(),
814        combined_values,
815        interleaved.nulls.clone(),
816    )))
817}
818
819/// Fallback implementation of interleave using [`MutableArrayData`]
820fn interleave_fallback(
821    values: &[&dyn Array],
822    indices: &[(usize, usize)],
823) -> Result<ArrayRef, ArrowError> {
824    let arrays: Vec<_> = values.iter().map(|x| x.to_data()).collect();
825    let arrays: Vec<_> = arrays.iter().collect();
826    let mut array_data = MutableArrayData::try_new(arrays, false, indices.len())?;
827
828    let mut cur_array = indices[0].0;
829    let mut start_row_idx = indices[0].1;
830    let mut end_row_idx = start_row_idx + 1;
831
832    for (array, row) in indices.iter().skip(1).copied() {
833        if array == cur_array && row == end_row_idx {
834            // subsequent row in same batch
835            end_row_idx += 1;
836            continue;
837        }
838
839        // emit current batch of rows for current buffer
840        array_data.try_extend(cur_array, start_row_idx, end_row_idx)?;
841
842        // start new batch of rows
843        cur_array = array;
844        start_row_idx = row;
845        end_row_idx = start_row_idx + 1;
846    }
847
848    // emit final batch of rows
849    array_data.try_extend(cur_array, start_row_idx, end_row_idx)?;
850    Ok(make_array(array_data.freeze()))
851}
852
853/// Fallback implementation for interleaving dictionaries when it was determined
854/// that the dictionary values should not be merged. This implementation concatenates
855/// the value slices and recomputes the resulting dictionary keys.
856///
857/// # Panics
858///
859/// This function assumes that the combined dictionary values will not overflow the
860/// key type. Callers must verify this condition [`should_merge_dictionary_values`]
861/// before calling this function.
862fn interleave_fallback_dictionary<K: ArrowDictionaryKeyType>(
863    dictionaries: &[&DictionaryArray<K>],
864    indices: &[(usize, usize)],
865) -> Result<ArrayRef, ArrowError> {
866    let relative_offsets: Vec<usize> = dictionaries
867        .iter()
868        .scan(0usize, |offset, dict| {
869            let current = *offset;
870            *offset += dict.values().len();
871            Some(current)
872        })
873        .collect();
874    let all_values: Vec<&dyn Array> = dictionaries.iter().map(|d| d.values().as_ref()).collect();
875    let concatenated_values = concat(&all_values)?;
876
877    let any_nulls = dictionaries.iter().any(|d| d.keys().nulls().is_some());
878    let (new_keys, nulls) = if any_nulls {
879        let mut has_nulls = false;
880        let new_keys: Vec<K::Native> = indices
881            .iter()
882            .map(|(array, row)| {
883                let old_keys = dictionaries[*array].keys();
884                if old_keys.is_valid(*row) {
885                    let old_key = old_keys.values()[*row].as_usize();
886                    K::Native::from_usize(relative_offsets[*array] + old_key)
887                        .expect("key overflow should be checked by caller")
888                } else {
889                    has_nulls = true;
890                    K::Native::ZERO
891                }
892            })
893            .collect();
894
895        let nulls = if has_nulls {
896            let null_buffer = BooleanBuffer::collect_bool(indices.len(), |i| {
897                let (array, row) = indices[i];
898                dictionaries[array].keys().is_valid(row)
899            });
900            Some(NullBuffer::new(null_buffer))
901        } else {
902            None
903        };
904        (new_keys, nulls)
905    } else {
906        let new_keys: Vec<K::Native> = indices
907            .iter()
908            .map(|(array, row)| {
909                let old_key = dictionaries[*array].keys().values()[*row].as_usize();
910                K::Native::from_usize(relative_offsets[*array] + old_key)
911                    .expect("key overflow should be checked by caller")
912            })
913            .collect();
914        (new_keys, None)
915    };
916
917    let keys_array = PrimitiveArray::<K>::new(new_keys.into(), nulls);
918    // SAFETY: keys_array is constructed from a valid set of keys.
919    let array = unsafe { DictionaryArray::new_unchecked(keys_array, concatenated_values) };
920    Ok(Arc::new(array))
921}
922
923/// Interleave rows by index from multiple [`RecordBatch`] instances and return a new [`RecordBatch`].
924///
925/// This function will call [`interleave`] on each array of the [`RecordBatch`] instances and assemble a new [`RecordBatch`].
926///
927/// # Example
928/// ```
929/// # use std::sync::Arc;
930/// # use arrow_array::{StringArray, Int32Array, RecordBatch, UInt32Array};
931/// # use arrow_schema::{DataType, Field, Schema};
932/// # use arrow_select::interleave::interleave_record_batch;
933///
934/// let schema = Arc::new(Schema::new(vec![
935///     Field::new("a", DataType::Int32, true),
936///     Field::new("b", DataType::Utf8, true),
937/// ]));
938///
939/// let batch1 = RecordBatch::try_new(
940///     schema.clone(),
941///     vec![
942///         Arc::new(Int32Array::from(vec![0, 1, 2])),
943///         Arc::new(StringArray::from(vec!["a", "b", "c"])),
944///     ],
945/// ).unwrap();
946///
947/// let batch2 = RecordBatch::try_new(
948///     schema.clone(),
949///     vec![
950///         Arc::new(Int32Array::from(vec![3, 4, 5])),
951///         Arc::new(StringArray::from(vec!["d", "e", "f"])),
952///     ],
953/// ).unwrap();
954///
955/// let indices = vec![(0, 1), (1, 2), (0, 0), (1, 1)];
956/// let interleaved = interleave_record_batch(&[&batch1, &batch2], &indices).unwrap();
957///
958/// let expected = RecordBatch::try_new(
959///     schema,
960///     vec![
961///         Arc::new(Int32Array::from(vec![1, 5, 0, 4])),
962///         Arc::new(StringArray::from(vec!["b", "f", "a", "e"])),
963///     ],
964/// ).unwrap();
965/// assert_eq!(interleaved, expected);
966/// ```
967pub fn interleave_record_batch(
968    record_batches: &[&RecordBatch],
969    indices: &[(usize, usize)],
970) -> Result<RecordBatch, ArrowError> {
971    let Some(first) = record_batches.first() else {
972        return Err(ArrowError::InvalidArgumentError(
973            "interleave_record_batch requires input of at least one record batch".to_string(),
974        ));
975    };
976    let schema = first.schema();
977    let columns = (0..schema.fields().len())
978        .map(|i| {
979            let column_values: Vec<&dyn Array> = record_batches
980                .iter()
981                .map(|batch| batch.column(i).as_ref())
982                .collect();
983            interleave(&column_values, indices)
984        })
985        .collect::<Result<Vec<_>, _>>()?;
986    // The row count is set explicitly so that a batch with no columns still
987    // produces a batch with `indices.len()` rows, like `take_record_batch`.
988    RecordBatch::try_new_with_options(
989        schema,
990        columns,
991        &RecordBatchOptions::new().with_row_count(Some(indices.len())),
992    )
993}
994
995#[cfg(test)]
996mod tests {
997    use super::*;
998    use arrow_array::Int32RunArray;
999    use arrow_array::builder::{
1000        GenericListBuilder, Int32Builder, PrimitiveBuilder, PrimitiveRunBuilder,
1001    };
1002    use arrow_array::types::{Decimal128Type, Int8Type, TimestampMicrosecondType};
1003    use arrow_buffer::ScalarBuffer;
1004    use arrow_schema::{Field, Schema, TimeUnit};
1005
1006    #[test]
1007    fn test_primitive() {
1008        let a = Int32Array::from_iter_values([1, 2, 3, 4]);
1009        let b = Int32Array::from_iter_values([5, 6, 7]);
1010        let c = Int32Array::from_iter_values([8, 9, 10]);
1011        let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap();
1012        let v = values.as_primitive::<Int32Type>();
1013        assert_eq!(v.values(), &[4, 4, 10, 8, 6]);
1014    }
1015
1016    #[test]
1017    fn test_primitive_nulls() {
1018        let a = Int32Array::from_iter_values([1, 2, 3, 4]);
1019        let b = Int32Array::from_iter([Some(1), Some(4), None]);
1020        let values = interleave(&[&a, &b], &[(0, 1), (1, 2), (1, 2), (0, 3), (0, 2)]).unwrap();
1021        let v: Vec<_> = values.as_primitive::<Int32Type>().into_iter().collect();
1022        assert_eq!(&v, &[Some(2), None, None, Some(4), Some(3)])
1023    }
1024
1025    #[test]
1026    fn test_interleave_boolean() {
1027        let a = BooleanArray::from(vec![Some(true), Some(false), None, Some(true)]);
1028        let b = BooleanArray::from(vec![Some(false), Some(true), Some(true), None]);
1029
1030        let values = interleave(&[&a, &b], &[(0, 0), (1, 1), (0, 2), (1, 3), (1, 2)]).unwrap();
1031
1032        let result = values.as_boolean();
1033
1034        assert_eq!(
1035            result.iter().collect::<Vec<_>>(),
1036            vec![Some(true), Some(true), None, None, Some(true),]
1037        );
1038    }
1039
1040    #[test]
1041    fn test_interleave_boolean_sliced() {
1042        let original = BooleanArray::from(vec![
1043            Some(false),
1044            Some(true),
1045            Some(false),
1046            Some(true),
1047            None,
1048            Some(false),
1049            Some(true),
1050        ]);
1051
1052        let sliced_array = original.slice(1, 5);
1053        let sliced = sliced_array
1054            .as_any()
1055            .downcast_ref::<BooleanArray>()
1056            .unwrap();
1057        let other = BooleanArray::from(vec![Some(false), Some(false), Some(true)]);
1058
1059        let values =
1060            interleave(&[sliced, &other], &[(0, 0), (1, 2), (0, 3), (0, 4), (1, 0)]).unwrap();
1061
1062        let result = values.as_boolean();
1063
1064        assert_eq!(
1065            result.iter().collect::<Vec<_>>(),
1066            vec![Some(true), Some(true), None, Some(false), Some(false),]
1067        );
1068    }
1069
1070    #[test]
1071    fn test_interleave_boolean_fallback_matches_specialized() {
1072        let a = BooleanArray::from_iter(
1073            (0..1024).map(|i| if i % 10 == 0 { None } else { Some(i % 2 == 0) }),
1074        );
1075        let b = BooleanArray::from_iter((0..1024).map(|i| Some(i % 3 == 0)));
1076
1077        let indices: Vec<_> = (0..1024).map(|i| (i % 2, (i * 37) % 1024)).collect();
1078
1079        let values: [&dyn Array; 2] = [&a, &b];
1080
1081        let specialized = interleave_boolean(&values, &indices).unwrap();
1082        let fallback = interleave_fallback(&values, &indices).unwrap();
1083
1084        assert_eq!(specialized.as_ref(), fallback.as_ref());
1085    }
1086
1087    #[test]
1088    fn test_primitive_empty() {
1089        let a = Int32Array::from_iter_values([1, 2, 3, 4]);
1090        let v = interleave(&[&a], &[]).unwrap();
1091        assert!(v.is_empty());
1092        assert_eq!(v.data_type(), &DataType::Int32);
1093    }
1094
1095    #[test]
1096    fn test_strings() {
1097        let a = StringArray::from_iter_values(["a", "b", "c"]);
1098        let b = StringArray::from_iter_values(["hello", "world", "foo"]);
1099        let values = interleave(&[&a, &b], &[(0, 2), (0, 2), (1, 0), (1, 1), (0, 1)]).unwrap();
1100        let v = values.as_string::<i32>();
1101        let values: Vec<_> = v.into_iter().collect();
1102        assert_eq!(
1103            &values,
1104            &[
1105                Some("c"),
1106                Some("c"),
1107                Some("hello"),
1108                Some("world"),
1109                Some("b")
1110            ]
1111        )
1112    }
1113
1114    #[test]
1115    fn test_interleave_dictionary() {
1116        let a = DictionaryArray::<Int32Type>::from_iter(["a", "b", "c", "a", "b"]);
1117        let b = DictionaryArray::<Int32Type>::from_iter(["a", "c", "a", "c", "a"]);
1118
1119        // Should not recompute dictionary
1120        let values =
1121            interleave(&[&a, &b], &[(0, 2), (0, 2), (0, 2), (1, 0), (1, 1), (0, 1)]).unwrap();
1122        let v = values.as_dictionary::<Int32Type>();
1123        assert_eq!(v.values().len(), 5);
1124
1125        let vc = v.downcast_dict::<StringArray>().unwrap();
1126        let collected: Vec<_> = vc.into_iter().map(Option::unwrap).collect();
1127        assert_eq!(&collected, &["c", "c", "c", "a", "c", "b"]);
1128
1129        // Should recompute dictionary
1130        let values = interleave(&[&a, &b], &[(0, 2), (0, 2), (1, 1)]).unwrap();
1131        let v = values.as_dictionary::<Int32Type>();
1132        assert_eq!(v.values().len(), 1);
1133
1134        let vc = v.downcast_dict::<StringArray>().unwrap();
1135        let collected: Vec<_> = vc.into_iter().map(Option::unwrap).collect();
1136        assert_eq!(&collected, &["c", "c", "c"]);
1137    }
1138
1139    #[test]
1140    fn test_interleave_sparse_primitive_dictionaries_with_nulls() {
1141        let values_len = 2 * crate::dictionary::MAX_DENSE_MAPPING_VALUES_PER_ROW + 1;
1142        let last = values_len as i32 - 1;
1143        let make_dictionary = |offset, null_value, keys| {
1144            let values = Int32Array::from_iter(
1145                (0..values_len).map(|idx| (idx != null_value).then_some(offset + idx as i32)),
1146            );
1147            DictionaryArray::<Int32Type>::new(Int32Array::from(keys), Arc::new(values))
1148        };
1149        let a = make_dictionary(0, 1, vec![Some(last), None, Some(1)]);
1150        let b = make_dictionary(10_000, 2, vec![Some(2), Some(last - 1), None]);
1151        let indices = [(0, 2), (1, 0), (0, 0), (1, 2)];
1152        let masks = [
1153            BooleanBuffer::from_iter([true, false, true]),
1154            BooleanBuffer::from_iter([true, false, true]),
1155        ];
1156
1157        let merged = merge_dictionary_values(&[&a, &b], Some(&masks), Some(&[2, 2])).unwrap();
1158        assert!(
1159            merged
1160                .key_mappings
1161                .iter()
1162                .all(|mapping| matches!(mapping, crate::dictionary::KeyMapping::Sparse(_)))
1163        );
1164
1165        let result = interleave(&[&a, &b], &indices).unwrap();
1166        let actual = result
1167            .as_dictionary::<Int32Type>()
1168            .downcast_dict::<Int32Array>()
1169            .unwrap()
1170            .into_iter()
1171            .collect::<Vec<_>>();
1172
1173        assert_eq!(actual, [None, None, Some(last), None]);
1174    }
1175
1176    #[test]
1177    fn test_interleave_dictionary_nulls() {
1178        let input_1_keys = Int32Array::from_iter_values([0, 2, 1, 3]);
1179        let input_1_values = StringArray::from(vec![Some("foo"), None, Some("bar"), Some("fiz")]);
1180        let input_1 = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
1181        let input_2: DictionaryArray<Int32Type> = vec![None].into_iter().collect();
1182
1183        let expected = vec![Some("fiz"), None, None, Some("foo")];
1184
1185        let values = interleave(
1186            &[&input_1 as _, &input_2 as _],
1187            &[(0, 3), (0, 2), (1, 0), (0, 0)],
1188        )
1189        .unwrap();
1190        let dictionary = values.as_dictionary::<Int32Type>();
1191        let actual: Vec<Option<&str>> = dictionary
1192            .downcast_dict::<StringArray>()
1193            .unwrap()
1194            .into_iter()
1195            .collect();
1196
1197        assert_eq!(actual, expected);
1198    }
1199
1200    #[test]
1201    fn test_interleave_dictionary_overflow_same_values() {
1202        let values: ArrayRef = Arc::new(StringArray::from_iter_values(
1203            (0..50).map(|i| format!("v{i}")),
1204        ));
1205
1206        // With 3 dictionaries of 50 values each, relative_offsets = [0, 50, 100]
1207        // Accessing key 49 from dict3 gives 100 + 49 = 149 which overflows Int8
1208        // (max 127).
1209        // This test case falls back to interleave_fallback because the
1210        // dictionaries share the same underlying values slice.
1211        let dict1 = DictionaryArray::<Int8Type>::new(
1212            Int8Array::from_iter_values([0, 1, 2]),
1213            values.clone(),
1214        );
1215        let dict2 = DictionaryArray::<Int8Type>::new(
1216            Int8Array::from_iter_values([0, 1, 2]),
1217            values.clone(),
1218        );
1219        let dict3 =
1220            DictionaryArray::<Int8Type>::new(Int8Array::from_iter_values([49]), values.clone());
1221
1222        let indices = &[(0, 0), (1, 0), (2, 0)];
1223        let result = interleave(&[&dict1, &dict2, &dict3], indices).unwrap();
1224
1225        let dict_result = result.as_dictionary::<Int8Type>();
1226        let string_result: Vec<_> = dict_result
1227            .downcast_dict::<StringArray>()
1228            .unwrap()
1229            .into_iter()
1230            .map(|x| x.unwrap())
1231            .collect();
1232        assert_eq!(string_result, vec!["v0", "v0", "v49"]);
1233    }
1234
1235    fn test_interleave_lists<O: OffsetSizeTrait>() {
1236        // [[1, 2], null, [3]]
1237        let mut a = GenericListBuilder::<O, _>::new(Int32Builder::new());
1238        a.values().append_value(1);
1239        a.values().append_value(2);
1240        a.append(true);
1241        a.append(false);
1242        a.values().append_value(3);
1243        a.append(true);
1244        let a = a.finish();
1245
1246        // [[4], null, [5, 6, null]]
1247        let mut b = GenericListBuilder::<O, _>::new(Int32Builder::new());
1248        b.values().append_value(4);
1249        b.append(true);
1250        b.append(false);
1251        b.values().append_value(5);
1252        b.values().append_value(6);
1253        b.values().append_null();
1254        b.append(true);
1255        let b = b.finish();
1256
1257        let values = interleave(&[&a, &b], &[(0, 2), (0, 1), (1, 0), (1, 2), (1, 1)]).unwrap();
1258        let v = values
1259            .as_any()
1260            .downcast_ref::<GenericListArray<O>>()
1261            .unwrap();
1262
1263        // [[3], null, [4], [5, 6, null], null]
1264        let mut expected = GenericListBuilder::<O, _>::new(Int32Builder::new());
1265        expected.values().append_value(3);
1266        expected.append(true);
1267        expected.append(false);
1268        expected.values().append_value(4);
1269        expected.append(true);
1270        expected.values().append_value(5);
1271        expected.values().append_value(6);
1272        expected.values().append_null();
1273        expected.append(true);
1274        expected.append(false);
1275        let expected = expected.finish();
1276
1277        assert_eq!(v, &expected);
1278    }
1279
1280    #[test]
1281    fn test_lists() {
1282        test_interleave_lists::<i32>();
1283    }
1284
1285    #[test]
1286    fn test_large_lists() {
1287        test_interleave_lists::<i64>();
1288    }
1289
1290    /// One list slot in a `List<Primitive>` fixture: `None` is a null slot,
1291    /// `Some(items)` is a list whose items may individually be null.
1292    type ListRow<T> = Option<Vec<Option<<T as ArrowPrimitiveType>::Native>>>;
1293
1294    /// Build a `List<Primitive>` from row fixtures. The primitive child carries
1295    /// `data_type` (e.g. its Decimal scale or timezone).
1296    fn list_of_primitive<O: OffsetSizeTrait, T: ArrowPrimitiveType>(
1297        data_type: &DataType,
1298        rows: &[ListRow<T>],
1299    ) -> GenericListArray<O> {
1300        let mut builder = GenericListBuilder::<O, _>::new(
1301            PrimitiveBuilder::<T>::new().with_data_type(data_type.clone()),
1302        );
1303        for row in rows {
1304            match row {
1305                Some(items) => {
1306                    items
1307                        .iter()
1308                        .for_each(|v| builder.values().append_option(*v));
1309                    builder.append(true);
1310                }
1311                None => builder.append(false),
1312            }
1313        }
1314        builder.finish()
1315    }
1316
1317    /// Interleave list fixtures and assert both the result and that the
1318    /// interleaved primitive child preserves the parameterized `data_type`.
1319    fn check_interleave_list_primitive<O: OffsetSizeTrait, T: ArrowPrimitiveType>(
1320        data_type: &DataType,
1321        inputs: &[&[ListRow<T>]],
1322        indices: &[(usize, usize)],
1323        expected: &[ListRow<T>],
1324    ) {
1325        let arrays: Vec<_> = inputs
1326            .iter()
1327            .map(|rows| list_of_primitive::<O, T>(data_type, rows))
1328            .collect();
1329        let refs: Vec<&dyn Array> = arrays.iter().map(|a| a as &dyn Array).collect();
1330
1331        let values = interleave(&refs, indices).unwrap();
1332        let v = values
1333            .as_any()
1334            .downcast_ref::<GenericListArray<O>>()
1335            .unwrap();
1336
1337        assert_eq!(v, &list_of_primitive::<O, T>(data_type, expected));
1338        // The child's logical type (Decimal precision/scale, Timestamp timezone)
1339        // must be preserved, not reset to the primitive's default.
1340        assert_eq!(v.values().data_type(), data_type);
1341    }
1342
1343    fn test_interleave_lists_decimal<O: OffsetSizeTrait>() {
1344        // List<Decimal128(20, 3)>, exercising child-element nulls and null slots.
1345        check_interleave_list_primitive::<O, Decimal128Type>(
1346            &DataType::Decimal128(20, 3),
1347            &[
1348                &[
1349                    Some(vec![Some(1), Some(2)]),
1350                    None,
1351                    Some(vec![Some(3), None]),
1352                ], // a
1353                &[Some(vec![Some(4)]), Some(vec![Some(5), Some(6)])], // b
1354            ],
1355            &[(0, 2), (0, 1), (1, 0), (1, 1)],
1356            &[
1357                Some(vec![Some(3), None]),
1358                None,
1359                Some(vec![Some(4)]),
1360                Some(vec![Some(5), Some(6)]),
1361            ],
1362        );
1363    }
1364
1365    #[test]
1366    fn test_lists_decimal() {
1367        test_interleave_lists_decimal::<i32>();
1368        test_interleave_lists_decimal::<i64>();
1369    }
1370
1371    fn test_interleave_lists_timestamp_tz<O: OffsetSizeTrait>() {
1372        // List<Timestamp(Microsecond, "+08:00")>, checking the timezone survives.
1373        check_interleave_list_primitive::<O, TimestampMicrosecondType>(
1374            &DataType::Timestamp(TimeUnit::Microsecond, Some("+08:00".into())),
1375            &[&[Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]],
1376            &[(0, 1), (0, 0)],
1377            &[Some(vec![Some(3)]), Some(vec![Some(1), Some(2)])],
1378        );
1379    }
1380
1381    #[test]
1382    fn test_lists_timestamp_tz() {
1383        test_interleave_lists_timestamp_tz::<i32>();
1384        test_interleave_lists_timestamp_tz::<i64>();
1385    }
1386
1387    fn test_interleave_list_views<O: OffsetSizeTrait>() {
1388        // [[1, 2], null, [3]]
1389        let mut a = GenericListBuilder::<O, _>::new(Int32Builder::new());
1390        a.values().append_value(1);
1391        a.values().append_value(2);
1392        a.append(true);
1393        a.append(false);
1394        a.values().append_value(3);
1395        a.append(true);
1396        let a: GenericListViewArray<O> = a.finish().into();
1397
1398        // [[4], null, [5, 6, null]]
1399        let mut b = GenericListBuilder::<O, _>::new(Int32Builder::new());
1400        b.values().append_value(4);
1401        b.append(true);
1402        b.append(false);
1403        b.values().append_value(5);
1404        b.values().append_value(6);
1405        b.values().append_null();
1406        b.append(true);
1407        let b: GenericListViewArray<O> = b.finish().into();
1408
1409        let values = interleave(&[&a, &b], &[(0, 2), (0, 1), (1, 0), (1, 2), (1, 1)]).unwrap();
1410        let v = values
1411            .as_any()
1412            .downcast_ref::<GenericListViewArray<O>>()
1413            .unwrap();
1414
1415        // [[3], null, [4], [5, 6, null], null]
1416        let mut expected = GenericListBuilder::<O, _>::new(Int32Builder::new());
1417        expected.values().append_value(3);
1418        expected.append(true);
1419        expected.append(false);
1420        expected.values().append_value(4);
1421        expected.append(true);
1422        expected.values().append_value(5);
1423        expected.values().append_value(6);
1424        expected.values().append_null();
1425        expected.append(true);
1426        expected.append(false);
1427        let expected: GenericListViewArray<O> = expected.finish().into();
1428
1429        assert_eq!(v, &expected);
1430    }
1431
1432    #[test]
1433    fn test_list_views() {
1434        test_interleave_list_views::<i32>();
1435    }
1436
1437    #[test]
1438    fn test_large_list_views() {
1439        test_interleave_list_views::<i64>();
1440    }
1441
1442    #[test]
1443    fn test_interleave_list_view_overlapping() {
1444        let field = Arc::new(Field::new_list_field(DataType::Int64, false));
1445
1446        // lv_a: 10 rows, two groups of 5 sharing the same backing elements.
1447        //   rows 0-4 → offset 0, size 5 → [0,1,2,3,4]
1448        //   rows 5-9 → offset 5, size 5 → [5,6,7,8,9]
1449        let lv_a = ListViewArray::new(
1450            Arc::clone(&field),
1451            ScalarBuffer::from(vec![0i32, 0, 0, 0, 0, 5, 5, 5, 5, 5]),
1452            ScalarBuffer::from(vec![5i32; 10]),
1453            Arc::new(Int64Array::from_iter_values(0..10)),
1454            None,
1455        );
1456
1457        // lv_b: 8 rows, two groups of 4 sharing the same backing elements.
1458        //   rows 0-3 → offset 0, size 3 → [100,101,102]
1459        //   rows 4-7 → offset 3, size 3 → [103,104,105]
1460        let lv_b = ListViewArray::new(
1461            Arc::clone(&field),
1462            ScalarBuffer::from(vec![0i32, 0, 0, 0, 3, 3, 3, 3]),
1463            ScalarBuffer::from(vec![3i32; 8]),
1464            Arc::new(Int64Array::from_iter_values(100..106)),
1465            None,
1466        );
1467
1468        let indices: Vec<(usize, usize)> = vec![
1469            (0, 0),
1470            (1, 0),
1471            (0, 5),
1472            (1, 4),
1473            (0, 1),
1474            (1, 1),
1475            (0, 6),
1476            (1, 5),
1477        ];
1478        let result = interleave(&[&lv_a as &dyn Array, &lv_b as &dyn Array], &indices).unwrap();
1479        result
1480            .to_data()
1481            .validate_full()
1482            .expect("result must be valid");
1483
1484        let result_lv = result.as_list_view::<i32>();
1485        assert_eq!(result_lv.len(), 8);
1486        assert_eq!(
1487            result_lv.value(0).as_primitive::<Int64Type>().values(),
1488            &[0, 1, 2, 3, 4]
1489        );
1490        assert_eq!(
1491            result_lv.value(1).as_primitive::<Int64Type>().values(),
1492            &[100, 101, 102]
1493        );
1494        assert_eq!(
1495            result_lv.value(2).as_primitive::<Int64Type>().values(),
1496            &[5, 6, 7, 8, 9]
1497        );
1498        assert_eq!(
1499            result_lv.value(3).as_primitive::<Int64Type>().values(),
1500            &[103, 104, 105]
1501        );
1502
1503        // Backing elements = sum of source arrays (10 + 6 = 16), not per-row
1504        // expansion (8 rows × avg ~4 = 32). Overlapping sharing is preserved.
1505        let total_input_elements = lv_a.values().len() + lv_b.values().len();
1506        assert_eq!(result_lv.values().len(), total_input_elements);
1507    }
1508
1509    #[test]
1510    fn test_struct_without_nulls() {
1511        let fields = Fields::from(vec![
1512            Field::new("number_col", DataType::Int32, false),
1513            Field::new("string_col", DataType::Utf8, false),
1514        ]);
1515        let a = {
1516            let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1517            let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1518
1519            StructArray::try_new(
1520                fields.clone(),
1521                vec![Arc::new(number_col), Arc::new(string_col)],
1522                None,
1523            )
1524            .unwrap()
1525        };
1526
1527        let b = {
1528            let number_col = Int32Array::from_iter_values([5, 6, 7]);
1529            let string_col = StringArray::from_iter_values(["hello", "world", "foo"]);
1530
1531            StructArray::try_new(
1532                fields.clone(),
1533                vec![Arc::new(number_col), Arc::new(string_col)],
1534                None,
1535            )
1536            .unwrap()
1537        };
1538
1539        let c = {
1540            let number_col = Int32Array::from_iter_values([8, 9, 10]);
1541            let string_col = StringArray::from_iter_values(["x", "y", "z"]);
1542
1543            StructArray::try_new(
1544                fields.clone(),
1545                vec![Arc::new(number_col), Arc::new(string_col)],
1546                None,
1547            )
1548            .unwrap()
1549        };
1550
1551        let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap();
1552        let values_struct = values.as_struct();
1553        assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
1554        assert_eq!(values_struct.null_count(), 0);
1555
1556        let values_number = values_struct.column(0).as_primitive::<Int32Type>();
1557        assert_eq!(values_number.values(), &[4, 4, 10, 8, 6]);
1558        let values_string = values_struct.column(1).as_string::<i32>();
1559        let values_string: Vec<_> = values_string.into_iter().collect();
1560        assert_eq!(
1561            &values_string,
1562            &[Some("d"), Some("d"), Some("z"), Some("x"), Some("world")]
1563        );
1564    }
1565
1566    #[test]
1567    fn test_struct_with_nulls_in_values() {
1568        let fields = Fields::from(vec![
1569            Field::new("number_col", DataType::Int32, true),
1570            Field::new("string_col", DataType::Utf8, true),
1571        ]);
1572        let a = {
1573            let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1574            let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1575
1576            StructArray::try_new(
1577                fields.clone(),
1578                vec![Arc::new(number_col), Arc::new(string_col)],
1579                None,
1580            )
1581            .unwrap()
1582        };
1583
1584        let b = {
1585            let number_col = Int32Array::from_iter([Some(1), Some(4), None]);
1586            let string_col = StringArray::from(vec![Some("hello"), None, Some("foo")]);
1587
1588            StructArray::try_new(
1589                fields.clone(),
1590                vec![Arc::new(number_col), Arc::new(string_col)],
1591                None,
1592            )
1593            .unwrap()
1594        };
1595
1596        let values = interleave(&[&a, &b], &[(0, 1), (1, 2), (1, 2), (0, 3), (1, 1)]).unwrap();
1597        let values_struct = values.as_struct();
1598        assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
1599
1600        // The struct itself has no nulls, but the values do
1601        assert_eq!(values_struct.null_count(), 0);
1602
1603        let values_number: Vec<_> = values_struct
1604            .column(0)
1605            .as_primitive::<Int32Type>()
1606            .into_iter()
1607            .collect();
1608        assert_eq!(values_number, &[Some(2), None, None, Some(4), Some(4)]);
1609
1610        let values_string = values_struct.column(1).as_string::<i32>();
1611        let values_string: Vec<_> = values_string.into_iter().collect();
1612        assert_eq!(
1613            &values_string,
1614            &[Some("b"), Some("foo"), Some("foo"), Some("d"), None]
1615        );
1616    }
1617
1618    #[test]
1619    fn test_struct_with_nulls() {
1620        let fields = Fields::from(vec![
1621            Field::new("number_col", DataType::Int32, false),
1622            Field::new("string_col", DataType::Utf8, false),
1623        ]);
1624        let a = {
1625            let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1626            let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1627
1628            StructArray::try_new(
1629                fields.clone(),
1630                vec![Arc::new(number_col), Arc::new(string_col)],
1631                None,
1632            )
1633            .unwrap()
1634        };
1635
1636        let b = {
1637            let number_col = Int32Array::from_iter_values([5, 6, 7]);
1638            let string_col = StringArray::from_iter_values(["hello", "world", "foo"]);
1639
1640            StructArray::try_new(
1641                fields.clone(),
1642                vec![Arc::new(number_col), Arc::new(string_col)],
1643                Some(NullBuffer::from(&[true, false, true])),
1644            )
1645            .unwrap()
1646        };
1647
1648        let c = {
1649            let number_col = Int32Array::from_iter_values([8, 9, 10]);
1650            let string_col = StringArray::from_iter_values(["x", "y", "z"]);
1651
1652            StructArray::try_new(
1653                fields.clone(),
1654                vec![Arc::new(number_col), Arc::new(string_col)],
1655                None,
1656            )
1657            .unwrap()
1658        };
1659
1660        let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (1, 1), (2, 0)]).unwrap();
1661        let values_struct = values.as_struct();
1662        assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
1663
1664        let validity: Vec<bool> = {
1665            let null_buffer = values_struct.nulls().expect("should_have_nulls");
1666
1667            null_buffer.iter().collect()
1668        };
1669        assert_eq!(validity, &[true, true, true, false, true]);
1670        let values_number = values_struct.column(0).as_primitive::<Int32Type>();
1671        assert_eq!(values_number.values(), &[4, 4, 10, 6, 8]);
1672        let values_string = values_struct.column(1).as_string::<i32>();
1673        let values_string: Vec<_> = values_string.into_iter().collect();
1674        assert_eq!(
1675            &values_string,
1676            &[Some("d"), Some("d"), Some("z"), Some("world"), Some("x"),]
1677        );
1678    }
1679
1680    #[test]
1681    fn test_struct_empty() {
1682        let fields = Fields::from(vec![
1683            Field::new("number_col", DataType::Int32, false),
1684            Field::new("string_col", DataType::Utf8, false),
1685        ]);
1686        let a = {
1687            let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1688            let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1689
1690            StructArray::try_new(
1691                fields.clone(),
1692                vec![Arc::new(number_col), Arc::new(string_col)],
1693                None,
1694            )
1695            .unwrap()
1696        };
1697        let v = interleave(&[&a], &[]).unwrap();
1698        assert!(v.is_empty());
1699        assert_eq!(v.data_type(), &DataType::Struct(fields));
1700    }
1701
1702    #[test]
1703    fn interleave_sparse_nulls() {
1704        let values = StringArray::from_iter_values((0..100).map(|x| x.to_string()));
1705        let keys = Int32Array::from_iter_values(0..10);
1706        let dict_a = DictionaryArray::new(keys, Arc::new(values));
1707        let values = StringArray::new_null(0);
1708        let keys = Int32Array::new_null(10);
1709        let dict_b = DictionaryArray::new(keys, Arc::new(values));
1710
1711        let indices = &[(0, 0), (0, 1), (0, 2), (1, 0)];
1712        let array = interleave(&[&dict_a, &dict_b], indices).unwrap();
1713
1714        let expected =
1715            DictionaryArray::<Int32Type>::from_iter(vec![Some("0"), Some("1"), Some("2"), None]);
1716        assert_eq!(array.as_ref(), &expected)
1717    }
1718
1719    #[test]
1720    fn test_interleave_views() {
1721        let values = StringArray::from_iter_values([
1722            "hello",
1723            "world_long_string_not_inlined",
1724            "foo",
1725            "bar",
1726            "baz",
1727        ]);
1728        let view_a = StringViewArray::from(&values);
1729
1730        let values = StringArray::from_iter_values([
1731            "test",
1732            "data",
1733            "more_long_string_not_inlined",
1734            "views",
1735            "here",
1736        ]);
1737        let view_b = StringViewArray::from(&values);
1738
1739        let indices = &[
1740            (0, 2), // "foo"
1741            (1, 0), // "test"
1742            (0, 4), // "baz"
1743            (1, 3), // "views"
1744            (0, 1), // "world_long_string_not_inlined"
1745        ];
1746
1747        // Test specialized implementation
1748        let values = interleave(&[&view_a, &view_b], indices).unwrap();
1749        let result = values.as_string_view();
1750        assert_eq!(result.data_buffers().len(), 1);
1751
1752        let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
1753        let fallback_result = fallback.as_string_view();
1754        // note that fallback_result has 2 buffers, but only one long enough string to warrant a buffer
1755        assert_eq!(fallback_result.data_buffers().len(), 2);
1756
1757        // Convert to strings for easier assertion
1758        let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
1759
1760        let fallback_collected: Vec<_> = fallback_result
1761            .iter()
1762            .map(|x| x.map(|s| s.to_string()))
1763            .collect();
1764
1765        assert_eq!(&collected, &fallback_collected);
1766
1767        assert_eq!(
1768            &collected,
1769            &[
1770                Some("foo".to_string()),
1771                Some("test".to_string()),
1772                Some("baz".to_string()),
1773                Some("views".to_string()),
1774                Some("world_long_string_not_inlined".to_string()),
1775            ]
1776        );
1777    }
1778
1779    #[test]
1780    fn test_interleave_views_with_nulls() {
1781        let values = StringArray::from_iter([
1782            Some("hello"),
1783            None,
1784            Some("foo_long_string_not_inlined"),
1785            Some("bar"),
1786            None,
1787        ]);
1788        let view_a = StringViewArray::from(&values);
1789
1790        let values = StringArray::from_iter([
1791            Some("test"),
1792            Some("data_long_string_not_inlined"),
1793            None,
1794            None,
1795            Some("here"),
1796        ]);
1797        let view_b = StringViewArray::from(&values);
1798
1799        let indices = &[
1800            (0, 1), // null
1801            (1, 2), // null
1802            (0, 2), // "foo_long_string_not_inlined"
1803            (1, 3), // null
1804            (0, 4), // null
1805        ];
1806
1807        // Test specialized implementation
1808        let values = interleave(&[&view_a, &view_b], indices).unwrap();
1809        let result = values.as_string_view();
1810        assert_eq!(result.data_buffers().len(), 1);
1811
1812        let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
1813        let fallback_result = fallback.as_string_view();
1814
1815        // Convert to strings for easier assertion
1816        let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
1817
1818        let fallback_collected: Vec<_> = fallback_result
1819            .iter()
1820            .map(|x| x.map(|s| s.to_string()))
1821            .collect();
1822
1823        assert_eq!(&collected, &fallback_collected);
1824
1825        assert_eq!(
1826            &collected,
1827            &[
1828                None,
1829                None,
1830                Some("foo_long_string_not_inlined".to_string()),
1831                None,
1832                None,
1833            ]
1834        );
1835    }
1836
1837    #[test]
1838    fn test_interleave_views_shared_buffers() {
1839        let long = |i: usize| format!("long_string_not_inlined_{i}");
1840        let base = StringViewArray::from_iter_values((0..10).map(long));
1841        assert_eq!(base.data_buffers().len(), 1);
1842
1843        // Slices share the same buffers
1844        let a = base.slice(0, 5);
1845        let b = base.slice(5, 5);
1846        // A distinct buffer list containing the same underlying buffer
1847        let c = StringViewArray::try_new(base.views().clone(), base.data_buffers().to_vec(), None)
1848            .unwrap();
1849
1850        let indices = &[(0, 1), (1, 2), (2, 3), (0, 4), (1, 0), (2, 9)];
1851        let values = interleave(&[&a, &b, &c], indices).unwrap();
1852        let result = values.as_string_view();
1853        assert_eq!(result.data_buffers().len(), 1);
1854
1855        let expected: Vec<_> = [1, 7, 3, 4, 5, 9].into_iter().map(long).collect();
1856        let actual: Vec<_> = result.iter().map(|x| x.unwrap().to_string()).collect();
1857        assert_eq!(actual, expected);
1858
1859        // Many slices, exceeding the linear scan threshold
1860        let slices: Vec<_> = (0..40).map(|i| base.slice(i % 10, 1)).collect();
1861        let arrays: Vec<&dyn Array> = slices.iter().map(|x| x as &dyn Array).collect();
1862        let indices: Vec<_> = (0..40).rev().map(|i| (i, 0)).collect();
1863        let values = interleave(&arrays, &indices).unwrap();
1864        let result = values.as_string_view();
1865        assert_eq!(result.data_buffers().len(), 1);
1866
1867        let expected: Vec<_> = (0..40).rev().map(|i| long(i % 10)).collect();
1868        let actual: Vec<_> = result.iter().map(|x| x.unwrap().to_string()).collect();
1869        assert_eq!(actual, expected);
1870    }
1871
1872    #[test]
1873    fn test_interleave_views_multiple_buffers() {
1874        let str1 = "very_long_string_from_first_buffer".as_bytes();
1875        let str2 = "very_long_string_from_second_buffer".as_bytes();
1876        let buffer1 = str1.to_vec().into();
1877        let buffer2 = str2.to_vec().into();
1878
1879        let view1 = ByteView::new(str1.len() as u32, &str1[..4])
1880            .with_buffer_index(0)
1881            .with_offset(0)
1882            .as_u128();
1883        let view2 = ByteView::new(str2.len() as u32, &str2[..4])
1884            .with_buffer_index(1)
1885            .with_offset(0)
1886            .as_u128();
1887        let view_a =
1888            StringViewArray::try_new(vec![view1, view2].into(), vec![buffer1, buffer2], None)
1889                .unwrap();
1890
1891        let str3 = "another_very_long_string_buffer_three".as_bytes();
1892        let str4 = "different_long_string_in_buffer_four".as_bytes();
1893        let buffer3 = str3.to_vec().into();
1894        let buffer4 = str4.to_vec().into();
1895
1896        let view3 = ByteView::new(str3.len() as u32, &str3[..4])
1897            .with_buffer_index(0)
1898            .with_offset(0)
1899            .as_u128();
1900        let view4 = ByteView::new(str4.len() as u32, &str4[..4])
1901            .with_buffer_index(1)
1902            .with_offset(0)
1903            .as_u128();
1904        let view_b =
1905            StringViewArray::try_new(vec![view3, view4].into(), vec![buffer3, buffer4], None)
1906                .unwrap();
1907
1908        let indices = &[
1909            (0, 0), // String from first buffer of array A
1910            (1, 0), // String from first buffer of array B
1911            (0, 1), // String from second buffer of array A
1912            (1, 1), // String from second buffer of array B
1913            (0, 0), // String from first buffer of array A again
1914            (1, 1), // String from second buffer of array B again
1915        ];
1916
1917        // Test interleave
1918        let values = interleave(&[&view_a, &view_b], indices).unwrap();
1919        let result = values.as_string_view();
1920
1921        assert_eq!(
1922            result.data_buffers().len(),
1923            4,
1924            "Expected four buffers (two from each input array)"
1925        );
1926
1927        let result_strings: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
1928        assert_eq!(
1929            result_strings,
1930            vec![
1931                Some("very_long_string_from_first_buffer".to_string()),
1932                Some("another_very_long_string_buffer_three".to_string()),
1933                Some("very_long_string_from_second_buffer".to_string()),
1934                Some("different_long_string_in_buffer_four".to_string()),
1935                Some("very_long_string_from_first_buffer".to_string()),
1936                Some("different_long_string_in_buffer_four".to_string()),
1937            ]
1938        );
1939
1940        let views = result.views();
1941        let buffer_indices: Vec<_> = views
1942            .iter()
1943            .map(|raw_view| ByteView::from(*raw_view).buffer_index)
1944            .collect();
1945
1946        assert_eq!(
1947            buffer_indices,
1948            vec![
1949                0, // First buffer from array A
1950                1, // First buffer from array B
1951                2, // Second buffer from array A
1952                3, // Second buffer from array B
1953                0, // First buffer from array A (reused)
1954                3, // Second buffer from array B (reused)
1955            ]
1956        );
1957    }
1958
1959    #[test]
1960    fn test_interleave_run_end_encoded_primitive() {
1961        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1962        builder.extend([1, 1, 2, 2, 2, 3].into_iter().map(Some));
1963        let a = builder.finish();
1964
1965        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1966        builder.extend([4, 5, 5, 6, 6, 6].into_iter().map(Some));
1967        let b = builder.finish();
1968
1969        let indices = &[(0, 1), (1, 0), (0, 4), (1, 2), (0, 5)];
1970        let result = interleave(&[&a, &b], indices).unwrap();
1971
1972        // The result should be a RunEndEncoded array
1973        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
1974
1975        // Cast to RunArray to access values
1976        let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
1977
1978        // Verify the logical values by accessing the logical array directly
1979        let expected = vec![1, 4, 2, 5, 3];
1980        let mut actual = Vec::new();
1981        for i in 0..result_run_array.len() {
1982            let physical_idx = result_run_array.get_physical_index(i);
1983            let value = result_run_array
1984                .values()
1985                .as_primitive::<Int32Type>()
1986                .value(physical_idx);
1987            actual.push(value);
1988        }
1989        assert_eq!(actual, expected);
1990    }
1991
1992    #[test]
1993    fn test_interleave_run_end_encoded_sliced() {
1994        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1995        builder.extend([1, 1, 2, 2, 2, 3].into_iter().map(Some));
1996        let a = builder.finish();
1997        let a = a.slice(2, 3); // [2, 2, 2]
1998
1999        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
2000        builder.extend([4, 5, 5, 6, 6, 6].into_iter().map(Some));
2001        let b = builder.finish();
2002        let b = b.slice(1, 3); // [5, 5, 6]
2003
2004        let indices = &[(0, 1), (1, 0), (0, 2), (1, 1), (1, 2)];
2005        let result = interleave(&[&a, &b], indices).unwrap();
2006
2007        let result = result.as_run::<Int32Type>();
2008        let result = result.downcast::<Int32Array>().unwrap();
2009
2010        let expected = vec![2, 5, 2, 5, 6];
2011        let actual = result.into_iter().flatten().collect::<Vec<_>>();
2012        assert_eq!(actual, expected);
2013    }
2014
2015    #[test]
2016    fn test_interleave_run_end_encoded_string() {
2017        let a: Int32RunArray = vec!["hello", "hello", "world", "world", "foo"]
2018            .into_iter()
2019            .collect();
2020        let b: Int32RunArray = vec!["bar", "baz", "baz", "qux"].into_iter().collect();
2021
2022        let indices = &[(0, 0), (1, 1), (0, 3), (1, 3), (0, 4)];
2023        let result = interleave(&[&a, &b], indices).unwrap();
2024
2025        // The result should be a RunEndEncoded array
2026        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
2027
2028        // Cast to RunArray to access values
2029        let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
2030
2031        // Verify the logical values by accessing the logical array directly
2032        let expected = vec!["hello", "baz", "world", "qux", "foo"];
2033        let mut actual = Vec::new();
2034        for i in 0..result_run_array.len() {
2035            let physical_idx = result_run_array.get_physical_index(i);
2036            let value = result_run_array
2037                .values()
2038                .as_string::<i32>()
2039                .value(physical_idx);
2040            actual.push(value);
2041        }
2042        assert_eq!(actual, expected);
2043    }
2044
2045    #[test]
2046    fn test_interleave_run_end_encoded_with_nulls() {
2047        let a: Int32RunArray = vec![Some("a"), Some("a"), None, None, Some("b")]
2048            .into_iter()
2049            .collect();
2050        let b: Int32RunArray = vec![None, Some("c"), Some("c"), Some("d")]
2051            .into_iter()
2052            .collect();
2053
2054        let indices = &[(0, 1), (1, 0), (0, 2), (1, 3), (0, 4)];
2055        let result = interleave(&[&a, &b], indices).unwrap();
2056
2057        // The result should be a RunEndEncoded array
2058        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
2059
2060        // Cast to RunArray to access values
2061        let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
2062
2063        // Verify the logical values by accessing the logical array directly
2064        let expected = vec![Some("a"), None, None, Some("d"), Some("b")];
2065        let mut actual = Vec::new();
2066        for i in 0..result_run_array.len() {
2067            let physical_idx = result_run_array.get_physical_index(i);
2068            if result_run_array.values().is_null(physical_idx) {
2069                actual.push(None);
2070            } else {
2071                let value = result_run_array
2072                    .values()
2073                    .as_string::<i32>()
2074                    .value(physical_idx);
2075                actual.push(Some(value));
2076            }
2077        }
2078        assert_eq!(actual, expected);
2079    }
2080
2081    #[test]
2082    fn test_interleave_run_end_encoded_different_run_types() {
2083        let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2084        builder.extend([1, 1, 2, 3, 3].into_iter().map(Some));
2085        let a = builder.finish();
2086
2087        let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2088        builder.extend([4, 5, 5, 6].into_iter().map(Some));
2089        let b = builder.finish();
2090
2091        let indices = &[(0, 0), (1, 1), (0, 3), (1, 3)];
2092        let result = interleave(&[&a, &b], indices).unwrap();
2093
2094        // The result should be a RunEndEncoded array
2095        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
2096
2097        // Cast to RunArray to access values
2098        let result_run_array: &RunArray<Int16Type> = result.as_any().downcast_ref().unwrap();
2099
2100        // Verify the logical values by accessing the logical array directly
2101        let expected = vec![1, 5, 3, 6];
2102        let mut actual = Vec::new();
2103        for i in 0..result_run_array.len() {
2104            let physical_idx = result_run_array.get_physical_index(i);
2105            let value = result_run_array
2106                .values()
2107                .as_primitive::<Int32Type>()
2108                .value(physical_idx);
2109            actual.push(value);
2110        }
2111        assert_eq!(actual, expected);
2112    }
2113
2114    #[test]
2115    fn test_interleave_run_end_encoded_mixed_run_lengths() {
2116        let mut builder = PrimitiveRunBuilder::<Int64Type, Int32Type>::new();
2117        builder.extend([1, 2, 2, 2, 2, 3, 3, 4].into_iter().map(Some));
2118        let a = builder.finish();
2119
2120        let mut builder = PrimitiveRunBuilder::<Int64Type, Int32Type>::new();
2121        builder.extend([5, 5, 5, 6, 7, 7, 8, 8].into_iter().map(Some));
2122        let b = builder.finish();
2123
2124        let indices = &[
2125            (0, 0), // 1
2126            (1, 2), // 5
2127            (0, 3), // 2
2128            (1, 3), // 6
2129            (0, 6), // 3
2130            (1, 6), // 8
2131            (0, 7), // 4
2132            (1, 4), // 7
2133        ];
2134        let result = interleave(&[&a, &b], indices).unwrap();
2135
2136        // The result should be a RunEndEncoded array
2137        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
2138
2139        // Cast to RunArray to access values
2140        let result_run_array: &RunArray<Int64Type> = result.as_any().downcast_ref().unwrap();
2141
2142        // Verify the logical values by accessing the logical array directly
2143        let expected = vec![1, 5, 2, 6, 3, 8, 4, 7];
2144        let mut actual = Vec::new();
2145        for i in 0..result_run_array.len() {
2146            let physical_idx = result_run_array.get_physical_index(i);
2147            let value = result_run_array
2148                .values()
2149                .as_primitive::<Int32Type>()
2150                .value(physical_idx);
2151            actual.push(value);
2152        }
2153        assert_eq!(actual, expected);
2154    }
2155
2156    #[test]
2157    fn test_interleave_run_end_encoded_empty_runs() {
2158        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
2159        builder.extend(std::iter::once(Some(1)));
2160        let a = builder.finish();
2161
2162        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
2163        builder.extend([2, 2, 2].into_iter().map(Some));
2164        let b = builder.finish();
2165
2166        let indices = &[(0, 0), (1, 1), (1, 2)];
2167        let result = interleave(&[&a, &b], indices).unwrap();
2168
2169        // The result should be a RunEndEncoded array
2170        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
2171
2172        // Cast to RunArray to access values
2173        let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
2174
2175        // Verify the logical values by accessing the logical array directly
2176        let expected = vec![1, 2, 2];
2177        let mut actual = Vec::new();
2178        for i in 0..result_run_array.len() {
2179            let physical_idx = result_run_array.get_physical_index(i);
2180            let value = result_run_array
2181                .values()
2182                .as_primitive::<Int32Type>()
2183                .value(physical_idx);
2184            actual.push(value);
2185        }
2186        assert_eq!(actual, expected);
2187    }
2188
2189    #[test]
2190    fn test_struct_no_fields() {
2191        let fields = Fields::empty();
2192        let a = StructArray::try_new_with_length(fields.clone(), vec![], None, 10).unwrap();
2193        let v = interleave(&[&a], &[(0, 0)]).unwrap();
2194        assert_eq!(v.len(), 1);
2195        assert_eq!(v.data_type(), &DataType::Struct(fields));
2196    }
2197
2198    #[test]
2199    fn test_interleave_fallback_dictionary_with_nulls() {
2200        let input_1_keys = Int32Array::from_iter([Some(0), None, Some(1)]);
2201        let input_1_values = StringArray::from_iter_values(["foo", "bar"]);
2202        let dict_a = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
2203
2204        let input_2_keys = Int32Array::from_iter([Some(0), Some(1), None]);
2205        let input_2_values = StringArray::from_iter_values(["baz", "qux"]);
2206        let dict_b = DictionaryArray::new(input_2_keys, Arc::new(input_2_values));
2207
2208        let indices = vec![
2209            (0, 0), // "foo"
2210            (0, 1), // null
2211            (1, 0), // "baz"
2212            (1, 2), // null
2213            (0, 2), // "bar"
2214            (1, 1), // "qux"
2215        ];
2216
2217        let result =
2218            interleave_fallback_dictionary::<Int32Type>(&[&dict_a, &dict_b], &indices).unwrap();
2219        let dict_result = result.as_dictionary::<Int32Type>();
2220
2221        let string_result = dict_result.downcast_dict::<StringArray>().unwrap();
2222        let collected: Vec<_> = string_result.into_iter().collect();
2223        assert_eq!(
2224            collected,
2225            vec![
2226                Some("foo"),
2227                None,
2228                Some("baz"),
2229                None,
2230                Some("bar"),
2231                Some("qux")
2232            ]
2233        );
2234    }
2235
2236    #[test]
2237    fn test_interleave_string_view_dictionary_overflow_returns_err() {
2238        // interleaving dictionaries which results in overflowing the key type should
2239        // surface an error not a panic
2240        let values_a: StringViewArray = (0..200).map(|i| Some(format!("a{i}"))).collect();
2241        let keys_a = UInt8Array::from_iter_values(0..200);
2242        let dict_a = DictionaryArray::<UInt8Type>::new(keys_a, Arc::new(values_a));
2243
2244        let values_b: StringViewArray = (0..200).map(|i| Some(format!("b{i}"))).collect();
2245        let keys_b = UInt8Array::from_iter_values(0..200);
2246        let dict_b = DictionaryArray::<UInt8Type>::new(keys_b, Arc::new(values_b));
2247
2248        let indices: Vec<_> = (0..200).flat_map(|i| [(0, i), (1, i)]).collect();
2249
2250        let err = interleave(&[&dict_a, &dict_b], &indices).unwrap_err();
2251        assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
2252    }
2253
2254    #[test]
2255    fn test_interleave_nested_dictionary_overflow_returns_err() {
2256        // same as above, but with the dictionary nested inside a FixedSizeList
2257        let field = Arc::new(arrow_schema::Field::new(
2258            "item",
2259            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8View)),
2260            false,
2261        ));
2262
2263        let values_a: StringViewArray = (0..200).map(|i| Some(format!("a{i}"))).collect();
2264        let keys_a = UInt8Array::from_iter_values(0..200);
2265        let dict_a = DictionaryArray::<UInt8Type>::new(keys_a, Arc::new(values_a));
2266        let list_a = FixedSizeListArray::new(field.clone(), 1, Arc::new(dict_a), None);
2267
2268        let values_b: StringViewArray = (0..200).map(|i| Some(format!("b{i}"))).collect();
2269        let keys_b = UInt8Array::from_iter_values(0..200);
2270        let dict_b = DictionaryArray::<UInt8Type>::new(keys_b, Arc::new(values_b));
2271        let list_b = FixedSizeListArray::new(field, 1, Arc::new(dict_b), None);
2272
2273        let indices: Vec<_> = (0..200).flat_map(|i| [(0, i), (1, i)]).collect();
2274
2275        let err = interleave(&[&list_a, &list_b], &indices).unwrap_err();
2276        assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
2277    }
2278
2279    #[test]
2280    #[cfg_attr(miri, ignore)] // Takes too long
2281    fn test_interleave_bytes_offset_overflow() {
2282        let indices: Vec<(usize, usize)> = vec![(0, 0); (i32::MAX >> 4) as usize];
2283        let text = ('a'..='z').collect::<String>();
2284        let values = StringArray::from(vec![Some(text)]);
2285        assert!(matches!(
2286            interleave(&[&values], &indices),
2287            Err(ArrowError::OffsetOverflowError(_))
2288        ));
2289    }
2290
2291    #[test]
2292    #[cfg_attr(miri, ignore)] // Takes too long
2293    fn test_interleave_list_offset_overflow() {
2294        // Build a ListArray<i32> with a single row containing many elements
2295        let mut builder = GenericListBuilder::<i32, _>::new(Int32Builder::new());
2296        for i in 0..32 {
2297            builder.values().append_value(i);
2298        }
2299        builder.append(true);
2300        let list = builder.finish();
2301
2302        // Interleave enough copies to overflow i32 offsets
2303        let indices: Vec<(usize, usize)> = vec![(0, 0); (i32::MAX as usize / 32) + 1];
2304        assert!(matches!(
2305            interleave(&[&list], &indices),
2306            Err(ArrowError::OffsetOverflowError(_))
2307        ));
2308    }
2309
2310    #[test]
2311    fn test_interleave_list_view() {
2312        // `interleave` for ListView falls through to `interleave_fallback`, which uses
2313        // `MutableArrayData`. `list_view::build_extend` copies offsets/sizes but never
2314        // extends the child array, so the result contains offsets/sizes that reference
2315        // positions in the now-absent original child arrays while the child is empty.
2316        //
2317        // lv_a: [[1, 2], [3]]   (values=[1,2,3], offsets=[0,2], sizes=[2,1])
2318        // lv_b: [[4, 5, 6]]     (values=[4,5,6], offsets=[0],   sizes=[3])
2319        // interleave at [(0,0), (1,0), (0,1)] should produce [[1, 2], [4, 5, 6], [3]]
2320        let field = Arc::new(Field::new_list_field(DataType::Int64, false));
2321
2322        let lv_a = ListViewArray::new(
2323            Arc::clone(&field),
2324            ScalarBuffer::from(vec![0i32, 2]),
2325            ScalarBuffer::from(vec![2i32, 1]),
2326            Arc::new(Int64Array::from(vec![1_i64, 2, 3])),
2327            None,
2328        );
2329        let lv_b = ListViewArray::new(
2330            field,
2331            ScalarBuffer::from(vec![0i32]),
2332            ScalarBuffer::from(vec![3i32]),
2333            Arc::new(Int64Array::from(vec![4_i64, 5, 6])),
2334            None,
2335        );
2336
2337        let result = interleave(
2338            &[&lv_a as &dyn Array, &lv_b as &dyn Array],
2339            &[(0, 0), (1, 0), (0, 1)],
2340        )
2341        .unwrap();
2342
2343        result
2344            .to_data()
2345            .validate_full()
2346            .expect("interleaved ListViewArray must be internally consistent");
2347
2348        let result_lv = result.as_list_view::<i32>();
2349        assert_eq!(result_lv.len(), 3);
2350        assert_eq!(
2351            result_lv.value(0).as_primitive::<Int64Type>().values(),
2352            &[1, 2]
2353        );
2354        assert_eq!(
2355            result_lv.value(1).as_primitive::<Int64Type>().values(),
2356            &[4, 5, 6]
2357        );
2358        assert_eq!(
2359            result_lv.value(2).as_primitive::<Int64Type>().values(),
2360            &[3]
2361        );
2362    }
2363
2364    #[test]
2365    fn test_interleave_fixed_size_list() {
2366        // a: [[1, 2], [3, 4], [5, 6]]
2367        let field = Arc::new(Field::new("item", DataType::Int32, false));
2368        let a = FixedSizeListArray::new(
2369            field.clone(),
2370            2,
2371            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])),
2372            None,
2373        );
2374        // b: [[7, 8], [9, 10]]
2375        let b = FixedSizeListArray::new(
2376            field.clone(),
2377            2,
2378            Arc::new(Int32Array::from(vec![7, 8, 9, 10])),
2379            None,
2380        );
2381
2382        let result = interleave(&[&a, &b], &[(0, 2), (1, 0), (0, 0), (1, 1), (0, 1)]).unwrap();
2383        let result = result.as_fixed_size_list();
2384        assert_eq!(result.len(), 5);
2385        assert_eq!(result.value_length(), 2);
2386
2387        let values = result.values().as_primitive::<Int32Type>();
2388        // [[5,6], [7,8], [1,2], [9,10], [3,4]]
2389        assert_eq!(values.values(), &[5, 6, 7, 8, 1, 2, 9, 10, 3, 4]);
2390    }
2391
2392    #[test]
2393    fn test_interleave_zero_sized_fixed_size_list() {
2394        let input = FixedSizeListArray::try_new_with_length(
2395            Field::new_list_field(DataType::Int32, true).into(),
2396            0,
2397            Arc::new(Int32Array::new_null(0)),
2398            None,
2399            3,
2400        )
2401        .unwrap();
2402
2403        let indices = [(0, 2), (0, 0)];
2404        let result = interleave(&[&input], &indices).unwrap();
2405
2406        assert_eq!(result.len(), 2);
2407    }
2408
2409    #[test]
2410    fn test_interleave_fixed_size_list_with_nulls() {
2411        let field = Arc::new(Field::new("item", DataType::Int32, true));
2412        // a: [[1, 2], null, [5, 6]]
2413        let a = FixedSizeListArray::new(
2414            field.clone(),
2415            2,
2416            Arc::new(Int32Array::from(vec![1, 2, 0, 0, 5, 6])),
2417            Some(NullBuffer::from(&[true, false, true])),
2418        );
2419        // b: [null, [9, 10]]
2420        let b = FixedSizeListArray::new(
2421            field.clone(),
2422            2,
2423            Arc::new(Int32Array::from(vec![0, 0, 9, 10])),
2424            Some(NullBuffer::from(&[false, true])),
2425        );
2426
2427        let result = interleave(&[&a, &b], &[(0, 0), (0, 1), (1, 0), (1, 1), (0, 2)]).unwrap();
2428        let result = result.as_fixed_size_list();
2429        assert_eq!(result.len(), 5);
2430
2431        let validity: Vec<bool> = result.nulls().unwrap().iter().collect();
2432        assert_eq!(validity, &[true, false, false, true, true]);
2433    }
2434
2435    fn run_fsl_string_child_test(
2436        child_data_type: DataType,
2437        create_child: impl Fn(Vec<&str>) -> ArrayRef,
2438        extract_values: impl Fn(&ArrayRef) -> Vec<String>,
2439    ) {
2440        let field = Arc::new(Field::new("item", child_data_type, false));
2441
2442        // a: [["a", "b"], ["c", "d"]]
2443        let a = FixedSizeListArray::new(
2444            field.clone(),
2445            2,
2446            create_child(vec!["a", "b", "c", "d"]),
2447            None,
2448        );
2449        // b: [["x", "y"], ["z", "w"]]
2450        let b = FixedSizeListArray::new(
2451            field.clone(),
2452            2,
2453            create_child(vec!["x", "y", "z", "w"]),
2454            None,
2455        );
2456
2457        let result = interleave(&[&a, &b], &[(0, 1), (1, 0), (0, 0)]).unwrap();
2458        let result = result.as_fixed_size_list();
2459        assert_eq!(result.len(), 3);
2460        assert_eq!(result.value_length(), 2);
2461
2462        // Expected: [[c,d], [x,y], [a,b]]
2463        let values = extract_values(result.values());
2464        assert_eq!(values, vec!["c", "d", "x", "y", "a", "b"]);
2465    }
2466
2467    #[test]
2468    fn test_interleave_fixed_size_list_string_child() {
2469        // FixedSizeList<Utf8> — exercises the non-primitive child path
2470        run_fsl_string_child_test(
2471            DataType::Utf8,
2472            |v| Arc::new(StringArray::from(v)),
2473            |a| {
2474                a.as_string::<i32>()
2475                    .iter()
2476                    .map(|s| s.unwrap().to_owned())
2477                    .collect()
2478            },
2479        );
2480    }
2481
2482    #[test]
2483    fn test_interleave_fixed_size_list_string_view_child() {
2484        // FixedSizeList<Utf8View> — exercises the non-primitive child path
2485        run_fsl_string_child_test(
2486            DataType::Utf8View,
2487            |v| Arc::new(StringViewArray::from(v)),
2488            |a| {
2489                a.as_string_view()
2490                    .iter()
2491                    .map(|s| s.unwrap().to_owned())
2492                    .collect()
2493            },
2494        );
2495    }
2496
2497    #[test]
2498    fn test_interleave_map() {
2499        use arrow_array::builder::MapBuilder;
2500        use arrow_array::builder::StringBuilder;
2501
2502        // a: [{k1: 1, k2: 2}, {k3: 3}]
2503        let mut a_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2504        a_builder.keys().append_value("k1");
2505        a_builder.values().append_value(1);
2506        a_builder.keys().append_value("k2");
2507        a_builder.values().append_value(2);
2508        a_builder.append(true).unwrap();
2509        a_builder.keys().append_value("k3");
2510        a_builder.values().append_value(3);
2511        a_builder.append(true).unwrap();
2512        let a = a_builder.finish();
2513
2514        // b: [{k4: 4}, {k5: 5, k6: 6, k7: 7}]
2515        let mut b_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2516        b_builder.keys().append_value("k4");
2517        b_builder.values().append_value(4);
2518        b_builder.append(true).unwrap();
2519        b_builder.keys().append_value("k5");
2520        b_builder.values().append_value(5);
2521        b_builder.keys().append_value("k6");
2522        b_builder.values().append_value(6);
2523        b_builder.keys().append_value("k7");
2524        b_builder.values().append_value(7);
2525        b_builder.append(true).unwrap();
2526        let b = b_builder.finish();
2527
2528        let result = interleave(&[&a, &b], &[(1, 0), (0, 0), (0, 1), (1, 1)]).unwrap();
2529        let result = result.as_map();
2530        assert_eq!(result.len(), 4);
2531
2532        // Row 0: {k4: 4}
2533        let row0 = result.value(0);
2534        assert_eq!(row0.len(), 1);
2535        assert_eq!(row0.column(0).as_string::<i32>().value(0), "k4");
2536        assert_eq!(row0.column(1).as_primitive::<Int32Type>().value(0), 4);
2537
2538        // Row 1: {k1: 1, k2: 2}
2539        let row1 = result.value(1);
2540        assert_eq!(row1.len(), 2);
2541
2542        // Row 2: {k3: 3}
2543        let row2 = result.value(2);
2544        assert_eq!(row2.len(), 1);
2545        assert_eq!(row2.column(0).as_string::<i32>().value(0), "k3");
2546
2547        // Row 3: {k5: 5, k6: 6, k7: 7}
2548        let row3 = result.value(3);
2549        assert_eq!(row3.len(), 3);
2550    }
2551
2552    #[test]
2553    fn test_interleave_map_with_nulls() {
2554        use arrow_array::builder::MapBuilder;
2555        use arrow_array::builder::StringBuilder;
2556
2557        // a: [{k1: 1}, null]
2558        let mut a_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2559        a_builder.keys().append_value("k1");
2560        a_builder.values().append_value(1);
2561        a_builder.append(true).unwrap();
2562        a_builder.append(false).unwrap();
2563        let a = a_builder.finish();
2564
2565        // b: [null, {k2: 2}]
2566        let mut b_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2567        b_builder.append(false).unwrap();
2568        b_builder.keys().append_value("k2");
2569        b_builder.values().append_value(2);
2570        b_builder.append(true).unwrap();
2571        let b = b_builder.finish();
2572
2573        let result = interleave(&[&a, &b], &[(0, 0), (1, 0), (0, 1), (1, 1)]).unwrap();
2574        let result = result.as_map();
2575        assert_eq!(result.len(), 4);
2576
2577        let validity: Vec<bool> = result.nulls().unwrap().iter().collect();
2578        assert_eq!(validity, &[true, false, false, true]);
2579
2580        assert_eq!(result.value(0).len(), 1);
2581        assert_eq!(result.value(3).len(), 1);
2582    }
2583
2584    #[test]
2585    fn test_interleave_record_batch_empty_input() {
2586        let err = interleave_record_batch(&[], &[]).unwrap_err();
2587        assert!(matches!(err, ArrowError::InvalidArgumentError(_)));
2588    }
2589
2590    #[test]
2591    fn test_interleave_record_batch_no_columns() {
2592        // A batch with no columns but a known row count interleaves to a
2593        // batch with `indices.len()` rows, matching `take_record_batch`.
2594        let options = RecordBatchOptions::new().with_row_count(Some(3));
2595        let batch =
2596            RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options).unwrap();
2597
2598        let result = interleave_record_batch(&[&batch], &[(0, 0), (0, 2)]).unwrap();
2599        assert_eq!(result.num_columns(), 0);
2600        assert_eq!(result.num_rows(), 2);
2601    }
2602}