arrow_data/transform/

mod.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.

//! Low-level array data abstractions.
//!
//! Provides utilities for creating, manipulating, and converting Arrow arrays
//! made of primitive types, strings, and nested types.

use super::{ArrayData, ArrayDataBuilder, ByteView, data::new_buffers};
use crate::bit_mask::set_bits;
use arrow_buffer::buffer::{BooleanBuffer, NullBuffer};
use arrow_buffer::{ArrowNativeType, Buffer, IntervalMonthDayNano, MutableBuffer, bit_util, i256};
use arrow_schema::{ArrowError, DataType, IntervalUnit, UnionMode};
use half::f16;
use num_integer::Integer;
use std::mem;

mod boolean;
mod fixed_binary;
mod fixed_size_list;
mod list;
mod list_view;
mod null;
mod primitive;
mod run;
mod structure;
mod union;
mod utils;
mod variable_size;

type ExtendNullBits<'a> = Box<dyn Fn(&mut _MutableArrayData, usize, usize) + 'a>;
// function that extends `[start..start+len]` to the mutable array.
// this is dynamic because different data_types influence how buffers and children are extended.
type Extend<'a> =
    Box<dyn Fn(&mut _MutableArrayData, usize, usize, usize) -> Result<(), ArrowError> + 'a>;

type ExtendNulls = Box<dyn Fn(&mut _MutableArrayData, usize) -> Result<(), ArrowError>>;

/// A mutable [ArrayData] that knows how to freeze itself into an [ArrayData].
/// This is just a data container.
#[derive(Debug)]
struct _MutableArrayData<'a> {
    pub data_type: DataType,
    pub null_count: usize,

    pub len: usize,
    pub null_buffer: Option<MutableBuffer>,

    // arrow specification only allows up to 3 buffers (2 ignoring the nulls above).
    // Thus, we place them in the stack to avoid bound checks and greater data locality.
    pub buffer1: MutableBuffer,
    pub buffer2: MutableBuffer,
    pub child_data: Vec<MutableArrayData<'a>>,
}

impl _MutableArrayData<'_> {
    fn null_buffer(&mut self) -> &mut MutableBuffer {
        self.null_buffer
            .as_mut()
            .expect("MutableArrayData not nullable")
    }
}

fn build_extend_null_bits(array: &ArrayData, use_nulls: bool) -> ExtendNullBits<'_> {
    if let Some(nulls) = array.nulls() {
        let bytes = nulls.validity();
        Box::new(move |mutable, start, len| {
            let mutable_len = mutable.len;
            let out = mutable.null_buffer();
            utils::resize_for_bits(out, mutable_len + len);
            mutable.null_count += set_bits(
                out.as_slice_mut(),
                bytes,
                mutable_len,
                nulls.offset() + start,
                len,
            );
        })
    } else if use_nulls {
        Box::new(|mutable, _, len| {
            let mutable_len = mutable.len;
            let out = mutable.null_buffer();
            utils::resize_for_bits(out, mutable_len + len);
            let write_data = out.as_slice_mut();
            (0..len).for_each(|i| {
                bit_util::set_bit(write_data, mutable_len + i);
            });
        })
    } else {
        Box::new(|_, _, _| {})
    }
}

/// Efficiently create an [ArrayData] from one or more existing [ArrayData]s by
/// copying chunks.
///
/// The main use case of this struct is to perform unary operations to arrays of
/// arbitrary types, such as `filter` and `take`.
///
/// # Example
/// ```
/// use arrow_buffer::Buffer;
/// use arrow_data::ArrayData;
/// use arrow_data::transform::MutableArrayData;
/// use arrow_schema::DataType;
/// fn i32_array(values: &[i32]) -> ArrayData {
///   ArrayData::try_new(DataType::Int32, values.len(), None, 0, vec![Buffer::from_slice_ref(values)], vec![]).unwrap()
/// }
/// let arr1  = i32_array(&[1, 2, 3, 4, 5]);
/// let arr2  = i32_array(&[6, 7, 8, 9, 10]);
/// // Create a mutable array for copying values from arr1 and arr2, with a capacity for 6 elements
/// let capacity = 3 * std::mem::size_of::<i32>();
/// let mut mutable = MutableArrayData::new(vec![&arr1, &arr2], false, 10);
/// // Copy the first 3 elements from arr1
/// mutable.extend(0, 0, 3);
/// // Copy the last 3 elements from arr2
/// mutable.extend(1, 2, 5);
/// // Complete the MutableArrayData into a new ArrayData
/// let frozen = mutable.freeze();
/// assert_eq!(frozen, i32_array(&[1, 2, 3, 8, 9, 10]));
/// ```
pub struct MutableArrayData<'a> {
    /// Input arrays: the data being read FROM.
    ///
    /// Note all actual reads of the arrays go through the closures for extending
    /// values and nulls; these references are only kept for bounds checking.
    arrays: Vec<&'a ArrayData>,

    /// In progress output array: The data being written TO
    ///
    /// Note these fields are in a separate struct, [_MutableArrayData], as they
    /// cannot be in [MutableArrayData] itself due to mutability invariants (interior
    /// mutability): [MutableArrayData] contains a function that can only mutate
    /// [_MutableArrayData], not [MutableArrayData] itself
    data: _MutableArrayData<'a>,

    /// The child data of the `Array` in Dictionary arrays.
    ///
    /// This is not stored in `_MutableArrayData` because these values are
    /// constant and only needed at the end, when freezing [_MutableArrayData].
    dictionary: Option<ArrayData>,

    /// Variadic data buffers referenced by views.
    ///
    /// Note this this is not stored in `_MutableArrayData` because these values
    /// are constant and only needed at the end, when freezing
    /// [_MutableArrayData]
    variadic_data_buffers: Vec<Buffer>,

    /// function used to extend output array with values from input arrays.
    ///
    /// This function's lifetime is bound to the input arrays because it reads
    /// values from them.
    extend_values: Vec<Extend<'a>>,

    /// function used to extend the output array with nulls from input arrays.
    ///
    /// This function's lifetime is bound to the input arrays because it reads
    /// nulls from it.
    extend_null_bits: Vec<ExtendNullBits<'a>>,

    /// function used to extend the output array with null elements.
    ///
    /// This function is independent of the arrays and therefore has no lifetime.
    extend_nulls: ExtendNulls,
}

