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

#![allow(clippy::enum_clike_unportable_variant)]

use crate::{make_array, Array, ArrayRef};

use arrow_buffer::bit_chunk_iterator::{BitChunkIterator, BitChunks};

use arrow_buffer::buffer::NullBuffer;

use arrow_buffer::{BooleanBuffer, MutableBuffer, ScalarBuffer};

use arrow_data::{ArrayData, ArrayDataBuilder};

use arrow_schema::{ArrowError, DataType, UnionFields, UnionMode};

/// Contains the `UnionArray` type.

///

use std::any::Any;

use std::collections::HashSet;

use std::sync::Arc;

/// An array of [values of varying types](https://arrow.apache.org/docs/format/Columnar.html#union-layout)

///

/// Each slot in a [UnionArray] can have a value chosen from a number

/// of types.  Each of the possible types are named like the fields of

/// a [`StructArray`](crate::StructArray).  A `UnionArray` can

/// have two possible memory layouts, "dense" or "sparse".  For more

/// information on please see the

/// [specification](https://arrow.apache.org/docs/format/Columnar.html#union-layout).

///

/// [UnionBuilder](crate::builder::UnionBuilder) can be used to

/// create [UnionArray]'s of primitive types. `UnionArray`'s of nested

/// types are also supported but not via `UnionBuilder`, see the tests

/// for examples.

///

/// # Examples

/// ## Create a dense UnionArray `[1, 3.2, 34]`

/// ```

/// use arrow_buffer::ScalarBuffer;

/// use arrow_schema::*;

/// use std::sync::Arc;

/// use arrow_array::{Array, Int32Array, Float64Array, UnionArray};

///

/// let int_array = Int32Array::from(vec![1, 34]);

/// let float_array = Float64Array::from(vec![3.2]);

/// let type_ids = [0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();

/// let offsets = [0, 0, 1].into_iter().collect::<ScalarBuffer<i32>>();

///

/// let union_fields = [

///     (0, Arc::new(Field::new("A", DataType::Int32, false))),

///     (1, Arc::new(Field::new("B", DataType::Float64, false))),

/// ].into_iter().collect::<UnionFields>();

///

/// let children = vec![

///     Arc::new(int_array) as Arc<dyn Array>,

///     Arc::new(float_array),

/// ];

///

/// let array = UnionArray::try_new(

///     union_fields,

///     type_ids,

///     Some(offsets),

///     children,

/// ).unwrap();

///

/// let value = array.value(0).as_any().downcast_ref::<Int32Array>().unwrap().value(0);

/// assert_eq!(1, value);

///

/// let value = array.value(1).as_any().downcast_ref::<Float64Array>().unwrap().value(0);

/// assert!(3.2 - value < f64::EPSILON);

///

/// let value = array.value(2).as_any().downcast_ref::<Int32Array>().unwrap().value(0);

/// assert_eq!(34, value);

/// ```

///

/// ## Create a sparse UnionArray `[1, 3.2, 34]`

/// ```

/// use arrow_buffer::ScalarBuffer;

/// use arrow_schema::*;

/// use std::sync::Arc;

/// use arrow_array::{Array, Int32Array, Float64Array, UnionArray};

///

/// let int_array = Int32Array::from(vec![Some(1), None, Some(34)]);

/// let float_array = Float64Array::from(vec![None, Some(3.2), None]);

/// let type_ids = [0_i8, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();

///

/// let union_fields = [

///     (0, Arc::new(Field::new("A", DataType::Int32, false))),

///     (1, Arc::new(Field::new("B", DataType::Float64, false))),

/// ].into_iter().collect::<UnionFields>();

///

/// let children = vec![

///     Arc::new(int_array) as Arc<dyn Array>,

///     Arc::new(float_array),

/// ];

///

/// let array = UnionArray::try_new(

///     union_fields,

///     type_ids,

///     None,

///     children,

/// ).unwrap();

///

/// let value = array.value(0).as_any().downcast_ref::<Int32Array>().unwrap().value(0);

/// assert_eq!(1, value);

///

/// let value = array.value(1).as_any().downcast_ref::<Float64Array>().unwrap().value(0);

