arrow-array 58.0.0

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

use crate::builder::buffer_builder::{Int8BufferBuilder, Int32BufferBuilder};
use crate::builder::{ArrayBuilder, BufferBuilder};
use crate::{ArrayRef, ArrowPrimitiveType, UnionArray, make_array};
use arrow_buffer::NullBufferBuilder;
use arrow_buffer::{ArrowNativeType, Buffer, ScalarBuffer};
use arrow_data::ArrayDataBuilder;
use arrow_schema::{ArrowError, DataType, Field};
use std::any::Any;
use std::collections::BTreeMap;
use std::sync::Arc;

/// `FieldData` is a helper struct to track the state of the fields in the `UnionBuilder`.
#[derive(Debug)]
struct FieldData {
    /// The type id for this field
    type_id: i8,
    /// The Arrow data type represented in the `values_buffer`, which is untyped
    data_type: DataType,
    /// A buffer containing the values for this field in raw bytes
    values_buffer: Box<dyn FieldDataValues>,
    ///  The number of array slots represented by the buffer
    slots: usize,
    /// A builder for the null bitmap
    null_buffer_builder: NullBufferBuilder,
}

/// A type-erased [`BufferBuilder`] used by [`FieldData`]
trait FieldDataValues: std::fmt::Debug + Send + Sync {
    fn as_mut_any(&mut self) -> &mut dyn Any;

    fn append_null(&mut self);

    fn finish(&mut self) -> Buffer;

    fn finish_cloned(&self) -> Buffer;
}

impl<T: ArrowNativeType> FieldDataValues for BufferBuilder<T> {
    fn as_mut_any(&mut self) -> &mut dyn Any {
        self
    }

    fn append_null(&mut self) {
        self.advance(1)
    }

    fn finish(&mut self) -> Buffer {
        self.finish()
    }

    fn finish_cloned(&self) -> Buffer {
        Buffer::from_slice_ref(self.as_slice())
    }
}

impl FieldData {
    /// Creates a new `FieldData`.
    fn new<T: ArrowPrimitiveType>(type_id: i8, data_type: DataType, capacity: usize) -> Self {
        Self {
            type_id,
            data_type,
            slots: 0,
            values_buffer: Box::new(BufferBuilder::<T::Native>::new(capacity)),
            null_buffer_builder: NullBufferBuilder::new(capacity),
        }
    }

    /// Appends a single value to this `FieldData`'s `values_buffer`.
    fn append_value<T: ArrowPrimitiveType>(&mut self, v: T::Native) {
        self.values_buffer
            .as_mut_any()
            .downcast_mut::<BufferBuilder<T::Native>>()
            .expect("Tried to append unexpected type")
            .append(v);

        self.null_buffer_builder.append(true);
        self.slots += 1;
    }

    /// Appends a null to this `FieldData`.
    fn append_null(&mut self) {
        self.values_buffer.append_null();
        self.null_buffer_builder.append(false);
        self.slots += 1;
    }
}