impl std::fmt::Debug for MutableArrayData<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        // ignores the closures.
        f.debug_struct("MutableArrayData")
            .field("data", &self.data)
            .finish()
    }
}

/// Builds an extend that adds `offset` to the source primitive
/// Additionally validates that `max` fits into the
/// the underlying primitive returning None if not
fn build_extend_dictionary(array: &ArrayData, offset: usize, max: usize) -> Option<Extend<'_>> {
    macro_rules! validate_and_build {
        ($dt: ty) => {{
            // `max` is the merged dictionary length; the largest key index is
            // `max - 1`, so the key type only needs to hold `max - 1` (e.g. 256
            // values use keys 0..=255, which fit in u8).
            let _: $dt = max.saturating_sub(1).try_into().ok()?;
            let offset: $dt = offset.try_into().ok()?;
            Some(primitive::build_extend_with_offset(array, offset))
        }};
    }
    match array.data_type() {
        DataType::Dictionary(child_data_type, _) => match child_data_type.as_ref() {
            DataType::UInt8 => validate_and_build!(u8),
            DataType::UInt16 => validate_and_build!(u16),
            DataType::UInt32 => validate_and_build!(u32),
            DataType::UInt64 => validate_and_build!(u64),
            DataType::Int8 => validate_and_build!(i8),
            DataType::Int16 => validate_and_build!(i16),
            DataType::Int32 => validate_and_build!(i32),
            DataType::Int64 => validate_and_build!(i64),
            _ => unreachable!(),
        },
        _ => None,
    }
}

/// Builds an extend that adds `buffer_offset` to any buffer indices encountered
fn build_extend_view(array: &ArrayData, buffer_offset: u32) -> Extend<'_> {
    let views = array.buffer::<u128>(0);
    Box::new(
        move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
            mutable
                .buffer1
                .extend(views[start..start + len].iter().map(|v| {
                    let len = *v as u32;
                    if len <= 12 {
                        return *v; // Stored inline
                    }
                    let mut view = ByteView::from(*v);
                    view.buffer_index += buffer_offset;
                    view.into()
                }));
            Ok(())
        },
    )
}

