add missing methods to decimal vectors

Adds `with_capacity` and other constructors for the mutable decimal
vectors, and additionally moves more things around.

I also removed the get method since it now returns a reference since you
can just call `.as_ref()` and call the native `slice::get` method.

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-vector/src/decimal/generic.rs

@@ -1,23 +1,34 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Definition and implementation of [`DVector<D>`].
+
 use vortex_buffer::Buffer;
 use vortex_dtype::{NativeDecimalType, PrecisionScale};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 use vortex_mask::Mask;
 
 use crate::{DVectorMut, VectorOps};
 
-/// A specifically typed decimal vector.
+/// An immutable vector of generic decimal values.
+///
+/// `DVector<D>` can be considered a borrowed / frozen  version of [`DVectorMut<D>`], which is
+/// created via the [`freeze`](crate::VectorMutOps::freeze) method.
+///
+/// See the documentation for [`DVectorMut<D>`] for more information.
 #[derive(Debug, Clone)]
 pub struct DVector<D> {
+    /// The precision and scale of each decimal in the decimal vector.
     pub(super) ps: PrecisionScale<D>,
+    /// The buffer representing the vector decimal elements.
     pub(super) elements: Buffer<D>,
+    /// The validity mask (where `true` represents an element is **not** null).
     pub(super) validity: Mask,
 }
 
 impl<D: NativeDecimalType> DVector<D> {
-    /// Try to create a new decimal vector from the given elements and validity.
+    /// Creates a new [`DVector<D>`] from the given [`PrecisionScale`], elements buffer, and
+    /// validity mask.
     ///
     /// # Panics
     ///
@@ -29,7 +40,8 @@ impl<D: NativeDecimalType> DVector<D> {
         Self::try_new(ps, elements, validity).vortex_expect("Failed to create `DVector`")
     }
 
-    /// Try to create a new decimal vector from the given elements and validity.
+    /// Tries to create a new [`DVector<D>`] from the given [`PrecisionScale`], elements buffer, and
+    /// validity mask.
     ///
     /// # Errors
     ///
@@ -66,7 +78,8 @@ impl<D: NativeDecimalType> DVector<D> {
         })
     }
 
-    /// Create a new decimal vector from the given elements and validity without validation.
+    /// Creates a new [`DVector<D>`] from the given [`PrecisionScale`], elements buffer, and
+    /// validity mask, _without_ validation.
     ///
     /// # Safety
     ///
