[Cpp] Pass path as output argument when BuildFileSystem (#78)

Signed-off-by: sunby <[email protected]>
milvus-io · Oct 18, 2023 · bef9bb3 · bef9bb3
1 parent 3b03ad8
commit bef9bb3
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 10 deletions.
diff --git a/cpp/include/milvus-storage/common/fs_util.h b/cpp/include/milvus-storage/common/fs_util.h
@@ -4,7 +4,7 @@
 #include "result.h"
 namespace milvus_storage {
 
-Result<std::shared_ptr<arrow::fs::FileSystem>> BuildFileSystem(const std::string& uri);
+Result<std::shared_ptr<arrow::fs::FileSystem>> BuildFileSystem(const std::string& uri, std::string* out_path = nullptr);
 
 std::string UriToPath(const std::string& uri);
 

diff --git a/cpp/src/common/fs_util.cpp b/cpp/src/common/fs_util.cpp
@@ -8,13 +8,15 @@
 #include "common/macro.h"
 namespace milvus_storage {
 
-Result<std::shared_ptr<arrow::fs::FileSystem>> BuildFileSystem(const std::string& uri) {
+Result<std::shared_ptr<arrow::fs::FileSystem>> BuildFileSystem(const std::string& uri, std::string* out_path) {
   arrow::internal::Uri uri_parser;
   RETURN_ARROW_NOT_OK(uri_parser.Parse(uri));
   auto schema = uri_parser.scheme();
   if (schema == "file") {
-    auto output_path = uri_parser.path();
-    ASSIGN_OR_RETURN_ARROW_NOT_OK(auto option, arrow::fs::LocalFileSystemOptions::FromUri(uri_parser, &output_path));
+    if (out_path == nullptr) {
+      return Status::InvalidArgument("out_path should not be nullptr if schema is file");
+    }
+    ASSIGN_OR_RETURN_ARROW_NOT_OK(auto option, arrow::fs::LocalFileSystemOptions::FromUri(uri_parser, out_path));
     return std::shared_ptr<arrow::fs::FileSystem>(new arrow::fs::LocalFileSystem(option));
   }
 
@@ -34,7 +36,7 @@ Result<std::shared_ptr<arrow::fs::FileSystem>> BuildFileSystem(const std::string
         }
       });
     }
-    ASSIGN_OR_RETURN_ARROW_NOT_OK(auto option, arrow::fs::S3Options::FromUri(uri_parser));
+    ASSIGN_OR_RETURN_ARROW_NOT_OK(auto option, arrow::fs::S3Options::FromUri(uri_parser, out_path));
     ASSIGN_OR_RETURN_ARROW_NOT_OK(auto fs, arrow::fs::S3FileSystem::Make(option));
 
     return std::shared_ptr<arrow::fs::FileSystem>(fs);

diff --git a/cpp/src/storage/space.cpp b/cpp/src/storage/space.cpp
@@ -232,11 +232,9 @@ Result<std::unique_ptr<Space>> Space::Open(const std::string& uri, Options optio
   std::string path;
   std::atomic_int64_t next_manifest_version = 1;
 
-  ASSIGN_OR_RETURN_NOT_OK(fs, BuildFileSystem(uri));
-  arrow::internal::Uri uri_parser;
-  RETURN_ARROW_NOT_OK(uri_parser.Parse(uri));
-  path = uri_parser.path();
+  ASSIGN_OR_RETURN_NOT_OK(fs, BuildFileSystem(uri, &path));
 
+  LOG_STORAGE_INFO_ << "Open space: " << path;
   RETURN_ARROW_NOT_OK(fs->CreateDir(GetManifestDir(path)));
   RETURN_ARROW_NOT_OK(fs->CreateDir(GetScalarDataDir(path)));
   RETURN_ARROW_NOT_OK(fs->CreateDir(GetVectorDataDir(path)));

diff --git a/cpp/test/multi_files_sequential_reader_test.cpp b/cpp/test/multi_files_sequential_reader_test.cpp
@@ -25,7 +25,8 @@ TEST(MultiFilesSeqReaderTest, ReadTest) {
   ASSERT_STATUS_OK(pk_builder.Finish(&pk_array));
   auto rec_batch = arrow::RecordBatch::Make(arrow_schema, 3, {pk_array});
 
-  ASSERT_AND_ASSIGN(auto fs, BuildFileSystem("file:///tmp/"));
+  std::string path;
+  ASSERT_AND_ASSIGN(auto fs, BuildFileSystem("file:///tmp/", &path));
   ASSERT_AND_ARROW_ASSIGN(auto f1, fs->OpenOutputStream("/tmp/file1"));
   ASSERT_AND_ARROW_ASSIGN(auto w1, parquet::arrow::FileWriter::Open(*arrow_schema, arrow::default_memory_pool(), f1));
   ASSERT_STATUS_OK(w1->WriteRecordBatch(*rec_batch));