Fix is_device_write_preferred in void_sink and user_sink_wrapper

cpp/benchmarks/io/cuio_common.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -41,7 +41,8 @@ std::string random_file_in_dir(std::string const& dir_path)
 cuio_source_sink_pair::cuio_source_sink_pair(io_type type)
   : type{type},
     d_buffer{0, cudf::get_default_stream()},
-    file_name{random_file_in_dir(tmpdir.path())}
+    file_name{random_file_in_dir(tmpdir.path())},
+    void_sink{cudf::io::data_sink::create()}
 {
 }
 
@@ -67,7 +68,7 @@ cudf::io::source_info cuio_source_sink_pair::make_source_info()
 cudf::io::sink_info cuio_source_sink_pair::make_sink_info()
 {
   switch (type) {
-    case io_type::VOID: return cudf::io::sink_info(&void_sink);
+    case io_type::VOID: return cudf::io::sink_info(void_sink.get());
     case io_type::FILEPATH: return cudf::io::sink_info(file_name);
     case io_type::HOST_BUFFER: [[fallthrough]];
     case io_type::DEVICE_BUFFER: return cudf::io::sink_info(&h_buffer);
@@ -78,7 +79,7 @@ cudf::io::sink_info cuio_source_sink_pair::make_sink_info()
 size_t cuio_source_sink_pair::size()
 {
   switch (type) {
-    case io_type::VOID: return void_sink.bytes_written();
+    case io_type::VOID: return void_sink->bytes_written();
     case io_type::FILEPATH:
       return static_cast<size_t>(
         std::ifstream(file_name, std::ifstream::ate | std::ifstream::binary).tellg());

cpp/benchmarks/io/cuio_common.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -32,15 +32,6 @@ std::string random_file_in_dir(std::string const& dir_path);
  * @brief Class to create a coupled `source_info` and `sink_info` of given type.
  */
 class cuio_source_sink_pair {
-  class bytes_written_only_sink : public cudf::io::data_sink {
-    size_t _bytes_written = 0;
-
-   public:
-    void host_write(void const* data, size_t size) override { _bytes_written += size; }
-    void flush() override {}
-    size_t bytes_written() override { return _bytes_written; }
-  };
-
  public:
   cuio_source_sink_pair(io_type type);
   ~cuio_source_sink_pair()
@@ -79,7 +70,7 @@ class cuio_source_sink_pair {
   std::vector<char> h_buffer;
   rmm::device_uvector<std::byte> d_buffer;
   std::string const file_name;
-  bytes_written_only_sink void_sink;
+  std::unique_ptr<cudf::io::data_sink> void_sink;
 };
 
 /**

cpp/src/io/utilities/data_sink.cpp

@@ -139,6 +139,8 @@ class void_sink : public data_sink {
 
   [[nodiscard]] bool supports_device_write() const override { return true; }
 
+  [[nodiscard]] bool is_device_write_preferred(size_t size) const override { return true; }
+
   void device_write(void const* gpu_data, size_t size, rmm::cuda_stream_view stream) override
   {
     _bytes_written += size;
@@ -189,6 +191,11 @@ class user_sink_wrapper : public data_sink {
     return user_sink->device_write_async(gpu_data, size, stream);
   }
 
+  [[nodiscard]] bool is_device_write_preferred(size_t size) const override
+  {
+    return user_sink->is_device_write_preferred(size);
+  }
+
   void flush() override { user_sink->flush(); }
 
   size_t bytes_written() override { return user_sink->bytes_written(); }