rust-ndarray · bluss · Apr 6, 2024 · Mar 9, 2023 · Mar 11, 2023 · Mar 11, 2023
diff --git a/benches/reserve.rs b/benches/reserve.rs
@@ -0,0 +1,31 @@
+#![feature(test)]
+
+extern crate test;
+use test::Bencher;
+
+use ndarray::prelude::*;
+
+#[bench]
+fn push_reserve(bench: &mut Bencher)
+{
+    let ones: Array<f32, _> = array![1f32];
+    bench.iter(|| {
+        let mut a: Array<f32, Ix1> = array![];
+        a.reserve(Axis(0), 100).unwrap();
+        for _ in 0..100 {
+            a.append(Axis(0), ones.view()).unwrap();
+        }
+    });
+}
+
+#[bench]
+fn push_no_reserve(bench: &mut Bencher)
+{
+    let ones: Array<f32, _> = array![1f32];
+    bench.iter(|| {
+        let mut a: Array<f32, Ix1> = array![];
+        for _ in 0..100 {
+            a.append(Axis(0), ones.view()).unwrap();
+        }
+    });
+}
diff --git a/src/impl_owned_array.rs b/src/impl_owned_array.rs
@@ -251,6 +251,52 @@ impl<A> Array<A, Ix2>
    {
        self.append(Axis(1), column.insert_axis(Axis(1)))
    }
+
+    /// Reserve capacity to grow array by at least `additional` rows.
+    ///
+    /// Existing elements of `array` are untouched and the backing storage is grown by
+    /// calling the underlying `reserve` method of the `OwnedRepr`.
+    ///
+    /// This is useful when pushing or appending repeatedly to an array to avoid multiple
+    /// allocations.
+    ///
+    /// ***Errors*** with a shape error if the resultant capacity is larger than the addressable
+    /// bounds; that is, the product of non-zero axis lengths once `axis` has been extended by
+    /// `additional` exceeds `isize::MAX`.
+    ///
+    /// ```rust
+    /// use ndarray::Array2;
+    /// let mut a = Array2::<i32>::zeros((2,4));
+    /// a.reserve_rows(1000).unwrap();
+    /// assert!(a.into_raw_vec().capacity() >= 4*1002);
+    /// ```
+    pub fn reserve_rows(&mut self, additional: usize) -> Result<(), ShapeError>
+    {
+        self.reserve(Axis(0), additional)
+    }
+
+    /// Reserve capacity to grow array by at least `additional` columns.
+    ///
+    /// Existing elements of `array` are untouched and the backing storage is grown by
+    /// calling the underlying `reserve` method of the `OwnedRepr`.
+    ///
+    /// This is useful when pushing or appending repeatedly to an array to avoid multiple
+    /// allocations.
+    ///
+    /// ***Errors*** with a shape error if the resultant capacity is larger than the addressable
+    /// bounds; that is, the product of non-zero axis lengths once `axis` has been extended by
+    /// `additional` exceeds `isize::MAX`.
+    ///
+    /// ```rust
+    /// use ndarray::Array2;
+    /// let mut a = Array2::<i32>::zeros((2,4));
+    /// a.reserve_columns(1000).unwrap();
+    /// assert!(a.into_raw_vec().capacity() >= 2*1002);
+    /// ```
+    pub fn reserve_columns(&mut self, additional: usize) -> Result<(), ShapeError>
+    {
+        self.reserve(Axis(1), additional)
+    }
 }

 impl<A, D> Array<A, D>
@@ -660,14 +706,10 @@ where D: Dimension
            self.strides.clone()
        };

