arrow_select/

interleave.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Interleave elements from multiple arrays

use crate::concat::concat;
use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values};
use arrow_array::builder::{BooleanBufferBuilder, PrimitiveBuilder};
use arrow_array::cast::AsArray;
use arrow_array::types::*;
use arrow_array::*;
use arrow_buffer::{ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer};
use arrow_data::ByteView;
use arrow_data::transform::MutableArrayData;
use arrow_schema::{ArrowError, DataType, FieldRef, Fields};
use std::sync::Arc;

macro_rules! primitive_helper {
    ($t:ty, $values:ident, $indices:ident, $data_type:ident) => {
        interleave_primitive::<$t>($values, $indices, $data_type)
    };
}

macro_rules! dict_helper {
    ($t:ty, $values:expr, $indices:expr) => {
        interleave_dictionaries::<$t>($values, $indices)
    };
}

///
/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values.
///
/// Each element in `indices` is a pair of `usize` with the first identifying the index
/// of the [`Array`] in `values`, and the second the index of the value within that [`Array`]
///
/// ```text
/// ┌─────────────────┐      ┌─────────┐                                  ┌─────────────────┐
/// │        A        │      │ (0, 0)  │        interleave(               │        A        │
/// ├─────────────────┤      ├─────────┤          [values0, values1],     ├─────────────────┤
/// │        D        │      │ (1, 0)  │          indices                 │        B        │
/// └─────────────────┘      ├─────────┤        )                         ├─────────────────┤
///   values array 0         │ (1, 1)  │      ─────────────────────────▶  │        C        │
///                          ├─────────┤                                  ├─────────────────┤
///                          │ (0, 1)  │                                  │        D        │
///                          └─────────┘                                  └─────────────────┘
/// ┌─────────────────┐       indices
/// │        B        │        array
/// ├─────────────────┤                                                    result
/// │        C        │
/// ├─────────────────┤
/// │        E        │
/// └─────────────────┘
///   values array 1
/// ```
///
/// For selecting values by index from a single array see [`crate::take`]
pub fn interleave(
    values: &[&dyn Array],
    indices: &[(usize, usize)],
) -> Result<ArrayRef, ArrowError> {
    if values.is_empty() {
        return Err(ArrowError::InvalidArgumentError(
            "interleave requires input of at least one array".to_string(),
        ));
    }
    let data_type = values[0].data_type();

    for array in values.iter().skip(1) {
        if array.data_type() != data_type {
            return Err(ArrowError::InvalidArgumentError(format!(
                "It is not possible to interleave arrays of different data types ({} and {})",
                data_type,
                array.data_type()
            )));
        }
    }

    if indices.is_empty() {
        return Ok(new_empty_array(data_type));
    }

    downcast_primitive! {
        data_type => (primitive_helper, values, indices, data_type),
        DataType::Utf8 => interleave_bytes::<Utf8Type>(values, indices),
        DataType::LargeUtf8 => interleave_bytes::<LargeUtf8Type>(values, indices),
        DataType::Binary => interleave_bytes::<BinaryType>(values, indices),
        DataType::LargeBinary => interleave_bytes::<LargeBinaryType>(values, indices),
        DataType::BinaryView => interleave_views::<BinaryViewType>(values, indices),
        DataType::Utf8View => interleave_views::<StringViewType>(values, indices),
        DataType::Dictionary(k, _) => downcast_integer! {
            k.as_ref() => (dict_helper, values, indices),
            _ => unreachable!("illegal dictionary key type {k}")
        },
        DataType::Struct(fields) => interleave_struct(fields, values, indices),
        DataType::List(field) => interleave_list::<i32>(values, indices, field),
        DataType::LargeList(field) => interleave_list::<i64>(values, indices, field),
        DataType::RunEndEncoded(r, _) => match r.data_type() {
            DataType::Int16 => interleave_run_end::<Int16Type>(values, indices),
            DataType::Int32 => interleave_run_end::<Int32Type>(values, indices),
            DataType::Int64 => interleave_run_end::<Int64Type>(values, indices),
            t => unreachable!("illegal run-end type {t}"),
        },
        DataType::ListView(field) => interleave_list_view::<i32>(values, indices, field),
        DataType::LargeListView(field) => interleave_list_view::<i64>(values, indices, field),
        _ => interleave_fallback(values, indices)
    }
}

/// Common functionality for interleaving arrays
///
/// T is the concrete Array type
struct Interleave<'a, T> {
    /// The input arrays downcast to T
    arrays: Vec<&'a T>,
    /// The null buffer of the interleaved output
    nulls: Option<NullBuffer>,
}

impl<'a, T: Array + 'static> Interleave<'a, T> {
    fn new(values: &[&'a dyn Array], indices: &'a [(usize, usize)]) -> Self {
        let mut has_nulls = false;
        let arrays: Vec<&T> = values
            .iter()
            .map(|x| {
                has_nulls = has_nulls || x.null_count() != 0;
                x.as_any().downcast_ref().unwrap()
            })
            .collect();

        let nulls = match has_nulls {
            true => {
                let nulls = BooleanBuffer::collect_bool(indices.len(), |i| {
                    let (a, b) = indices[i];
                    arrays[a].is_valid(b)
                });
                Some(nulls.into())
            }
            false => None,
        };

        Self { arrays, nulls }
    }
}

fn interleave_primitive<T: ArrowPrimitiveType>(
    values: &[&dyn Array],
    indices: &[(usize, usize)],
    data_type: &DataType,
) -> Result<ArrayRef, ArrowError> {
    let interleaved = Interleave::<'_, PrimitiveArray<T>>::new(values, indices);
    let arrays = &interleaved.arrays;
    let len = indices.len();

    let mut output = Vec::with_capacity(len);
    let dst: *mut T::Native = output.as_mut_ptr();
    let mut base = 0;

    // Process 8 elements at a time to issue multiple independent loads
    // and increase memory-level parallelism for random access patterns.
    let chunks = indices.chunks_exact(8);
    let remainder = chunks.remainder();
    for chunk in chunks {
        let v0 = arrays[chunk[0].0].value(chunk[0].1);
        let v1 = arrays[chunk[1].0].value(chunk[1].1);
        let v2 = arrays[chunk[2].0].value(chunk[2].1);
        let v3 = arrays[chunk[3].0].value(chunk[3].1);
        let v4 = arrays[chunk[4].0].value(chunk[4].1);
        let v5 = arrays[chunk[5].0].value(chunk[5].1);
        let v6 = arrays[chunk[6].0].value(chunk[6].1);
        let v7 = arrays[chunk[7].0].value(chunk[7].1);

        // SAFETY: base+7 < len == output capacity
        debug_assert!(base + 7 < len);
        unsafe {
            dst.add(base).write(v0);
            dst.add(base + 1).write(v1);
            dst.add(base + 2).write(v2);
            dst.add(base + 3).write(v3);
            dst.add(base + 4).write(v4);
            dst.add(base + 5).write(v5);
            dst.add(base + 6).write(v6);
            dst.add(base + 7).write(v7);
        }
        base += 8;
    }

    for idx in remainder {
        // SAFETY: base < len == output capacity
        debug_assert!(base < len);
        unsafe { dst.add(base).write(arrays[idx.0].value(idx.1)) };
        base += 1;
    }

    // SAFETY: all `len` elements have been initialized
    debug_assert!(base == len);
    unsafe { output.set_len(len) };

    let array = PrimitiveArray::<T>::try_new(output.into(), interleaved.nulls)?;
    Ok(Arc::new(array.with_data_type(data_type.clone())))
}

