[arrow-rs] branch master updated: InMemory append API (#4153)

[tu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 521fdb993 InMemory append API (#4153)
521fdb993 is described below

commit 521fdb99374cfe856fdb03e84bb56f09b020d785
Author: Berkay Şahin <12...@users.noreply.github.com>
AuthorDate: Fri Apr 28 21:15:09 2023 +0300

    InMemory append API (#4153)
    
    * ready to review
    
    * clippy fix
    
    * Refactor code to remove byte duplication
    
    * simplify shutdown
    
    Co-authored-by: Raphael Taylor-Davies <17...@users.noreply.github.com>
    
    ---------
    
    Co-authored-by: metesynnada <10...@users.noreply.github.com>
    Co-authored-by: Raphael Taylor-Davies <17...@users.noreply.github.com>
---
 object_store/src/memory.rs | 104 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 104 insertions(+)

diff --git a/object_store/src/memory.rs b/object_store/src/memory.rs
index 057a260f7..b01ffbb02 100644
--- a/object_store/src/memory.rs
+++ b/object_store/src/memory.rs
@@ -117,6 +117,17 @@ impl ObjectStore for InMemory {
         Ok(())
     }
 
+    async fn append(
+        &self,
+        location: &Path,
+    ) -> Result<Box<dyn AsyncWrite + Unpin + Send>> {
+        Ok(Box::new(InMemoryAppend {
+            location: location.clone(),
+            data: Vec::<u8>::new(),
+            storage: StorageType::clone(&self.storage),
+        }))
+    }
+
     async fn get(&self, location: &Path) -> Result<GetResult> {
         let data = self.entry(location).await?;
 
@@ -329,8 +340,55 @@ impl AsyncWrite for InMemoryUpload {
     }
 }
 
+struct InMemoryAppend {
+    location: Path,
+    data: Vec<u8>,
+    storage: StorageType,
+}
+
+impl AsyncWrite for InMemoryAppend {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        _cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> std::task::Poll<Result<usize, io::Error>> {
+        self.data.extend_from_slice(buf);
+        Poll::Ready(Ok(buf.len()))
+    }
+
+    fn poll_flush(
+        mut self: Pin<&mut Self>,
+        _cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), io::Error>> {
+        let storage = StorageType::clone(&self.storage);
+
+        let mut writer = storage.write();
+
+        if let Some((bytes, _)) = writer.remove(&self.location) {
+            let buf = std::mem::take(&mut self.data);
+            let concat = Bytes::from_iter(bytes.into_iter().chain(buf.into_iter()));
+            writer.insert(self.location.clone(), (concat, Utc::now()));
+        } else {
+            writer.insert(
+                self.location.clone(),
+                (Bytes::from(std::mem::take(&mut self.data)), Utc::now()),
+            );
+        };
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), io::Error>> {
+        self.poll_flush(cx)
+    }
+}
+
 #[cfg(test)]
 mod tests {
+    use tokio::io::AsyncWriteExt;
+
     use super::*;
 
     use crate::{
@@ -396,4 +454,50 @@ mod tests {
             panic!("unexpected error type: {err:?}");
         }
     }
+
+    #[tokio::test]
+    async fn test_append_new() {
+        let in_memory = InMemory::new();
+        let location = Path::from("some_file");
+        let data = Bytes::from("arbitrary data");
+        let expected_data = data.clone();
+
+        let mut writer = in_memory.append(&location).await.unwrap();
+        writer.write_all(&data).await.unwrap();
+        writer.flush().await.unwrap();
+
+        let read_data = in_memory
+            .get(&location)
+            .await
+            .unwrap()
+            .bytes()
+            .await
+            .unwrap();
+        assert_eq!(&*read_data, expected_data);
+    }
+
+    #[tokio::test]
+    async fn test_append_existing() {
+        let in_memory = InMemory::new();
+        let location = Path::from("some_file");
+        let data = Bytes::from("arbitrary");
+        let data_appended = Bytes::from(" data");
+        let expected_data = Bytes::from("arbitrary data");
+
+        let mut writer = in_memory.append(&location).await.unwrap();
+        writer.write_all(&data).await.unwrap();
+        writer.flush().await.unwrap();
+
+        writer.write_all(&data_appended).await.unwrap();
+        writer.flush().await.unwrap();
+
+        let read_data = in_memory
+            .get(&location)
+            .await
+            .unwrap()
+            .bytes()
+            .await
+            .unwrap();
+        assert_eq!(&*read_data, expected_data);
+    }
 }