-        unsafe {
-            // grow backing storage and update head ptr
-            let offset_from_alloc_to_logical = self.offset_from_alloc_to_logical_ptr().unwrap_or(0);
-            self.ptr = self
-                .data
-                .reserve(len_to_append)
-                .add(offset_from_alloc_to_logical);
+        // grow backing storage and update head ptr
+        self.reserve(axis, array_dim[axis.index()])?;

+        unsafe {
            // clone elements from view to the array now
            //
            // To be robust for panics and drop the right elements, we want
@@ -751,6 +793,70 @@ where D: Dimension

        Ok(())
    }
+
+    /// Reserve capacity to grow array along `axis` by at least `additional` elements.
+    ///
+    /// The axis should be in the range `Axis(` 0 .. *n* `)` where *n* is the
+    /// number of dimensions (axes) of the array.
+    ///
+    /// Existing elements of `array` are untouched and the backing storage is grown by
+    /// calling the underlying `reserve` method of the `OwnedRepr`.
+    ///
+    /// This is useful when pushing or appending repeatedly to an array to avoid multiple
+    /// allocations.
+    ///
+    /// ***Panics*** if the axis is out of bounds.
+    ///
+    /// ***Errors*** with a shape error if the resultant capacity is larger than the addressable
+    /// bounds; that is, the product of non-zero axis lengths once `axis` has been extended by
+    /// `additional` exceeds `isize::MAX`.
+    ///
+    /// ```rust
+    /// use ndarray::{Array3, Axis};
+    /// let mut a = Array3::<i32>::zeros((0,2,4));
+    /// a.reserve(Axis(0), 1000).unwrap();
+    /// assert!(a.into_raw_vec().capacity() >= 2*4*1000);
+    /// ```
+    ///
+    pub fn reserve(&mut self, axis: Axis, additional: usize) -> Result<(), ShapeError>
+    where D: RemoveAxis
+    {
+        debug_assert!(axis.index() < self.ndim());
+        let self_dim = self.raw_dim();
+        let remaining_shape = self_dim.remove_axis(axis);
+
+        // Make sure added capacity doesn't overflow usize::MAX
+        let len_to_append = remaining_shape
+            .size()
+            .checked_mul(additional)
+            .ok_or(ShapeError::from_kind(ErrorKind::Overflow))?;
+
+        // Make sure new capacity is still in bounds
+        let mut res_dim = self_dim;
+        res_dim[axis.index()] += additional;
+        let new_len = dimension::size_of_shape_checked(&res_dim)?;
+
+        // Check whether len_to_append would cause an overflow
+        debug_assert_eq!(self.len().checked_add(len_to_append).unwrap(), new_len);
+
+        unsafe {
+            // grow backing storage and update head ptr
+            let data_to_array_offset = if std::mem::size_of::<A>() != 0 {
+                self.as_ptr().offset_from(self.data.as_ptr())
+            } else {
+                0
+            };
+            debug_assert!(data_to_array_offset >= 0);
+            self.ptr = self
+                .data
+                .reserve(len_to_append)
+                .offset(data_to_array_offset);
+        }
+
+        debug_assert!(self.pointer_is_inbounds());
+
+        Ok(())
+    }
 }

 /// This drops all "unreachable" elements in `self_` given the data pointer and data length.

diff --git a/tests/reserve.rs b/tests/reserve.rs
@@ -0,0 +1,70 @@
+use ndarray::prelude::*;
+
+fn into_raw_vec_capacity<T, D: Dimension>(a: Array<T, D>) -> usize
+{
+    a.into_raw_vec_and_offset().0.capacity()
+}
+
+#[test]
+fn reserve_1d()
+{
+    let mut a = Array1::<i32>::zeros((4,));
+    a.reserve(Axis(0), 1000).unwrap();
+    assert_eq!(a.shape(), &[4]);
+    assert!(into_raw_vec_capacity(a) >= 1004);
+}
+
+#[test]
+fn reserve_3d()
+{
+    let mut a = Array3::<i32>::zeros((0, 4, 8));
+    a.reserve(Axis(0), 10).unwrap();
+    assert_eq!(a.shape(), &[0, 4, 8]);
+    assert!(into_raw_vec_capacity(a) >= 4 * 8 * 10);
+}
+
+#[test]
+fn reserve_empty_3d()
+{
+    let mut a = Array3::<i32>::zeros((0, 0, 0));
+    a.reserve(Axis(0), 10).unwrap();
+}
+
+#[test]
+fn reserve_3d_axis1()
+{
+    let mut a = Array3::<i32>::zeros((2, 4, 8));
+    a.reserve(Axis(1), 10).unwrap();
+    assert!(into_raw_vec_capacity(a) >= 2 * 8 * 10);
+}
+
+#[test]
+fn reserve_3d_repeat()
+{
+    let mut a = Array3::<i32>::zeros((2, 4, 8));
+    a.reserve(Axis(1), 10).unwrap();
+    a.reserve(Axis(2), 30).unwrap();
+    assert!(into_raw_vec_capacity(a) >= 2 * 4 * 30);
+}
+
+#[test]
+fn reserve_2d_with_data()
+{
+    let mut a = array![[1, 2], [3, 4], [5, 6]];
+    a.reserve(Axis(1), 100).unwrap();
+    assert_eq!(a, array![[1, 2], [3, 4], [5, 6]]);
+    assert!(into_raw_vec_capacity(a) >= 3 * 100);
+}
+
+#[test]
+fn reserve_2d_inverted_with_data()
+{
+    let mut a = array![[1, 2], [3, 4], [5, 6]];
+    a.invert_axis(Axis(1));
+    assert_eq!(a, array![[2, 1], [4, 3], [6, 5]]);
+    a.reserve(Axis(1), 100).unwrap();
+    assert_eq!(a, array![[2, 1], [4, 3], [6, 5]]);
+    let (raw_vec, offset) = a.into_raw_vec_and_offset();
+    assert!(raw_vec.capacity() >= 3 * 100);
+    assert_eq!(offset, Some(1));
+}