fn interleave_bytes<T: ByteArrayType>(
    values: &[&dyn Array],
    indices: &[(usize, usize)],
) -> Result<ArrayRef, ArrowError> {
    let interleaved = Interleave::<'_, GenericByteArray<T>>::new(values, indices);

    let mut capacity = 0;
    let mut offsets = Vec::with_capacity(indices.len() + 1);
    offsets.push(T::Offset::from_usize(0).unwrap());
    for (a, b) in indices {
        let o = interleaved.arrays[*a].value_offsets();
        let element_len = o[*b + 1].as_usize() - o[*b].as_usize();
        capacity += element_len;
        offsets.push(
            T::Offset::from_usize(capacity)
                .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
        );
    }

    let mut values = Vec::with_capacity(capacity);
    for (a, b) in indices {
        values.extend_from_slice(interleaved.arrays[*a].value(*b).as_ref());
    }

    // Safety: safe by construction
    let array = unsafe {
        let offsets = OffsetBuffer::new_unchecked(offsets.into());
        GenericByteArray::<T>::new_unchecked(offsets, values.into(), interleaved.nulls)
    };
    Ok(Arc::new(array))
}

fn interleave_dictionaries<K: ArrowDictionaryKeyType>(
    arrays: &[&dyn Array],
    indices: &[(usize, usize)],
) -> Result<ArrayRef, ArrowError> {
    let dictionaries: Vec<_> = arrays.iter().map(|x| x.as_dictionary::<K>()).collect();
    let (should_merge, has_overflow) =
        should_merge_dictionary_values::<K>(&dictionaries, indices.len());
    if !should_merge {
        return if has_overflow {
            interleave_fallback(arrays, indices)
        } else {
            interleave_fallback_dictionary::<K>(&dictionaries, indices)
        };
    }

    let masks: Vec<_> = dictionaries
        .iter()
        .enumerate()
        .map(|(a_idx, dictionary)| {
            let mut key_mask = BooleanBufferBuilder::new_from_buffer(
                MutableBuffer::new_null(dictionary.len()),
                dictionary.len(),
            );

            for (_, key_idx) in indices.iter().filter(|(a, _)| *a == a_idx) {
                key_mask.set_bit(*key_idx, true);
            }
            key_mask.finish()
        })
        .collect();

    let merged = merge_dictionary_values(&dictionaries, Some(&masks))?;

    // Recompute keys
    let mut keys = PrimitiveBuilder::<K>::with_capacity(indices.len());
    for (a, b) in indices {
        let old_keys: &PrimitiveArray<K> = dictionaries[*a].keys();
        match old_keys.is_valid(*b) {
            true => {
                let old_key = old_keys.values()[*b];
                keys.append_value(merged.key_mappings[*a][old_key.as_usize()])
            }
            false => keys.append_null(),
        }
    }
    let array = unsafe { DictionaryArray::new_unchecked(keys.finish(), merged.values) };
    Ok(Arc::new(array))
}

fn interleave_views<T: ByteViewType>(
    values: &[&dyn Array],
    indices: &[(usize, usize)],
) -> Result<ArrayRef, ArrowError> {
    let interleaved = Interleave::<'_, GenericByteViewArray<T>>::new(values, indices);
    let mut buffers = Vec::new();

    // Contains the offsets of start buffer in `buffer_to_new_index`
    let mut offsets = Vec::with_capacity(interleaved.arrays.len() + 1);
    offsets.push(0);
    let mut total_buffers = 0;
    for a in interleaved.arrays.iter() {
        total_buffers += a.data_buffers().len();
        offsets.push(total_buffers);
    }

    // contains the mapping from old buffer index to new buffer index
    let mut buffer_to_new_index = vec![None; total_buffers];

    let views: Vec<u128> = indices
        .iter()
        .map(|(array_idx, value_idx)| {
            let array = interleaved.arrays[*array_idx];
            let view = array.views().get(*value_idx).unwrap();
            let view_len = *view as u32;
            if view_len <= 12 {
                return *view;
            }
            // value is big enough to be in a variadic buffer
            let view = ByteView::from(*view);
            let buffer_to_new_idx = offsets[*array_idx] + view.buffer_index as usize;
            let new_buffer_idx: u32 =
                *buffer_to_new_index[buffer_to_new_idx].get_or_insert_with(|| {
                    buffers.push(array.data_buffers()[view.buffer_index as usize].clone());
                    (buffers.len() - 1) as u32
                });
            view.with_buffer_index(new_buffer_idx).as_u128()
        })
        .collect();

    let array = unsafe {
        GenericByteViewArray::<T>::new_unchecked(views.into(), buffers, interleaved.nulls)
    };
    Ok(Arc::new(array))
}

fn interleave_struct(
    fields: &Fields,
    values: &[&dyn Array],
    indices: &[(usize, usize)],
) -> Result<ArrayRef, ArrowError> {
    let interleaved = Interleave::<'_, StructArray>::new(values, indices);

    if fields.is_empty() {
        let array = StructArray::try_new_with_length(
            fields.clone(),
            vec![],
            interleaved.nulls,
            indices.len(),
        )?;
        return Ok(Arc::new(array));
    }

    let struct_fields_array: Result<Vec<_>, _> = (0..fields.len())
        .map(|i| {
            let field_values: Vec<&dyn Array> = interleaved
                .arrays
                .iter()
                .map(|x| x.column(i).as_ref())
                .collect();
            interleave(&field_values, indices)
        })
        .collect();

    let struct_array =
        StructArray::try_new(fields.clone(), struct_fields_array?, interleaved.nulls)?;
    Ok(Arc::new(struct_array))
}

fn interleave_list<O: OffsetSizeTrait>(
    values: &[&dyn Array],
    indices: &[(usize, usize)],
    field: &FieldRef,
) -> Result<ArrayRef, ArrowError> {
    let interleaved = Interleave::<'_, GenericListArray<O>>::new(values, indices);

    let mut capacity = 0usize;
    let mut offsets = Vec::with_capacity(indices.len() + 1);
    offsets.push(O::from_usize(0).unwrap());
    for (array, row) in indices {
        let o = interleaved.arrays[*array].value_offsets();
        let element_len = o[*row + 1].as_usize() - o[*row].as_usize();
        capacity += element_len;
        offsets.push(
            O::from_usize(capacity).ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
        );
    }

    let mut child_indices = Vec::with_capacity(capacity);
    for (array, row) in indices {
        let list = interleaved.arrays[*array];
        let start = list.value_offsets()[*row].as_usize();
        let end = list.value_offsets()[*row + 1].as_usize();
        child_indices.extend((start..end).map(|i| (*array, i)));
    }

    let child_arrays: Vec<&dyn Array> = interleaved
        .arrays
        .iter()
        .map(|list| list.values().as_ref())
        .collect();

    let interleaved_values = interleave(&child_arrays, &child_indices)?;

    let offsets = OffsetBuffer::new(offsets.into());
    let list_array = GenericListArray::<O>::new(
        field.clone(),
        offsets,
        interleaved_values,
        interleaved.nulls,
    );

    Ok(Arc::new(list_array))
}

