vortex-array 0.84.0

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! Builders for Vortex arrays.
//!
//! Every logical type in Vortex has a canonical (uncompressed) in-memory encoding. This module
//! provides pre-allocated builders to construct new canonical arrays.
//!
//! ## Example:
//!
//! ```
//! use vortex_array::builders::{builder_with_capacity, ArrayBuilder};
//! use vortex_array::dtype::{DType, Nullability};
//! use vortex_array::{VortexSessionExecute, array_session};
//!
//! // Create a new builder for string data.
//! let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4);
//!
//! builder.append_scalar(&"a".into()).unwrap();
//! builder.append_scalar(&"b".into()).unwrap();
//! builder.append_scalar(&"c".into()).unwrap();
//! builder.append_scalar(&"d".into()).unwrap();
//!
//! let strings = builder.finish();
//! let mut ctx = array_session().create_execution_ctx();
//!
//! assert_eq!(strings.execute_scalar(0, &mut ctx).unwrap(), "a".into());
//! assert_eq!(strings.execute_scalar(1, &mut ctx).unwrap(), "b".into());
//! assert_eq!(strings.execute_scalar(2, &mut ctx).unwrap(), "c".into());
//! assert_eq!(strings.execute_scalar(3, &mut ctx).unwrap(), "d".into());
//! ```

use std::any::Any;
use std::sync::Arc;

use vortex_error::VortexResult;
use vortex_mask::Mask;

use crate::ArrayRef;
use crate::ExecutionCtx;
use crate::canonical::Canonical;
use crate::dtype::DType;
use crate::match_each_decimal_value_type;
use crate::match_each_native_ptype;
use crate::memory::HostAllocatorRef;
use crate::scalar::Scalar;

mod lazy_null_builder;
pub(crate) use lazy_null_builder::LazyBitBufferBuilder;

mod bool;
mod decimal;
pub mod dict;
mod extension;
mod fixed_size_list;
mod list;
mod listview;
mod map;
mod null;
mod primitive;
mod struct_;
mod varbinview;

pub use bool::*;
pub use decimal::*;
pub use extension::*;
pub use fixed_size_list::*;
pub use list::*;
pub use listview::*;
pub use map::*;
pub use null::*;
pub use primitive::*;
pub use struct_::*;
pub use varbinview::*;

pub use crate::arrays::varbin::builder::VarBinBuilder;

#[cfg(test)]
mod tests;

/// The default capacity for builders.
///
/// This is equal to the default capacity for Arrow Arrays.
pub const DEFAULT_BUILDER_CAPACITY: usize = 1024;

pub trait ArrayBuilder: Send {
    fn as_any(&self) -> &dyn Any;

    fn as_any_mut(&mut self) -> &mut dyn Any;

    fn dtype(&self) -> &DType;

    fn len(&self) -> usize;

    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Append a "zero" value to the array.
    ///
    /// Zero values are generally determined by [`Scalar::default_value`].
    fn append_zero(&mut self) {
        self.append_zeros(1)
    }

    /// Appends n "zero" values to the array.
    ///
    /// Zero values are generally determined by [`Scalar::default_value`].
    fn append_zeros(&mut self, n: usize);

    /// Append a "null" value to the array.
    ///
    /// Implementors should panic if this method is called on a non-nullable [`ArrayBuilder`].
    fn append_null(&mut self) {
        self.append_nulls(1)
    }

    /// The inner part of `append_nulls`.
    ///
    /// # Safety
    ///
    /// The array builder must be nullable.
    unsafe fn append_nulls_unchecked(&mut self, n: usize);

    /// Appends n "null" values to the array.
    ///
    /// Implementors should panic if this method is called on a non-nullable [`ArrayBuilder`].
    fn append_nulls(&mut self, n: usize) {
        assert!(
            self.dtype().is_nullable(),
            "tried to append {n} nulls to a non-nullable array builder"
        );

        // SAFETY: We check above that the array builder is nullable.
        unsafe {
            self.append_nulls_unchecked(n);
        }
    }

    /// Appends a default value to the array.
    fn append_default(&mut self) {
        self.append_defaults(1)
    }

    /// Appends n default values to the array.
    ///
    /// If the array builder is nullable, then this has the behavior of `self.append_nulls(n)`.
    /// If the array builder is non-nullable, then it has the behavior of `self.append_zeros(n)`.
    fn append_defaults(&mut self, n: usize) {
        if self.dtype().is_nullable() {
            self.append_nulls(n);
        } else {
            self.append_zeros(n);
        }
    }

    /// A generic function to append a scalar to the builder.
    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()>;

    /// Allocate space for extra `additional` items
    fn reserve_exact(&mut self, additional: usize);

    /// Override builders validity with the one provided.
    ///
    /// Note that this will have no effect on the final array if the array builder is non-nullable.
    fn set_validity(&mut self, validity: Mask) {
        if !self.dtype().is_nullable() {
            return;
        }
        assert_eq!(self.len(), validity.len());
        unsafe { self.set_validity_unchecked(validity) }
    }