/// Builder for [`UnionArray`]
///
/// Example: **Dense Memory Layout**
///
/// ```
/// # use arrow_array::builder::UnionBuilder;
/// # use arrow_array::types::{Float64Type, Int32Type};
///
/// let mut builder = UnionBuilder::new_dense();
/// builder.append::<Int32Type>("a", 1).unwrap();
/// builder.append::<Float64Type>("b", 3.0).unwrap();
/// builder.append::<Int32Type>("a", 4).unwrap();
/// let union = builder.build().unwrap();
///
/// assert_eq!(union.type_id(0), 0);
/// assert_eq!(union.type_id(1), 1);
/// assert_eq!(union.type_id(2), 0);
///
/// assert_eq!(union.value_offset(0), 0);
/// assert_eq!(union.value_offset(1), 0);
/// assert_eq!(union.value_offset(2), 1);
/// ```
///
/// Example: **Sparse Memory Layout**
/// ```
/// # use arrow_array::builder::UnionBuilder;
/// # use arrow_array::types::{Float64Type, Int32Type};
///
/// let mut builder = UnionBuilder::new_sparse();
/// builder.append::<Int32Type>("a", 1).unwrap();
/// builder.append::<Float64Type>("b", 3.0).unwrap();
/// builder.append::<Int32Type>("a", 4).unwrap();
/// let union = builder.build().unwrap();
///
/// assert_eq!(union.type_id(0), 0);
/// assert_eq!(union.type_id(1), 1);
/// assert_eq!(union.type_id(2), 0);
///
/// assert_eq!(union.value_offset(0), 0);
/// assert_eq!(union.value_offset(1), 1);
/// assert_eq!(union.value_offset(2), 2);
/// ```
#[derive(Debug, Default)]
pub struct UnionBuilder {
    /// The current number of slots in the array
    len: usize,
    /// Maps field names to `FieldData` instances which track the builders for that field
    fields: BTreeMap<String, FieldData>,
    /// Builder to keep track of type ids
    type_id_builder: Int8BufferBuilder,
    /// Builder to keep track of offsets (`None` for sparse unions)
    value_offset_builder: Option<Int32BufferBuilder>,
    initial_capacity: usize,
}

impl UnionBuilder {
    /// Creates a new dense array builder.
    pub fn new_dense() -> Self {
        Self::with_capacity_dense(1024)
    }

    /// Creates a new sparse array builder.
    pub fn new_sparse() -> Self {
        Self::with_capacity_sparse(1024)
    }

    /// Creates a new dense array builder with capacity.
    pub fn with_capacity_dense(capacity: usize) -> Self {
        Self {
            len: 0,
            fields: Default::default(),
            type_id_builder: Int8BufferBuilder::new(capacity),
            value_offset_builder: Some(Int32BufferBuilder::new(capacity)),
            initial_capacity: capacity,
        }
    }

    /// Creates a new sparse array builder  with capacity.
    pub fn with_capacity_sparse(capacity: usize) -> Self {
        Self {
            len: 0,
            fields: Default::default(),
            type_id_builder: Int8BufferBuilder::new(capacity),
            value_offset_builder: None,
            initial_capacity: capacity,
        }
    }

    /// Appends a null to this builder, encoding the null in the array
    /// of the `type_name` child / field.
    ///
    /// Since `UnionArray` encodes nulls as an entry in its children
    /// (it doesn't have a validity bitmap itself), and where the null
    /// is part of the final array, appending a NULL requires
    /// specifying which field (child) to use.
    #[inline]
    pub fn append_null<T: ArrowPrimitiveType>(
        &mut self,
        type_name: &str,
    ) -> Result<(), ArrowError> {
        self.append_option::<T>(type_name, None)
    }

    /// Appends a value to this builder.
    #[inline]
    pub fn append<T: ArrowPrimitiveType>(
        &mut self,
        type_name: &str,
        v: T::Native,
    ) -> Result<(), ArrowError> {
        self.append_option::<T>(type_name, Some(v))
    }