fn build_extend(array: &ArrayData) -> Extend<'_> {
    match array.data_type() {
        DataType::Null => null::build_extend(array),
        DataType::Boolean => boolean::build_extend(array),
        DataType::UInt8 => primitive::build_extend::<u8>(array),
        DataType::UInt16 => primitive::build_extend::<u16>(array),
        DataType::UInt32 => primitive::build_extend::<u32>(array),
        DataType::UInt64 => primitive::build_extend::<u64>(array),
        DataType::Int8 => primitive::build_extend::<i8>(array),
        DataType::Int16 => primitive::build_extend::<i16>(array),
        DataType::Int32 => primitive::build_extend::<i32>(array),
        DataType::Int64 => primitive::build_extend::<i64>(array),
        DataType::Float32 => primitive::build_extend::<f32>(array),
        DataType::Float64 => primitive::build_extend::<f64>(array),
        DataType::Date32 | DataType::Time32(_) | DataType::Interval(IntervalUnit::YearMonth) => {
            primitive::build_extend::<i32>(array)
        }
        DataType::Date64
        | DataType::Time64(_)
        | DataType::Timestamp(_, _)
        | DataType::Duration(_)
        | DataType::Interval(IntervalUnit::DayTime) => primitive::build_extend::<i64>(array),
        DataType::Interval(IntervalUnit::MonthDayNano) => {
            primitive::build_extend::<IntervalMonthDayNano>(array)
        }
        DataType::Decimal32(_, _) => primitive::build_extend::<i32>(array),
        DataType::Decimal64(_, _) => primitive::build_extend::<i64>(array),
        DataType::Decimal128(_, _) => primitive::build_extend::<i128>(array),
        DataType::Decimal256(_, _) => primitive::build_extend::<i256>(array),
        DataType::Utf8 | DataType::Binary => variable_size::build_extend::<i32>(array),
        DataType::LargeUtf8 | DataType::LargeBinary => variable_size::build_extend::<i64>(array),
        DataType::BinaryView | DataType::Utf8View => unreachable!("should use build_extend_view"),
        DataType::Map(_, _) | DataType::List(_) => list::build_extend::<i32>(array),
        DataType::LargeList(_) => list::build_extend::<i64>(array),
        DataType::ListView(_) => list_view::build_extend::<i32>(array),
        DataType::LargeListView(_) => list_view::build_extend::<i64>(array),
        DataType::Dictionary(_, _) => unreachable!("should use build_extend_dictionary"),
        DataType::Struct(_) => structure::build_extend(array),
        DataType::FixedSizeBinary(_) => fixed_binary::build_extend(array),
        DataType::Float16 => primitive::build_extend::<f16>(array),
        DataType::FixedSizeList(_, _) => fixed_size_list::build_extend(array),
        DataType::Union(_, mode) => match mode {
            UnionMode::Sparse => union::build_extend_sparse(array),
            UnionMode::Dense => union::build_extend_dense(array),
        },
        DataType::RunEndEncoded(_, _) => run::build_extend(array),
    }
}

fn build_extend_nulls(data_type: &DataType) -> ExtendNulls {
    Box::new(match data_type {
        DataType::Null => null::extend_nulls,
        DataType::Boolean => boolean::extend_nulls,
        DataType::UInt8 => primitive::extend_nulls::<u8>,
        DataType::UInt16 => primitive::extend_nulls::<u16>,
        DataType::UInt32 => primitive::extend_nulls::<u32>,
        DataType::UInt64 => primitive::extend_nulls::<u64>,
        DataType::Int8 => primitive::extend_nulls::<i8>,
        DataType::Int16 => primitive::extend_nulls::<i16>,
        DataType::Int32 => primitive::extend_nulls::<i32>,
        DataType::Int64 => primitive::extend_nulls::<i64>,
        DataType::Float32 => primitive::extend_nulls::<f32>,
        DataType::Float64 => primitive::extend_nulls::<f64>,
        DataType::Date32 | DataType::Time32(_) | DataType::Interval(IntervalUnit::YearMonth) => {
            primitive::extend_nulls::<i32>
        }
        DataType::Date64
        | DataType::Time64(_)
        | DataType::Timestamp(_, _)
        | DataType::Duration(_)
        | DataType::Interval(IntervalUnit::DayTime) => primitive::extend_nulls::<i64>,
        DataType::Interval(IntervalUnit::MonthDayNano) => {
            primitive::extend_nulls::<IntervalMonthDayNano>
        }
        DataType::Decimal32(_, _) => primitive::extend_nulls::<i32>,
        DataType::Decimal64(_, _) => primitive::extend_nulls::<i64>,
        DataType::Decimal128(_, _) => primitive::extend_nulls::<i128>,
        DataType::Decimal256(_, _) => primitive::extend_nulls::<i256>,
        DataType::Utf8 | DataType::Binary => variable_size::extend_nulls::<i32>,
        DataType::LargeUtf8 | DataType::LargeBinary => variable_size::extend_nulls::<i64>,
        DataType::BinaryView | DataType::Utf8View => primitive::extend_nulls::<u128>,
        DataType::Map(_, _) | DataType::List(_) => list::extend_nulls::<i32>,
        DataType::LargeList(_) => list::extend_nulls::<i64>,
        DataType::ListView(_) => list_view::extend_nulls::<i32>,
        DataType::LargeListView(_) => list_view::extend_nulls::<i64>,
        DataType::Dictionary(child_data_type, _) => match child_data_type.as_ref() {
            DataType::UInt8 => primitive::extend_nulls::<u8>,
            DataType::UInt16 => primitive::extend_nulls::<u16>,
            DataType::UInt32 => primitive::extend_nulls::<u32>,
            DataType::UInt64 => primitive::extend_nulls::<u64>,
            DataType::Int8 => primitive::extend_nulls::<i8>,
            DataType::Int16 => primitive::extend_nulls::<i16>,
            DataType::Int32 => primitive::extend_nulls::<i32>,
            DataType::Int64 => primitive::extend_nulls::<i64>,
            _ => unreachable!(),
        },
        DataType::Struct(_) => structure::extend_nulls,
        DataType::FixedSizeBinary(_) => fixed_binary::extend_nulls,
        DataType::Float16 => primitive::extend_nulls::<f16>,
        DataType::FixedSizeList(_, _) => fixed_size_list::extend_nulls,
        DataType::Union(_, mode) => match mode {
            UnionMode::Sparse => union::extend_nulls_sparse,
            UnionMode::Dense => union::extend_nulls_dense,
        },
        DataType::RunEndEncoded(_, _) => run::extend_nulls,
    })
}

