io/json.hpp

前往此文件的文档。

/*
 * 版权所有 (c) 2020-2024, NVIDIA CORPORATION。
 *
 * 根据 Apache 许可证 2.0 版（“许可证”）的规定获得许可；
 * 除非遵守许可证，否则不得使用此文件。
 * 您可以在以下位置获取许可证副本
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，软件
 * 根据许可证分发，不提供任何形式的明示或暗示的保证或条件。
 * 请参阅许可证中关于特定语言管理权限和
 * 许可证下的限制。
 */

#pragma once

#include "types.hpp"

#include <cudf/detail/utilities/visitor_overload.hpp>
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/error.hpp>
#include <cudf/utilities/memory_resource.hpp>

#include <map>
#include <string>
#include <utility>
#include <variant>
#include <vector>

namespace CUDF_EXPORT cudf {
namespace io {
// JSON 读取器选项构建器
class json_reader_options_builder;

struct schema_element {
 data_type type;

 std::map<std::string, schema_element> child_types;

 std::optional<std::vector<std::string>> column_order;
};

// JSON 读取器恢复模式类型
enum class json_recovery_mode_t {
 FAIL, // 在无效 JSON 行失败
 RECOVER_WITH_NULL // 在无效 JSON 行时以 null 恢复
};

// JSON 读取器选项
class json_reader_options {
 public
 // 用于指定数据类型的可能变体
 using dtype_variant =
 std::variant<std::vector<data_type>,
 std::map<std::string, data_type>,
 std::map<std::string, schema_element>,
 schema_element>;

 private
 source_info _source; // 源信息

 // 列的数据类型；为空则推断数据类型
 dtype_variant _dtypes;
 // 指定源文件的压缩格式或从文件扩展名推断
 compression_type _compression = compression_type::AUTO;

 // 将文件按行读取为 JSON 对象
 bool _lines = false;
 // 将混合类型解析为字符串列
 bool _mixed_types_as_string = false;
 // JSON 行中分隔记录的分隔符
 char _delimiter = '\n';
 // 读取时修剪列，根据 _dtypes 选项选择
 bool _prune_columns = false;
 // 实验性功能：新的列树构建
 bool _experimental = false;

 // 从头跳过的字节数
 size_t _byte_range_offset = 0;
 // 读取的字节数；总是读取完整的行
 size_t _byte_range_size = 0;

 // 是否将日期解析为 DD/MM 而非 MM/DD
 bool _dayfirst = false;

 // 是否保留字符串值的引号字符
 bool _keep_quotes = false;

 // 规范化单引号
 bool _normalize_single_quotes = false;

 // 规范化未加引号的空格和制表符
 bool _normalize_whitespace = false;

 // 在无效 JSON 行后是否恢复
 json_recovery_mode_t _recovery_mode = json_recovery_mode_t::FAIL;

 // 用于 spark 的验证检查
 // JSON 验证是否应严格
 // 注意：严格验证强制执行 JSON 规范 https://www.json.org/json-en.html
 bool _strict_validation = false;
 // 允许数值带有前导零。
 bool _allow_numeric_leading_zeros = true;
 // 允许非数字：NaN, +INF, -INF, +Infinity, Infinity, -Infinity
 bool _allow_nonnumeric_numbers = true;
 // 允许未加引号的控制字符
 bool _allow_unquoted_control_chars = true;
 // 额外识别为 null 值的数值
 std::vector<std::string> _na_values;

 // 供构建器使用的私有构造函数
 explicit json_reader_options(source_info src) : _source{std::move(src)} {}

 friend json_reader_options_builder;

 public
 // 构造具有默认值的新选项
 explicit json_reader_options() = default;

 // 创建新的构建器实例
 static json_reader_options_builder builder(source_info src);

 // 返回源信息
 [[nodiscard]] source_info const& get_source() const { return _source; }

 // 返回数据类型
 [[nodiscard]] dtype_variant const& get_dtypes() const { return _dtypes; }

 // 返回压缩类型
 [[nodiscard]] compression_type get_compression() const { return _compression; }

