use kvikio::remote_handle.hpp

rapidsai · Oct 29, 2024 · a875da8 · a875da8
1 parent 007c651
commit a875da8
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/cpp/src/io/utilities/datasource.cpp b/cpp/src/io/utilities/datasource.cpp
@@ -26,6 +26,7 @@
 #include <cudf/utilities/span.hpp>
 
 #include <kvikio/file_handle.hpp>
+#include <kvikio/remote_handle.hpp>
 
 #include <rmm/device_buffer.hpp>
 
@@ -395,6 +396,12 @@ class user_datasource_wrapper : public datasource {
  */
 class remote_file_source {
  public:
+  static std::unique_ptr<kvikio::S3Endpoint> create_s3_endpoint(std::string const& filepath)
+  {
+    auto [bucket_name, bucket_object] = kvikio::S3Endpoint::parse_s3_url(filepath);
+    return std::make_unique<kvikio::S3Endpoint>(bucket_name, bucket_object);
+  }
+
   /**
    * @brief Is `url` referring to a remote file supported by KvikIO?
    *
@@ -422,14 +429,16 @@ std::unique_ptr<datasource> datasource::create(std::string const& filepath,
 
     CUDF_FAIL("Invalid LIBCUDF_MMAP_ENABLED value: " + policy);
   }();
-  CUDF_EXPECTS(!remote_file_source::is_supported_remote_url(filepath),
-               "didn't expect a remote URL");
-  if (use_memory_mapping) {
+  if (remote_file_source::is_supported_remote_url(filepath)) {
+    remote_file_source::create_s3_endpoint(filepath);
+    CUDF_FAIL("is_supported_remote_url(): didn't expect a s3 url!");
+  } else if (use_memory_mapping) {
     return std::make_unique<memory_mapped_source>(filepath.c_str(), offset, max_size_estimate);
   } else {
     // `file_source` reads the file directly, without memory mapping
     return std::make_unique<file_source>(filepath.c_str());
   }
+  return std::make_unique<file_source>(filepath.c_str());
 }
 
 std::unique_ptr<datasource> datasource::create(host_buffer const& buffer)