fn preallocate_offset_and_binary_buffer<Offset: ArrowNativeType + Integer>(
    capacity: usize,
    binary_size: usize,
) -> [MutableBuffer; 2] {
    // offsets
    let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<Offset>());
    // safety: `unsafe` code assumes that this buffer is initialized with one element
    buffer.push(Offset::zero());

    [
        buffer,
        MutableBuffer::new(binary_size * mem::size_of::<u8>()),
    ]
}

/// Define capacities to pre-allocate for child data or data buffers.
#[derive(Debug, Clone)]
pub enum Capacities {
    /// Binary, Utf8 and LargeUtf8 data types
    ///
    /// Defines
    /// * the capacity of the array offsets
    /// * the capacity of the binary/ str buffer
    Binary(usize, Option<usize>),
    /// List, LargeList and Map data types
    ///
    /// Defines
    /// * the capacity of the array offsets
    /// * the capacity of the child data
    ///
    /// For Map the child data is the entries [`DataType::Struct`], so the child
    /// capacity is a [`Capacities::Struct`] holding the key and value capacities.
    List(usize, Option<Box<Capacities>>),
    /// Struct type
    ///
    /// Defines
    /// * the capacity of the array
    /// * the capacities of the fields
    Struct(usize, Option<Vec<Capacities>>),
    /// Dictionary type
    ///
    /// Defines
    /// * the capacity of the array/keys
    /// * the capacity of the values
    Dictionary(usize, Option<Box<Capacities>>),
    /// Don't preallocate inner buffers and rely on array growth strategy
    Array(usize),
}

impl<'a> MutableArrayData<'a> {
    /// Returns a new [MutableArrayData] with capacity to `capacity` slots and
    /// specialized to create an [ArrayData] from multiple `arrays`.
    ///
    /// # Arguments
    /// * `arrays` - the source arrays to copy from
    /// * `use_nulls` - a flag indicating whether the caller intends to call `extend_nulls`.
    ///   Note: null-handling is enabled automatically if any source array contains nulls.
    /// * `capacity` - the preallocated capacity of the output array, in slots (number of elements)
    ///
    /// if `use_nulls` is `false` and no source arrays contains nulls, calling
    /// [MutableArrayData::extend_nulls] or [MutableArrayData::try_extend_nulls] will panic.
    pub fn new(arrays: Vec<&'a ArrayData>, use_nulls: bool, capacity: usize) -> Self {
        Self::with_capacities(arrays, use_nulls, Capacities::Array(capacity))
    }

    /// Fallible variant of [MutableArrayData::new].
    ///
    /// Unlike [MutableArrayData::new], this does not panic when merging dictionary
    /// arrays whose combined values would overflow the dictionary key type. Instead,
    /// it returns an error, letting callers (e.g. [`interleave`](crate) / `concat`)
    /// surface it as a normal error.
    pub fn try_new(
        arrays: Vec<&'a ArrayData>,
        use_nulls: bool,
        capacity: usize,
    ) -> Result<Self, ArrowError> {
        Self::try_with_capacities(arrays, use_nulls, Capacities::Array(capacity))
    }

    /// Similar to [MutableArrayData::new], but lets users define the
    /// preallocated capacities of the array with more granularity.
    ///
    /// See [MutableArrayData::new] for more information on the arguments.
    ///
    /// # Panics
    ///
    /// * if the given `capacities` don't match the data type of `arrays`
    /// * if a [Capacities] variant is not yet supported
    /// * when merging dictionary arrays whose combined values overflow the
    ///   dictionary key type — see [MutableArrayData::try_with_capacities] for a
    ///   fallible variant
    pub fn with_capacities(
        arrays: Vec<&'a ArrayData>,
        use_nulls: bool,
        capacities: Capacities,
    ) -> Self {
        Self::try_with_capacities(arrays, use_nulls, capacities)
            .expect("MutableArrayData::new is infallible")
    }

