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