/// assert!(3.2 - value < f64::EPSILON);

///

/// let value = array.value(2).as_any().downcast_ref::<Int32Array>().unwrap().value(0);

/// assert_eq!(34, value);

/// ```

#[derive(Clone)]

pub struct UnionArray {

    data_type: DataType,

    type_ids: ScalarBuffer<i8>,

    offsets: Option<ScalarBuffer<i32>>,

    fields: Vec<Option<ArrayRef>>,

}

impl UnionArray {

    /// Creates a new `UnionArray`.

    ///

    /// Accepts type ids, child arrays and optionally offsets (for dense unions) to create

    /// a new `UnionArray`.  This method makes no attempt to validate the data provided by the

    /// caller and assumes that each of the components are correct and consistent with each other.

    /// See `try_new` for an alternative that validates the data provided.

    ///

    /// # Safety

    ///

    /// The `type_ids` values should be positive and must match one of the type ids of the fields provided in `fields`.

    /// These values are used to index into the `children` arrays.

    ///

    /// The `offsets` is provided in the case of a dense union, sparse unions should use `None`.

    /// If provided the `offsets` values should be positive and must be less than the length of the

    /// corresponding array.

    ///

    /// In both cases above we use signed integer types to maintain compatibility with other

    /// Arrow implementations.

    pub unsafe fn new_unchecked(

        fields: UnionFields,

        type_ids: ScalarBuffer<i8>,

        offsets: Option<ScalarBuffer<i32>>,

        children: Vec<ArrayRef>,

    ) -> Self {

        let mode = if offsets.is_some() {

            UnionMode::Dense

        } else {

            UnionMode::Sparse

        };

        let len = type_ids.len();

        let builder = ArrayData::builder(DataType::Union(fields, mode))

            .add_buffer(type_ids.into_inner())

            .child_data(children.into_iter().map(Array::into_data).collect())

            .len(len);

        let data = match offsets {

            Some(offsets) => builder.add_buffer(offsets.into_inner()).build_unchecked(),

            None => builder.build_unchecked(),

        };

        Self::from(data)

    }

    /// Attempts to create a new `UnionArray`, validating the inputs provided.

    ///

    /// The order of child arrays child array order must match the fields order

    pub fn try_new(

        fields: UnionFields,

        type_ids: ScalarBuffer<i8>,

        offsets: Option<ScalarBuffer<i32>>,

        children: Vec<ArrayRef>,

    ) -> Result<Self, ArrowError> {

        // There must be a child array for every field.

        if fields.len() != children.len() {

            return Err(ArrowError::InvalidArgumentError(

                "Union fields length must match child arrays length".to_string(),

            ));

        }

        if let Some(offsets) = &offsets {

            // There must be an offset value for every type id value.

            if offsets.len() != type_ids.len() {

                return Err(ArrowError::InvalidArgumentError(

                    "Type Ids and Offsets lengths must match".to_string(),

                ));

            }

        } else {

            // Sparse union child arrays must be equal in length to the length of the union

            for child in &children {

                if child.len() != type_ids.len() {

                    return Err(ArrowError::InvalidArgumentError(

                        "Sparse union child arrays must be equal in length to the length of the union".to_string(),

                    ));

                }

            }

        }

        // Create mapping from type id to array lengths.

        let max_id = fields.iter().map(|(i, _)| i).max().unwrap_or_default() as usize;

        let mut array_lens = vec![i32::MIN; max_id + 1];

        for (cd, (field_id, _)) in children.iter().zip(fields.iter()) {

            array_lens[field_id as usize] = cd.len() as i32;

        }

        // Type id values must match one of the fields.

        for id in &type_ids {

            match array_lens.get(*id as usize) {

                Some(x) if *x != i32::MIN => {}

                _ => {

                    return Err(ArrowError::InvalidArgumentError(

                        "Type Ids values must match one of the field type ids".to_owned(),

                    ))

                }

            }

        }

        // Check the value offsets are in bounds.

        if let Some(offsets) = &offsets {

            let mut iter = type_ids.iter().zip(offsets.iter());

            if iter.any(|(type_id, &offset)| offset < 0 || offset >= array_lens[*type_id as usize])

            {

                return Err(ArrowError::InvalidArgumentError(

                    "Offsets must be positive and within the length of the Array".to_owned(),

                ));

            }

        }

        // Safety:

        // - Arguments validated above.

        let union_array = unsafe { Self::new_unchecked(fields, type_ids, offsets, children) };

        Ok(union_array)

    }

