csv.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 <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;
 
  // Read settings
 
  // Specify the compression format of the source or infer from file extension
  compression_type _compression = compression_type::AUTO;
  // Bytes to skip from the source start
  std::size_t _byte_range_offset = 0;
  // Bytes to read; always reads complete rows
  std::size_t _byte_range_size = 0;
  // Names of all the columns; if empty then names are auto-generated
  std::vector<std::string> _names;
  // If there is no header or names, prepend this to the column ID as the name
  std::string _prefix;
  // Whether to rename duplicate column names
  bool _mangle_dupe_cols = true;
 
  // Filter settings
 
  // Names of columns to read; empty is all columns
  std::vector<std::string> _use_cols_names;
  // Indexes of columns to read; empty is all columns
  std::vector<int> _use_cols_indexes;
  // Rows to read; -1 is all
  size_type _nrows = -1;
  // Rows to skip from the start
  size_type _skiprows = 0;
  // Rows to skip from the end
  size_type _skipfooter = 0;
  // Header row index
  size_type _header = 0;
 
  // Parsing settings
 
  // Line terminator
  char _lineterminator = '\n';
  // Field delimiter
  char _delimiter = ',';
  // Numeric data thousands separator; cannot match delimiter
  char _thousands = '\0';
  // Decimal point character; cannot match delimiter
  char _decimal = '.';
  // Comment line start character
  char _comment                = '\0';
  bool _windowslinetermination = false;
  // Treat whitespace as field delimiter; overrides character delimiter
  bool _delim_whitespace = false;
  // Skip whitespace after the delimiter
  bool _skipinitialspace = false;
  // Ignore empty lines or parse line values as invalid
  bool _skip_blank_lines = true;
  // Treatment of quoting behavior
  quote_style _quoting = quote_style::MINIMAL;
  // Quoting character (if `quoting` is true)
  char _quotechar = '"';
  // Whether a quote inside a value is double-quoted
  bool _doublequote = true;
  // Whether to detect quotes surrounded by spaces e.g. `   "data"   `. This flag has no effect when
  // _doublequote is true
  bool _detect_whitespace_around_quotes = false;
  // Names of columns to read as datetime
  std::vector<std::string> _parse_dates_names;
  // Indexes of columns to read as datetime
  std::vector<int> _parse_dates_indexes;
  // Names of columns to parse as hexadecimal
  std::vector<std::string> _parse_hex_names;
  // Indexes of columns to parse as hexadecimal
  std::vector<int> _parse_hex_indexes;
 
  // Conversion settings
 
  // Per-column types; disables type inference on those columns
  std::variant<std::vector<data_type>, std::map<std::string, data_type>> _dtypes;
  // Additional values to recognize as boolean true values
  std::vector<std::string> _true_values{"True", "TRUE", "true"};
  // Additional values to recognize as boolean false values
  std::vector<std::string> _false_values{"False", "FALSE", "false"};
  // Additional values to recognize as null values
  std::vector<std::string> _na_values;
  // Whether to keep the built-in default NA values
  bool _keep_default_na = true;
  // Whether to disable null filter; disabling can improve performance
  bool _na_filter = true;
  // Whether to parse dates as DD/MM versus MM/DD
  bool _dayfirst = false;
  // Cast timestamp columns to a specific type
  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) {
      // 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]] 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(
        "When there is valid value in skiprows or skipfooter or nrows, offset can't have non-zero "
        "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(
        "If the value of any of skiprows, skipfooter or nrows is valid, range size cannot be "
        "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), "Cannot use both `nrows` and `skipfooter`");
    if ((nrows != -1) and ((_byte_range_offset != 0) or (_byte_range_size != 0))) {
      CUDF_FAIL(
        "nrows can't be a non negative value if range offset and/or range size has been set");
    }
 
    _nrows = nrows;
  }
 
  void set_skiprows(size_type skiprows)
  {
    if ((skiprows != 0) and ((_byte_range_offset != 0) or (_byte_range_size != 0))) {
      CUDF_FAIL("skiprows must be zero if range offset or range size has been set",
                std::invalid_argument);
    }
    _skiprows = skiprows;
  }
 
  void set_skipfooter(size_type skipfooter)
  {
    CUDF_EXPECTS((skipfooter == 0) or (_nrows == -1),
                 "Cannot use both `nrows` and `skipfooter`",
                 std::invalid_argument);
    if ((skipfooter != 0) and ((_byte_range_offset != 0) or (_byte_range_size != 0))) {
      CUDF_FAIL("skipfooter must be zero if range offset or range size has been set",
                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,
                 "Only MINIMAL and NONE are supported for quoting.");
    _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("Can't set na_values when na_filtering is disabled");
    }
 
    _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 = 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());
  // end of group
class csv_writer_options_builder;
 
class csv_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 write headers to csv
  bool _include_header = true;
  // maximum number of rows to write in each chunk (limits memory use)
  size_type _rows_per_chunk = std::numeric_limits<size_type>::max();
  // character to use for separating lines (default "\n")
  std::string _line_terminator = "\n";
  // character to use for separating column values (default ",")
  char _inter_column_delimiter = ',';
  // 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::vector<std::string> _names;
  // Quote style. Currently only MINIMAL and NONE are supported.
  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 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 get_line_terminator() const { return _line_terminator; }
 
  [[nodiscard]] char get_inter_column_delimiter() const { return _inter_column_delimiter; }
 
  [[nodiscard]] std::string get_true_value() const { return _true_value; }
 
  [[nodiscard]] std::string get_false_value() const { return _false_value; }
 
  [[nodiscard]] quote_style get_quoting() const { return _quoting; }
 
  // Setter
  void set_names(std::vector<std::string> names) { _names = std::move(names); }
 
  void set_na_rep(std::string val) { _na_rep = 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 = term; }
 
  void set_inter_column_delimiter(char delim) { _inter_column_delimiter = delim; }
 
  void set_true_value(std::string val) { _true_value = val; }
 
  void set_false_value(std::string val) { _false_value = 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());
  // end of group
}  // namespace io
}  // namespace CUDF_EXPORT cudf