csv.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/table/table_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/error.hpp>
#include <cudf/utilities/memory_resource.hpp>

#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <variant>
#include <vector>

namespace CUDF_EXPORT cudf {
namespace io {

class csv_reader_options_builder;

class csv_reader_options {
 source_info _source;

 // 读取设置

 // 指定源的压缩格式或从文件扩展名推断
 compression_type _compression = compression_type::AUTO;
 // 从源开始跳过的字节数
 std::size_t _byte_range_offset = 0;
 // 要读取的字节数；始终读取完整的行
 std::size_t _byte_range_size = 0;
 // 所有列的名称；如果为空，则自动生成名称
 std::vector<std::string> _names;
 // 如果没有头或名称，将其作为名称前缀添加到列 ID
 std::string _prefix;
 // 是否重命名重复的列名
 bool _mangle_dupe_cols = true;

 // 过滤设置

 // 要读取的列的名称；空表示所有列
 std::vector<std::string> _use_cols_names;
 // 要读取的列的索引；空表示所有列
 std::vector<int> _use_cols_indexes;
 // 要读取的行数；-1 表示所有行
 size_type _nrows = -1;
 // 从开始跳过的行数
 size_type _skiprows = 0;
 // 从末尾跳过的行数
 size_type _skipfooter = 0;
 // 头部行索引
 size_type _header = 0;

 // 解析设置

 // 行终止符
 char _lineterminator = '\n';
 // 字段分隔符
 char _delimiter = ',';
 // 数字数据的千位分隔符；不能与分隔符相同
 char _thousands = '\0';
 // 小数点字符；不能与分隔符相同
 char _decimal = '.';
 // 注释行开始字符
 char _comment = '\0';
 bool _windowslinetermination = false;
 // 将空格视为字段分隔符；覆盖字符分隔符
 bool _delim_whitespace = false;
 // 跳过分隔符后的空格
 bool _skipinitialspace = false;
 // 忽略空行或将行值解析为无效
 bool _skip_blank_lines = true;
 // 引号处理行为
 quote_style _quoting = quote_style::MINIMAL;
 // 引号字符（如果 `quoting` 为 true）
 char _quotechar = '"';
 // 值内部的引号是否双重引用
 bool _doublequote = true;
 // 是否检测引号周围的空格，例如 ` "data" `。当 _doublequote 为 true 时，此标志无效
 // _doublequote is true
 bool _detect_whitespace_around_quotes = false;
 // 要读取为日期时间的列的名称
 std::vector<std::string> _parse_dates_names;
 // 要读取为日期时间的列的索引
 std::vector<int> _parse_dates_indexes;
 // 要解析为十六进制的列的名称
 std::vector<std::string> _parse_hex_names;
 // 要解析为十六进制的列的索引
 std::vector<int> _parse_hex_indexes;

 // 转换设置

 // 每列类型；禁用这些列上的类型推断
 std::variant<std::vector<data_type>, std::map<std::string, data_type>> _dtypes;
 // 识别为布尔真值的附加值
 std::vector<std::string> _true_values{"True", "TRUE", "true"};
 // 识别为布尔假值的附加值
 std::vector<std::string> _false_values{"False", "FALSE", "false"};
 // 识别为空值的附加值
 std::vector<std::string> _na_values;
 // 是否保留内置的默认 NA 值
 bool _keep_default_na = true;
 // 是否禁用空值过滤；禁用可以提高性能
 bool _na_filter = true;
 // 是否将日期解析为 DD/MM 而不是 MM/DD
 bool _dayfirst = false;
 // 将时间戳列强制转换为特定类型
 data_type _timestamp_type{type_id::EMPTY};

 explicit csv_reader_options(source_info src) : _source{std::move(src)} {}

 friend csv_reader_options_builder;

 public
 csv_reader_options() = default;

 static csv_reader_options_builder builder(source_info src);

 [[nodiscard]] source_info const& get_source() const { return _source; }

 [[nodiscard]] compression_type get_compression() const { return _compression; }