 // 返回要读取的字节范围的偏移量
 [[nodiscard]] size_t get_byte_range_offset() const { return _byte_range_offset; }

 // 返回要读取的字节范围的大小
 [[nodiscard]] size_t get_byte_range_size() const { return _byte_range_size; }

 // 返回要读取的字节范围的大小，包括填充
 [[nodiscard]] 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]] size_t get_byte_range_padding() const
 {
 auto const num_columns =
 std::visit(cudf::detail::visitor_overload{
 [](auto const& dtypes) { return dtypes.size(); },
 [](schema_element const& dtypes) { return dtypes.child_types.size(); }},
 _dtypes);

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

 if (num_columns == 0) {
 // Use flat size if the number of columns is not known
 return max_row_bytes;
 } else {
 // Expand the size based on the number of columns, if available
 return base_padding + num_columns * column_bytes;
 }
 }

 // 返回用于在 JSON 行中分隔记录的分隔符
 [[nodiscard]] char get_delimiter() const { return _delimiter; }

 // 返回是否启用 JSON 行读取
 [[nodiscard]] bool is_enabled_lines() const { return _lines; }

 // 返回是否将混合类型解析为字符串
 [[nodiscard]] bool is_enabled_mixed_types_as_string() const { return _mixed_types_as_string; }

 // 返回是否根据数据类型修剪列
 [[nodiscard]] bool is_enabled_prune_columns() const { return _prune_columns; }

 // 返回是否启用实验性功能
 [[nodiscard]] bool is_enabled_experimental() const { return _experimental; }

 // 返回是否将日期解析为 DD/MM 而非 MM/DD
 [[nodiscard]] bool is_enabled_dayfirst() const { return _dayfirst; }

 // 返回是否保留字符串值的引号字符
 [[nodiscard]] bool is_enabled_keep_quotes() const { return _keep_quotes; }

 // 返回是否规范化单引号
 [[nodiscard]] bool is_enabled_normalize_single_quotes() const { return _normalize_single_quotes; }

 // 返回是否规范化未加引号的空格和制表符
 [[nodiscard]] bool is_enabled_normalize_whitespace() const { return _normalize_whitespace; }

 // 返回无效 JSON 行的恢复模式
 [[nodiscard]] json_recovery_mode_t recovery_mode() const { return _recovery_mode; }

 // 返回是否启用严格 JSON 验证
 [[nodiscard]] bool is_strict_validation() const { return _strict_validation; }

 // 返回是否允许带有前导零的数值。
 [[nodiscard]] bool is_allowed_numeric_leading_zeros() const
 {
 return _allow_numeric_leading_zeros;
 }

 // 返回是否允许非数字。
 [[nodiscard]] bool is_allowed_nonnumeric_numbers() const { return _allow_nonnumeric_numbers; }

 // 返回是否允许未加引号的控制字符。
 [[nodiscard]] bool is_allowed_unquoted_control_chars() const
 {
 return _allow_unquoted_control_chars;
 }

 // 返回额外识别为 null 值的数值
 [[nodiscard]] std::vector<std::string> const& get_na_values() const { return _na_values; }

 // 通过 std::vector 设置数据类型
 void set_dtypes(std::vector<data_type> types) { _dtypes = std::move(types); }

 // 通过 std::map 设置数据类型
 void set_dtypes(std::map<std::string, data_type> types) { _dtypes = std::move(types); }

 // 通过 std::map 设置数据类型
 void set_dtypes(std::map<std::string, schema_element> types) { _dtypes = std::move(types); }

 // 通过 schema_element 设置数据类型。
 void set_dtypes(schema_element types);

 // 设置压缩类型
 void set_compression(compression_type comp_type) { _compression = comp_type; }

 // 设置要读取的字节范围的偏移量
 void set_byte_range_offset(size_t offset) { _byte_range_offset = offset; }

 // 设置要读取的字节范围的大小
 void set_byte_range_size(size_t size) { _byte_range_size = size; }

