data_sink.hpp

Go to the documentation of this file.

/*
 * 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.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <cudf/types.hpp>
#include <cudf/utilities/error.hpp>
 
#include <rmm/cuda_stream_view.hpp>
 
#include <algorithm>
#include <future>
#include <memory>
#include <string>
#include <vector>
 
namespace CUDF_EXPORT cudf {
namespace io {
 
class data_sink {
 public:
  static std::unique_ptr<data_sink> create(std::string const& filepath);
 
  static std::unique_ptr<data_sink> create(std::vector<char>* buffer);
 
  static std::unique_ptr<data_sink> create();
 
  static std::unique_ptr<data_sink> create(cudf::io::data_sink* const user_sink);
 
  template <typename T>
  static std::vector<std::unique_ptr<data_sink>> create(std::vector<T> const& args)
  {
    std::vector<std::unique_ptr<data_sink>> sinks;
    sinks.reserve(args.size());
    std::transform(args.cbegin(), args.cend(), std::back_inserter(sinks), [](auto const& arg) {
      return data_sink::create(arg);
    });
    return sinks;
  }
 
  virtual ~data_sink(){};
 
  virtual void host_write(void const* data, size_t size) = 0;
 
  [[nodiscard]] virtual bool supports_device_write() const { return false; }
 
  [[nodiscard]] virtual bool is_device_write_preferred(size_t size) const
  {
    return supports_device_write();
  }
 
  virtual void device_write(void const* gpu_data, size_t size, rmm::cuda_stream_view stream)
  {
    CUDF_FAIL("data_sink classes that support device_write must override it.");
  }
 
  virtual std::future<void> device_write_async(void const* gpu_data,
                                               size_t size,
                                               rmm::cuda_stream_view stream)
  {
    CUDF_FAIL("data_sink classes that support device_write_async must override it.");
  }
 
  virtual void flush() = 0;
 
  virtual size_t bytes_written() = 0;
};
  // end of group
}  // namespace io
}  // namespace CUDF_EXPORT cudf