 [[nodiscard]] std::size_t get_byte_range_offset() const { return _byte_range_offset; }

 [[nodiscard]] std::size_t get_byte_range_size() const { return _byte_range_size; }

 [[nodiscard]] std::size_t get_byte_range_size_with_padding() const
 {
 if (_byte_range_size == 0) {
 return 0;
 } else {
 return _byte_range_size + get_byte_range_padding();
 }
 }

 [[nodiscard]] std::size_t get_byte_range_padding() const
 {
 auto const num_names = _names.size();
 auto const num_dtypes = std::visit([](auto const& dtypes) { return dtypes.size(); }, _dtypes);
 auto const num_columns = std::max(num_dtypes, num_names);

 auto const max_row_bytes = 16 * 1024; // 16KB
 auto const column_bytes = 64;
 auto const base_padding = 1024; // 1KB

 if (num_columns == 0) {
 // 如果列数未知，使用固定大小
 return max_row_bytes;
 }

 // 如果可用，根据列数扩展大小
 return base_padding + num_columns * column_bytes;
 }
 }
 [[nodiscard]] std::vector<std::string> const& get_names() const { return _names; }

 [[nodiscard]] std::string get_prefix() const { return _prefix; }

 [[nodiscard]] bool is_enabled_mangle_dupe_cols() const { return _mangle_dupe_cols; }

 [[nodiscard]] std::vector<std::string> const& get_use_cols_names() const
 {
 return _use_cols_names;
 }

 [[nodiscard]] std::vector<int> const& get_use_cols_indexes() const { return _use_cols_indexes; }

 [[nodiscard]] size_type get_nrows() const { return _nrows; }

 [[nodiscard]] size_type get_skiprows() const { return _skiprows; }

 [[nodiscard]] size_type get_skipfooter() const { return _skipfooter; }

 [[nodiscard]] size_type get_header() const { return _header; }

 [[nodiscard]] char get_lineterminator() const { return _lineterminator; }

 [[nodiscard]] char get_delimiter() const { return _delimiter; }

 [[nodiscard]] char get_thousands() const { return _thousands; }

 [[nodiscard]] char get_decimal() const { return _decimal; }

 [[nodiscard]] char get_comment() const { return _comment; }

 [[nodiscard]] bool is_enabled_windowslinetermination() const { return _windowslinetermination; }

 [[nodiscard]] bool is_enabled_delim_whitespace() const { return _delim_whitespace; }

 [[nodiscard]] bool is_enabled_skipinitialspace() const { return _skipinitialspace; }

 [[nodiscard]] bool is_enabled_skip_blank_lines() const { return _skip_blank_lines; }

 [[nodiscard]] quote_style get_quoting() const { return _quoting; }

 [[nodiscard]] char get_quotechar() const { return _quotechar; }

 [[nodiscard]] bool is_enabled_doublequote() const { return _doublequote; }

 [[nodiscard]] bool is_enabled_detect_whitespace_around_quotes() const
 {
 return _detect_whitespace_around_quotes;
 }

 [[nodiscard]] std::vector<std::string> const& get_parse_dates_names() const
 {
 return _parse_dates_names;
 }

 [[nodiscard]] std::vector<int> const& get_parse_dates_indexes() const
 {
 return _parse_dates_indexes;
 }

 [[nodiscard]] std::vector<std::string> const& get_parse_hex_names() const
 {
 return _parse_hex_names;
 }

 [[nodiscard]] std::vector<int> const& get_parse_hex_indexes() const { return _parse_hex_indexes; }

 [[nodiscard]] std::variant<std::vector<data_type>, std::map<std::string, data_type>> const&
 get_dtypes() const
 {
 return _dtypes;
 }

 [[nodiscard]] std::vector<std::string> const& get_true_values() const { return _true_values; }

 [[nodiscard]] std::vector<std::string> const& get_false_values() const { return _false_values; }

 [[nodiscard]] std::vector<std::string> const& get_na_values() const { return _na_values; }

 [[nodiscard]] bool is_enabled_keep_default_na() const { return _keep_default_na; }

