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/table/table_view.hpp>
#include <cudf/types.hpp>
 
#include <rmm/mr/device/per_device_resource.hpp>
 
#include <map>
#include <string>
#include <variant>
#include <vector>
 
namespace cudf {
namespace io {
class json_reader_options_builder;
 
struct schema_element {
  data_type type;
 
  std::map<std::string, schema_element> child_types;
};
 
enum class json_recovery_mode_t {
  FAIL,              
  RECOVER_WITH_NULL  
};
 
class json_reader_options {
  source_info _source;
 
  // Data types of the column; empty to infer dtypes
  std::variant<std::vector<data_type>,
               std::map<std::string, data_type>,
               std::map<std::string, schema_element>>
    _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;
 
  // 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 use the legacy reader
  bool _legacy = false;
 
  // Whether to keep the quote characters of string values
  bool _keep_quotes = false;
 
  // Normalize single quotes
  bool _normalize_single_quotes = false;
 
  // Whether to recover after an invalid JSON line
  json_recovery_mode_t _recovery_mode = json_recovery_mode_t::FAIL;
 
  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; }
 
  std::variant<std::vector<data_type>,
               std::map<std::string, data_type>,
               std::map<std::string, schema_element>> const&
  get_dtypes() const
  {
    return _dtypes;
  }
 
  compression_type get_compression() const { return _compression; }
 
  size_t get_byte_range_offset() const { return _byte_range_offset; }
 
  size_t get_byte_range_size() const { return _byte_range_size; }
 
  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();
    }
  }
 
  size_t get_byte_range_padding() const
  {
    auto const num_columns = std::visit([](auto const& dtypes) { return dtypes.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;
  }
 
  bool is_enabled_lines() const { return _lines; }
 
  bool is_enabled_mixed_types_as_string() const { return _mixed_types_as_string; }
 
  bool is_enabled_dayfirst() const { return _dayfirst; }
 
  bool is_enabled_legacy() const { return _legacy; }
 
  bool is_enabled_keep_quotes() const { return _keep_quotes; }
 
  bool is_enabled_normalize_single_quotes() const { return _normalize_single_quotes; }
 
  json_recovery_mode_t recovery_mode() const { return _recovery_mode; }
 
  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_compression(compression_type comp_type) { _compression = comp_type; }
 
  void set_byte_range_offset(size_type offset) { _byte_range_offset = offset; }
 
  void set_byte_range_size(size_type size) { _byte_range_size = size; }
 
  void enable_lines(bool val) { _lines = val; }
 
  void enable_mixed_types_as_string(bool val) { _mixed_types_as_string = val; }
 
  void enable_dayfirst(bool val) { _dayfirst = val; }
 
  void enable_legacy(bool val) { _legacy = val; }
 
  void enable_keep_quotes(bool val) { _keep_quotes = val; }
 
  void enable_normalize_single_quotes(bool val) { _normalize_single_quotes = val; }
 
  void set_recovery_mode(json_recovery_mode_t val) { _recovery_mode = val; }
};
 
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& 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& 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& dayfirst(bool val)
  {
    options._dayfirst = val;
    return *this;
  }
 
  json_reader_options_builder& legacy(bool val)
  {
    options._legacy = 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& recovery_mode(json_recovery_mode_t val)
  {
    options._recovery_mode = val;
    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::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
  // end of group
 
class json_writer_options_builder;
 
class json_writer_options {
  // Specify the sink to use for writer output
  sink_info _sink;
  // 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;
  // maximum number of rows to write in each chunk (limits memory use)
  size_type _rows_per_chunk = std::numeric_limits<size_type>::max();
  // 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 const& sink, table_view const& table)
    : _sink(sink), _table(table), _rows_per_chunk(table.num_rows())
  {
  }
 
  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(),
                rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
  // end of group
}  // namespace io
}  // namespace cudf