data_sink.hpp

前往此文件的文档。

/*
 * Copyright (c) 2020-2025, 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
 *
 * https://apache.ac.cn/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