 [[nodiscard]] bool is_enabled_na_filter() const { return _na_filter; }

 [[nodiscard]] bool is_enabled_dayfirst() const { return _dayfirst; }

 [[nodiscard]] data_type get_timestamp_type() const { return _timestamp_type; }
};
 void set_compression(compression_type comp) { _compression = comp; }

 void set_byte_range_offset(std::size_t offset)
 {
 if ((offset != 0) and ((_skiprows != 0) or (_skipfooter != 0) or (_nrows != -1))) {
 CUDF_FAIL(
 "当 skiprows、skipfooter 或 nrows 中有有效值时，offset 不能为非零值"
 "value");
 }
 _byte_range_offset = offset;
 }

 void set_byte_range_size(std::size_t size)
 {
 if ((size != 0) and ((_skiprows != 0) or (_skipfooter != 0) or (_nrows != -1))) {
 CUDF_FAIL(
 "如果 skiprows、skipfooter 或 nrows 中的任何值有效，则 range size 不能为非零。"
 "non-zero.");
 }
 _byte_range_size = size;
 }

 void set_names(std::vector<std::string> col_names) { _names = std::move(col_names); }

 void set_prefix(std::string pfx) { _prefix = pfx; }

 void enable_mangle_dupe_cols(bool val) { _mangle_dupe_cols = val; }

 void set_use_cols_names(std::vector<std::string> col_names)
 {
 _use_cols_names = std::move(col_names);
 }

 void set_use_cols_indexes(std::vector<int> col_indices)
 {
 _use_cols_indexes = std::move(col_indices);
 }

 void set_nrows(size_type nrows)
 {
 CUDF_EXPECTS((nrows == 0) or (_skipfooter == 0), "不能同时使用 `nrows` 和 `skipfooter`");
 if ((nrows != -1) and ((_byte_range_offset != 0) or (_byte_range_size != 0))) {
 CUDF_FAIL(
 "如果已设置 range offset 和/或 range size，nrows 不能为非负值");
 }

 _nrows = nrows;
 }

 void set_skiprows(size_type skiprows)
 {
 if ((skiprows != 0) and ((_byte_range_offset != 0) or (_byte_range_size != 0))) {
 CUDF_FAIL("如果已设置 range offset 或 range size，skiprows 必须为零",
 std::invalid_argument);
 }
 _skiprows = skiprows;
 }

 void set_skipfooter(size_type skipfooter)
 {
 CUDF_EXPECTS((skipfooter == 0) or (_nrows == -1),
 "不能同时使用 `nrows` 和 `skipfooter`",
 std::invalid_argument);
 if ((skipfooter != 0) and ((_byte_range_offset != 0) or (_byte_range_size != 0))) {
 CUDF_FAIL("如果已设置 range offset 或 range size，skipfooter 必须为零",
 std::invalid_argument);
 }

 _skipfooter = skipfooter;
 }

 void set_header(size_type hdr) { _header = hdr; }

 void set_lineterminator(char term) { _lineterminator = term; }

 void set_delimiter(char delim) { _delimiter = delim; }

 void set_thousands(char val) { _thousands = val; }

 void set_decimal(char val) { _decimal = val; }

 void set_comment(char val) { _comment = val; }

 void enable_windowslinetermination(bool val) { _windowslinetermination = val; }

 void enable_delim_whitespace(bool val) { _delim_whitespace = val; }

 void enable_skipinitialspace(bool val) { _skipinitialspace = val; }

 void enable_skip_blank_lines(bool val) { _skip_blank_lines = val; }

 void set_quoting(quote_style quoting)
 {
 CUDF_EXPECTS(quoting == quote_style::MINIMAL || quoting == quote_style::NONE,
 "引号只支持 MINIMAL 和 NONE。");
 _quoting = quoting;
 }

 void set_quotechar(char ch) { _quotechar = ch; }

 void enable_doublequote(bool val) { _doublequote = val; }

 void enable_detect_whitespace_around_quotes(bool val) { _detect_whitespace_around_quotes = val; }

