datasource.hpp

前往此文件的文档。

/*
 * 版权所有 (c) 2020-2025, NVIDIA CORPORATION.
 *
 * 根据 Apache 许可，版本 2.0 (以下简称“许可”) 获得许可；
 * 除非遵守许可，否则不得使用此文件。
 * 您可以获取一份许可的副本，地址如下：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则根据许可分发
 * 的软件是按“原样”分发的，
 * 不附带任何明示或默示的担保或条件。
 * 请参阅许可了解特定语言的管理权限和
 * 限制。
 */

#pragma once

#include <cudf/io/types.hpp>
#include <cudf/utilities/error.hpp>
#include <cudf/utilities/export.hpp>
#include <cudf/utilities/span.hpp>

#include <rmm/cuda_stream_view.hpp>

#include <future>
#include <memory>

namespace CUDF_EXPORT cudf {
namespace io {

/**
 * @brief 为读取器提供输入数据的接口类。
 */
 template <typename Container>
 class owning_buffer; // 前向声明
 class buffer {
 public
 /**
 * @brief 返回缓冲区的字节大小。
 */ [[nodiscard]] virtual size_t size() const = 0;

 /**
 * @brief 返回缓冲区中数据的地址。
 */ [[nodiscard]] virtual uint8_t const* data() const = 0;
 virtual ~buffer() = default;

 /**
 * @brief 从容器构造 datasource 缓冲区对象的工厂函数。
 template <typename Container>
 static std::unique_ptr<buffer> create(Container&& data_owner)
 {
 return std::make_unique<owning_buffer<Container>>(std::forward<Container>(data_owner));
 }
 static std::unique_ptr<datasource> create(std::string const& filepath,
 size_t offset = 0,
 size_t max_size_estimate = 0);

 /**
 * @brief 从主机内存缓冲区创建 source。
 static std::unique_ptr<datasource> create(host_buffer const& buffer);

 /**
 * @brief 从主机内存缓冲区创建 source。
 */ static std::unique_ptr<datasource> create(cudf::host_span<std::byte const> buffer);

 /**
 * @brief 从设备内存缓冲区创建 source。
 */ static std::unique_ptr<datasource> create(cudf::device_span<std::byte const> buffer);

 /**
 * @brief 从用户实现的 datasource 对象创建 source。
 */ static std::unique_ptr<datasource> create(datasource* source);
 template <typename T>
 /**
 * @brief 创建 datasource 向量，输入向量中的每个元素对应一个 datasource。
 */ static std::vector<std::unique_ptr<datasource>> create(std::vector<T> const& args)
 {
 std::vector<std::unique_ptr<datasource>> sources;
 sources.reserve(args.size());
 std::transform(args.cbegin(), args.cend(), std::back_inserter(sources), [](auto const& arg) {
 return datasource::create(arg);
 });
 return sources;
 }
 /**
 * @brief 返回包含 source 部分数据的缓冲区。
 */ virtual ~datasource() = default;

 /**
 * @brief 从 datasource 读取指定部分的数据。
 */ virtual std::unique_ptr<datasource::buffer> host_read(size_t offset, size_t size) = 0;

 /**
 * @brief 从 datasource 异步读取指定部分的数据。
 */ virtual std::future<std::unique_ptr<datasource::buffer>> host_read_async(size_t offset,
 size_t size);

 /**
 * @brief 将选定范围的数据读取到预分配的缓冲区中。
 */ virtual size_t host_read(size_t offset, size_t size, uint8_t* dst) = 0;

 /**
 * @brief 将数据从 source 异步读取到提供的主机内存缓冲区中。
 */ virtual std::future<size_t> host_read_async(size_t offset, size_t size, uint8_t* dst);

 /**
 * @brief 此 source 是否支持直接读取到设备内存中。
 */ [[nodiscard]] virtual bool supports_device_read() const { return false; }
 /**
 * @brief 估计对于给定大小，直接设备读取是否更优。
 */ [[nodiscard]] virtual bool is_device_read_preferred(size_t size) const
 {
 return supports_device_read();
 }

 /**
 * @brief 返回包含 source 部分数据的设备缓冲区。
 */ virtual std::unique_ptr<datasource::buffer> device_read(size_t offset,
 size_t size,
 rmm::cuda_stream_view stream)
 {
 CUDF_FAIL("datasource classes that support device_read must override it.");
 }

 /**
 * @brief 将选定范围的数据读取到预分配的设备缓冲区中。
 */ virtual size_t device_read(size_t offset, size_t size, uint8_t* dst, rmm::cuda_stream_view stream)
 {
 CUDF_FAIL("datasource classes that support device_read must override it.");
 }
 /**
 * @brief 将选定范围的数据异步读取到预分配的设备缓冲区中。
 */ virtual std::future<size_t> device_read_async(size_t offset,
 size_t size,
 uint8_t* dst,
 rmm::cuda_stream_view stream)
 {
 CUDF_FAIL("datasource classes that support device_read_async must override it.");
 }

 /**
 * @brief 返回 source 中数据的大小。
 */ [[nodiscard]] virtual size_t size() const = 0;

 /**
 * @brief 返回 source 是否包含任何数据。
 */ [[nodiscard]] virtual bool is_empty() const { return size() == 0; }

 /**
 * @brief 非拥有缓冲区的实现，datasource 在销毁前持有该缓冲区。
 */ class non_owning_buffer : public buffer {
 public
 non_owning_buffer() = default;

 /**
 * @brief 构造一个新的非拥有缓冲区对象。
 */ non_owning_buffer(uint8_t const* data, size_t size) : _data(data), _size(size) {}

 /**
 * @brief 返回缓冲区的大小。
 */ [[nodiscard]] size_t size() const override { return _size; }

 /**