    fn append_option<T: ArrowPrimitiveType>(
        &mut self,
        type_name: &str,
        v: Option<T::Native>,
    ) -> Result<(), ArrowError> {
        let type_name = type_name.to_string();

        let mut field_data = match self.fields.remove(&type_name) {
            Some(data) => {
                if data.data_type != T::DATA_TYPE {
                    return Err(ArrowError::InvalidArgumentError(format!(
                        "Attempt to write col \"{}\" with type {} doesn't match existing type {}",
                        type_name,
                        T::DATA_TYPE,
                        data.data_type
                    )));
                }
                data
            }
            None => match self.value_offset_builder {
                Some(_) => FieldData::new::<T>(
                    self.fields.len() as i8,
                    T::DATA_TYPE,
                    self.initial_capacity,
                ),
                // In the case of a sparse union, we should pass the maximum of the currently length and the capacity.
                None => {
                    let mut fd = FieldData::new::<T>(
                        self.fields.len() as i8,
                        T::DATA_TYPE,
                        self.len.max(self.initial_capacity),
                    );
                    for _ in 0..self.len {
                        fd.append_null();
                    }
                    fd
                }
            },
        };
        self.type_id_builder.append(field_data.type_id);

        match &mut self.value_offset_builder {
            // Dense Union
            Some(offset_builder) => {
                offset_builder.append(field_data.slots as i32);
            }
            // Sparse Union
            None => {
                for (_, fd) in self.fields.iter_mut() {
                    // Append to all bar the FieldData currently being appended to
                    fd.append_null();
                }
            }
        }

        match v {
            Some(v) => field_data.append_value::<T>(v),
            None => field_data.append_null(),
        }

        self.fields.insert(type_name, field_data);
        self.len += 1;
        Ok(())
    }

    /// Builds this builder creating a new `UnionArray`.
    pub fn build(self) -> Result<UnionArray, ArrowError> {
        let mut children = Vec::with_capacity(self.fields.len());
        let union_fields = self
            .fields
            .into_iter()
            .map(
                |(
                    name,
                    FieldData {
                        type_id,
                        data_type,
                        mut values_buffer,
                        slots,
                        mut null_buffer_builder,
                    },
                )| {
                    let array_ref = make_array(unsafe {
                        ArrayDataBuilder::new(data_type.clone())
                            .add_buffer(values_buffer.finish())
                            .len(slots)
                            .nulls(null_buffer_builder.finish())
                            .build_unchecked()
                    });
                    children.push(array_ref);
                    (type_id, Arc::new(Field::new(name, data_type, false)))
                },
            )
            .collect();
        UnionArray::try_new(
            union_fields,
            self.type_id_builder.into(),
            self.value_offset_builder.map(Into::into),
            children,
        )
    }

    /// Builds this builder creating a new `UnionArray` without consuming the builder.
    ///
    /// This is used for the `finish_cloned` implementation in `ArrayBuilder`.
    fn build_cloned(&self) -> Result<UnionArray, ArrowError> {
        let mut children = Vec::with_capacity(self.fields.len());
        let union_fields: Vec<_> = self
            .fields
            .iter()
            .map(|(name, field_data)| {
                let FieldData {
                    type_id,
                    data_type,
                    values_buffer,
                    slots,
                    null_buffer_builder,
                } = field_data;

                let array_ref = make_array(unsafe {
                    ArrayDataBuilder::new(data_type.clone())
                        .add_buffer(values_buffer.finish_cloned())
                        .len(*slots)
                        .nulls(null_buffer_builder.finish_cloned())
                        .build_unchecked()
                });
                children.push(array_ref);
                (
                    *type_id,
                    Arc::new(Field::new(name.clone(), data_type.clone(), false)),
                )
            })
            .collect();
        UnionArray::try_new(
            union_fields.into_iter().collect(),
            ScalarBuffer::from(self.type_id_builder.as_slice().to_vec()),
            self.value_offset_builder
                .as_ref()
                .map(|builder| ScalarBuffer::from(builder.as_slice().to_vec())),
            children,
        )
    }
}

impl ArrayBuilder for UnionBuilder {
    /// Returns the number of array slots in the builder
    fn len(&self) -> usize {
        self.len
    }

    /// Builds the array
    fn finish(&mut self) -> ArrayRef {
        // Even simpler - just move the builder using mem::take and replace with default
        let builder = std::mem::take(self);

        // Since UnionBuilder controls all invariants, this should never fail
        Arc::new(builder.build().unwrap())
    }

    /// Builds the array without resetting the underlying builder
    fn finish_cloned(&self) -> ArrayRef {
        // We construct the UnionArray carefully to ensure try_new cannot fail.
        // Since UnionBuilder controls all the invariants, this should never panic.
        Arc::new(self.build_cloned().unwrap_or_else(|err| {
            panic!("UnionBuilder::build_cloned failed unexpectedly: {}", err)
        }))
    }