 // 设置用于在 JSON 行中分隔记录的分隔符。
 void set_delimiter(char delimiter)
 {
 switch (delimiter) {
 case '{'
 case '['
 case '}'
 case ']'
 case ','
 case ':'
 case '"'
 case '\''
 case '\\'
 case ' '
 case '\t'
 case '\r': CUDF_FAIL("不支持的分隔符字符。", std::invalid_argument); break;
 }
 _delimiter = delimiter;
 }

 // 启用/禁用 JSON 行读取
 void enable_lines(bool val) { _lines = val; }

 // 启用/禁用将混合类型解析为字符串
 void enable_mixed_types_as_string(bool val) { _mixed_types_as_string = val; }

 // 启用/禁用根据数据类型修剪列。
 void enable_prune_columns(bool val) { _prune_columns = val; }

 // 启用/禁用实验性功能。
 void enable_experimental(bool val) { _experimental = val; }

 // 启用/禁用将日期解析为 DD/MM 而非 MM/DD
 void enable_dayfirst(bool val) { _dayfirst = val; }

 // 启用/禁用保留字符串值的引号字符
 void enable_keep_quotes(bool val) { _keep_quotes = val; }

 // 启用/禁用规范化单引号
 void enable_normalize_single_quotes(bool val) { _normalize_single_quotes = val; }

 // 启用/禁用规范化未加引号的空格和制表符
 void enable_normalize_whitespace(bool val) { _normalize_whitespace = val; }

 // 设置无效 JSON 行的恢复模式
 void set_recovery_mode(json_recovery_mode_t val) { _recovery_mode = val; }

 // 设置是否启用严格 JSON 验证
 void set_strict_validation(bool val) { _strict_validation = val; }

 // 设置是否允许带有前导零的数值。
 void allow_numeric_leading_zeros(bool val)
 {
 CUDF_EXPECTS(_strict_validation, "必须启用严格验证才能使其工作。");
 _allow_numeric_leading_zeros = val;
 }

 // 设置是否允许非数字。
 void allow_nonnumeric_numbers(bool val)
 {
 CUDF_EXPECTS(_strict_validation, "必须启用严格验证才能使其工作。");
 _allow_nonnumeric_numbers = val;
 }

 // 设置是否允许未加引号的控制字符。
 void allow_unquoted_control_chars(bool val)
 {
 CUDF_EXPECTS(_strict_validation, "必须启用严格验证才能使其工作。");
 _allow_unquoted_control_chars = val;
 }

 // 设置额外识别为 null 值的数值
 void set_na_values(std::vector<std::string> vals) { _na_values = std::move(vals); }
};

// JSON 读取器选项构建器
class json_reader_options_builder {
 json_reader_options options;

 public
 // 构造新的构建器实例
 explicit json_reader_options_builder() = default;

 // 从源构造新的构建器实例
 explicit json_reader_options_builder(source_info src) : options{std::move(src)} {}

 // 设置数据类型
 json_reader_options_builder& dtypes(std::vector<data_type> types)
 {
 options._dtypes = std::move(types);
 return *this;
 }

 // 设置数据类型
 json_reader_options_builder& dtypes(std::map<std::string, data_type> types)
 {
 options._dtypes = std::move(types);
 return *this;
 }

 // 设置数据类型
 json_reader_options_builder& dtypes(std::map<std::string, schema_element> types)
 {
 options._dtypes = std::move(types);
 return *this;
 }

 // 设置数据类型
 json_reader_options_builder& dtypes(schema_element types)
 {
 options.set_dtypes(std::move(types));
 return *this;
 }

 // 设置压缩类型
 json_reader_options_builder& compression(compression_type comp_type)
 {
 options._compression = comp_type;
 return *this;
 }

 // 设置要读取的字节范围的偏移量
 json_reader_options_builder& byte_range_offset(size_type offset)
 {
 options._byte_range_offset = offset;
 return *this;
 }

 // 设置要读取的字节范围的大小
 json_reader_options_builder& byte_range_size(size_type size)
 {
 options._byte_range_size = size;
 return *this;
 }

