From 984c91e2f82cdc2fc43434e400e3d4ae9c3bf000 Mon Sep 17 00:00:00 2001
From: ZENOTME <43447882+ZENOTME@users.noreply.github.com>
Date: Thu, 26 Sep 2024 17:56:02 +0800
Subject: [PATCH] avoid to create memory schema operator every time (#635)

Co-authored-by: ZENOTME <st810918843@gmail.com>
---
 crates/iceberg/src/io/file_io.rs | 19 +++++++++++++++++++
 crates/iceberg/src/io/storage.rs | 12 +++++-------
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/crates/iceberg/src/io/file_io.rs b/crates/iceberg/src/io/file_io.rs
index 5e30968f5..e0d83a08b 100644
--- a/crates/iceberg/src/io/file_io.rs
+++ b/crates/iceberg/src/io/file_io.rs
@@ -368,6 +368,7 @@ mod tests {
     use std::io::Write;
     use std::path::Path;
 
+    use bytes::Bytes;
     use futures::io::AllowStdIo;
     use futures::AsyncReadExt;
     use tempfile::TempDir;
@@ -490,4 +491,22 @@ mod tests {
         let io = FileIO::from_path("tmp/||c");
         assert!(io.is_err());
     }
+
+    #[tokio::test]
+    async fn test_memory_io() {
+        let io = FileIOBuilder::new("memory").build().unwrap();
+
+        let path = format!("{}/1.txt", TempDir::new().unwrap().path().to_str().unwrap());
+
+        let output_file = io.new_output(&path).unwrap();
+        output_file.write("test".into()).await.unwrap();
+
+        assert!(io.is_exist(&path.clone()).await.unwrap());
+        let input_file = io.new_input(&path).unwrap();
+        let content = input_file.read().await.unwrap();
+        assert_eq!(content, Bytes::from("test"));
+
+        io.delete(&path).await.unwrap();
+        assert!(!io.is_exist(&path).await.unwrap());
+    }
 }
diff --git a/crates/iceberg/src/io/storage.rs b/crates/iceberg/src/io/storage.rs
index 890104448..cd9b54d54 100644
--- a/crates/iceberg/src/io/storage.rs
+++ b/crates/iceberg/src/io/storage.rs
@@ -30,7 +30,7 @@ use crate::{Error, ErrorKind};
 #[derive(Debug)]
 pub(crate) enum Storage {
     #[cfg(feature = "storage-memory")]
-    Memory,
+    Memory(Operator),
     #[cfg(feature = "storage-fs")]
     LocalFs,
     #[cfg(feature = "storage-s3")]
@@ -56,7 +56,7 @@ impl Storage {
 
         match scheme {
             #[cfg(feature = "storage-memory")]
-            Scheme::Memory => Ok(Self::Memory),
+            Scheme::Memory => Ok(Self::Memory(super::memory_config_build()?)),
             #[cfg(feature = "storage-fs")]
             Scheme::Fs => Ok(Self::LocalFs),
             #[cfg(feature = "storage-s3")]
@@ -96,13 +96,11 @@ impl Storage {
         let path = path.as_ref();
         match self {
             #[cfg(feature = "storage-memory")]
-            Storage::Memory => {
-                let op = super::memory_config_build()?;
-
+            Storage::Memory(op) => {
                 if let Some(stripped) = path.strip_prefix("memory:/") {
-                    Ok((op, stripped))
+                    Ok((op.clone(), stripped))
                 } else {
-                    Ok((op, &path[1..]))
+                    Ok((op.clone(), &path[1..]))
                 }
             }
             #[cfg(feature = "storage-fs")]