    /// Fallible variant of [MutableArrayData::with_capacities].
    ///
    /// Returns an error instead of panicking when merging dictionary arrays whose
    /// combined values would overflow the dictionary key type. Still panics for
    /// other unsupported combinations (inconsistent input types, unsupported
    /// `Capacities` variants) as documented on [MutableArrayData::with_capacities].
    pub fn try_with_capacities(
        arrays: Vec<&'a ArrayData>,
        use_nulls: bool,
        capacities: Capacities,
    ) -> Result<Self, ArrowError> {
        let data_type = arrays[0].data_type();

        for a in arrays.iter().skip(1) {
            assert_eq!(
                data_type,
                a.data_type(),
                "Arrays with inconsistent types passed to MutableArrayData"
            )
        }

        // if any of the arrays has nulls, insertions from any array requires setting bits
        // as there is at least one array with nulls.
        let use_nulls = use_nulls | arrays.iter().any(|array| array.null_count() > 0);

        let mut array_capacity;

        let [buffer1, buffer2] = match (data_type, &capacities) {
            (
                DataType::LargeUtf8 | DataType::LargeBinary,
                Capacities::Binary(capacity, Some(value_cap)),
            ) => {
                array_capacity = *capacity;
                preallocate_offset_and_binary_buffer::<i64>(*capacity, *value_cap)
            }
            (DataType::Utf8 | DataType::Binary, Capacities::Binary(capacity, Some(value_cap))) => {
                array_capacity = *capacity;
                preallocate_offset_and_binary_buffer::<i32>(*capacity, *value_cap)
            }
            (_, Capacities::Array(capacity)) => {
                array_capacity = *capacity;
                new_buffers(data_type, *capacity)
            }
            (
                DataType::List(_)
                | DataType::LargeList(_)
                | DataType::ListView(_)
                | DataType::LargeListView(_)
                | DataType::FixedSizeList(_, _)
                | DataType::Map(_, _),
                Capacities::List(capacity, _),
            ) => {
                array_capacity = *capacity;
                new_buffers(data_type, *capacity)
            }
            (DataType::Struct(_), Capacities::Struct(capacity, _)) => {
                array_capacity = *capacity;
                new_buffers(data_type, *capacity)
            }
            _ => panic!("Capacities: {capacities:?} not yet supported"),
        };

        let child_data = match &data_type {
            DataType::Decimal32(_, _)
            | DataType::Decimal64(_, _)
            | DataType::Decimal128(_, _)
            | DataType::Decimal256(_, _)
            | DataType::Null
            | DataType::Boolean
            | DataType::UInt8
            | DataType::UInt16
            | DataType::UInt32
            | DataType::UInt64
            | DataType::Int8
            | DataType::Int16
            | DataType::Int32
            | DataType::Int64
            | DataType::Float16
            | DataType::Float32
            | DataType::Float64
            | DataType::Date32
            | DataType::Date64
            | DataType::Time32(_)
            | DataType::Time64(_)
            | DataType::Duration(_)
            | DataType::Timestamp(_, _)
            | DataType::Utf8
            | DataType::Binary
            | DataType::LargeUtf8
            | DataType::LargeBinary
            | DataType::BinaryView
            | DataType::Utf8View
            | DataType::Interval(_)
            | DataType::FixedSizeBinary(_) => vec![],
            DataType::Map(_, _)
            | DataType::List(_)
            | DataType::LargeList(_)
            | DataType::ListView(_)
            | DataType::LargeListView(_) => {
                let children = arrays
                    .iter()
                    .map(|array| &array.child_data()[0])
                    .collect::<Vec<_>>();

                let capacities =
                    if let Capacities::List(capacity, ref child_capacities) = capacities {
                        child_capacities
                            .clone()
                            .map(|c| *c)
                            .unwrap_or(Capacities::Array(capacity))
                    } else {
                        Capacities::Array(array_capacity)
                    };

                vec![MutableArrayData::try_with_capacities(
                    children, use_nulls, capacities,
                )?]
            }
            // the dictionary type just appends keys and clones the values.
            DataType::Dictionary(_, _) => vec![],
            DataType::Struct(fields) => match capacities {
                Capacities::Struct(capacity, Some(ref child_capacities)) => {
                    array_capacity = capacity;
                    (0..fields.len())
                        .zip(child_capacities)
                        .map(|(i, child_cap)| {
                            let child_arrays = arrays
                                .iter()
                                .map(|array| &array.child_data()[i])
                                .collect::<Vec<_>>();
                            MutableArrayData::try_with_capacities(
                                child_arrays,
                                use_nulls,
                                child_cap.clone(),
                            )
                        })
                        .collect::<Result<Vec<_>, _>>()?
                }
                Capacities::Struct(capacity, None) => {
                    array_capacity = capacity;
                    (0..fields.len())
                        .map(|i| {
                            let child_arrays = arrays
                                .iter()
                                .map(|array| &array.child_data()[i])
                                .collect::<Vec<_>>();
                            MutableArrayData::try_new(child_arrays, use_nulls, capacity)
                        })
                        .collect::<Result<Vec<_>, _>>()?
                }
                _ => (0..fields.len())
                    .map(|i| {
                        let child_arrays = arrays
                            .iter()
                            .map(|array| &array.child_data()[i])
                            .collect::<Vec<_>>();
                        MutableArrayData::try_new(child_arrays, use_nulls, array_capacity)
                    })
                    .collect::<Result<Vec<_>, _>>()?,
            },
            DataType::RunEndEncoded(_, _) => {
                let run_ends_child = arrays
                    .iter()
                    .map(|array| &array.child_data()[0])
                    .collect::<Vec<_>>();
                let value_child = arrays
                    .iter()
                    .map(|array| &array.child_data()[1])
                    .collect::<Vec<_>>();
                vec![
                    MutableArrayData::try_new(run_ends_child, false, array_capacity)?,
                    MutableArrayData::try_new(value_child, use_nulls, array_capacity)?,
                ]
            }
            DataType::FixedSizeList(_, size) => {
                let children = arrays
                    .iter()
                    .map(|array| &array.child_data()[0])
                    .collect::<Vec<_>>();
                let capacities =
                    if let Capacities::List(capacity, ref child_capacities) = capacities {
                        child_capacities
                            .clone()
                            .map(|c| *c)
                            .unwrap_or(Capacities::Array(capacity * *size as usize))
                    } else {
                        Capacities::Array(array_capacity * *size as usize)
                    };
                vec![MutableArrayData::try_with_capacities(
                    children, use_nulls, capacities,
                )?]
            }
            DataType::Union(fields, _) => (0..fields.len())
                .map(|i| {
                    let child_arrays = arrays
                        .iter()
                        .map(|array| &array.child_data()[i])
                        .collect::<Vec<_>>();
                    MutableArrayData::try_new(child_arrays, use_nulls, array_capacity)
                })
                .collect::<Result<Vec<_>, _>>()?,
        };

        // Get the dictionary if any, and if it is a concatenation of multiple
        let (dictionary, dict_concat) = match &data_type {
            DataType::Dictionary(_, _) => {
                // If more than one dictionary, concatenate dictionaries together
                let dict_concat = !arrays
                    .windows(2)
                    .all(|a| a[0].child_data()[0].ptr_eq(&a[1].child_data()[0]));

                match dict_concat {
                    false => (Some(arrays[0].child_data()[0].clone()), false),
                    true => {
                        if let Capacities::Dictionary(_, _) = capacities {
                            panic!("dictionary capacity not yet supported")
                        }
                        let dictionaries: Vec<_> =
                            arrays.iter().map(|array| &array.child_data()[0]).collect();
                        let lengths: Vec<_> = dictionaries
                            .iter()
                            .map(|dictionary| dictionary.len())
                            .collect();
                        let capacity = lengths.iter().sum();

                        let mut mutable = MutableArrayData::new(dictionaries, false, capacity);

                        for (i, len) in lengths.iter().enumerate() {
                            mutable.try_extend(i, 0, *len).expect(
                                "extend failed while building dictionary; \
                                 this is a bug in MutableArrayData",
                            )
                        }

                        (Some(mutable.freeze()), true)
                    }
                }
            }
            _ => (None, false),
        };

        let variadic_data_buffers = match &data_type {
            DataType::BinaryView | DataType::Utf8View => arrays
                .iter()
                .flat_map(|x| x.buffers().iter().skip(1))
                .map(Buffer::clone)
                .collect(),
            _ => vec![],
        };

        let extend_nulls = build_extend_nulls(data_type);

        let extend_null_bits = arrays
            .iter()
            .map(|array| build_extend_null_bits(array, use_nulls))
            .collect();

        let null_buffer = use_nulls.then(|| {
            let null_bytes = bit_util::ceil(array_capacity, 8);
            MutableBuffer::from_len_zeroed(null_bytes)
        });

        let extend_values = match &data_type {
            DataType::Dictionary(_, _) => {
                let mut next_offset = 0;
                let extend_values: Result<Vec<_>, _> = arrays
                    .iter()
                    .map(|array| {
                        let offset = next_offset;
                        let dict_len = array.child_data()[0].len();

                        if dict_concat {
                            next_offset += dict_len;
                        }

                        build_extend_dictionary(array, offset, offset + dict_len)
                            .ok_or(ArrowError::DictionaryKeyOverflowError)
                    })
                    .collect();

                extend_values?
            }
            DataType::BinaryView | DataType::Utf8View => {
                let mut next_offset = 0u32;
                arrays
                    .iter()
                    .map(|arr| {
                        let num_data_buffers = (arr.buffers().len() - 1) as u32;
                        let offset = next_offset;
                        next_offset = next_offset
                            .checked_add(num_data_buffers)
                            .expect("view buffer index overflow");
                        build_extend_view(arr, offset)
                    })
                    .collect()
            }
            _ => arrays.iter().map(|array| build_extend(array)).collect(),
        };

        let data = _MutableArrayData {
            data_type: data_type.clone(),
            len: 0,
            null_count: 0,
            null_buffer,
            buffer1,
            buffer2,
            child_data,
        };
        Ok(Self {
            arrays,
            data,
            dictionary,
            variadic_data_buffers,
            extend_values,
            extend_null_bits,
            extend_nulls,
        })
    }

