diff --git a/vortex-vector/src/binaryview/vector.rs b/vortex-vector/src/binaryview/vector.rs index 3b07ba08e36..1038df94fb7 100644 --- a/vortex-vector/src/binaryview/vector.rs +++ b/vortex-vector/src/binaryview/vector.rs @@ -203,6 +203,21 @@ impl VectorOps for BinaryViewVector { )) } } + + fn into_mut(self) -> BinaryViewVectorMut { + 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)] diff --git a/vortex-vector/src/bool/vector.rs b/vortex-vector/src/bool/vector.rs index 53ece89eeb9..89d4808e1a5 100644 --- a/vortex-vector/src/bool/vector.rs +++ b/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(), + } + } } diff --git a/vortex-vector/src/decimal/generic.rs b/vortex-vector/src/decimal/generic.rs index 6f529d1134a..b73459f6d74 100644 --- a/vortex-vector/src/decimal/generic.rs +++ b/vortex-vector/src/decimal/generic.rs @@ -180,4 +180,12 @@ impl VectorOps for DVector { }), } } + + fn into_mut(self) -> DVectorMut { + DVectorMut { + ps: self.ps, + elements: self.elements.into_mut(), + validity: self.validity.into_mut(), + } + } } diff --git a/vortex-vector/src/decimal/vector.rs b/vortex-vector/src/decimal/vector.rs index e424f7a159d..247663d420b 100644 --- a/vortex-vector/src/decimal/vector.rs +++ b/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 { diff --git a/vortex-vector/src/fixed_size_list/vector.rs b/vortex-vector/src/fixed_size_list/vector.rs index b274984a304..77d9e437a7e 100644 --- a/vortex-vector/src/fixed_size_list/vector.rs +++ b/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)] diff --git a/vortex-vector/src/listview/vector.rs b/vortex-vector/src/listview/vector.rs index 134ec416dc8..c6fb5c190da 100644 --- a/vortex-vector/src/listview/vector.rs +++ b/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, diff --git a/vortex-vector/src/null/vector.rs b/vortex-vector/src/null/vector.rs index c45a3b8818c..c505bca8298 100644 --- a/vortex-vector/src/null/vector.rs +++ b/vortex-vector/src/null/vector.rs @@ -47,4 +47,8 @@ impl VectorOps for NullVector { fn try_into_mut(self) -> Result { Ok(NullVectorMut::new(self.len)) } + + fn into_mut(self) -> NullVectorMut { + NullVectorMut::new(self.len) + } } diff --git a/vortex-vector/src/ops.rs b/vortex-vector/src/ops.rs index 3c5a996dc18..8f82ae8ac28 100644 --- a/vortex-vector/src/ops.rs +++ b/vortex-vector/src/ops.rs @@ -39,6 +39,19 @@ pub trait VectorOps: private::Sealed + Into + Sized { /// /// If `self` is not unique, this will fail and return `self` back to the caller. fn try_into_mut(self) -> Result; + + /// 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`]). diff --git a/vortex-vector/src/primitive/generic.rs b/vortex-vector/src/primitive/generic.rs index 36d6935918b..ee5d30925d7 100644 --- a/vortex-vector/src/primitive/generic.rs +++ b/vortex-vector/src/primitive/generic.rs @@ -142,4 +142,11 @@ impl VectorOps for PVector { }), } } + + fn into_mut(self) -> PVectorMut { + let elements = self.elements.into_mut(); + let validity = self.validity.into_mut(); + + PVectorMut { elements, validity } + } } diff --git a/vortex-vector/src/primitive/vector.rs b/vortex-vector/src/primitive/vector.rs index 21a5e02dd92..b616a861e1e 100644 --- a/vortex-vector/src/primitive/vector.rs +++ b/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 { diff --git a/vortex-vector/src/struct_/vector.rs b/vortex-vector/src/struct_/vector.rs index 2b1f8cd01ed..014788c6691 100644 --- a/vortex-vector/src/struct_/vector.rs +++ b/vortex-vector/src/struct_/vector.rs @@ -186,4 +186,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, + } + } } diff --git a/vortex-vector/src/vector.rs b/vortex-vector/src/vector.rs index 19e20e7a702..de7961aa238 100644 --- a/vortex-vector/src/vector.rs +++ b/vortex-vector/src/vector.rs @@ -75,6 +75,10 @@ impl VectorOps for Vector { v.try_into_mut().map(VectorMut::from).map_err(Vector::from) }) } + + fn into_mut(self) -> VectorMut { + match_each_vector!(self, |v| { VectorMut::from(v.into_mut()) }) + } } impl Vector {