    /// Accesses the child array for `type_id`.

    ///

    /// # Panics

    ///

    /// Panics if the `type_id` provided is not present in the array's DataType

    /// in the `Union`.

    pub fn child(&self, type_id: i8) -> &ArrayRef {

        assert!((type_id as usize) < self.fields.len());

        let boxed = &self.fields[type_id as usize];

        boxed.as_ref().expect("invalid type id")

    }

    /// Returns the `type_id` for the array slot at `index`.

    ///

    /// # Panics

    ///

    /// Panics if `index` is greater than or equal to the number of child arrays

    pub fn type_id(&self, index: usize) -> i8 {

        assert!(index < self.type_ids.len());

        self.type_ids[index]

    }

    /// Returns the `type_ids` buffer for this array

    pub fn type_ids(&self) -> &ScalarBuffer<i8> {

        &self.type_ids

    }

    /// Returns the `offsets` buffer if this is a dense array

    pub fn offsets(&self) -> Option<&ScalarBuffer<i32>> {

        self.offsets.as_ref()

    }

    /// Returns the offset into the underlying values array for the array slot at `index`.

    ///

    /// # Panics

    ///

    /// Panics if `index` is greater than or equal the length of the array.

    pub fn value_offset(&self, index: usize) -> usize {

        assert!(index < self.len());

        match &self.offsets {

            Some(offsets) => offsets[index] as usize,

            None => self.offset() + index,

        }

    }

    /// Returns the array's value at index `i`.

    ///

    /// Note: This method does not check for nulls and the value is arbitrary

    /// (but still well-defined) if [`is_null`](Self::is_null) returns true for the index.

    ///

    /// # Panics

    /// Panics if index `i` is out of bounds

    pub fn value(&self, i: usize) -> ArrayRef {

        let type_id = self.type_id(i);

        let value_offset = self.value_offset(i);

        let child = self.child(type_id);

        child.slice(value_offset, 1)

    }

    /// Returns the names of the types in the union.

    pub fn type_names(&self) -> Vec<&str> {

        match self.data_type() {

            DataType::Union(fields, _) => fields

                .iter()

                .map(|(_, f)| f.name().as_str())

                .collect::<Vec<&str>>(),

            _ => unreachable!("Union array's data type is not a union!"),

        }

    }

    /// Returns whether the `UnionArray` is dense (or sparse if `false`).

    fn is_dense(&self) -> bool {

        match self.data_type() {

            DataType::Union(_, mode) => mode == &UnionMode::Dense,

            _ => unreachable!("Union array's data type is not a union!"),

        }

    }

    /// Returns a zero-copy slice of this array with the indicated offset and length.

    pub fn slice(&self, offset: usize, length: usize) -> Self {

        let (offsets, fields) = match self.offsets.as_ref() {

            // If dense union, slice offsets

            Some(offsets) => (Some(offsets.slice(offset, length)), self.fields.clone()),

            // Otherwise need to slice sparse children

            None => {

                let fields = self

                    .fields

                    .iter()

                    .map(|x| x.as_ref().map(|x| x.slice(offset, length)))

                    .collect();

                (None, fields)

            }

        };

        Self {

            data_type: self.data_type.clone(),

            type_ids: self.type_ids.slice(offset, length),

            offsets,

            fields,

        }

    }

    /// Deconstruct this array into its constituent parts

    ///

    /// # Example

    ///

    /// ```

    /// # use arrow_array::array::UnionArray;

    /// # use arrow_array::types::Int32Type;

    /// # use arrow_array::builder::UnionBuilder;

    /// # use arrow_buffer::ScalarBuffer;

    /// # fn main() -> Result<(), arrow_schema::ArrowError> {

    /// let mut builder = UnionBuilder::new_dense();

    /// builder.append::<Int32Type>("a", 1).unwrap();

