Feature: add into_mut to vectors (#5221)

Revived from #5215 since that
was supposed to merge _after_
#5214

---------

Signed-off-by: Connor Tsui <[email protected]>
Co-authored-by: Nicholas Gates <[email protected]>

Author: connortsui20

vortex-vector/src/binaryview/vector.rs

@@ -254,6 +254,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

@@ -121,4 +121,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

@@ -203,4 +203,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

@@ -70,6 +70,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

@@ -199,6 +199,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

@@ -285,6 +285,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

@@ -60,4 +60,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/primitive/generic.rs

@@ -156,4 +156,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

@@ -93,6 +93,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 {

vortex-vector/src/struct_/vector.rs

@@ -197,4 +197,25 @@ impl VectorOps for StructVector {
             validity,
         })
     }
+
+    fn into_mut(self) -> StructVectorMut {
+        let len = self.len;
+        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 fields = Arc::try_unwrap(self.fields).unwrap_or_else(|arc| (*arc).clone());
+
+        let mutable_fields: Box<[_]> = fields
+            .into_vec()
+            .into_iter()
+            .map(|field| field.into_mut())
+            .collect();
+
+        StructVectorMut {
+            fields: mutable_fields,
+            len,
+            validity,
+        }
+    }
 }