Add take method to SegmentContainer

Author: fabubaker

db4-storage/src/segments/graph/segment.rs

@@ -51,6 +51,18 @@ impl MemGraphSegment {
         &self.layers
     }
 
+    /// Replaces this segment with an empty instance, returning the old segment
+    /// with its data.
+    ///
+    /// The new segment will have the same number of layers as the original.
+    pub fn take(&mut self) -> Self {
+        let layers = self.layers.iter_mut().map(|layer| layer.take()).collect();
+
+        Self {
+            layers,
+        }
+    }
+
     pub fn add_properties(&mut self, t: TimeIndexEntry, props: impl IntoIterator<Item = (usize, Prop)>) {
         let layer = &mut self.layers[Self::LAYER];
 

db4-storage/src/segments/mod.rs

@@ -184,6 +184,15 @@ impl<T: HasRow> SegmentContainer<T> {
         }
     }
 
+    /// Replaces this container with an empty instance, returning the
+    /// old container with its data.
+    pub fn take(&mut self) -> Self {
+        std::mem::replace(
+            self,
+            Self::new(self.segment_id, self.max_page_len, self.meta.clone()),
+        )
+    }
+
     #[inline]
     pub fn est_size(&self) -> usize {
         // TODO: this is a rough estimate and should be improved