 void set_parse_dates(std::vector<std::string> col_names)
 {
 _parse_dates_names = std::move(col_names);
 }

 void set_parse_dates(std::vector<int> col_indices)
 {
 _parse_dates_indexes = std::move(col_indices);
 }

 void set_parse_hex(std::vector<std::string> col_names)
 {
 _parse_hex_names = std::move(col_names);
 }

 void set_parse_hex(std::vector<int> col_indices) { _parse_hex_indexes = std::move(col_indices); }

 void set_dtypes(std::map<std::string, data_type> types) { _dtypes = std::move(types); }

 void set_dtypes(std::vector<data_type> types) { _dtypes = std::move(types); }

 void set_true_values(std::vector<std::string> vals)
 {
 _true_values.insert(_true_values.end(), vals.begin(), vals.end());
 }

 void set_false_values(std::vector<std::string> vals)
 {
 _false_values.insert(_false_values.end(), vals.begin(), vals.end());
 }

 void set_na_values(std::vector<std::string> vals)
 {
 if ((!vals.empty()) and (!_na_filter)) {
 CUDF_FAIL("禁用 na_filtering 时不能设置 na_values");
 }

 _na_values = std::move(vals);
 }

 void enable_keep_default_na(bool val) { _keep_default_na = val; }

 void enable_na_filter(bool val)
 {
 if (!val) { _na_values.clear(); }
 _na_filter = val;
 }

 void enable_dayfirst(bool val) { _dayfirst = val; }

 void set_timestamp_type(data_type type) { _timestamp_type = type; }
};

class csv_reader_options_builder {
 csv_reader_options options;

 public
 csv_reader_options_builder() = default;

 csv_reader_options_builder(source_info src) : options{std::move(src)} {}

 csv_reader_options_builder& compression(compression_type comp)
 {
 options._compression = comp;
 return *this;
 }

 csv_reader_options_builder& byte_range_offset(std::size_t offset)
 {
 options.set_byte_range_offset(offset);
 return *this;
 }

 csv_reader_options_builder& byte_range_size(std::size_t size)
 {
 options.set_byte_range_size(size);
 return *this;
 }

 csv_reader_options_builder& names(std::vector<std::string> col_names)
 {
 options._names = std::move(col_names);
 return *this;
 }

 csv_reader_options_builder& prefix(std::string pfx)
 {
 options._prefix = std::move(pfx);
 return *this;
 }

 csv_reader_options_builder& mangle_dupe_cols(bool val)
 {
 options._mangle_dupe_cols = val;
 return *this;
 }

 csv_reader_options_builder& use_cols_names(std::vector<std::string> col_names)
 {
 options._use_cols_names = std::move(col_names);
 return *this;
 }

 csv_reader_options_builder& use_cols_indexes(std::vector<int> col_indices)
 {
 options._use_cols_indexes = std::move(col_indices);
 return *this;
 }

 csv_reader_options_builder& nrows(size_type rows)
 {
 options.set_nrows(rows);
 return *this;
 }

 csv_reader_options_builder& skiprows(size_type skip)
 {
 options.set_skiprows(skip);
 return *this;
 }

 csv_reader_options_builder& skipfooter(size_type skip)
 {
 options.set_skipfooter(skip);
 return *this;
 }

 csv_reader_options_builder& header(size_type hdr)
 {
 options._header = hdr;
 return *this;
 }

 csv_reader_options_builder& lineterminator(char term)
 {
 options._lineterminator = term;
 return *this;
 }

 csv_reader_options_builder& delimiter(char delim)
 {
 options._delimiter = delim;
 return *this;
 }

 csv_reader_options_builder& thousands(char val)
 {
 options._thousands = val;
 return *this;
 }

 csv_reader_options_builder& decimal(char val)
 {
 options._decimal = val;
 return *this;
 }

 csv_reader_options_builder& comment(char val)
 {
 options._comment = val;
 return *this;
 }

 csv_reader_options_builder& windowslinetermination(bool val)
 {
 options._windowslinetermination = val;
 return *this;
 }

