vortex-array 0.54.0

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

use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer};
use vortex_dtype::Nullability;
use vortex_dtype::Nullability::{NonNullable, Nullable};
use vortex_error::{VortexExpect, vortex_panic};
use vortex_mask::Mask;

use crate::validity::Validity;

/// This is borrowed from arrow's null buffer builder, however we expose a `append_buffer`
/// method to append a boolean buffer directly.
pub struct LazyNullBufferBuilder {
    inner: Option<BooleanBufferBuilder>,
    len: usize,
    capacity: usize,
}

impl LazyNullBufferBuilder {
    /// Creates a new empty builder.
    /// `capacity` is the number of bits in the null buffer.
    pub fn new(capacity: usize) -> Self {
        Self {
            inner: None,
            len: 0,
            capacity,
        }
    }

    /// Appends `n` non-null values to the builder.
    #[inline]
    pub fn append_n_non_nulls(&mut self, n: usize) {
        if let Some(buf) = self.inner.as_mut() {
            buf.append_n(n, true)
        } else {
            self.len += n;
        }
    }

    /// Appends a single non-null value to the builder.
    #[inline]
    pub fn append_non_null(&mut self) {
        if let Some(buf) = self.inner.as_mut() {
            buf.append(true)
        } else {
            self.len += 1;
        }
    }

    /// Appends `n` null values to the builder.
    #[inline]
    pub fn append_n_nulls(&mut self, n: usize) {
        self.materialize_if_needed();
        self.inner
            .as_mut()
            .vortex_expect("cannot append null to non-nullable builder")
            .append_n(n, false);
    }

    /// Appends a single null value to the builder.
    #[inline]
    pub fn append_null(&mut self) {
        self.materialize_if_needed();
        self.inner
            .as_mut()
            .vortex_expect("cannot append null to non-nullable builder")
            .append(false);
    }

    /// Appends values from a boolean buffer where `true` indicates non-null.
    #[inline]
    pub fn append_buffer(&mut self, bool_buffer: &BooleanBuffer) {
        self.materialize_if_needed();
        self.inner
            .as_mut()
            .vortex_expect("buffer just materialized")
            .append_buffer(bool_buffer);
    }

    /// Appends values from a validity mask.
    pub fn append_validity_mask(&mut self, validity_mask: Mask) {
        match validity_mask {
            Mask::AllTrue(len) => self.append_n_non_nulls(len),
            Mask::AllFalse(len) => self.append_n_nulls(len),
            Mask::Values(is_valid) => self.append_buffer(is_valid.boolean_buffer()),
        }
    }

    /// Sets the validity bit at the given index.
    pub fn set_bit(&mut self, index: usize, v: bool) {
        self.materialize_if_needed();
        self.inner
            .as_mut()
            .vortex_expect("buffer just materialized")
            .set_bit(index, v);
    }

    /// Returns the current length of the builder.
    pub fn len(&self) -> usize {
        // self.len is the length of the builder if the inner buffer is not materialized
        self.inner.as_ref().map(|i| i.len()).unwrap_or(self.len)
    }

    fn finish(&mut self) -> Option<NullBuffer> {
        self.len = 0;
        Some(NullBuffer::new(self.inner.take()?.finish()))
    }

    /// Finishes the builder and returns a `Validity` based on the given nullability.
    pub fn finish_with_nullability(&mut self, nullability: Nullability) -> Validity {
        let nulls = self.finish();

        match (nullability, nulls) {
            (NonNullable, None) => Validity::NonNullable,
            (Nullable, None) => Validity::AllValid,
            (Nullable, Some(arr)) => Validity::from(arr),
            _ => vortex_panic!("Invalid nullability/nulls combination"),
        }
    }

    /// Ensures the builder can hold `additional` extra values.
    pub fn reserve_exact(&mut self, additional: usize) {
        if self.inner.is_none() {
            self.capacity += additional;
        } else {
            self.inner
                .as_mut()
                .vortex_expect("buffer just materialized")
                .reserve(additional);
        }
    }

    #[inline]
    fn materialize_if_needed(&mut self) {
        if self.inner.is_none() {
            self.materialize()
        }
    }

    // This only happens once per builder
    #[cold]
    #[inline(never)]
    fn materialize(&mut self) {
        if self.inner.is_none() {
            let mut b = BooleanBufferBuilder::new(self.len.max(self.capacity));
            b.append_n(self.len, true);
            self.inner = Some(b);
        }
    }
}