    /// override validity with the one provided, without checking lengths
    ///
    /// # Safety
    ///
    /// Given validity must have an equal length to [`self.len()`](Self::len).
    unsafe fn set_validity_unchecked(&mut self, validity: Mask);

    /// Constructs an Array from the builder components.
    ///
    /// # Panics
    ///
    /// This function may panic if the builder's methods are called with invalid arguments. If only
    /// the methods on this interface are used, the builder should not panic. However, specific
    /// builders have interfaces that may be misused. For example, if the number of values in a
    /// [PrimitiveBuilder]'s [vortex_buffer::BufferMut] does not match the number of validity bits,
    /// the PrimitiveBuilder's [Self::finish] will panic.
    fn finish(&mut self) -> ArrayRef;

    /// Constructs a canonical array directly from the builder.
    ///
    /// This method provides a default implementation that creates an [`ArrayRef`] via `finish` and
    /// then converts it to canonical form. Specific builders can override this with optimized
    /// implementations that avoid the intermediate [`ArrayRef`] creation.
    fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical;
}

/// Matches a `&mut dyn ArrayBuilder` against every concrete list builder type, i.e. every
/// [`ListBuilder`]`<O>` and [`ListViewBuilder`]`<O, S>` instantiation over the
/// [`OffsetBuilderPType`](crate::dtype::OffsetBuilderPType) offset/size types (`u32`, `u64`, `i32`,
/// `i64`).
///
/// Binds the downcast builder as `$builder` and evaluates `$body` with it, yielding
/// `Some($body)`; yields `None` when the builder is not a list builder. List encodings dispatch
/// through this matcher because the concrete list builders are generic over their offset/size
/// integer types, which cannot be named through a `dyn ArrayBuilder`. The matcher is exhaustive
/// because `OffsetBuilderPType` is sealed, so no other instantiations can be constructed.
#[macro_export]
macro_rules! match_each_list_builder {
    ($dyn_builder:expr, | $builder:ident | $body:expr) => {{
        let __dyn_builder: &mut dyn $crate::builders::ArrayBuilder = $dyn_builder;
        match $crate::__match_each_list_builder!(
            __dyn_builder,
            $builder,
            $body,
            [u32, u64, i32, i64]
        ) {
            ::core::option::Option::Some(__result) => ::core::option::Option::Some(__result),
            ::core::option::Option::None => $crate::__match_each_listview_builder!(
                __dyn_builder,
                $builder,
                $body,
                [u32, u64, i32, i64]
            ),
        }
    }};
}

#[doc(hidden)]
#[macro_export]
macro_rules! __match_each_list_builder {
    ($target:ident, $builder:ident, $body:expr, []) => {
        ::core::option::Option::None
    };
    ($target:ident, $builder:ident, $body:expr, [$offset:ty $(, $rest:ty)*]) => {
        if let ::core::option::Option::Some($builder) =
            $crate::builders::ArrayBuilder::as_any_mut($target)
                .downcast_mut::<$crate::builders::ListBuilder<$offset>>()
        {
            ::core::option::Option::Some($body)
        } else {
            $crate::__match_each_list_builder!($target, $builder, $body, [$($rest),*])
        }
    };
}

#[doc(hidden)]
#[macro_export]
macro_rules! __match_each_listview_builder {
    ($target:ident, $builder:ident, $body:expr, []) => {
        ::core::option::Option::None
    };
    ($target:ident, $builder:ident, $body:expr, [$offset:ty $(, $rest:ty)*]) => {
        match $crate::__match_each_listview_builder_size!(
            $target,
            $builder,
            $body,
            $offset,
            [u32, u64, i32, i64]
        ) {
            ::core::option::Option::Some(__result) => ::core::option::Option::Some(__result),
            ::core::option::Option::None => $crate::__match_each_listview_builder!(
                $target,
                $builder,
                $body,
                [$($rest),*]
            ),
        }
    };
}

#[doc(hidden)]
#[macro_export]
macro_rules! __match_each_listview_builder_size {
    ($target:ident, $builder:ident, $body:expr, $offset:ty, []) => {
        ::core::option::Option::None
    };
    ($target:ident, $builder:ident, $body:expr, $offset:ty, [$size:ty $(, $rest:ty)*]) => {
        if let ::core::option::Option::Some($builder) =
            $crate::builders::ArrayBuilder::as_any_mut($target)
                .downcast_mut::<$crate::builders::ListViewBuilder<$offset, $size>>()
        {
            ::core::option::Option::Some($body)
        } else {
            $crate::__match_each_listview_builder_size!(
                $target, $builder, $body, $offset, [$($rest),*]
            )
        }
    };
}

/// Matches a `&mut dyn ArrayBuilder` against every concrete map builder type.
///
/// Binds the downcast builder as `$builder` and evaluates `$body` with it, yielding
/// `Some($body)`; yields `None` when the builder is not a map builder.
#[macro_export]
macro_rules! match_each_map_builder {
    ($dyn_builder:expr, | $builder:ident | $body:expr) => {{
        let __dyn_builder: &mut dyn $crate::builders::ArrayBuilder = $dyn_builder;
        $crate::__match_each_map_builder!(__dyn_builder, $builder, $body, [u32, u64, i32, i64])
    }};
}