 csv_reader_options_builder& delim_whitespace(bool val)
 {
 options._delim_whitespace = val;
 return *this;
 }

 csv_reader_options_builder& skipinitialspace(bool val)
 {
 options._skipinitialspace = val;
 return *this;
 }

 csv_reader_options_builder& skip_blank_lines(bool val)
 {
 options._skip_blank_lines = val;
 return *this;
 }

 csv_reader_options_builder& quoting(quote_style style)
 {
 options._quoting = style;
 return *this;
 }

 csv_reader_options_builder& quotechar(char ch)
 {
 options._quotechar = ch;
 return *this;
 }

 csv_reader_options_builder& doublequote(bool val)
 {
 options._doublequote = val;
 return *this;
 }

 csv_reader_options_builder& detect_whitespace_around_quotes(bool val)
 {
 options._detect_whitespace_around_quotes = val;
 return *this;
 }

 csv_reader_options_builder& parse_dates(std::vector<std::string> col_names)
 {
 options._parse_dates_names = std::move(col_names);
 return *this;
 }

 csv_reader_options_builder& parse_dates(std::vector<int> col_indices)
 {
 options._parse_dates_indexes = std::move(col_indices);
 return *this;
 }

 csv_reader_options_builder& parse_hex(std::vector<std::string> col_names)
 {
 options._parse_hex_names = std::move(col_names);
 return *this;
 }

 csv_reader_options_builder& parse_hex(std::vector<int> col_indices)
 {
 options._parse_hex_indexes = std::move(col_indices);
 return *this;
 }

 csv_reader_options_builder& dtypes(std::map<std::string, data_type> types)
 {
 options._dtypes = std::move(types);
 return *this;
 }

 csv_reader_options_builder& dtypes(std::vector<data_type> types)
 {
 options._dtypes = std::move(types);
 return *this;
 }

 csv_reader_options_builder& true_values(std::vector<std::string> vals)
 {
 options._true_values.insert(options._true_values.end(), vals.begin(), vals.end());
 return *this;
 }

 csv_reader_options_builder& false_values(std::vector<std::string> vals)
 {
 options._false_values.insert(options._false_values.end(), vals.begin(), vals.end());
 return *this;
 }

 csv_reader_options_builder& na_values(std::vector<std::string> vals)
 {
 options.set_na_values(std::move(vals));
 return *this;
 }

 csv_reader_options_builder& keep_default_na(bool val)
 {
 options.enable_keep_default_na(val);
 return *this;
 }

 csv_reader_options_builder& na_filter(bool val)
 {
 options.enable_na_filter(val);
 return *this;
 }

 csv_reader_options_builder& dayfirst(bool val)
 {
 options._dayfirst = val;
 return *this;
 }

 csv_reader_options_builder& timestamp_type(data_type type)
 {
 options._timestamp_type = type;
 return *this;
 }

 operator csv_reader_options&&() { return std::move(options); }

 csv_reader_options&& build() { return std::move(options); }
};

