vortex_vector/primitive/

generic.rs

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

//! Definition and implementation of [`PVector<T>`].

use std::fmt::Debug;
use std::ops::BitAnd;
use std::ops::RangeBounds;

use vortex_buffer::Buffer;
use vortex_dtype::NativePType;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_ensure;
use vortex_mask::Mask;

use crate::VectorOps;
use crate::primitive::PScalar;
use crate::primitive::PVectorMut;

/// An immutable vector of generic primitive values.
///
/// `T` is expected to be bound by [`NativePType`], which templates an internal [`Buffer<T>`] that
/// stores the elements of the vector.
#[derive(Default, Debug, Clone)]
pub struct PVector<T> {
    /// The buffer representing the vector elements.
    pub(super) elements: Buffer<T>,
    /// The validity mask (where `true` represents an element is **not** null).
    pub(super) validity: Mask,
}

impl<T: NativePType + PartialEq> PartialEq for PVector<T> {
    fn eq(&self, other: &Self) -> bool {
        if self.len() != other.len() {
            return false;
        }
        // Validity patterns must match
        if self.validity != other.validity {
            return false;
        }
        // Compare all elements, OR with !validity to ignore invalid positions
        self.elements
            .iter()
            .zip(other.elements.iter())
            .enumerate()
            .all(|(i, (a, b))| !self.validity.value(i) | (a == b))
    }
}

impl<T: NativePType + Eq> Eq for PVector<T> {}

impl<T> PVector<T> {
    /// Creates a new [`PVector<T>`] from the given elements buffer and validity mask.
    ///
    /// # Panics
    ///
    /// Panics if the length of the validity mask does not match the length of the elements buffer.
    pub fn new(elements: Buffer<T>, validity: Mask) -> Self {
        Self::try_new(elements, validity).vortex_expect("Failed to create `PVector`")
    }

    /// Tries to create a new [`PVector<T>`] from the given elements buffer and validity mask.
    ///
    /// # Errors
    ///
    /// Returns an error if the length of the validity mask does not match the length of the
    /// elements buffer.
    pub fn try_new(elements: Buffer<T>, validity: Mask) -> VortexResult<Self> {
        vortex_ensure!(
            validity.len() == elements.len(),
            "`PVector` validity mask must have the same length as elements"
        );
        Ok(Self { elements, validity })
    }

    /// Creates a new [`PVector<T>`] from the given elements buffer and validity mask without
    /// validation.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the validity mask has the same length as the elements buffer.
    pub unsafe fn new_unchecked(elements: Buffer<T>, validity: Mask) -> Self {
        if cfg!(debug_assertions) {
            Self::new(elements, validity)
        } else {
            Self { elements, validity }
        }
    }

    /// Decomposes the primitive vector into its constituent parts (buffer and validity).
    pub fn into_parts(self) -> (Buffer<T>, Mask) {
        (self.elements, self.validity)
    }

    /// Decomposes the primitive vector into its internal buffer, consuming the vector.
    ///
    /// # Panics
    ///
    /// Panics if there are any null values in the vector.
    pub fn into_nonnull_buffer(self) -> Buffer<T> {
        assert!(
            self.validity.all_true(),
            "Cannot convert to buffer: vector contains null values"
        );
        self.elements
    }

    /// Decomposes the primitive vector into its constituent parts by mutable reference.
    ///
    /// # Safety
    ///
    /// The caller must ensure that no other references to the internal parts exist while mutable
    pub unsafe fn as_parts_mut(&mut self) -> (&mut Buffer<T>, &mut Mask) {
        (&mut self.elements, &mut self.validity)
    }

    /// Gets a nullable element at the given index, panicking on out-of-bounds.
    ///
    /// If the element at the given index is null, returns `None`. Otherwise, returns `Some(x)`,
    /// where `x: T`.
    ///
    /// Note that this `get` method is different from the standard library [`slice::get`], which
    /// returns `None` if the index is out of bounds. This method will panic if the index is out of
    /// bounds, and return `None` if the elements is null.
    ///
    /// # Panics
    ///
    /// Panics if the index is out of bounds.
    pub fn get(&self, index: usize) -> Option<&T> {
        self.validity.value(index).then(|| &self.elements[index])
    }