    /// let union_array = builder.build()?;

    ///

    /// // Deconstruct into parts

    /// let (union_fields, type_ids, offsets, children) = union_array.into_parts();

    ///

    /// // Reconstruct from parts

    /// let union_array = UnionArray::try_new(

    ///     union_fields,

    ///     type_ids,

    ///     offsets,

    ///     children,

    /// );

    /// # Ok(())

    /// # }

    /// ```

    #[allow(clippy::type_complexity)]

    pub fn into_parts(

        self,

    ) -> (

        UnionFields,

        ScalarBuffer<i8>,

        Option<ScalarBuffer<i32>>,

        Vec<ArrayRef>,

    ) {

        let Self {

            data_type,

            type_ids,

            offsets,

            mut fields,

        } = self;

        match data_type {

            DataType::Union(union_fields, _) => {

                let children = union_fields

                    .iter()

                    .map(|(type_id, _)| fields[type_id as usize].take().unwrap())

                    .collect();

                (union_fields, type_ids, offsets, children)

            }

            _ => unreachable!(),

        }

    }

    /// Computes the logical nulls for a sparse union, optimized for when there's a lot of fields without nulls

    fn mask_sparse_skip_without_nulls(&self, nulls: Vec<(i8, NullBuffer)>) -> BooleanBuffer {

        // Example logic for a union with 5 fields, a, b & c with nulls, d & e without nulls:

        // let [a_nulls, b_nulls, c_nulls] = nulls;

        // let [is_a, is_b, is_c] = masks;

        // let is_d_or_e = !(is_a | is_b | is_c)

        // let union_chunk_nulls = is_d_or_e  | (is_a & a_nulls) | (is_b & b_nulls) | (is_c & c_nulls)

        let fold = |(with_nulls_selected, union_nulls), (is_field, field_nulls)| {

            (

                with_nulls_selected | is_field,

                union_nulls | (is_field & field_nulls),

            )

        };

        self.mask_sparse_helper(

            nulls,

            |type_ids_chunk_array, nulls_masks_iters| {

                let (with_nulls_selected, union_nulls) = nulls_masks_iters

                    .iter_mut()

                    .map(|(field_type_id, field_nulls)| {

                        let field_nulls = field_nulls.next().unwrap();

                        let is_field = selection_mask(type_ids_chunk_array, *field_type_id);

                        (is_field, field_nulls)

                    })

                    .fold((0, 0), fold);

                // In the example above, this is the is_d_or_e = !(is_a | is_b) part

                let without_nulls_selected = !with_nulls_selected;

                // if a field without nulls is selected, the value is always true(set bit)

                // otherwise, the true/set bits have been computed above

                without_nulls_selected | union_nulls

            },

            |type_ids_remainder, bit_chunks| {

                let (with_nulls_selected, union_nulls) = bit_chunks

                    .iter()

                    .map(|(field_type_id, field_bit_chunks)| {

                        let field_nulls = field_bit_chunks.remainder_bits();

                        let is_field = selection_mask(type_ids_remainder, *field_type_id);

                        (is_field, field_nulls)

                    })

                    .fold((0, 0), fold);

                let without_nulls_selected = !with_nulls_selected;

                without_nulls_selected | union_nulls

            },

        )

    }

    /// Computes the logical nulls for a sparse union, optimized for when there's a lot of fields fully null