    /// Returns the builder as a non-mutable `Any` reference
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Returns the builder as a mutable `Any` reference
    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    /// Returns the boxed builder as a box of `Any`
    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::array::Array;
    use crate::cast::AsArray;
    use crate::types::{Float64Type, Int32Type};

    #[test]
    fn test_union_builder_array_builder_trait() {
        // Test that UnionBuilder implements ArrayBuilder trait
        let mut builder = UnionBuilder::new_dense();

        // Add some data
        builder.append::<Int32Type>("a", 1).unwrap();
        builder.append::<Float64Type>("b", 3.0).unwrap();
        builder.append::<Int32Type>("a", 4).unwrap();

        assert_eq!(builder.len(), 3);

        // Test finish_cloned (non-destructive)
        let array1 = builder.finish_cloned();
        assert_eq!(array1.len(), 3);

        // Verify values in cloned array
        let union1 = array1.as_any().downcast_ref::<UnionArray>().unwrap();
        assert_eq!(union1.type_ids(), &[0, 1, 0]);
        assert_eq!(union1.offsets().unwrap().as_ref(), &[0, 0, 1]);
        let int_array1 = union1.child(0).as_primitive::<Int32Type>();
        let float_array1 = union1.child(1).as_primitive::<Float64Type>();
        assert_eq!(int_array1.value(0), 1);
        assert_eq!(int_array1.value(1), 4);
        assert_eq!(float_array1.value(0), 3.0);

        // Builder should still be usable after finish_cloned
        builder.append::<Float64Type>("b", 5.0).unwrap();
        assert_eq!(builder.len(), 4);

        // Test finish (destructive)
        let array2 = builder.finish();
        assert_eq!(array2.len(), 4);

        // Verify values in final array
        let union2 = array2.as_any().downcast_ref::<UnionArray>().unwrap();
        assert_eq!(union2.type_ids(), &[0, 1, 0, 1]);
        assert_eq!(union2.offsets().unwrap().as_ref(), &[0, 0, 1, 1]);
        let int_array2 = union2.child(0).as_primitive::<Int32Type>();
        let float_array2 = union2.child(1).as_primitive::<Float64Type>();
        assert_eq!(int_array2.value(0), 1);
        assert_eq!(int_array2.value(1), 4);
        assert_eq!(float_array2.value(0), 3.0);
        assert_eq!(float_array2.value(1), 5.0);
    }

    #[test]
    fn test_union_builder_type_erased() {
        // Test type-erased usage with Box<dyn ArrayBuilder>
        let mut builders: Vec<Box<dyn ArrayBuilder>> = vec![Box::new(UnionBuilder::new_sparse())];

        // Downcast and use
        let union_builder = builders[0]
            .as_any_mut()
            .downcast_mut::<UnionBuilder>()
            .unwrap();
        union_builder.append::<Int32Type>("x", 10).unwrap();
        union_builder.append::<Float64Type>("y", 20.0).unwrap();

        assert_eq!(builders[0].len(), 2);

        let result = builders
            .into_iter()
            .map(|mut b| b.finish())
            .collect::<Vec<_>>();
        assert_eq!(result[0].len(), 2);

        // Verify sparse union values
        let union = result[0].as_any().downcast_ref::<UnionArray>().unwrap();
        assert_eq!(union.type_ids(), &[0, 1]);
        assert!(union.offsets().is_none()); // Sparse union has no offsets
        let int_array = union.child(0).as_primitive::<Int32Type>();
        let float_array = union.child(1).as_primitive::<Float64Type>();
        assert_eq!(int_array.value(0), 10);
        assert!(int_array.is_null(1)); // Null in sparse layout
        assert!(float_array.is_null(0)); // Null in sparse layout
        assert_eq!(float_array.value(1), 20.0);
    }
}