#[doc(hidden)]
#[macro_export]
macro_rules! __match_each_map_builder {
    ($target:ident, $builder:ident, $body:expr, []) => {
        ::core::option::Option::None
    };
    ($target:ident, $builder:ident, $body:expr, [$offset:ty $(, $rest:ty)*]) => {
        match $crate::__match_each_map_builder_size!(
            $target,
            $builder,
            $body,
            $offset,
            [u32, u64, i32, i64]
        ) {
            ::core::option::Option::Some(__result) => ::core::option::Option::Some(__result),
            ::core::option::Option::None => $crate::__match_each_map_builder!(
                $target,
                $builder,
                $body,
                [$($rest),*]
            ),
        }
    };
}

#[doc(hidden)]
#[macro_export]
macro_rules! __match_each_map_builder_size {
    ($target:ident, $builder:ident, $body:expr, $offset:ty, []) => {
        ::core::option::Option::None
    };
    ($target:ident, $builder:ident, $body:expr, $offset:ty, [$size:ty $(, $rest:ty)*]) => {
        if let ::core::option::Option::Some($builder) =
            $crate::builders::ArrayBuilder::as_any_mut($target)
                .downcast_mut::<$crate::builders::MapBuilder<$offset, $size>>()
        {
            ::core::option::Option::Some($body)
        } else {
            $crate::__match_each_map_builder_size!(
                $target, $builder, $body, $offset, [$($rest),*]
            )
        }
    };
}

/// Construct a new canonical builder for the given [`DType`].
///
///
/// # Example
///
/// ```
/// use vortex_array::builders::{builder_with_capacity, ArrayBuilder};
/// use vortex_array::dtype::{DType, Nullability};
/// use vortex_array::{VortexSessionExecute, array_session};
///
/// // Create a new builder for string data.
/// let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4);
///
/// builder.append_scalar(&"a".into()).unwrap();
/// builder.append_scalar(&"b".into()).unwrap();
/// builder.append_scalar(&"c".into()).unwrap();
/// builder.append_scalar(&"d".into()).unwrap();
///
/// let strings = builder.finish();
/// let mut ctx = array_session().create_execution_ctx();
///
/// assert_eq!(strings.execute_scalar(0, &mut ctx).unwrap(), "a".into());
/// assert_eq!(strings.execute_scalar(1, &mut ctx).unwrap(), "b".into());
/// assert_eq!(strings.execute_scalar(2, &mut ctx).unwrap(), "c".into());
/// assert_eq!(strings.execute_scalar(3, &mut ctx).unwrap(), "d".into());
/// ```
pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box<dyn ArrayBuilder> {
    match dtype {
        DType::Null => Box::new(NullBuilder::new()),
        DType::Bool(n) => Box::new(BoolBuilder::with_capacity(*n, capacity)),
        DType::Primitive(ptype, n) => {
            match_each_native_ptype!(ptype, |P| {
                Box::new(PrimitiveBuilder::<P>::with_capacity(*n, capacity))
            })
        }
        DType::Decimal(decimal_type, n) => {
            match_each_decimal_value_type!(
                DecimalType::smallest_decimal_value_type(decimal_type),
                |D| {
                    Box::new(DecimalBuilder::with_capacity::<D>(
                        capacity,
                        *decimal_type,
                        *n,
                    ))
                }
            )
        }
        DType::Utf8(n) => Box::new(VarBinViewBuilder::with_capacity(DType::Utf8(*n), capacity)),
        DType::Binary(n) => Box::new(VarBinViewBuilder::with_capacity(
            DType::Binary(*n),
            capacity,
        )),
        DType::List(dtype, n) => Box::new(ListViewBuilder::<u64, u64>::with_capacity(
            Arc::clone(dtype),
            *n,
            2 * capacity, // Arbitrarily choose 2 times the `offsets` capacity here.
            capacity,
        )),
        DType::Map(map_dtype, nullability) => Box::new(MapBuilder::<u64, u64>::with_capacity(
            map_dtype.clone(),
            *nullability,
            capacity,
        )),
        DType::FixedSizeList(elem_dtype, list_size, null) => {
            Box::new(FixedSizeListBuilder::with_capacity(
                Arc::clone(elem_dtype),
                *list_size,
                *null,
                capacity,
            ))
        }
        DType::Struct(struct_dtype, n) => Box::new(StructBuilder::with_capacity(
            struct_dtype.clone(),
            *n,
            capacity,
        )),
        DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"),
        DType::Variant(_) => {
            unimplemented!()
        }
        DType::Extension(ext_dtype) => {
            Box::new(ExtensionBuilder::with_capacity(ext_dtype.clone(), capacity))
        }
    }
}

/// Construct a new canonical builder for the given [`DType`] using a host
/// [`crate::memory::HostAllocator`].
pub fn builder_with_capacity_in(
    allocator: HostAllocatorRef,
    dtype: &DType,
    capacity: usize,
) -> Box<dyn ArrayBuilder> {
    let _allocator = allocator;
    builder_with_capacity(dtype, capacity)
}