    fn mask_sparse_skip_fully_null(&self, mut nulls: Vec<(i8, NullBuffer)>) -> BooleanBuffer {

        let fields = match self.data_type() {

            DataType::Union(fields, _) => fields,

            _ => unreachable!("Union array's data type is not a union!"),

        };

        let type_ids = fields.iter().map(|(id, _)| id).collect::<HashSet<_>>();

        let with_nulls = nulls.iter().map(|(id, _)| *id).collect::<HashSet<_>>();

        let without_nulls_ids = type_ids

            .difference(&with_nulls)

            .copied()

            .collect::<Vec<_>>();

        nulls.retain(|(_, nulls)| nulls.null_count() < nulls.len());

        // Example logic for a union with 6 fields, a, b & c with nulls, d & e without nulls, and f fully_null:

        // let [a_nulls, b_nulls, c_nulls] = nulls;

        // let [is_a, is_b, is_c, is_d, is_e] = masks;

        // let union_chunk_nulls = is_d | is_e | (is_a & a_nulls) | (is_b & b_nulls) | (is_c & c_nulls)

        self.mask_sparse_helper(

            nulls,

            |type_ids_chunk_array, nulls_masks_iters| {

                let union_nulls = nulls_masks_iters.iter_mut().fold(

                    0,

                    |union_nulls, (field_type_id, nulls_iter)| {

                        let field_nulls = nulls_iter.next().unwrap();

                        if field_nulls == 0 {

                            union_nulls

                        } else {

                            let is_field = selection_mask(type_ids_chunk_array, *field_type_id);

                            union_nulls | (is_field & field_nulls)

                        }

                    },

                );

                // Given the example above, this is the is_d_or_e = (is_d | is_e) part

                let without_nulls_selected =

                    without_nulls_selected(type_ids_chunk_array, &without_nulls_ids);

                // if a field without nulls is selected, the value is always true(set bit)

                // otherwise, the true/set bits have been computed above

                union_nulls | without_nulls_selected

            },

            |type_ids_remainder, bit_chunks| {

                let union_nulls =

                    bit_chunks

                        .iter()

                        .fold(0, |union_nulls, (field_type_id, field_bit_chunks)| {

                            let is_field = selection_mask(type_ids_remainder, *field_type_id);

                            let field_nulls = field_bit_chunks.remainder_bits();

                            union_nulls | is_field & field_nulls

                        });

                union_nulls | without_nulls_selected(type_ids_remainder, &without_nulls_ids)

            },

        )

    }

    /// Computes the logical nulls for a sparse union, optimized for when all fields contains nulls

    fn mask_sparse_all_with_nulls_skip_one(&self, nulls: Vec<(i8, NullBuffer)>) -> BooleanBuffer {

        // Example logic for a union with 3 fields, a, b & c, all containing nulls:

        // let [a_nulls, b_nulls, c_nulls] = nulls;

        // We can skip the first field: it's selection mask is the negation of all others selection mask

        // let [is_b, is_c] = selection_masks;

        // let is_a = !(is_b | is_c)

        // let union_chunk_nulls = (is_a & a_nulls) | (is_b & b_nulls) | (is_c & c_nulls)

        self.mask_sparse_helper(

            nulls,

            |type_ids_chunk_array, nulls_masks_iters| {

                let (is_not_first, union_nulls) = nulls_masks_iters[1..] // skip first

                    .iter_mut()

                    .fold(

                        (0, 0),

                        |(is_not_first, union_nulls), (field_type_id, nulls_iter)| {

                            let field_nulls = nulls_iter.next().unwrap();

                            let is_field = selection_mask(type_ids_chunk_array, *field_type_id);

                            (

                                is_not_first | is_field,

                                union_nulls | (is_field & field_nulls),

                            )

                        },

                    );

                let is_first = !is_not_first;

                let first_nulls = nulls_masks_iters[0].1.next().unwrap();

                (is_first & first_nulls) | union_nulls

            },

            |type_ids_remainder, bit_chunks| {

                bit_chunks

                    .iter()

                    .fold(0, |union_nulls, (field_type_id, field_bit_chunks)| {

                        let field_nulls = field_bit_chunks.remainder_bits();

                        // The same logic as above, except that since this runs at most once,

                        // it doesn't make difference to speed-up the first selection mask

                        let is_field = selection_mask(type_ids_remainder, *field_type_id);

                        union_nulls | (is_field & field_nulls)

                    })

            },

        )

    }

    /// Maps `nulls` to `BitChunk's` and then to `BitChunkIterator's`, then divides `self.type_ids` into exact chunks of 64 values,

    /// calling `mask_chunk` for every exact chunk, and `mask_remainder` for the remainder, if any, collecting the result in a `BooleanBuffer`

