io/json.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 "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 {
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;
};
 
enum class json_recovery_mode_t {
  FAIL,              
  RECOVER_WITH_NULL  
};
 
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;
 
  // Data types of the column; empty to infer dtypes
  dtype_variant _dtypes;
  // Specify the compression format of the source or infer from file extension
  compression_type _compression = compression_type::AUTO;
 
  // Read the file as a json object per line
  bool _lines = false;
  // Parse mixed types as a string column
  bool _mixed_types_as_string = false;
  // Delimiter separating records in JSON lines
  char _delimiter = '\n';
  // Prune columns on read, selected based on the _dtypes option
  bool _prune_columns = false;
  // Experimental features: new column tree construction
  bool _experimental = false;
 
  // Bytes to skip from the start
  size_t _byte_range_offset = 0;
  // Bytes to read; always reads complete rows
  size_t _byte_range_size = 0;
 
  // Whether to parse dates as DD/MM versus MM/DD
  bool _dayfirst = false;
 
  // Whether to keep the quote characters of string values
  bool _keep_quotes = false;
 
  // Normalize single quotes
  bool _normalize_single_quotes = false;
 
  // Normalize unquoted spaces and tabs
  bool _normalize_whitespace = false;
 
  // Whether to recover after an invalid JSON line
  json_recovery_mode_t _recovery_mode = json_recovery_mode_t::FAIL;
 
  // Validation checks for spark
  // Should the json validation be strict or not
  // Note: strict validation enforces the JSON specification https://www.json.org/json-en.html
  bool _strict_validation = false;
  // Allow leading zeros for numeric values.
  bool _allow_numeric_leading_zeros = true;
  // Allow non-numeric numbers: NaN, +INF, -INF, +Infinity, Infinity, -Infinity
  bool _allow_nonnumeric_numbers = true;
  // Allow unquoted control characters
  bool _allow_unquoted_control_chars = true;
  // Additional values to recognize as null values
  std::vector<std::string> _na_values;
 
  explicit json_reader_options(source_info src) : _source{std::move(src)} {}
 
  friend json_reader_options_builder;
 
 public:
  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;
    }
 
    // Expand the size based on the number of columns, if available
    return base_padding + num_columns * column_bytes;
  }
 
  [[nodiscard]] char get_delimiter() const { return _delimiter; }
 
  [[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; }
 
  [[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; }
 
  [[nodiscard]] json_recovery_mode_t recovery_mode() const { return _recovery_mode; }
 
  [[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;
  }
 
  [[nodiscard]] std::vector<std::string> const& get_na_values() const { return _na_values; }
 
  void set_dtypes(std::vector<data_type> types) { _dtypes = std::move(types); }
 
  void set_dtypes(std::map<std::string, data_type> types) { _dtypes = std::move(types); }
 
  void set_dtypes(std::map<std::string, schema_element> types) { _dtypes = std::move(types); }
 
  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; }
 
  void set_delimiter(char delimiter)
  {
    switch (delimiter) {
      case '{':
      case '[':
      case '}':
      case ']':
      case ',':
      case ':':
      case '"':
      case '\'':
      case '\\':
      case ' ':
      case '\t':
      case '\r': CUDF_FAIL("Unsupported delimiter character.", std::invalid_argument); break;
    }
    _delimiter = delimiter;
  }
 
  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; }
 
  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; }
 
  void set_recovery_mode(json_recovery_mode_t val) { _recovery_mode = val; }
 
  void set_strict_validation(bool val) { _strict_validation = val; }
 
  void allow_numeric_leading_zeros(bool val)
  {
    CUDF_EXPECTS(_strict_validation, "Strict validation must be enabled for this to work.");
    _allow_numeric_leading_zeros = val;
  }
 
  void allow_nonnumeric_numbers(bool val)
  {
    CUDF_EXPECTS(_strict_validation, "Strict validation must be enabled for this to work.");
    _allow_nonnumeric_numbers = val;
  }
 
  void allow_unquoted_control_chars(bool val)
  {
    CUDF_EXPECTS(_strict_validation, "Strict validation must be enabled for this to work.");
    _allow_unquoted_control_chars = val;
  }
 
  void set_na_values(std::vector<std::string> vals) { _na_values = std::move(vals); }
};
 
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_reader_options_builder& delimiter(char delimiter)
  {
    options.set_delimiter(delimiter);
    return *this;
  }
 
  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;
  }
 
  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_reader_options_builder& recovery_mode(json_recovery_mode_t val)
  {
    options._recovery_mode = val;
    return *this;
  }
 
  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;
  }
 
  json_reader_options_builder& na_values(std::vector<std::string> vals)
  {
    options.set_na_values(std::move(vals));
    return *this;
  }
 
  operator json_reader_options&&() { return std::move(options); }
 
  json_reader_options&& build() { return std::move(options); }
};
 
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());
  // end of group
 
class json_writer_options_builder;
 
class json_writer_options {
  // Specify the sink to use for writer output
  sink_info _sink;
  // maximum number of rows to write in each chunk (limits memory use)
  size_type _rows_per_chunk = std::numeric_limits<size_type>::max();
  // Set of columns to output
  table_view _table;
  // string to use for null entries
  std::string _na_rep = "";
  // Indicates whether to output nulls as 'null' or exclude the field
  bool _include_nulls = false;
  // Indicates whether to use JSON lines for records format
  bool _lines = false;
  // string to use for values != 0 in INT8 types (default 'true')
  std::string _true_value = std::string{"true"};
  // string to use for values == 0 in INT8 types (default 'false')
  std::string _false_value = std::string{"false"};
  // Names of all columns; if empty, writer will generate column names
  std::optional<table_metadata> _metadata;  // Optional column names
 
  explicit json_writer_options(sink_info sink, table_view table)
    : _sink(std::move(sink)), _rows_per_chunk(table.num_rows()), _table(std::move(table))
  {
  }
 
  friend json_writer_options_builder;
 
 public:
  explicit json_writer_options() = default;
 
  static json_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::optional<table_metadata> const& get_metadata() const { return _metadata; }
 
  [[nodiscard]] std::string const& get_na_rep() const { return _na_rep; }
 
  [[nodiscard]] bool is_enabled_include_nulls() const { return _include_nulls; }
 
  [[nodiscard]] bool is_enabled_lines() const { return _lines; }
 
  [[nodiscard]] size_type get_rows_per_chunk() const { return _rows_per_chunk; }
 
  [[nodiscard]] std::string const& get_true_value() const { return _true_value; }
 
  [[nodiscard]] std::string const& get_false_value() const { return _false_value; }
 
  // Setter
 
  void set_table(table_view tbl) { _table = tbl; }
 
  void set_metadata(table_metadata metadata) { _metadata = std::move(metadata); }
 
  void set_na_rep(std::string val) { _na_rep = std::move(val); }
 
  void enable_include_nulls(bool val) { _include_nulls = val; }
 
  void enable_lines(bool val) { _lines = val; }
 
  void set_rows_per_chunk(size_type val) { _rows_per_chunk = val; }
 
  void set_true_value(std::string val) { _true_value = std::move(val); }
 
  void set_false_value(std::string val) { _false_value = std::move(val); }
};
 
class json_writer_options_builder {
  json_writer_options options;  
 
 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& 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