/// Specialized [`interleave`] for [`RunArray`].
fn interleave_run_end<R: RunEndIndexType>(
    values: &[&dyn Array],
    indices: &[(usize, usize)],
) -> Result<ArrayRef, ArrowError> {
    if indices.is_empty() {
        return Ok(new_empty_array(values[0].data_type()));
    }

    let n = indices.len();
    R::Native::from_usize(n).ok_or_else(|| {
        ArrowError::ComputeError(format!(
            "interleave_run_end: output length {n} does not fit run-end type"
        ))
    })?;

    let runs: Vec<&RunArray<R>> = values.iter().map(|a| a.as_run::<R>()).collect();
    let value_arrays: Vec<&dyn Array> = runs.iter().map(|r| r.values().as_ref()).collect();

    // Resolve each (array, logical_row) to (array, physical_row), so we can
    // lookup physical indices by batch.
    let mut phys_pairs: Vec<(usize, usize)> = vec![(0, 0); n];
    let mut grouped: Vec<(Vec<R::Native>, Vec<usize>)> =
        (0..runs.len()).map(|_| (Vec::new(), Vec::new())).collect();
    for (out_pos, &(arr, row)) in indices.iter().enumerate() {
        let row = R::Native::from_usize(row).ok_or_else(|| {
            ArrowError::InvalidArgumentError(format!(
                "interleave_run_end: row index {row} not representable as run-end type {}",
                R::DATA_TYPE
            ))
        })?;
        grouped[arr].0.push(row);
        grouped[arr].1.push(out_pos);
    }
    for (arr_idx, (logical_rows, out_positions)) in grouped.into_iter().enumerate() {
        let phys = runs[arr_idx].get_physical_indices(&logical_rows)?;
        for (p, out_pos) in phys.iter().zip(out_positions.iter()) {
            phys_pairs[*out_pos] = (arr_idx, *p);
        }
    }

    // Coalesce by physical-pair equality only: emit a new run when the
    // (array_idx, physical_idx) pair changes between adjacent output rows.
    // TODO: We could perform an equality check across sources to extend the
    // output run, but we can't call make_comparator from this crate.
    let mut run_ends_buf: Vec<R::Native> = Vec::with_capacity(n);
    let mut dedup_pairs: Vec<(usize, usize)> = Vec::with_capacity(n);
    dedup_pairs.push(phys_pairs[0]);
    for i in 1..n {
        if phys_pairs[i] != phys_pairs[i - 1] {
            run_ends_buf.push(R::Native::from_usize(i).unwrap());
            dedup_pairs.push(phys_pairs[i]);
        }
    }
    run_ends_buf.push(R::Native::from_usize(n).unwrap());

    let taken_values = interleave(&value_arrays, &dedup_pairs)?;
    let run_ends = PrimitiveArray::<R>::from_iter_values(run_ends_buf);

    Ok(Arc::new(RunArray::<R>::try_new(
        &run_ends,
        taken_values.as_ref(),
    )?))
}

fn interleave_list_view<O: OffsetSizeTrait>(
    values: &[&dyn Array],
    indices: &[(usize, usize)],
    field: &FieldRef,
) -> Result<ArrayRef, ArrowError> {
    let interleaved = Interleave::<'_, GenericListViewArray<O>>::new(values, indices);

    // Pick whichever strategy produces fewer child elements:
    // - Per-row copy: total = sum of selected sizes. Better for sparse selections.
    // - Concat + offset adjustment: total = sum of source backing array lengths.
    //   Better when rows share backing elements via overlapping offset/size ranges.
    let concat_cost: usize = interleaved.arrays.iter().map(|lv| lv.values().len()).sum();
    let per_row_cost: usize = indices
        .iter()
        .map(|&(a, r)| interleaved.arrays[a].sizes()[r].as_usize())
        .sum();

    if per_row_cost <= concat_cost {
        interleave_list_view_copy::<O>(&interleaved, indices, field)
    } else {
        interleave_list_view_concat::<O>(&interleaved, indices, field)
    }
}

/// Per-row copy: copies each selected row's child elements into a new flat array.
fn interleave_list_view_copy<O: OffsetSizeTrait>(
    interleaved: &Interleave<'_, GenericListViewArray<O>>,
    indices: &[(usize, usize)],
    field: &FieldRef,
) -> Result<ArrayRef, ArrowError> {
    let mut capacity = 0usize;
    let mut offsets = Vec::with_capacity(indices.len());
    let mut sizes = Vec::with_capacity(indices.len());
    for &(array_idx, row_idx) in indices {
        let list = interleaved.arrays[array_idx];
        let size = list.sizes()[row_idx].as_usize();
        offsets.push(
            O::from_usize(capacity).ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
        );
        sizes.push(O::from_usize(size).ok_or_else(|| ArrowError::OffsetOverflowError(size))?);
        capacity += size;
    }

    let child_data: Vec<_> = interleaved
        .arrays
        .iter()
        .map(|list| list.values().to_data())
        .collect();
    let child_data_refs: Vec<_> = child_data.iter().collect();
    let mut mutable_child = MutableArrayData::new(child_data_refs, false, capacity);
    for &(array_idx, row_idx) in indices {
        let list = interleaved.arrays[array_idx];
        let start = list.offsets()[row_idx].as_usize();
        let size = list.sizes()[row_idx].as_usize();
        if size > 0 {
            mutable_child.extend(array_idx, start, start + size);
        }
    }

    Ok(Arc::new(GenericListViewArray::<O>::new(
        field.clone(),
        offsets.into(),
        sizes.into(),
        make_array(mutable_child.freeze()),
        interleaved.nulls.clone(),
    )))
}

/// Concat backing arrays: concatenates all source value arrays and adjusts offsets.
/// Preserves within-source element sharing.
fn interleave_list_view_concat<O: OffsetSizeTrait>(
    interleaved: &Interleave<'_, GenericListViewArray<O>>,
    indices: &[(usize, usize)],
    field: &FieldRef,
) -> Result<ArrayRef, ArrowError> {
    let child_arrays: Vec<&dyn Array> = interleaved
        .arrays
        .iter()
        .map(|lv| lv.values().as_ref())
        .collect();
    let mut base_offsets = Vec::with_capacity(interleaved.arrays.len());
    let mut running = 0usize;
    for lv in &interleaved.arrays {
        base_offsets.push(running);
        running += lv.values().len();
    }
    let combined_values = concat(&child_arrays)?;

    let mut new_offsets = Vec::with_capacity(indices.len());
    let mut new_sizes = Vec::with_capacity(indices.len());
    for &(array_idx, row_idx) in indices {
        let lv = interleaved.arrays[array_idx];
        let adjusted = lv.offsets()[row_idx].as_usize() + base_offsets[array_idx];
        new_offsets.push(
            O::from_usize(adjusted).ok_or_else(|| ArrowError::OffsetOverflowError(adjusted))?,
        );
        new_sizes.push(lv.sizes()[row_idx]);
    }

    Ok(Arc::new(GenericListViewArray::<O>::new(
        field.clone(),
        new_offsets.into(),
        new_sizes.into(),
        combined_values,
        interleaved.nulls.clone(),
    )))
}

/// Fallback implementation of interleave using [`MutableArrayData`]
fn interleave_fallback(
    values: &[&dyn Array],
    indices: &[(usize, usize)],
) -> Result<ArrayRef, ArrowError> {
    let arrays: Vec<_> = values.iter().map(|x| x.to_data()).collect();
    let arrays: Vec<_> = arrays.iter().collect();
    let mut array_data = MutableArrayData::new(arrays, false, indices.len());

    let mut cur_array = indices[0].0;
    let mut start_row_idx = indices[0].1;
    let mut end_row_idx = start_row_idx + 1;

    for (array, row) in indices.iter().skip(1).copied() {
        if array == cur_array && row == end_row_idx {
            // subsequent row in same batch
            end_row_idx += 1;
            continue;
        }

        // emit current batch of rows for current buffer
        array_data.extend(cur_array, start_row_idx, end_row_idx);

        // start new batch of rows
        cur_array = array;
        start_row_idx = row;
        end_row_idx = start_row_idx + 1;
    }

    // emit final batch of rows
    array_data.extend(cur_array, start_row_idx, end_row_idx);
    Ok(make_array(array_data.freeze()))
}