    /// Extends the in progress array with a region of the input arrays, returning an error on
    /// overflow.
    ///
    /// # Arguments
    /// * `index` - the index of array that you want to copy values from
    /// * `start` - the start index of the chunk (inclusive)
    /// * `end` - the end index of the chunk (exclusive)
    ///
    /// # Errors
    /// Returns an error if
    /// * `index` >= the number of source arrays,
    /// * `start..end` is not a valid range within the `index`th array, or
    /// * offset arithmetic overflows the underlying integer type.
    pub fn try_extend(&mut self, index: usize, start: usize, end: usize) -> Result<(), ArrowError> {
        let Some(array_len) = self.arrays.get(index).map(|array| array.len()) else {
            return Err(ArrowError::InvalidArgumentError(format!(
                "Source array index {index} is out of bounds: there are {} source arrays",
                self.arrays.len()
            )));
        };
        if end < start || array_len < end {
            return Err(ArrowError::InvalidArgumentError(format!(
                "Invalid range {start}..{end} for source array {index} of length {array_len}"
            )));
        }

        let len = end - start;
        (self.extend_null_bits[index])(&mut self.data, start, len);
        // Snapshot buffer lengths before attempting the extend so we can roll
        // back to a consistent state if it fails.
        let buf1_len = self.data.buffer1.len();
        let buf2_len = self.data.buffer2.len();
        if let Err(e) = (self.extend_values[index])(&mut self.data, index, start, len) {
            // Restore buffers to their pre-call lengths so the array remains
            // in a valid state for the caller to inspect or retry.
            self.data.buffer1.truncate(buf1_len);
            self.data.buffer2.truncate(buf2_len);
            return Err(e);
        }
        self.data.len += len;
        Ok(())
    }