@@ -100,6 +113,33 @@ impl<D: NativeDecimalType> DVector<D> {
     pub fn precision_scale(&self) -> PrecisionScale<D> {
         self.ps
     }
+
+    /// Returns a reference to the underlying elements buffer containing the decimal data.
+    pub fn elements(&self) -> &Buffer<D> {
+        &self.elements
+    }
+
+    /// Gets a nullable element at the given index, **WITHOUT** bounds checking.
+    ///
+    /// If the element at the given index is null, returns `None`. Otherwise, returns `Some(x)`,
+    /// where `x: D`.
+    ///
+    /// 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<&D> {
+        self.validity.value(index).then(|| &self.elements[index])
+    }
+}
+
+impl<D: NativeDecimalType> AsRef<[D]> for DVector<D> {
+    fn as_ref(&self) -> &[D] {
+        &self.elements
+    }
 }
 
 impl<D: NativeDecimalType> VectorOps for DVector<D> {

vortex-vector/src/decimal/generic_mut.rs

@@ -1,28 +1,151 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Definition and implementation of [`DVectorMut<D>`].
+
 use vortex_buffer::BufferMut;
-use vortex_dtype::{NativeDecimalType, PrecisionScale};
-use vortex_error::{VortexResult, vortex_bail};
+use vortex_dtype::{DecimalDType, NativeDecimalType, PrecisionScale};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 use vortex_mask::MaskMut;
 
 use crate::{DVector, VectorMutOps, VectorOps};
 
 /// A specifically typed mutable decimal vector.
 #[derive(Debug, Clone)]
 pub struct DVectorMut<D> {
+    /// The precision and scale of each decimal in the decimal vector.
     pub(super) ps: PrecisionScale<D>,
+    /// The mutable buffer representing the vector decimal elements.
     pub(super) elements: BufferMut<D>,
+    /// The validity mask (where `true` represents an element is **not** null).
     pub(super) validity: MaskMut,
 }
 
 impl<D: NativeDecimalType> DVectorMut<D> {
+    /// Creates a new [`DVectorMut<D>`] from the given [`PrecisionScale`], elements buffer, and
+    /// validity mask.
+    ///
+    /// # Panics
+    ///
+    /// Panics if:
+    ///
+    /// - The lengths of the `elements` and `validity` do not match.
+    /// - Any of the elements are out of bounds for the given [`PrecisionScale`].
+    pub fn new(ps: PrecisionScale<D>, elements: BufferMut<D>, validity: MaskMut) -> Self {
+        Self::try_new(ps, elements, validity).vortex_expect("Failed to create `DVector`")
+    }
+
+    /// Tries to create a new [`DVectorMut<D>`] from the given [`PrecisionScale`], elements buffer,
+    /// and validity mask.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    ///
+    /// - The lengths of the `elements` and `validity` do not match.
+    /// - Any of the elements are out of bounds for the given [`PrecisionScale`].
+    pub fn try_new(
+        ps: PrecisionScale<D>,
+        elements: BufferMut<D>,
+        validity: MaskMut,
+    ) -> VortexResult<Self> {
+        if elements.len() != validity.len() {
+            vortex_bail!(
+                "Elements length {} does not match validity length {}",
+                elements.len(),
+                validity.len()
+            );
+        }
+
+        // We assert that each element is within bounds for the given precision/scale.
+        if !elements.iter().all(|e| ps.is_valid(*e)) {
+            vortex_bail!(
+                "One or more elements are out of bounds for precision {} and scale {}",
+                ps.precision(),
+                ps.scale()
+            );
+        }
+
+        Ok(Self {
+            ps,
+            elements,
+            validity,
+        })
+    }
+
+    /// Creates a new [`DVectorMut<D>`] from the given [`PrecisionScale`], elements buffer, and
+    /// validity mask, _without_ validation.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure:
+    ///
+    /// - The lengths of the elements and validity are equal.
+    /// - All elements are in bounds for the given [`PrecisionScale`].
+    pub unsafe fn new_unchecked(
+        ps: PrecisionScale<D>,
+        elements: BufferMut<D>,
+        validity: MaskMut,
+    ) -> Self {
+        if cfg!(debug_assertions) {
+            Self::try_new(ps, elements, validity).vortex_expect("Failed to create `DVectorMut`")
+        } else {
+            Self {
+                ps,
+                elements,
+                validity,
+            }
+        }
+    }
+
+    /// Create a new mutable primitive vector with the given capacity.
+    pub fn with_capacity(decimal_dtype: &DecimalDType, capacity: usize) -> Self {
+        Self {
+            ps: PrecisionScale::try_from(decimal_dtype)
+                .vortex_expect("TODO(someone): This definitely should not be fallible"),
+            elements: BufferMut::with_capacity(capacity),
+            validity: MaskMut::with_capacity(capacity),
+        }
+    }
+
+    /// Decomposes the decimal vector into its constituent parts ([`PrecisionScale`], decimal
+    /// buffer, and validity).
+    pub fn into_parts(self) -> (PrecisionScale<D>, BufferMut<D>, MaskMut) {
+        (self.ps, self.elements, self.validity)
+    }
+
     /// Get the precision/scale of the decimal vector.
     pub fn precision_scale(&self) -> PrecisionScale<D> {
         self.ps
     }
 
-    /// Get a nullable element at the given index.
+    /// Returns a reference to the underlying elements buffer containing the decimal data.
+    pub fn elements(&self) -> &BufferMut<D> {
+        &self.elements
+    }
+
+    /// Returns a mutable reference to the underlying elements buffer containing the decimal data.
+    ///
+    /// # Safety
+    ///
+    /// Modifying the elements buffer directly may violate the precision/scale constraints.
+    /// The caller must ensure that any modifications maintain these invariants.
+    pub unsafe fn elements_mut(&mut self) -> &mut BufferMut<D> {
+        &mut self.elements
+    }
+
+    /// Gets a nullable element at the given index, **WITHOUT** bounds checking.
+    ///
+    /// If the element at the given index is null, returns `None`. Otherwise, returns `Some(x)`,
+    /// where `x: D`.
+    ///
+    /// 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<&D> {
         self.validity.value(index).then(|| &self.elements[index])
     }
@@ -41,16 +164,6 @@ impl<D: NativeDecimalType> DVectorMut<D> {
         self.validity.append_n(true, 1);
         Ok(())
     }
-
-    /// Returns a mutable reference to the underlying elements buffer.
-    ///
-    /// # Safety
-    ///
-    /// Modifying the elements buffer directly may violate the precision/scale constraints.
-    /// The caller must ensure that any modifications maintain these invariants.
-    pub unsafe fn elements_mut(&mut self) -> &mut [D] {
-        &mut self.elements
-    }
 }
 
 impl<D: NativeDecimalType> AsRef<[D]> for DVectorMut<D> {

vortex-vector/src/decimal/mod.rs

@@ -1,6 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Definitions and implementations of decimal vector types.
+//!
+//! The types that hold data are [`DVector`] and [`DVectorMut`], which are generic over types `D`
+//! that implement [`NativeDecimalType`].
+//!
+//! [`DecimalVector`] and [`DecimalVectorMut`] are enums that wrap all of the different possible
+//! [`DVector`]s. There are several macros defined in this crate to make working with these
+//! primitive vector types easier.
+
 mod generic;
 pub use generic::DVector;
 

vortex-vector/src/decimal/vector.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Definition and implementation of [`DecimalVector`].
+
 use vortex_dtype::{DecimalTypeDowncast, DecimalTypeUpcast, NativeDecimalType, i256};
 use vortex_error::vortex_panic;
 use vortex_mask::Mask;

vortex-vector/src/decimal/vector_mut.rs

@@ -1,7 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_dtype::{DecimalTypeDowncast, DecimalTypeUpcast, NativeDecimalType, i256};
+//! Definition and implementation of [`DecimalVectorMut`].
+
+use vortex_dtype::{
+    DecimalDType, DecimalType, DecimalTypeDowncast, DecimalTypeUpcast, NativeDecimalType, i256,
+};
 use vortex_error::vortex_panic;
 
 use crate::decimal::DVectorMut;
@@ -24,6 +28,34 @@ pub enum DecimalVectorMut {
     D256(DVectorMut<i256>),
 }
 
+impl DecimalVectorMut {
+    /// Create a new mutable decimal vector with the given primitive type and capacity.
+    pub fn with_capacity(decimal_dtype: &DecimalDType, capacity: usize) -> Self {
+        let decimal_kind = DecimalType::smallest_decimal_value_type(decimal_dtype);
+
+        match decimal_kind {
+            DecimalType::I8 => {
+                DecimalVectorMut::D8(DVectorMut::<i8>::with_capacity(decimal_dtype, capacity))
+            }
+            DecimalType::I16 => {
+                DecimalVectorMut::D16(DVectorMut::<i16>::with_capacity(decimal_dtype, capacity))
+            }
+            DecimalType::I32 => {
+                DecimalVectorMut::D32(DVectorMut::<i32>::with_capacity(decimal_dtype, capacity))
+            }
+            DecimalType::I64 => {
+                DecimalVectorMut::D64(DVectorMut::<i64>::with_capacity(decimal_dtype, capacity))
+            }
+            DecimalType::I128 => {
+                DecimalVectorMut::D128(DVectorMut::<i128>::with_capacity(decimal_dtype, capacity))
+            }
+            DecimalType::I256 => {
+                DecimalVectorMut::D256(DVectorMut::<i256>::with_capacity(decimal_dtype, capacity))
+            }
+        }
+    }
+}
+
 impl VectorMutOps for DecimalVectorMut {
     type Immutable = DecimalVector;
 

vortex-vector/src/primitive/generic.rs

@@ -71,11 +71,15 @@ impl<T> PVector<T> {
         (self.elements, self.validity)
     }
 
-    /// Gets a nullable element at the given index.
+    /// Gets a nullable element at the given index, **WITHOUT** bounds checking.
     ///
     /// 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.

vortex-vector/src/primitive/mod.rs

@@ -11,7 +11,6 @@
 //! [`PVector`]s. There are several macros defined in this crate to make working with these
 //! primitive vector types easier.
 //!
-//! [`NativePType`]: vortex_dtype::NativePType
 //! [`f16`]: vortex_dtype::half::f16
 
 mod generic;

vortex-vector/src/primitive/vector_mut.rs

@@ -60,9 +60,7 @@ impl PrimitiveVectorMut {
             PrimitiveVectorMut::F64(_) => PType::F64,
         }
     }
-}
 
-impl PrimitiveVectorMut {
     /// Create a new mutable primitive vector with the given primitive type and capacity.
     pub fn with_capacity(ptype: PType, capacity: usize) -> Self {
         match ptype {