    /// Gets an element at the given index and converts it to type `U`.
    ///
    /// Returns `None` if:
    /// - The element at the given index is null
    /// - The conversion from `T` to `U` fails
    ///
    /// # Panics
    ///
    /// Panics if the index is out of bounds.
    pub fn get_as<U>(&self, index: usize) -> Option<U>
    where
        U: TryFrom<T>,
        T: Copy,
    {
        self.get(index).and_then(|&v| U::try_from(v).ok())
    }

    /// Returns the internal [`Buffer`] of the [`PVector`].
    ///
    /// Note that the internal buffer may hold garbage data in place of nulls. That information is
    /// tracked by the [`validity()`](Self::validity).
    #[inline]
    pub fn elements(&self) -> &Buffer<T> {
        &self.elements
    }

    /// Transmute a `PVector<T>` into a `PVector<U>`.
    ///
    /// # Safety
    ///
    /// The caller must ensure that all values of type `T` in this vector are valid as type `U`.
    /// See [`std::mem::transmute`] for more details.
    ///
    /// # Panics
    ///
    /// Panics if the type `U` does not have the same size and alignment as `T`.
    pub unsafe fn transmute<U: NativePType>(self) -> PVector<U> {
        let (buffer, mask) = self.into_parts();

        // SAFETY: same guarantees as this function.
        let buffer = unsafe { buffer.transmute::<U>() };

        PVector::new(buffer, mask)
    }
}

impl<T: NativePType> AsRef<[T]> for PVector<T> {
    /// Returns an immutable slice over the internal buffer with elements of type `T`.
    ///
    /// Note that this slice may contain garbage data where the [`validity()`] mask states that an
    /// element is invalid.
    ///
    /// The caller should check the [`validity()`] before performing any operations.
    ///
    /// [`validity()`]: crate::VectorOps::validity
    #[inline]
    fn as_ref(&self) -> &[T] {
        self.elements.as_slice()
    }
}

impl<T: NativePType> VectorOps for PVector<T> {
    type Mutable = PVectorMut<T>;
    type Scalar = PScalar<T>;

    fn len(&self) -> usize {
        self.elements.len()
    }

    fn validity(&self) -> &Mask {
        &self.validity
    }

    fn mask_validity(&mut self, mask: &Mask) {
        self.validity = self.validity.bitand(mask);
    }

    fn scalar_at(&self, index: usize) -> PScalar<T> {
        assert!(index < self.len(), "Index out of bounds in `PVector`");
        PScalar::<T>::new(self.validity.value(index).then(|| self.elements[index]))
    }

    fn slice(&self, range: impl RangeBounds<usize> + Clone + Debug) -> Self {
        let elements = self.elements.slice(range.clone());
        let validity = self.validity.slice(range);
        Self::new(elements, validity)
    }

    fn clear(&mut self) {
        self.elements.clear();
        self.validity.clear();
    }

    /// Try to convert self into a mutable vector.
    fn try_into_mut(self) -> Result<PVectorMut<T>, Self> {
        let elements = match self.elements.try_into_mut() {
            Ok(elements) => elements,
            Err(elements) => {
                return Err(Self {
                    elements,
                    validity: self.validity,
                });
            }
        };

        match self.validity.try_into_mut() {
            Ok(validity_mut) => Ok(PVectorMut {
                elements,
                validity: validity_mut,
            }),
            Err(validity) => Err(Self {
                elements: elements.freeze(),
                validity,
            }),
        }
    }

    fn into_mut(self) -> PVectorMut<T> {
        let elements = self.elements.into_mut();
        let validity = self.validity.into_mut();

        PVectorMut { elements, validity }
    }
}

impl<T: NativePType> From<Buffer<T>> for PVector<T> {
    /// Creates a new [`PVector<T>`] from the given elements buffer, with an all valid validity.
    fn from(value: Buffer<T>) -> Self {
        let len = value.len();
        Self {
            elements: value,
            validity: Mask::new_true(len),
        }
    }
}