    fn mask_sparse_helper(

        &self,

        nulls: Vec<(i8, NullBuffer)>,

        mut mask_chunk: impl FnMut(&[i8; 64], &mut [(i8, BitChunkIterator)]) -> u64,

        mask_remainder: impl FnOnce(&[i8], &[(i8, BitChunks)]) -> u64,

    ) -> BooleanBuffer {

        let bit_chunks = nulls

            .iter()

            .map(|(type_id, nulls)| (*type_id, nulls.inner().bit_chunks()))

            .collect::<Vec<_>>();

        let mut nulls_masks_iter = bit_chunks

            .iter()

            .map(|(type_id, bit_chunks)| (*type_id, bit_chunks.iter()))

            .collect::<Vec<_>>();

        let chunks_exact = self.type_ids.chunks_exact(64);

        let remainder = chunks_exact.remainder();

        let chunks = chunks_exact.map(|type_ids_chunk| {

            let type_ids_chunk_array = <&[i8; 64]>::try_from(type_ids_chunk).unwrap();

            mask_chunk(type_ids_chunk_array, &mut nulls_masks_iter)

        });

        // SAFETY:

        // chunks is a ChunksExact iterator, which implements TrustedLen, and correctly reports its length

        let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks) };

        if !remainder.is_empty() {

            buffer.push(mask_remainder(remainder, &bit_chunks));

        }

        BooleanBuffer::new(buffer.into(), 0, self.type_ids.len())

    }

    /// Computes the logical nulls for a sparse or dense union, by gathering individual bits from the null buffer of the selected field

    fn gather_nulls(&self, nulls: Vec<(i8, NullBuffer)>) -> BooleanBuffer {

        let one_null = NullBuffer::new_null(1);

        let one_valid = NullBuffer::new_valid(1);

        // Unsafe code below depend on it:

        // To remove one branch from the loop, if the a type_id is not utilized, or it's logical_nulls is None/all set,

        // we use a null buffer of len 1 and a index_mask of 0, or the true null buffer and usize::MAX otherwise.

        // We then unconditionally access the null buffer with index & index_mask,

        // which always return 0 for the 1-len buffer, or the true index unchanged otherwise

        // We also use a 256 array, so llvm knows that `type_id as u8 as usize` is always in bounds

        let mut logical_nulls_array = [(&one_valid, Mask::Zero); 256];

        for (type_id, nulls) in &nulls {

            if nulls.null_count() == nulls.len() {

                // Similarly, if all values are null, use a 1-null null-buffer to reduce cache pressure a bit

                logical_nulls_array[*type_id as u8 as usize] = (&one_null, Mask::Zero);

            } else {

                logical_nulls_array[*type_id as u8 as usize] = (nulls, Mask::Max);

            }

        }

        match &self.offsets {

            Some(offsets) => {

                assert_eq!(self.type_ids.len(), offsets.len());

                BooleanBuffer::collect_bool(self.type_ids.len(), |i| unsafe {

                    // SAFETY: BooleanBuffer::collect_bool calls us 0..self.type_ids.len()

                    let type_id = *self.type_ids.get_unchecked(i);

                    // SAFETY: We asserted that offsets len and self.type_ids len are equal

                    let offset = *offsets.get_unchecked(i);

                    let (nulls, offset_mask) = &logical_nulls_array[type_id as u8 as usize];

                    // SAFETY:

                    // If offset_mask is Max

                    // 1. Offset validity is checked at union creation

                    // 2. If the null buffer len equals it's array len is checked at array creation

                    // If offset_mask is Zero, the null buffer len is 1

                    nulls

                        .inner()

                        .value_unchecked(offset as usize & *offset_mask as usize)

                })

            }

            None => {

                BooleanBuffer::collect_bool(self.type_ids.len(), |index| unsafe {

                    // SAFETY: BooleanBuffer::collect_bool calls us 0..self.type_ids.len()

                    let type_id = *self.type_ids.get_unchecked(index);

                    let (nulls, index_mask) = &logical_nulls_array[type_id as u8 as usize];

                    // SAFETY:

                    // If index_mask is Max

                    // 1. On sparse union, every child len match it's parent, this is checked at union creation

                    // 2. If the null buffer len equals it's array len is checked at array creation

                    // If index_mask is Zero, the null buffer len is 1

                    nulls.inner().value_unchecked(index & *index_mask as usize)

                })

            }

        }

    }

    /// Returns a vector of tuples containing each field's type_id and its logical null buffer.

    /// Only fields with non-zero null counts are included.

    fn fields_logical_nulls(&self) -> Vec<(i8, NullBuffer)> {

        self.fields

            .iter()

            .enumerate()

            .filter_map(|(type_id, field)| Some((type_id as i8, field.as_ref()?.logical_nulls()?)))

            .filter(|(_, nulls)| nulls.null_count() > 0)

            .collect()

    }

}