 // 设置用于在 JSON 行中分隔记录的分隔符。
 json_reader_options_builder& delimiter(char delimiter)
 {
 options.set_delimiter(delimiter);
 return *this;
 }

 // 启用/禁用 JSON 行读取
 json_reader_options_builder& lines(bool val)
 {
 options._lines = val;
 return *this;
 }

 // 启用/禁用将混合类型解析为字符串
 json_reader_options_builder& mixed_types_as_string(bool val)
 {
 options._mixed_types_as_string = val;
 return *this;
 }

 // 启用/禁用根据数据类型修剪列。
 json_reader_options_builder& prune_columns(bool val)
 {
 options._prune_columns = val;
 return *this;
 }

 // 启用/禁用实验性功能。
 json_reader_options_builder& experimental(bool val)
 {
 options._experimental = val;
 return *this;
 }

 // 启用/禁用将日期解析为 DD/MM 而非 MM/DD
 json_reader_options_builder& dayfirst(bool val)
 {
 options._dayfirst = val;
 return *this;
 }

 // 启用/禁用保留字符串值的引号字符
 json_reader_options_builder& keep_quotes(bool val)
 {
 options._keep_quotes = val;
 return *this;
 }

 // 启用/禁用规范化单引号
 json_reader_options_builder& normalize_single_quotes(bool val)
 {
 options._normalize_single_quotes = val;
 return *this;
 }

 // 启用/禁用规范化未加引号的空格和制表符
 json_reader_options_builder& normalize_whitespace(bool val)
 {
 options._normalize_whitespace = val;
 return *this;
 }

 // 设置无效 JSON 行的恢复模式
 json_reader_options_builder& recovery_mode(json_recovery_mode_t val)
 {
 options._recovery_mode = val;
 return *this;
 }

 // 设置是否启用严格 JSON 验证
 json_reader_options_builder& strict_validation(bool val)
 {
 options.set_strict_validation(val);
 return *this;
 }

 // 设置是否允许带有前导零的数值。
 json_reader_options_builder& numeric_leading_zeros(bool val)
 {
 options.allow_numeric_leading_zeros(val);
 return *this;
 }

 // 设置是否允许非数字。
 json_reader_options_builder& nonnumeric_numbers(bool val)
 {
 options.allow_nonnumeric_numbers(val);
 return *this;
 }

 // 设置是否允许未加引号的控制字符。
 json_reader_options_builder& unquoted_control_chars(bool val)
 {
 options.allow_unquoted_control_chars(val);
 return *this;
 }

 // 设置额外识别为 null 值的数值
 json_reader_options_builder& na_values(std::vector<std::string> vals)
 {
 options.set_na_values(std::move(vals));
 return *this;
 }

 // 构建 json_reader_options
 operator json_reader_options&&() { return std::move(options); }

 // 构建 json_reader_options
 json_reader_options&& build() { return std::move(options); }
};

// 将 JSON 数据读取到表中。
table_with_metadata read_json(
 json_reader_options options,
 rmm::cuda_stream_view stream = cudf::get_default_stream(),
 rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 // 组结束

// JSON 写入器选项构建器

 public
 explicit json_writer_options_builder() = default;

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

 json_writer_options_builder& table(table_view tbl)
 {
 options._table = tbl;
 return *this;
 }

 json_writer_options_builder& compression(compression_type comptype)
 {
 options._compression = comptype;
 return *this;
 }

 json_writer_options_builder& metadata(table_metadata metadata)
 {
 options._metadata = std::move(metadata);
 return *this;
 }

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

 json_writer_options_builder& include_nulls(bool val)
 {
 options._include_nulls = val;
 return *this;
 }

 json_writer_options_builder& lines(bool val)
 {
 options._lines = val;
 return *this;
 }

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

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

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

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

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

void write_json(json_writer_options const& options,
 rmm::cuda_stream_view stream = cudf::get_default_stream());
 // end of group
} // namespace io
} // namespace CUDF_EXPORT cudf