table_with_metadata read_csv(
 csv_reader_options options,
 rmm::cuda_stream_view stream = cudf::get_default_stream(),
 rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 // 组结束
class csv_writer_options_builder;

class csv_writer_options {
 // 指定写入器输出使用的接收器
 sink_info _sink;
 // 要输出的列集合
 table_view _table;
 // 用于表示空值的字符串
 std::string _na_rep = "";
 // 指示是否将头部写入CSV
 bool _include_header = true;
 // 每个块写入的最大行数（限制内存使用）
 size_type _rows_per_chunk = std::numeric_limits<size_type>::max();
 // 用于分隔行的字符（默认为 "\n"）
 std::string _line_terminator = "\n";
 // 用于分隔列值的字符（默认为 ","）
 char _inter_column_delimiter = ',';
 // INT8 类型中非零值使用的字符串（默认为 'true'）
 std::string _true_value = std::string{"true"};
 // INT8 类型中值为零使用的字符串（默认为 'false'）
 std::string _false_value = std::string{"false"};
 // 所有列的名称；如果为空，写入器将生成列名
 std::vector<std::string> _names;
 // 引用风格。目前仅支持 MINIMAL 和 NONE。
 quote_style _quoting = quote_style::MINIMAL;

 explicit csv_writer_options(sink_info sink, table_view const& table)
 : _sink(std::move(sink)), _table(table), _rows_per_chunk(table.num_rows())
 {
 }

 friend csv_writer_options_builder;

 public
 explicit csv_writer_options() = default;

 static csv_writer_options_builder builder(sink_info const& sink, table_view const& table);

 [[nodiscard]] sink_info const& get_sink() const { return _sink; }

 [[nodiscard]] table_view const& get_table() const { return _table; }

 [[nodiscard]] std::vector<std::string> const& get_names() const { return _names; }

 [[nodiscard]] std::string const& get_na_rep() const { return _na_rep; }

 [[nodiscard]] bool is_enabled_include_header() const { return _include_header; }

 [[nodiscard]] size_type get_rows_per_chunk() const { return _rows_per_chunk; }

 [[nodiscard]] std::string const& get_line_terminator() const { return _line_terminator; }

 [[nodiscard]] char get_inter_column_delimiter() const { return _inter_column_delimiter; }

 [[nodiscard]] std::string const& get_true_value() const { return _true_value; }

 [[nodiscard]] std::string const& get_false_value() const { return _false_value; }

 [[nodiscard]] quote_style get_quoting() const { return _quoting; }

 // 设置器
 void set_names(std::vector<std::string> names) { _names = std::move(names); }

 void set_na_rep(std::string val) { _na_rep = std::move(val); }

 void enable_include_header(bool val) { _include_header = val; }

 void set_rows_per_chunk(size_type val) { _rows_per_chunk = val; }

 void set_line_terminator(std::string term) { _line_terminator = std::move(term); }

 void set_inter_column_delimiter(char delim) { _inter_column_delimiter = delim; }

 void set_true_value(std::string val) { _true_value = std::move(val); }

 void set_false_value(std::string val) { _false_value = std::move(val); }

 void set_table(table_view const& table) { _table = table; }

 void set_quoting(quote_style quoting)
 {
 CUDF_EXPECTS(quoting == quote_style::MINIMAL || quoting == quote_style::NONE,
 "Only MINIMAL and NONE are supported for quoting.");
 _quoting = quoting;
 }
};

class csv_writer_options_builder {
 csv_writer_options options;

 public
 explicit csv_writer_options_builder() = default;

 explicit csv_writer_options_builder(sink_info const& sink, table_view const& table)
 : options{sink, table}
 {
 }

 csv_writer_options_builder& names(std::vector<std::string> names)
 {
 options._names = names;
 return *this;
 }

 csv_writer_options_builder& na_rep(std::string val)
 {
 options._na_rep = val;
 return *this;
 };

 csv_writer_options_builder& include_header(bool val)
 {
 options._include_header = val;
 return *this;
 }

 csv_writer_options_builder& rows_per_chunk(int val)
 {
 options._rows_per_chunk = val;
 return *this;
 }

 csv_writer_options_builder& line_terminator(std::string term)
 {
 options._line_terminator = term;
 return *this;
 }

 csv_writer_options_builder& inter_column_delimiter(char delim)
 {
 options._inter_column_delimiter = delim;
 return *this;
 }

 csv_writer_options_builder& true_value(std::string val)
 {
 options._true_value = val;
 return *this;
 }

 csv_writer_options_builder& false_value(std::string val)
 {
 options._false_value = val;
 return *this;
 }

 csv_writer_options_builder& quoting(quote_style quoting)
 {
 options.set_quoting(quoting);
 return *this;
 }

 operator csv_writer_options&&() { return std::move(options); }

 csv_writer_options&& build() { return std::move(options); }
};

void write_csv(csv_writer_options const& options,
 rmm::cuda_stream_view stream = cudf::get_default_stream());
 // 组结束
} // 命名空间 io
} // 命名空间 CUDF_EXPORT cudf