impl From<ArrayData> for UnionArray {

    fn from(data: ArrayData) -> Self {

        let (fields, mode) = match data.data_type() {

            DataType::Union(fields, mode) => (fields, *mode),

            d => panic!("UnionArray expected ArrayData with type Union got {d}"),

        };

        let (type_ids, offsets) = match mode {

            UnionMode::Sparse => (

                ScalarBuffer::new(data.buffers()[0].clone(), data.offset(), data.len()),

                None,

            ),

            UnionMode::Dense => (

                ScalarBuffer::new(data.buffers()[0].clone(), data.offset(), data.len()),

                Some(ScalarBuffer::new(

                    data.buffers()[1].clone(),

                    data.offset(),

                    data.len(),

                )),

            ),

        };

        let max_id = fields.iter().map(|(i, _)| i).max().unwrap_or_default() as usize;

        let mut boxed_fields = vec![None; max_id + 1];

        for (cd, (field_id, _)) in data.child_data().iter().zip(fields.iter()) {

            boxed_fields[field_id as usize] = Some(make_array(cd.clone()));

        }

        Self {

            data_type: data.data_type().clone(),

            type_ids,

            offsets,

            fields: boxed_fields,

        }

    }

}

impl From<UnionArray> for ArrayData {

    fn from(array: UnionArray) -> Self {

        let len = array.len();

        let f = match &array.data_type {

            DataType::Union(f, _) => f,

            _ => unreachable!(),

        };

        let buffers = match array.offsets {

            Some(o) => vec![array.type_ids.into_inner(), o.into_inner()],

            None => vec![array.type_ids.into_inner()],

        };

        let child = f

            .iter()

            .map(|(i, _)| array.fields[i as usize].as_ref().unwrap().to_data())

            .collect();

        let builder = ArrayDataBuilder::new(array.data_type)

            .len(len)

            .buffers(buffers)

            .child_data(child);

        unsafe { builder.build_unchecked() }

    }

}

impl Array for UnionArray {

    fn as_any(&self) -> &dyn Any {

        self

    }

    fn to_data(&self) -> ArrayData {

        self.clone().into()

    }

    fn into_data(self) -> ArrayData {

        self.into()

    }

    fn data_type(&self) -> &DataType {

        &self.data_type

    }

    fn slice(&self, offset: usize, length: usize) -> ArrayRef {

        Arc::new(self.slice(offset, length))

    }

    fn len(&self) -> usize {

        self.type_ids.len()

    }

    fn is_empty(&self) -> bool {

        self.type_ids.is_empty()

    }

    fn shrink_to_fit(&mut self) {

        self.type_ids.shrink_to_fit();

        if let Some(offsets) = &mut self.offsets {

            offsets.shrink_to_fit();

        }

        for array in self.fields.iter_mut().flatten() {

            array.shrink_to_fit();

        }

        self.fields.shrink_to_fit();

    }

    fn offset(&self) -> usize {

        0

    }

    fn nulls(&self) -> Option<&NullBuffer> {

        None

    }