/// Fallback implementation for interleaving dictionaries when it was determined
/// that the dictionary values should not be merged. This implementation concatenates
/// the value slices and recomputes the resulting dictionary keys.
///
/// # Panics
///
/// This function assumes that the combined dictionary values will not overflow the
/// key type. Callers must verify this condition [`should_merge_dictionary_values`]
/// before calling this function.
fn interleave_fallback_dictionary<K: ArrowDictionaryKeyType>(
    dictionaries: &[&DictionaryArray<K>],
    indices: &[(usize, usize)],
) -> Result<ArrayRef, ArrowError> {
    let relative_offsets: Vec<usize> = dictionaries
        .iter()
        .scan(0usize, |offset, dict| {
            let current = *offset;
            *offset += dict.values().len();
            Some(current)
        })
        .collect();
    let all_values: Vec<&dyn Array> = dictionaries.iter().map(|d| d.values().as_ref()).collect();
    let concatenated_values = concat(&all_values)?;

    let any_nulls = dictionaries.iter().any(|d| d.keys().nulls().is_some());
    let (new_keys, nulls) = if any_nulls {
        let mut has_nulls = false;
        let new_keys: Vec<K::Native> = indices
            .iter()
            .map(|(array, row)| {
                let old_keys = dictionaries[*array].keys();
                if old_keys.is_valid(*row) {
                    let old_key = old_keys.values()[*row].as_usize();
                    K::Native::from_usize(relative_offsets[*array] + old_key)
                        .expect("key overflow should be checked by caller")
                } else {
                    has_nulls = true;
                    K::Native::ZERO
                }
            })
            .collect();

        let nulls = if has_nulls {
            let null_buffer = BooleanBuffer::collect_bool(indices.len(), |i| {
                let (array, row) = indices[i];
                dictionaries[array].keys().is_valid(row)
            });
            Some(NullBuffer::new(null_buffer))
        } else {
            None
        };
        (new_keys, nulls)
    } else {
        let new_keys: Vec<K::Native> = indices
            .iter()
            .map(|(array, row)| {
                let old_key = dictionaries[*array].keys().values()[*row].as_usize();
                K::Native::from_usize(relative_offsets[*array] + old_key)
                    .expect("key overflow should be checked by caller")
            })
            .collect();
        (new_keys, None)
    };

    let keys_array = PrimitiveArray::<K>::new(new_keys.into(), nulls);
    // SAFETY: keys_array is constructed from a valid set of keys.
    let array = unsafe { DictionaryArray::new_unchecked(keys_array, concatenated_values) };
    Ok(Arc::new(array))
}

