Feature: add into_mut to vectors (#5215)

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

Author: connortsui20

vortex-vector/src/binaryview/vector.rs

@@ -203,6 +203,21 @@ impl<T: BinaryViewType> VectorOps for BinaryViewVector<T> {
             ))
         }
     }
+
+    fn into_mut(self) -> BinaryViewVectorMut<T> {
+        let views_mut = self.views.into_mut();
+        let validity_mut = self.validity.into_mut();
+
+        // If someone else has a strong reference to the `Arc`, clone the underlying data (which is
+        // just a **different** reference count increment).
+        let buffers_mut = Arc::try_unwrap(self.buffers)
+            .unwrap_or_else(|arc| (*arc).clone())
+            .into_vec();
+
+        // SAFETY: The BinaryViewVector maintains the exact same invariants as the immutable
+        // version, so all invariants are still upheld.
+        unsafe { BinaryViewVectorMut::new_unchecked(views_mut, validity_mut, buffers_mut) }
+    }
 }
 
 #[cfg(test)]

vortex-vector/src/bool/vector.rs

@@ -103,4 +103,11 @@ impl VectorOps for BoolVector {
             }),
         }
     }
+
+    fn into_mut(self) -> BoolVectorMut {
+        BoolVectorMut {
+            bits: self.bits.into_mut(),
+            validity: self.validity.into_mut(),
+        }
+    }
 }

vortex-vector/src/decimal/generic.rs

@@ -180,4 +180,12 @@ impl<D: NativeDecimalType> VectorOps for DVector<D> {
             }),
         }
     }
+
+    fn into_mut(self) -> DVectorMut<D> {
+        DVectorMut {
+            ps: self.ps,
+            elements: self.elements.into_mut(),
+            validity: self.validity.into_mut(),
+        }
+    }
 }

vortex-vector/src/decimal/vector.rs

@@ -45,6 +45,10 @@ impl VectorOps for DecimalVector {
                 .map_err(Self::from)
         })
     }
+
+    fn into_mut(self) -> DecimalVectorMut {
+        match_each_dvector!(self, |v| { DecimalVectorMut::from(v.into_mut()) })
+    }
 }
 
 impl DecimalTypeDowncast for DecimalVector {

vortex-vector/src/fixed_size_list/vector.rs

@@ -188,6 +188,23 @@ impl VectorOps for FixedSizeListVector {
             }),
         }
     }
+
+    fn into_mut(self) -> FixedSizeListVectorMut {
+        let len = self.len;
+        let list_size = self.list_size;
+        let validity = self.validity.into_mut();
+
+        // If someone else has a strong reference to the `Arc`, clone the underlying data (which is
+        // just a **different** reference count increment).
+        let elements = Arc::try_unwrap(self.elements).unwrap_or_else(|arc| (*arc).clone());
+
+        FixedSizeListVectorMut {
+            elements: Box::new(elements.into_mut()),
+            list_size,
+            validity,
+            len,
+        }
+    }
 }
 
 #[cfg(test)]

vortex-vector/src/listview/vector.rs

@@ -274,6 +274,25 @@ impl VectorOps for ListViewVector {
             }),
         }
     }
+
+    fn into_mut(self) -> ListViewVectorMut {
+        let len = self.len;
+        let validity = self.validity.into_mut();
+        let offsets = self.offsets.into_mut();
+        let sizes = self.sizes.into_mut();
+
+        // If someone else has a strong reference to the `Arc`, clone the underlying data (which is
+        // just a **different** reference count increment).
+        let elements = Arc::try_unwrap(self.elements).unwrap_or_else(|arc| (*arc).clone());
+
+        ListViewVectorMut {
+            offsets,
+            sizes,
+            elements: Box::new(elements.into_mut()),
+            validity,
+            len,
+        }
+    }
 }
 
 // TODO(connor): It would be better to separate everything inside the macros into its own function,

vortex-vector/src/null/vector.rs

@@ -47,4 +47,8 @@ impl VectorOps for NullVector {
     fn try_into_mut(self) -> Result<NullVectorMut, Self> {
         Ok(NullVectorMut::new(self.len))
     }
+
+    fn into_mut(self) -> NullVectorMut {
+        NullVectorMut::new(self.len)
+    }
 }

vortex-vector/src/ops.rs

@@ -39,6 +39,19 @@ pub trait VectorOps: private::Sealed + Into<Vector> + Sized {
     ///
     /// If `self` is not unique, this will fail and return `self` back to the caller.
     fn try_into_mut(self) -> Result<Self::Mutable, Self>;
+
+    /// Converts `self` into a mutable vector (implementing [`VectorMutOps`]).
+    ///
+    /// This method uses "clone-on-write" semantics, meaning it will clone any underlying data that
+    /// has multiple references (preventing mutable access). `into_mut` can be more efficient than
+    /// [`try_into_mut()`] when mutations are infrequent.
+    ///
+    /// The semantics of `into_mut` are somewhat similar to that of [`Arc::make_mut()`], but instead
+    /// of working with references, this works with owned immutable / mutable types.
+    ///
+    /// [`try_into_mut()`]: Self::try_into_mut
+    /// [`Arc::make_mut()`]: std::sync::Arc::make_mut
+    fn into_mut(self) -> Self::Mutable;
 }
 
 /// Common operations for mutable vectors (all the variants of [`VectorMut`]).

vortex-vector/src/primitive/generic.rs

@@ -142,4 +142,11 @@ impl<T: NativePType> VectorOps for PVector<T> {
             }),
         }
     }
+
+    fn into_mut(self) -> PVectorMut<T> {
+        let elements = self.elements.into_mut();
+        let validity = self.validity.into_mut();
+
+        PVectorMut { elements, validity }
+    }
 }

vortex-vector/src/primitive/vector.rs

@@ -82,6 +82,10 @@ impl VectorOps for PrimitiveVector {
                 .map_err(Self::from)
         })
     }
+
+    fn into_mut(self) -> PrimitiveVectorMut {
+        match_each_pvector!(self, |v| { v.into_mut().into() })
+    }
 }
 
 impl PTypeUpcast for PrimitiveVector {