Skip to main content

arrow_select/
interleave.rs

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