Provide data_chunk_source wrapper for datasource

cpp/include/cudf/io/text/data_chunk_source_factories.hpp

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <cudf/io/datasource.hpp>
 #include <cudf/io/text/data_chunk_source.hpp>
 #include <cudf/scalar/scalar.hpp>
 #include <cudf/utilities/span.hpp>
@@ -25,6 +26,14 @@
 
 namespace cudf::io::text {
 
+/**
+ * @brief Creates a data source capable of producing device-buffered views of a datasource.
+ * @param data the datasource to be exposed as a data chunk source
+ * @return the data chunk source for the provided datasource. It must not outlive the datasource
+ *         used to construct it.
+ */
+std::unique_ptr<data_chunk_source> make_source(datasource& data);
+
 /**
  * @brief Creates a data source capable of producing device-buffered views of the given string.
  * @param data the host data to be exposed as a data chunk source. Its lifetime must be at least as

cpp/src/io/text/data_chunk_source_factories.cpp

@@ -30,6 +30,88 @@ namespace cudf::io::text {
 
 namespace {
 
+/**
+ * @brief A reader which produces owning chunks of device memory which contain a copy of the data
+ * from an istream.
+ */
+class datasource_chunk_reader : public data_chunk_reader {
+  struct host_ticket {
+    cudaEvent_t event;
+    thrust::host_vector<char, thrust::system::cuda::experimental::pinned_allocator<char>> buffer;
+  };
+
+ public:
+  datasource_chunk_reader(datasource* source) : _source(source), _tickets(2)
+  {
+    // create an event to track the completion of the last device-to-host copy.
+    for (auto& ticket : _tickets) {
+      CUDF_CUDA_TRY(cudaEventCreate(&(ticket.event)));
+    }
+  }
+
+  ~datasource_chunk_reader() override
+  {
+    for (auto& ticket : _tickets) {
+      CUDF_CUDA_TRY(cudaEventDestroy(ticket.event));
+    }
+  }
+
+  void skip_bytes(std::size_t size) override
+  {
+    _offset += std::min(_source->size() - _offset, size);
+  };
+
+  std::unique_ptr<device_data_chunk> get_next_chunk(std::size_t read_size,
+                                                    rmm::cuda_stream_view stream) override
+  {
+    CUDF_FUNC_RANGE();
+
+    read_size = std::min(_source->size() - _offset, read_size);
+
+    // get a device buffer containing read data on the device.
+    auto chunk = rmm::device_uvector<char>(read_size, stream);
+
+    if (_source->supports_device_read() && _source->is_device_read_preferred(read_size)) {
+      _source->device_read_async(
+        _offset, read_size, reinterpret_cast<uint8_t*>(chunk.data()), stream);
+    } else {
+      auto& h_ticket = _tickets[_next_ticket_idx];
+
+      _next_ticket_idx = (_next_ticket_idx + 1) % _tickets.size();
+
+      // synchronize on the last host-to-device copy, so we don't clobber the host buffer.
+      CUDF_CUDA_TRY(cudaEventSynchronize(h_ticket.event));
+
+      // resize the host buffer as necessary to contain the requested number of bytes
+      if (h_ticket.buffer.size() < read_size) { h_ticket.buffer.resize(read_size); }
+
+      _source->host_read(_offset, read_size, reinterpret_cast<uint8_t*>(h_ticket.buffer.data()));
+
+      // copy the host-pinned data on to device
+      CUDF_CUDA_TRY(cudaMemcpyAsync(  //
+        chunk.data(),
+        h_ticket.buffer.data(),
+        read_size,
+        cudaMemcpyHostToDevice,
+        stream.value()));
+
+      // record the host-to-device copy.
+      CUDF_CUDA_TRY(cudaEventRecord(h_ticket.event, stream.value()));
+    }
+
+    _offset += read_size;
+
+    // return the device buffer so it can be processed.
+    return std::make_unique<device_uvector_data_chunk>(std::move(chunk));
+  }
+
+ private:
+  std::size_t _offset          = 0;
+  std::size_t _next_ticket_idx = 0;
+  datasource* _source;
+  std::vector<host_ticket> _tickets;
+};
+
 /**
  * @brief A reader which produces owning chunks of device memory which contain a copy of the data
  * from an istream.
@@ -180,6 +262,21 @@ class device_span_data_chunk_reader : public data_chunk_reader {
   uint64_t _position = 0;
 };
 
+/**
+ * @brief A datasource-based data chunk source which creates a datasource_chunk_reader.
+ */
+class datasource_chunk_source : public data_chunk_source {
+ public:
+  datasource_chunk_source(datasource& source) : _source(&source) {}
+  [[nodiscard]] std::unique_ptr<data_chunk_reader> create_reader() const override
+  {
+    return std::make_unique<datasource_chunk_reader>(_source);
+  }
+
+ private:
+  datasource* _source;
+};
+
 /**
  * @brief A file data source which creates an istream_data_chunk_reader.
  */
@@ -228,6 +325,11 @@ class device_span_data_chunk_source : public data_chunk_source {
 
 }  // namespace
 
+std::unique_ptr<data_chunk_source> make_source(datasource& data)
+{
+  return std::make_unique<datasource_chunk_source>(data);
+}
+
 std::unique_ptr<data_chunk_source> make_source(host_span<const char> data)
 {
   return std::make_unique<host_span_data_chunk_source>(data);

cpp/tests/io/text/data_chunk_source_test.cpp

@@ -96,6 +96,34 @@ void test_source(const std::string& content, const cudf::io::text::data_chunk_so
   }
 }
 
+TEST_F(DataChunkSourceTest, DataSourceHost)
+{
+  std::string const content = "host buffer source";
+  auto const datasource =
+    cudf::io::datasource::create(cudf::io::host_buffer{content.data(), content.size()});
+  auto const source = cudf::io::text::make_source(*datasource);
+
+  test_source(content, *source);
+}
+
+TEST_F(DataChunkSourceTest, DataSourceFile)
+{
+  std::string content = "file datasource";
+  // make it big enought to have is_device_read_preferred return true
+  for (int i = 0; i < 20; i++) {
+    content = content + content;
+  }
+  auto const filename = temp_env->get_temp_filepath("file_source");
+  {
+    std::ofstream file{filename};
+    file << content;
+  }
+  auto const datasource = cudf::io::datasource::create(filename);
+  auto const source     = cudf::io::text::make_source(*datasource);
+
+  test_source(content, *source);
+}
+
 TEST_F(DataChunkSourceTest, Device)
 {
   std::string const content = "device buffer source";