    /// Extends the in progress array with a region of the input arrays.
    ///
    /// # Panics
    /// This function panics if
    /// * `index` >= the number of source arrays,
    /// * `start..end` is not a valid range within the `index`th array, or
    /// * the offset type overflows (e.g. more than 2 GiB in a `StringArray`).
    #[deprecated(
        since = "59.0.0",
        note = "Use `try_extend` which returns an error on overflow instead of panicking"
    )]
    pub fn extend(&mut self, index: usize, start: usize, end: usize) {
        self.try_extend(index, start, end).expect("extend failed")
    }

    /// Extends the in progress array with null elements, ignoring the input arrays, returning an
    /// error on overflow.
    ///
    /// Prefer this over [`extend_nulls`](Self::extend_nulls) to handle cases where the run-end
    /// counter overflows (relevant for `RunEndEncoded` arrays).
    ///
    /// # Errors
    ///
    /// Returns an error if this [`MutableArrayData`] was not created with `use_nulls` and none
    /// of the source arrays are nullable, or if the run-end counter overflows.
    pub fn try_extend_nulls(&mut self, len: usize) -> Result<(), ArrowError> {
        if self.data.null_buffer.is_none() {
            return Err(ArrowError::InvalidArgumentError(
                "MutableArrayData cannot be extended with nulls: it was created with `use_nulls` \
                 set to false and no source array is nullable"
                    .to_owned(),
            ));
        }

        self.data.len += len;
        let bit_len = bit_util::ceil(self.data.len, 8);
        let nulls = self.data.null_buffer();
        nulls
            .try_resize(bit_len, 0)
            .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
        self.data.null_count += len;
        (self.extend_nulls)(&mut self.data, len)?;
        Ok(())
    }

    /// Extends the in progress array with null elements, ignoring the input arrays.
    ///
    /// # Panics
    ///
    /// Panics if this [`MutableArrayData`] was not created with `use_nulls` and none of the
    /// source arrays are nullable, or if the run-end counter overflows.
    #[deprecated(
        since = "59.0.0",
        note = "Use `try_extend_nulls` which returns an error on overflow instead of panicking"
    )]
    pub fn extend_nulls(&mut self, len: usize) {
        self.try_extend_nulls(len).expect("extend_nulls failed")
    }

    /// Returns the current length
    #[inline]
    pub fn len(&self) -> usize {
        self.data.len
    }

    /// Returns true if len is 0
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.data.len == 0
    }

    /// Returns the current null count
    #[inline]
    pub fn null_count(&self) -> usize {
        self.data.null_count
    }

    /// Creates a [ArrayData] from the in progress array, consuming `self`.
    pub fn freeze(self) -> ArrayData {
        unsafe { self.into_builder().build_unchecked() }
    }

    /// Consume self and returns the in progress array as [`ArrayDataBuilder`].
    ///
    /// This is useful for extending the default behavior of MutableArrayData.
    pub fn into_builder(self) -> ArrayDataBuilder {
        let data = self.data;

        let buffers = match data.data_type {
            DataType::Null
            | DataType::Struct(_)
            | DataType::FixedSizeList(_, _)
            | DataType::RunEndEncoded(_, _) => {
                vec![]
            }
            DataType::BinaryView | DataType::Utf8View => {
                let mut b = self.variadic_data_buffers;
                b.insert(0, data.buffer1.into());
                b
            }
            DataType::Utf8
            | DataType::Binary
            | DataType::LargeUtf8
            | DataType::LargeBinary
            | DataType::ListView(_)
            | DataType::LargeListView(_) => {
                vec![data.buffer1.into(), data.buffer2.into()]
            }
            DataType::Union(_, mode) => {
                match mode {
                    // Based on Union's DataTypeLayout
                    UnionMode::Sparse => vec![data.buffer1.into()],
                    UnionMode::Dense => vec![data.buffer1.into(), data.buffer2.into()],
                }
            }
            _ => vec![data.buffer1.into()],
        };

        let child_data = match data.data_type {
            DataType::Dictionary(_, _) => vec![self.dictionary.unwrap()],
            _ => data.child_data.into_iter().map(|x| x.freeze()).collect(),
        };

        let nulls = match data.data_type {
            // RunEndEncoded, Null, and Union arrays cannot have top-level null bitmasks
            DataType::RunEndEncoded(_, _) | DataType::Null | DataType::Union(_, _) => None,
            _ => data
                .null_buffer
                .map(|nulls| {
                    let bools = BooleanBuffer::new(nulls.into(), 0, data.len);
                    unsafe { NullBuffer::new_unchecked(bools, data.null_count) }
                })
                .filter(|n| n.null_count() > 0),
        };

        ArrayDataBuilder::new(data.data_type)
            .offset(0)
            .len(data.len)
            .nulls(nulls)
            .buffers(buffers)
            .child_data(child_data)
    }
}