/// Interleave rows by index from multiple [`RecordBatch`] instances and return a new [`RecordBatch`].
///
/// This function will call [`interleave`] on each array of the [`RecordBatch`] instances and assemble a new [`RecordBatch`].
///
/// # Example
/// ```
/// # use std::sync::Arc;
/// # use arrow_array::{StringArray, Int32Array, RecordBatch, UInt32Array};
/// # use arrow_schema::{DataType, Field, Schema};
/// # use arrow_select::interleave::interleave_record_batch;
///
/// let schema = Arc::new(Schema::new(vec![
///     Field::new("a", DataType::Int32, true),
///     Field::new("b", DataType::Utf8, true),
/// ]));
///
/// let batch1 = RecordBatch::try_new(
///     schema.clone(),
///     vec![
///         Arc::new(Int32Array::from(vec![0, 1, 2])),
///         Arc::new(StringArray::from(vec!["a", "b", "c"])),
///     ],
/// ).unwrap();
///
/// let batch2 = RecordBatch::try_new(
///     schema.clone(),
///     vec![
///         Arc::new(Int32Array::from(vec![3, 4, 5])),
///         Arc::new(StringArray::from(vec!["d", "e", "f"])),
///     ],
/// ).unwrap();
///
/// let indices = vec![(0, 1), (1, 2), (0, 0), (1, 1)];
/// let interleaved = interleave_record_batch(&[&batch1, &batch2], &indices).unwrap();
///
/// let expected = RecordBatch::try_new(
///     schema,
///     vec![
///         Arc::new(Int32Array::from(vec![1, 5, 0, 4])),
///         Arc::new(StringArray::from(vec!["b", "f", "a", "e"])),
///     ],
/// ).unwrap();
/// assert_eq!(interleaved, expected);
/// ```
pub fn interleave_record_batch(
    record_batches: &[&RecordBatch],
    indices: &[(usize, usize)],
) -> Result<RecordBatch, ArrowError> {
    let schema = record_batches[0].schema();
    let columns = (0..schema.fields().len())
        .map(|i| {
            let column_values: Vec<&dyn Array> = record_batches
                .iter()
                .map(|batch| batch.column(i).as_ref())
                .collect();
            interleave(&column_values, indices)
        })
        .collect::<Result<Vec<_>, _>>()?;
    RecordBatch::try_new(schema, columns)
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow_array::Int32RunArray;
    use arrow_array::builder::{GenericListBuilder, Int32Builder, PrimitiveRunBuilder};
    use arrow_array::types::Int8Type;
    use arrow_buffer::ScalarBuffer;
    use arrow_schema::Field;

    #[test]
    fn test_primitive() {
        let a = Int32Array::from_iter_values([1, 2, 3, 4]);
        let b = Int32Array::from_iter_values([5, 6, 7]);
        let c = Int32Array::from_iter_values([8, 9, 10]);
        let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap();
        let v = values.as_primitive::<Int32Type>();
        assert_eq!(v.values(), &[4, 4, 10, 8, 6]);
    }

    #[test]
    fn test_primitive_nulls() {
        let a = Int32Array::from_iter_values([1, 2, 3, 4]);
        let b = Int32Array::from_iter([Some(1), Some(4), None]);
        let values = interleave(&[&a, &b], &[(0, 1), (1, 2), (1, 2), (0, 3), (0, 2)]).unwrap();
        let v: Vec<_> = values.as_primitive::<Int32Type>().into_iter().collect();
        assert_eq!(&v, &[Some(2), None, None, Some(4), Some(3)])
    }

    #[test]
    fn test_primitive_empty() {
        let a = Int32Array::from_iter_values([1, 2, 3, 4]);
        let v = interleave(&[&a], &[]).unwrap();
        assert!(v.is_empty());
        assert_eq!(v.data_type(), &DataType::Int32);
    }

    #[test]
    fn test_strings() {
        let a = StringArray::from_iter_values(["a", "b", "c"]);
        let b = StringArray::from_iter_values(["hello", "world", "foo"]);
        let values = interleave(&[&a, &b], &[(0, 2), (0, 2), (1, 0), (1, 1), (0, 1)]).unwrap();
        let v = values.as_string::<i32>();
        let values: Vec<_> = v.into_iter().collect();
        assert_eq!(
            &values,
            &[
                Some("c"),
                Some("c"),
                Some("hello"),
                Some("world"),
                Some("b")
            ]
        )
    }

    #[test]
    fn test_interleave_dictionary() {
        let a = DictionaryArray::<Int32Type>::from_iter(["a", "b", "c", "a", "b"]);
        let b = DictionaryArray::<Int32Type>::from_iter(["a", "c", "a", "c", "a"]);

        // Should not recompute dictionary
        let values =
            interleave(&[&a, &b], &[(0, 2), (0, 2), (0, 2), (1, 0), (1, 1), (0, 1)]).unwrap();
        let v = values.as_dictionary::<Int32Type>();
        assert_eq!(v.values().len(), 5);

        let vc = v.downcast_dict::<StringArray>().unwrap();
        let collected: Vec<_> = vc.into_iter().map(Option::unwrap).collect();
        assert_eq!(&collected, &["c", "c", "c", "a", "c", "b"]);

        // Should recompute dictionary
        let values = interleave(&[&a, &b], &[(0, 2), (0, 2), (1, 1)]).unwrap();
        let v = values.as_dictionary::<Int32Type>();
        assert_eq!(v.values().len(), 1);

        let vc = v.downcast_dict::<StringArray>().unwrap();
        let collected: Vec<_> = vc.into_iter().map(Option::unwrap).collect();
        assert_eq!(&collected, &["c", "c", "c"]);
    }

    #[test]
    fn test_interleave_dictionary_nulls() {
        let input_1_keys = Int32Array::from_iter_values([0, 2, 1, 3]);
        let input_1_values = StringArray::from(vec![Some("foo"), None, Some("bar"), Some("fiz")]);
        let input_1 = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
        let input_2: DictionaryArray<Int32Type> = vec![None].into_iter().collect();

        let expected = vec![Some("fiz"), None, None, Some("foo")];

        let values = interleave(
            &[&input_1 as _, &input_2 as _],
            &[(0, 3), (0, 2), (1, 0), (0, 0)],
        )
        .unwrap();
        let dictionary = values.as_dictionary::<Int32Type>();
        let actual: Vec<Option<&str>> = dictionary
            .downcast_dict::<StringArray>()
            .unwrap()
            .into_iter()
            .collect();

        assert_eq!(actual, expected);
    }

    #[test]
    fn test_interleave_dictionary_overflow_same_values() {
        let values: ArrayRef = Arc::new(StringArray::from_iter_values(
            (0..50).map(|i| format!("v{i}")),
        ));

        // With 3 dictionaries of 50 values each, relative_offsets = [0, 50, 100]
        // Accessing key 49 from dict3 gives 100 + 49 = 149 which overflows Int8
        // (max 127).
        // This test case falls back to interleave_fallback because the
        // dictionaries share the same underlying values slice.
        let dict1 = DictionaryArray::<Int8Type>::new(
            Int8Array::from_iter_values([0, 1, 2]),
            values.clone(),
        );
        let dict2 = DictionaryArray::<Int8Type>::new(
            Int8Array::from_iter_values([0, 1, 2]),
            values.clone(),
        );
        let dict3 =
            DictionaryArray::<Int8Type>::new(Int8Array::from_iter_values([49]), values.clone());

        let indices = &[(0, 0), (1, 0), (2, 0)];
        let result = interleave(&[&dict1, &dict2, &dict3], indices).unwrap();

        let dict_result = result.as_dictionary::<Int8Type>();
        let string_result: Vec<_> = dict_result
            .downcast_dict::<StringArray>()
            .unwrap()
            .into_iter()
            .map(|x| x.unwrap())
            .collect();
        assert_eq!(string_result, vec!["v0", "v0", "v49"]);
    }

    fn test_interleave_lists<O: OffsetSizeTrait>() {
        // [[1, 2], null, [3]]
        let mut a = GenericListBuilder::<O, _>::new(Int32Builder::new());
        a.values().append_value(1);
        a.values().append_value(2);
        a.append(true);
        a.append(false);
        a.values().append_value(3);
        a.append(true);
        let a = a.finish();

        // [[4], null, [5, 6, null]]
        let mut b = GenericListBuilder::<O, _>::new(Int32Builder::new());
        b.values().append_value(4);
        b.append(true);
        b.append(false);
        b.values().append_value(5);
        b.values().append_value(6);
        b.values().append_null();
        b.append(true);
        let b = b.finish();

        let values = interleave(&[&a, &b], &[(0, 2), (0, 1), (1, 0), (1, 2), (1, 1)]).unwrap();
        let v = values
            .as_any()
            .downcast_ref::<GenericListArray<O>>()
            .unwrap();

        // [[3], null, [4], [5, 6, null], null]
        let mut expected = GenericListBuilder::<O, _>::new(Int32Builder::new());
        expected.values().append_value(3);
        expected.append(true);
        expected.append(false);
        expected.values().append_value(4);
        expected.append(true);
        expected.values().append_value(5);
        expected.values().append_value(6);
        expected.values().append_null();
        expected.append(true);
        expected.append(false);
        let expected = expected.finish();

        assert_eq!(v, &expected);
    }

    #[test]
    fn test_lists() {
        test_interleave_lists::<i32>();
    }

    #[test]
    fn test_large_lists() {
        test_interleave_lists::<i64>();
    }

    fn test_interleave_list_views<O: OffsetSizeTrait>() {
        // [[1, 2], null, [3]]
        let mut a = GenericListBuilder::<O, _>::new(Int32Builder::new());
        a.values().append_value(1);
        a.values().append_value(2);
        a.append(true);
        a.append(false);
        a.values().append_value(3);
        a.append(true);
        let a: GenericListViewArray<O> = a.finish().into();

        // [[4], null, [5, 6, null]]
        let mut b = GenericListBuilder::<O, _>::new(Int32Builder::new());
        b.values().append_value(4);
        b.append(true);
        b.append(false);
        b.values().append_value(5);
        b.values().append_value(6);
        b.values().append_null();
        b.append(true);
        let b: GenericListViewArray<O> = b.finish().into();

        let values = interleave(&[&a, &b], &[(0, 2), (0, 1), (1, 0), (1, 2), (1, 1)]).unwrap();
        let v = values
            .as_any()
            .downcast_ref::<GenericListViewArray<O>>()
            .unwrap();

        // [[3], null, [4], [5, 6, null], null]
        let mut expected = GenericListBuilder::<O, _>::new(Int32Builder::new());
        expected.values().append_value(3);
        expected.append(true);
        expected.append(false);
        expected.values().append_value(4);
        expected.append(true);
        expected.values().append_value(5);
        expected.values().append_value(6);
        expected.values().append_null();
        expected.append(true);
        expected.append(false);
        let expected: GenericListViewArray<O> = expected.finish().into();

        assert_eq!(v, &expected);
    }

    #[test]
    fn test_list_views() {
        test_interleave_list_views::<i32>();
    }

    #[test]
    fn test_large_list_views() {
        test_interleave_list_views::<i64>();
    }

    #[test]
    fn test_interleave_list_view_overlapping() {
        let field = Arc::new(Field::new_list_field(DataType::Int64, false));

        // lv_a: 10 rows, two groups of 5 sharing the same backing elements.
        //   rows 0-4 → offset 0, size 5 → [0,1,2,3,4]
        //   rows 5-9 → offset 5, size 5 → [5,6,7,8,9]
        let lv_a = ListViewArray::new(
            Arc::clone(&field),
            ScalarBuffer::from(vec![0i32, 0, 0, 0, 0, 5, 5, 5, 5, 5]),
            ScalarBuffer::from(vec![5i32; 10]),
            Arc::new(Int64Array::from_iter_values(0..10)),
            None,
        );

        // lv_b: 8 rows, two groups of 4 sharing the same backing elements.
        //   rows 0-3 → offset 0, size 3 → [100,101,102]
        //   rows 4-7 → offset 3, size 3 → [103,104,105]
        let lv_b = ListViewArray::new(
            Arc::clone(&field),
            ScalarBuffer::from(vec![0i32, 0, 0, 0, 3, 3, 3, 3]),
            ScalarBuffer::from(vec![3i32; 8]),
            Arc::new(Int64Array::from_iter_values(100..106)),
            None,
        );

        let indices: Vec<(usize, usize)> = vec![
            (0, 0),
            (1, 0),
            (0, 5),
            (1, 4),
            (0, 1),
            (1, 1),
            (0, 6),
            (1, 5),
        ];
        let result = interleave(&[&lv_a as &dyn Array, &lv_b as &dyn Array], &indices).unwrap();
        result
            .to_data()
            .validate_full()
            .expect("result must be valid");

        let result_lv = result.as_list_view::<i32>();
        assert_eq!(result_lv.len(), 8);
        assert_eq!(
            result_lv.value(0).as_primitive::<Int64Type>().values(),
            &[0, 1, 2, 3, 4]
        );
        assert_eq!(
            result_lv.value(1).as_primitive::<Int64Type>().values(),
            &[100, 101, 102]
        );
        assert_eq!(
            result_lv.value(2).as_primitive::<Int64Type>().values(),
            &[5, 6, 7, 8, 9]
        );
        assert_eq!(
            result_lv.value(3).as_primitive::<Int64Type>().values(),
            &[103, 104, 105]
        );

        // Backing elements = sum of source arrays (10 + 6 = 16), not per-row
        // expansion (8 rows × avg ~4 = 32). Overlapping sharing is preserved.
        let total_input_elements = lv_a.values().len() + lv_b.values().len();
        assert_eq!(result_lv.values().len(), total_input_elements);
    }

    #[test]
    fn test_struct_without_nulls() {
        let fields = Fields::from(vec![
            Field::new("number_col", DataType::Int32, false),
            Field::new("string_col", DataType::Utf8, false),
        ]);
        let a = {
            let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
            let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);

            StructArray::try_new(
                fields.clone(),
                vec![Arc::new(number_col), Arc::new(string_col)],
                None,
            )
            .unwrap()
        };

        let b = {
            let number_col = Int32Array::from_iter_values([5, 6, 7]);
            let string_col = StringArray::from_iter_values(["hello", "world", "foo"]);

            StructArray::try_new(
                fields.clone(),
                vec![Arc::new(number_col), Arc::new(string_col)],
                None,
            )
            .unwrap()
        };

        let c = {
            let number_col = Int32Array::from_iter_values([8, 9, 10]);
            let string_col = StringArray::from_iter_values(["x", "y", "z"]);

            StructArray::try_new(
                fields.clone(),
                vec![Arc::new(number_col), Arc::new(string_col)],
                None,
            )
            .unwrap()
        };

        let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap();
        let values_struct = values.as_struct();
        assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
        assert_eq!(values_struct.null_count(), 0);

        let values_number = values_struct.column(0).as_primitive::<Int32Type>();
        assert_eq!(values_number.values(), &[4, 4, 10, 8, 6]);
        let values_string = values_struct.column(1).as_string::<i32>();
        let values_string: Vec<_> = values_string.into_iter().collect();
        assert_eq!(
            &values_string,
            &[Some("d"), Some("d"), Some("z"), Some("x"), Some("world")]
        );
    }

    #[test]
    fn test_struct_with_nulls_in_values() {
        let fields = Fields::from(vec![
            Field::new("number_col", DataType::Int32, true),
            Field::new("string_col", DataType::Utf8, true),
        ]);
        let a = {
            let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
            let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);

            StructArray::try_new(
                fields.clone(),
                vec![Arc::new(number_col), Arc::new(string_col)],
                None,
            )
            .unwrap()
        };

        let b = {
            let number_col = Int32Array::from_iter([Some(1), Some(4), None]);
            let string_col = StringArray::from(vec![Some("hello"), None, Some("foo")]);

            StructArray::try_new(
                fields.clone(),
                vec![Arc::new(number_col), Arc::new(string_col)],
                None,
            )
            .unwrap()
        };

        let values = interleave(&[&a, &b], &[(0, 1), (1, 2), (1, 2), (0, 3), (1, 1)]).unwrap();
        let values_struct = values.as_struct();
        assert_eq!(values_struct.data_type(), &DataType::Struct(fields));

        // The struct itself has no nulls, but the values do
        assert_eq!(values_struct.null_count(), 0);

        let values_number: Vec<_> = values_struct
            .column(0)
            .as_primitive::<Int32Type>()
            .into_iter()
            .collect();
        assert_eq!(values_number, &[Some(2), None, None, Some(4), Some(4)]);

        let values_string = values_struct.column(1).as_string::<i32>();
        let values_string: Vec<_> = values_string.into_iter().collect();
        assert_eq!(
            &values_string,
            &[Some("b"), Some("foo"), Some("foo"), Some("d"), None]
        );
    }

    #[test]
    fn test_struct_with_nulls() {
        let fields = Fields::from(vec![
            Field::new("number_col", DataType::Int32, false),
            Field::new("string_col", DataType::Utf8, false),
        ]);
        let a = {
            let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
            let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);

            StructArray::try_new(
                fields.clone(),
                vec![Arc::new(number_col), Arc::new(string_col)],
                None,
            )
            .unwrap()
        };

        let b = {
            let number_col = Int32Array::from_iter_values([5, 6, 7]);
            let string_col = StringArray::from_iter_values(["hello", "world", "foo"]);

            StructArray::try_new(
                fields.clone(),
                vec![Arc::new(number_col), Arc::new(string_col)],
                Some(NullBuffer::from(&[true, false, true])),
            )
            .unwrap()
        };

        let c = {
            let number_col = Int32Array::from_iter_values([8, 9, 10]);
            let string_col = StringArray::from_iter_values(["x", "y", "z"]);

            StructArray::try_new(
                fields.clone(),
                vec![Arc::new(number_col), Arc::new(string_col)],
                None,
            )
            .unwrap()
        };

        let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (1, 1), (2, 0)]).unwrap();
        let values_struct = values.as_struct();
        assert_eq!(values_struct.data_type(), &DataType::Struct(fields));

        let validity: Vec<bool> = {
            let null_buffer = values_struct.nulls().expect("should_have_nulls");

            null_buffer.iter().collect()
        };
        assert_eq!(validity, &[true, true, true, false, true]);
        let values_number = values_struct.column(0).as_primitive::<Int32Type>();
        assert_eq!(values_number.values(), &[4, 4, 10, 6, 8]);
        let values_string = values_struct.column(1).as_string::<i32>();
        let values_string: Vec<_> = values_string.into_iter().collect();
        assert_eq!(
            &values_string,
            &[Some("d"), Some("d"), Some("z"), Some("world"), Some("x"),]
        );
    }

    #[test]
    fn test_struct_empty() {
        let fields = Fields::from(vec![
            Field::new("number_col", DataType::Int32, false),
            Field::new("string_col", DataType::Utf8, false),
        ]);
        let a = {
            let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
            let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);

            StructArray::try_new(
                fields.clone(),
                vec![Arc::new(number_col), Arc::new(string_col)],
                None,
            )
            .unwrap()
        };
        let v = interleave(&[&a], &[]).unwrap();
        assert!(v.is_empty());
        assert_eq!(v.data_type(), &DataType::Struct(fields));
    }

    #[test]
    fn interleave_sparse_nulls() {
        let values = StringArray::from_iter_values((0..100).map(|x| x.to_string()));
        let keys = Int32Array::from_iter_values(0..10);
        let dict_a = DictionaryArray::new(keys, Arc::new(values));
        let values = StringArray::new_null(0);
        let keys = Int32Array::new_null(10);
        let dict_b = DictionaryArray::new(keys, Arc::new(values));

        let indices = &[(0, 0), (0, 1), (0, 2), (1, 0)];
        let array = interleave(&[&dict_a, &dict_b], indices).unwrap();

        let expected =
            DictionaryArray::<Int32Type>::from_iter(vec![Some("0"), Some("1"), Some("2"), None]);
        assert_eq!(array.as_ref(), &expected)
    }

    #[test]
    fn test_interleave_views() {
        let values = StringArray::from_iter_values([
            "hello",
            "world_long_string_not_inlined",
            "foo",
            "bar",
            "baz",
        ]);
        let view_a = StringViewArray::from(&values);

        let values = StringArray::from_iter_values([
            "test",
            "data",
            "more_long_string_not_inlined",
            "views",
            "here",
        ]);
        let view_b = StringViewArray::from(&values);

        let indices = &[
            (0, 2), // "foo"
            (1, 0), // "test"
            (0, 4), // "baz"
            (1, 3), // "views"
            (0, 1), // "world_long_string_not_inlined"
        ];

        // Test specialized implementation
        let values = interleave(&[&view_a, &view_b], indices).unwrap();
        let result = values.as_string_view();
        assert_eq!(result.data_buffers().len(), 1);

        let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
        let fallback_result = fallback.as_string_view();
        // note that fallback_result has 2 buffers, but only one long enough string to warrant a buffer
        assert_eq!(fallback_result.data_buffers().len(), 2);

        // Convert to strings for easier assertion
        let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();

        let fallback_collected: Vec<_> = fallback_result
            .iter()
            .map(|x| x.map(|s| s.to_string()))
            .collect();

        assert_eq!(&collected, &fallback_collected);

        assert_eq!(
            &collected,
            &[
                Some("foo".to_string()),
                Some("test".to_string()),
                Some("baz".to_string()),
                Some("views".to_string()),
                Some("world_long_string_not_inlined".to_string()),
            ]
        );
    }

    #[test]
    fn test_interleave_views_with_nulls() {
        let values = StringArray::from_iter([
            Some("hello"),
            None,
            Some("foo_long_string_not_inlined"),
            Some("bar"),
            None,
        ]);
        let view_a = StringViewArray::from(&values);

        let values = StringArray::from_iter([
            Some("test"),
            Some("data_long_string_not_inlined"),
            None,
            None,
            Some("here"),
        ]);
        let view_b = StringViewArray::from(&values);

        let indices = &[
            (0, 1), // null
            (1, 2), // null
            (0, 2), // "foo_long_string_not_inlined"
            (1, 3), // null
            (0, 4), // null
        ];

        // Test specialized implementation
        let values = interleave(&[&view_a, &view_b], indices).unwrap();
        let result = values.as_string_view();
        assert_eq!(result.data_buffers().len(), 1);

        let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
        let fallback_result = fallback.as_string_view();

        // Convert to strings for easier assertion
        let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();

        let fallback_collected: Vec<_> = fallback_result
            .iter()
            .map(|x| x.map(|s| s.to_string()))
            .collect();

        assert_eq!(&collected, &fallback_collected);

        assert_eq!(
            &collected,
            &[
                None,
                None,
                Some("foo_long_string_not_inlined".to_string()),
                None,
                None,
            ]
        );
    }

    #[test]
    fn test_interleave_views_multiple_buffers() {
        let str1 = "very_long_string_from_first_buffer".as_bytes();
        let str2 = "very_long_string_from_second_buffer".as_bytes();
        let buffer1 = str1.to_vec().into();
        let buffer2 = str2.to_vec().into();

        let view1 = ByteView::new(str1.len() as u32, &str1[..4])
            .with_buffer_index(0)
            .with_offset(0)
            .as_u128();
        let view2 = ByteView::new(str2.len() as u32, &str2[..4])
            .with_buffer_index(1)
            .with_offset(0)
            .as_u128();
        let view_a =
            StringViewArray::try_new(vec![view1, view2].into(), vec![buffer1, buffer2], None)
                .unwrap();

        let str3 = "another_very_long_string_buffer_three".as_bytes();
        let str4 = "different_long_string_in_buffer_four".as_bytes();
        let buffer3 = str3.to_vec().into();
        let buffer4 = str4.to_vec().into();

        let view3 = ByteView::new(str3.len() as u32, &str3[..4])
            .with_buffer_index(0)
            .with_offset(0)
            .as_u128();
        let view4 = ByteView::new(str4.len() as u32, &str4[..4])
            .with_buffer_index(1)
            .with_offset(0)
            .as_u128();
        let view_b =
            StringViewArray::try_new(vec![view3, view4].into(), vec![buffer3, buffer4], None)
                .unwrap();

        let indices = &[
            (0, 0), // String from first buffer of array A
            (1, 0), // String from first buffer of array B
            (0, 1), // String from second buffer of array A
            (1, 1), // String from second buffer of array B
            (0, 0), // String from first buffer of array A again
            (1, 1), // String from second buffer of array B again
        ];

        // Test interleave
        let values = interleave(&[&view_a, &view_b], indices).unwrap();
        let result = values.as_string_view();

        assert_eq!(
            result.data_buffers().len(),
            4,
            "Expected four buffers (two from each input array)"
        );

        let result_strings: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
        assert_eq!(
            result_strings,
            vec![
                Some("very_long_string_from_first_buffer".to_string()),
                Some("another_very_long_string_buffer_three".to_string()),
                Some("very_long_string_from_second_buffer".to_string()),
                Some("different_long_string_in_buffer_four".to_string()),
                Some("very_long_string_from_first_buffer".to_string()),
                Some("different_long_string_in_buffer_four".to_string()),
            ]
        );

        let views = result.views();
        let buffer_indices: Vec<_> = views
            .iter()
            .map(|raw_view| ByteView::from(*raw_view).buffer_index)
            .collect();

        assert_eq!(
            buffer_indices,
            vec![
                0, // First buffer from array A
                1, // First buffer from array B
                2, // Second buffer from array A
                3, // Second buffer from array B
                0, // First buffer from array A (reused)
                3, // Second buffer from array B (reused)
            ]
        );
    }

    #[test]
    fn test_interleave_run_end_encoded_primitive() {
        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
        builder.extend([1, 1, 2, 2, 2, 3].into_iter().map(Some));
        let a = builder.finish();

        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
        builder.extend([4, 5, 5, 6, 6, 6].into_iter().map(Some));
        let b = builder.finish();

        let indices = &[(0, 1), (1, 0), (0, 4), (1, 2), (0, 5)];
        let result = interleave(&[&a, &b], indices).unwrap();

        // The result should be a RunEndEncoded array
        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));

        // Cast to RunArray to access values
        let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();

        // Verify the logical values by accessing the logical array directly
        let expected = vec![1, 4, 2, 5, 3];
        let mut actual = Vec::new();
        for i in 0..result_run_array.len() {
            let physical_idx = result_run_array.get_physical_index(i);
            let value = result_run_array
                .values()
                .as_primitive::<Int32Type>()
                .value(physical_idx);
            actual.push(value);
        }
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_interleave_run_end_encoded_sliced() {
        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
        builder.extend([1, 1, 2, 2, 2, 3].into_iter().map(Some));
        let a = builder.finish();
        let a = a.slice(2, 3); // [2, 2, 2]

        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
        builder.extend([4, 5, 5, 6, 6, 6].into_iter().map(Some));
        let b = builder.finish();
        let b = b.slice(1, 3); // [5, 5, 6]

        let indices = &[(0, 1), (1, 0), (0, 2), (1, 1), (1, 2)];
        let result = interleave(&[&a, &b], indices).unwrap();

        let result = result.as_run::<Int32Type>();
        let result = result.downcast::<Int32Array>().unwrap();

        let expected = vec![2, 5, 2, 5, 6];
        let actual = result.into_iter().flatten().collect::<Vec<_>>();
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_interleave_run_end_encoded_string() {
        let a: Int32RunArray = vec!["hello", "hello", "world", "world", "foo"]
            .into_iter()
            .collect();
        let b: Int32RunArray = vec!["bar", "baz", "baz", "qux"].into_iter().collect();

        let indices = &[(0, 0), (1, 1), (0, 3), (1, 3), (0, 4)];
        let result = interleave(&[&a, &b], indices).unwrap();

        // The result should be a RunEndEncoded array
        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));

        // Cast to RunArray to access values
        let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();

        // Verify the logical values by accessing the logical array directly
        let expected = vec!["hello", "baz", "world", "qux", "foo"];
        let mut actual = Vec::new();
        for i in 0..result_run_array.len() {
            let physical_idx = result_run_array.get_physical_index(i);
            let value = result_run_array
                .values()
                .as_string::<i32>()
                .value(physical_idx);
            actual.push(value);
        }
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_interleave_run_end_encoded_with_nulls() {
        let a: Int32RunArray = vec![Some("a"), Some("a"), None, None, Some("b")]
            .into_iter()
            .collect();
        let b: Int32RunArray = vec![None, Some("c"), Some("c"), Some("d")]
            .into_iter()
            .collect();

        let indices = &[(0, 1), (1, 0), (0, 2), (1, 3), (0, 4)];
        let result = interleave(&[&a, &b], indices).unwrap();

        // The result should be a RunEndEncoded array
        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));

        // Cast to RunArray to access values
        let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();

        // Verify the logical values by accessing the logical array directly
        let expected = vec![Some("a"), None, None, Some("d"), Some("b")];
        let mut actual = Vec::new();
        for i in 0..result_run_array.len() {
            let physical_idx = result_run_array.get_physical_index(i);
            if result_run_array.values().is_null(physical_idx) {
                actual.push(None);
            } else {
                let value = result_run_array
                    .values()
                    .as_string::<i32>()
                    .value(physical_idx);
                actual.push(Some(value));
            }
        }
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_interleave_run_end_encoded_different_run_types() {
        let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
        builder.extend([1, 1, 2, 3, 3].into_iter().map(Some));
        let a = builder.finish();

        let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
        builder.extend([4, 5, 5, 6].into_iter().map(Some));
        let b = builder.finish();

        let indices = &[(0, 0), (1, 1), (0, 3), (1, 3)];
        let result = interleave(&[&a, &b], indices).unwrap();

        // The result should be a RunEndEncoded array
        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));

        // Cast to RunArray to access values
        let result_run_array: &RunArray<Int16Type> = result.as_any().downcast_ref().unwrap();

        // Verify the logical values by accessing the logical array directly
        let expected = vec![1, 5, 3, 6];
        let mut actual = Vec::new();
        for i in 0..result_run_array.len() {
            let physical_idx = result_run_array.get_physical_index(i);
            let value = result_run_array
                .values()
                .as_primitive::<Int32Type>()
                .value(physical_idx);
            actual.push(value);
        }
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_interleave_run_end_encoded_mixed_run_lengths() {
        let mut builder = PrimitiveRunBuilder::<Int64Type, Int32Type>::new();
        builder.extend([1, 2, 2, 2, 2, 3, 3, 4].into_iter().map(Some));
        let a = builder.finish();

        let mut builder = PrimitiveRunBuilder::<Int64Type, Int32Type>::new();
        builder.extend([5, 5, 5, 6, 7, 7, 8, 8].into_iter().map(Some));
        let b = builder.finish();

        let indices = &[
            (0, 0), // 1
            (1, 2), // 5
            (0, 3), // 2
            (1, 3), // 6
            (0, 6), // 3
            (1, 6), // 8
            (0, 7), // 4
            (1, 4), // 7
        ];
        let result = interleave(&[&a, &b], indices).unwrap();

        // The result should be a RunEndEncoded array
        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));

        // Cast to RunArray to access values
        let result_run_array: &RunArray<Int64Type> = result.as_any().downcast_ref().unwrap();

        // Verify the logical values by accessing the logical array directly
        let expected = vec![1, 5, 2, 6, 3, 8, 4, 7];
        let mut actual = Vec::new();
        for i in 0..result_run_array.len() {
            let physical_idx = result_run_array.get_physical_index(i);
            let value = result_run_array
                .values()
                .as_primitive::<Int32Type>()
                .value(physical_idx);
            actual.push(value);
        }
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_interleave_run_end_encoded_empty_runs() {
        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
        builder.extend([1].into_iter().map(Some));
        let a = builder.finish();

        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
        builder.extend([2, 2, 2].into_iter().map(Some));
        let b = builder.finish();

        let indices = &[(0, 0), (1, 1), (1, 2)];
        let result = interleave(&[&a, &b], indices).unwrap();

        // The result should be a RunEndEncoded array
        assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));

        // Cast to RunArray to access values
        let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();

        // Verify the logical values by accessing the logical array directly
        let expected = vec![1, 2, 2];
        let mut actual = Vec::new();
        for i in 0..result_run_array.len() {
            let physical_idx = result_run_array.get_physical_index(i);
            let value = result_run_array
                .values()
                .as_primitive::<Int32Type>()
                .value(physical_idx);
            actual.push(value);
        }
        assert_eq!(actual, expected);
    }

    #[test]
    fn test_struct_no_fields() {
        let fields = Fields::empty();
        let a = StructArray::try_new_with_length(fields.clone(), vec![], None, 10).unwrap();
        let v = interleave(&[&a], &[(0, 0)]).unwrap();
        assert_eq!(v.len(), 1);
        assert_eq!(v.data_type(), &DataType::Struct(fields));
    }

    #[test]
    fn test_interleave_fallback_dictionary_with_nulls() {
        let input_1_keys = Int32Array::from_iter([Some(0), None, Some(1)]);
        let input_1_values = StringArray::from_iter_values(["foo", "bar"]);
        let dict_a = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));

        let input_2_keys = Int32Array::from_iter([Some(0), Some(1), None]);
        let input_2_values = StringArray::from_iter_values(["baz", "qux"]);
        let dict_b = DictionaryArray::new(input_2_keys, Arc::new(input_2_values));

        let indices = vec![
            (0, 0), // "foo"
            (0, 1), // null
            (1, 0), // "baz"
            (1, 2), // null
            (0, 2), // "bar"
            (1, 1), // "qux"
        ];

        let result =
            interleave_fallback_dictionary::<Int32Type>(&[&dict_a, &dict_b], &indices).unwrap();
        let dict_result = result.as_dictionary::<Int32Type>();

        let string_result = dict_result.downcast_dict::<StringArray>().unwrap();
        let collected: Vec<_> = string_result.into_iter().collect();
        assert_eq!(
            collected,
            vec![
                Some("foo"),
                None,
                Some("baz"),
                None,
                Some("bar"),
                Some("qux")
            ]
        );
    }

    #[test]
    fn test_interleave_bytes_offset_overflow() {
        let indices: Vec<(usize, usize)> = vec![(0, 0); (i32::MAX >> 4) as usize];
        let text = ('a'..='z').collect::<String>();
        let values = StringArray::from(vec![Some(text)]);
        assert!(matches!(
            interleave(&[&values], &indices),
            Err(ArrowError::OffsetOverflowError(_))
        ));
    }

    #[test]
    fn test_interleave_list_offset_overflow() {
        // Build a ListArray<i32> with a single row containing many elements
        let mut builder = GenericListBuilder::<i32, _>::new(Int32Builder::new());
        for i in 0..32 {
            builder.values().append_value(i);
        }
        builder.append(true);
        let list = builder.finish();

        // Interleave enough copies to overflow i32 offsets
        let indices: Vec<(usize, usize)> = vec![(0, 0); (i32::MAX as usize / 32) + 1];
        assert!(matches!(
            interleave(&[&list], &indices),
            Err(ArrowError::OffsetOverflowError(_))
        ));
    }

    #[test]
    fn test_interleave_list_view() {
        // `interleave` for ListView falls through to `interleave_fallback`, which uses
        // `MutableArrayData`. `list_view::build_extend` copies offsets/sizes but never
        // extends the child array, so the result contains offsets/sizes that reference
        // positions in the now-absent original child arrays while the child is empty.
        //
        // lv_a: [[1, 2], [3]]   (values=[1,2,3], offsets=[0,2], sizes=[2,1])
        // lv_b: [[4, 5, 6]]     (values=[4,5,6], offsets=[0],   sizes=[3])
        // interleave at [(0,0), (1,0), (0,1)] should produce [[1, 2], [4, 5, 6], [3]]
        let field = Arc::new(Field::new_list_field(DataType::Int64, false));

        let lv_a = ListViewArray::new(
            Arc::clone(&field),
            ScalarBuffer::from(vec![0i32, 2]),
            ScalarBuffer::from(vec![2i32, 1]),
            Arc::new(Int64Array::from(vec![1_i64, 2, 3])),
            None,
        );
        let lv_b = ListViewArray::new(
            field,
            ScalarBuffer::from(vec![0i32]),
            ScalarBuffer::from(vec![3i32]),
            Arc::new(Int64Array::from(vec![4_i64, 5, 6])),
            None,
        );

        let result = interleave(
            &[&lv_a as &dyn Array, &lv_b as &dyn Array],
            &[(0, 0), (1, 0), (0, 1)],
        )
        .unwrap();

        result
            .to_data()
            .validate_full()
            .expect("interleaved ListViewArray must be internally consistent");

        let result_lv = result.as_list_view::<i32>();
        assert_eq!(result_lv.len(), 3);
        assert_eq!(
            result_lv.value(0).as_primitive::<Int64Type>().values(),
            &[1, 2]
        );
        assert_eq!(
            result_lv.value(1).as_primitive::<Int64Type>().values(),
            &[4, 5, 6]
        );
        assert_eq!(
            result_lv.value(2).as_primitive::<Int64Type>().values(),
            &[3]
        );
    }
}