    fn logical_nulls(&self) -> Option<NullBuffer> {

        let fields = match self.data_type() {

            DataType::Union(fields, _) => fields,

            _ => unreachable!(),

        };

        if fields.len() <= 1 {

            return self.fields.iter().find_map(|field_opt| {

                field_opt

                    .as_ref()

                    .and_then(|field| field.logical_nulls())

                    .map(|logical_nulls| {

                        if self.is_dense() {

                            self.gather_nulls(vec![(0, logical_nulls)]).into()

                        } else {

                            logical_nulls

                        }

                    })

            });

        }

        let logical_nulls = self.fields_logical_nulls();

        if logical_nulls.is_empty() {

            return None;

        }

        let fully_null_count = logical_nulls

            .iter()

            .filter(|(_, nulls)| nulls.null_count() == nulls.len())

            .count();

        if fully_null_count == fields.len() {

            if let Some((_, exactly_sized)) = logical_nulls

                .iter()

                .find(|(_, nulls)| nulls.len() == self.len())

            {

                return Some(exactly_sized.clone());

            }

            if let Some((_, bigger)) = logical_nulls

                .iter()

                .find(|(_, nulls)| nulls.len() > self.len())

            {

                return Some(bigger.slice(0, self.len()));

            }

            return Some(NullBuffer::new_null(self.len()));

        }

        let boolean_buffer = match &self.offsets {

            Some(_) => self.gather_nulls(logical_nulls),

            None => {

                // Choose the fastest way to compute the logical nulls

                // Gather computes one null per iteration, while the others work on 64 nulls chunks,

                // but must also compute selection masks, which is expensive,

                // so it's cost is the number of selection masks computed per chunk

                // Since computing the selection mask gets auto-vectorized, it's performance depends on which simd feature is enabled

                // For gather, the cost is the threshold where masking becomes slower than gather, which is determined with benchmarks

                // TODO: bench on avx512f(feature is still unstable)

                let gather_relative_cost = if cfg!(target_feature = "avx2") {

                    10

                } else if cfg!(target_feature = "sse4.1") {

                    3

                } else if cfg!(target_arch = "x86") || cfg!(target_arch = "x86_64") {

                    // x86 baseline includes sse2

                    2

                } else {

                    // TODO: bench on non x86

                    // Always use gather on non benchmarked archs because even though it may slower on some cases,

                    // it's performance depends only on the union length, without being affected by the number of fields

                    0

                };

                let strategies = [

                    (SparseStrategy::Gather, gather_relative_cost, true),

                    (

                        SparseStrategy::MaskAllFieldsWithNullsSkipOne,

                        fields.len() - 1,

                        fields.len() == logical_nulls.len(),

                    ),

                    (

                        SparseStrategy::MaskSkipWithoutNulls,

                        logical_nulls.len(),

                        true,

                    ),

                    (

                        SparseStrategy::MaskSkipFullyNull,

                        fields.len() - fully_null_count,

                        true,

                    ),

                ];

                let (strategy, _, _) = strategies

                    .iter()

                    .filter(|(_, _, applicable)| *applicable)

                    .min_by_key(|(_, cost, _)| cost)

                    .unwrap();

                match strategy {

                    SparseStrategy::Gather => self.gather_nulls(logical_nulls),

                    SparseStrategy::MaskAllFieldsWithNullsSkipOne => {

                        self.mask_sparse_all_with_nulls_skip_one(logical_nulls)

                    }

                    SparseStrategy::MaskSkipWithoutNulls => {

                        self.mask_sparse_skip_without_nulls(logical_nulls)

                    }

                    SparseStrategy::MaskSkipFullyNull => {

                        self.mask_sparse_skip_fully_null(logical_nulls)

                    }

                }

            }

        };

        let null_buffer = NullBuffer::from(boolean_buffer);

        if null_buffer.null_count() > 0 {

            Some(null_buffer)

        } else {

            None

        }

    }

    fn is_nullable(&self) -> bool {

        self.fields

            .iter()

            .flatten()

            .any(|field| field.is_nullable())

    }

    fn get_buffer_memory_size(&self) -> usize {

        let mut sum = self.type_ids.inner().capacity();

        if let Some(o) = self.offsets.as_ref() {

            sum += o.inner().capacity()

        }

        self.fields

            .iter()

            .flat_map(|x| x.as_ref().map(|x| x.get_buffer_memory_size()))

            .sum::<usize>()

            + sum

    }

    fn get_array_memory_size(&self) -> usize {

        let mut sum = self.type_ids.inner().capacity();

        if let Some(o) = self.offsets.as_ref() {

            sum += o.inner().capacity()

        }

        std::mem::size_of::<Self>()

            + self