// See arrow/tests/array_transform.rs for tests of transform functionality

#[cfg(test)]
mod test {
    use super::*;
    use arrow_schema::Field;
    use std::sync::Arc;

    fn int64_array_data(values: Vec<i64>) -> ArrayData {
        let len = values.len();
        ArrayData::try_new(
            DataType::Int64,
            len,
            None,
            0,
            vec![arrow_buffer::Buffer::from_slice_ref(&values)],
            vec![],
        )
        .unwrap()
    }

    #[test]
    fn test_try_extend_invalid_index_and_range() {
        let array = int64_array_data(vec![1, 2, 3]);
        let mut mutable = MutableArrayData::new(vec![&array], false, 3);

        let err = mutable.try_extend(1, 0, 1).unwrap_err();
        assert_eq!(
            err.to_string(),
            "Invalid argument error: Source array index 1 is out of bounds: there are 1 source arrays"
        );

        let err = mutable.try_extend(0, 0, 4).unwrap_err();
        assert_eq!(
            err.to_string(),
            "Invalid argument error: Invalid range 0..4 for source array 0 of length 3"
        );

        // `end < start` used to underflow:
        let err = mutable.try_extend(0, 2, 1).unwrap_err();
        assert_eq!(
            err.to_string(),
            "Invalid argument error: Invalid range 2..1 for source array 0 of length 3"
        );

        // The bounds are inclusive of the full array:
        mutable.try_extend(0, 3, 3).unwrap();
        mutable.try_extend(0, 0, 3).unwrap();
        assert_eq!(mutable.len(), 3);
    }

    #[test]
    fn test_try_extend_nulls_without_null_buffer() {
        let array = int64_array_data(vec![1, 2, 3]);
        let mut mutable = MutableArrayData::new(vec![&array], false, 3);
        let err = mutable.try_extend_nulls(1).unwrap_err();
        assert!(
            err.to_string().contains("cannot be extended with nulls"),
            "unexpected error: {err}"
        );

        let mut mutable = MutableArrayData::new(vec![&array], true, 3);
        mutable.try_extend_nulls(1).unwrap();
        assert_eq!(mutable.len(), 1);
    }

    #[test]
    fn test_list_append_with_capacities() {
        let array = ArrayData::new_empty(&DataType::List(Arc::new(Field::new(
            "element",
            DataType::Int64,
            false,
        ))));

        let mutable = MutableArrayData::with_capacities(
            vec![&array],
            false,
            Capacities::List(6, Some(Box::new(Capacities::Array(17)))),
        );

        // capacities are rounded up to multiples of 64 by MutableBuffer
        assert_eq!(mutable.data.buffer1.capacity(), 64);
        assert_eq!(mutable.data.child_data[0].data.buffer1.capacity(), 192);
    }

    #[test]
    fn test_map_append_with_capacities() {
        let entries = Arc::new(Field::new(
            "entries",
            DataType::Struct(
                vec![
                    Field::new("keys", DataType::Int64, false),
                    Field::new("values", DataType::Int64, true),
                ]
                .into(),
            ),
            false,
        ));
        let array = ArrayData::new_empty(&DataType::Map(entries, false));

        let mutable = MutableArrayData::with_capacities(
            vec![&array],
            false,
            Capacities::List(
                6,
                Some(Box::new(Capacities::Struct(
                    17,
                    Some(vec![Capacities::Array(17), Capacities::Array(17)]),
                ))),
            ),
        );

        // capacities are rounded up to multiples of 64 by MutableBuffer
        // the map offsets buffer holds `1 + 6` i32s
        assert_eq!(mutable.data.buffer1.capacity(), 64);

        // the entries struct itself has no buffers of its own
        let entries = &mutable.data.child_data[0];
        assert_eq!(entries.data.buffer1.capacity(), 0);

        // both key and value buffers hold 17 i64s
        assert_eq!(entries.data.child_data[0].data.buffer1.capacity(), 192);
        assert_eq!(entries.data.child_data[1].data.buffer1.capacity(), 192);
    }
}