Expand CSV and JSON reader APIs to accept dtypes as a vector or map of data_type objects

cpp/include/cudf/io/csv.hpp

@@ -104,14 +104,19 @@ class csv_reader_options {
   // Whether a quote inside a value is double-quoted
   bool _doublequote = true;
   // Names of columns to read as datetime
-  std::vector<std::string> _infer_date_names;
+  std::vector<std::string> _parse_dates_names;
   // Indexes of columns to read as datetime
-  std::vector<int> _infer_date_indexes;
+  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<std::string>, std::vector<data_type>> _dtypes;
+  std::variant<std::vector<std::string>, 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
@@ -280,17 +285,30 @@ class csv_reader_options {
   /**
    * @brief Returns names of columns to read as datetime.
    */
-  std::vector<std::string> const& get_infer_date_names() const { return _infer_date_names; }
+  std::vector<std::string> const& get_parse_dates_names() const { return _parse_dates_names; }
 
   /**
    * @brief Returns indexes of columns to read as datetime.
    */
-  std::vector<int> const& get_infer_date_indexes() const { return _infer_date_indexes; }
+  std::vector<int> const& get_parse_dates_indexes() const { return _parse_dates_indexes; }
+
+  /**
+   * @brief Returns names of columns to read as datetime.
+   */
+  std::vector<std::string> const& get_parse_hex_names() const { return _parse_hex_names; }
+
+  /**
+   * @brief Returns indexes of columns to read as datetime.
+   */
+  std::vector<int> const& get_parse_hex_indexes() const { return _parse_hex_indexes; }
 
   /**
    * @brief Returns per-column types.
    */
-  std::variant<std::vector<std::string>, std::vector<data_type>> const& get_dtypes() const
+  std::variant<std::vector<std::string>,
+               std::vector<data_type>,
+               std::map<std::string, data_type>> const&
+  get_dtypes() const
   {
     return _dtypes;
   }
@@ -547,21 +565,42 @@ class csv_reader_options {
    *
    * @param col_names Vector of column names to infer as datetime.
    */
-  void set_infer_date_names(std::vector<std::string> col_names)
+  void set_parse_dates(std::vector<std::string> col_names)
   {
-    _infer_date_names = std::move(col_names);
+    _parse_dates_names = std::move(col_names);
   }
 
   /**
    * @brief Sets indexes of columns to read as datetime.
    *
    * @param col_names Vector of column indices to infer as datetime.
    */
-  void set_infer_date_indexes(std::vector<int> col_ind)
+  void set_parse_dates(std::vector<int> col_ind) { _parse_dates_indexes = std::move(col_ind); }
+
+  /**
+   * @brief Sets names of columns to parse as hexadecimal
+   *
+   * @param col_names Vector of column names to parse as hexadecimal
+   */
+  void set_parse_hex(std::vector<std::string> col_names)
   {
-    _infer_date_indexes = std::move(col_ind);
+    _parse_hex_names = std::move(col_names);
   }
 
+  /**
+   * @brief Sets indexes of columns to parse as hexadecimal
+   *
+   * @param col_names Vector of column indices to parse as hexadecimal
+   */
+  void set_parse_hex(std::vector<int> col_ind) { _parse_hex_indexes = std::move(col_ind); }
+
+  /**
+   * @brief Sets per-column types
+   *
+   * @param types Column name -> data type map specifying the columns' target data types
+   */
+  void set_dtypes(std::map<std::string, data_type> types) { _dtypes = std::move(types); }
+
   /**
    * @brief Sets per-column types
    *
@@ -958,24 +997,60 @@ class csv_reader_options_builder {
   /**
    * @brief Sets names of columns to read as datetime.
    *
-   * @param col_names Vector of column names to infer as datetime.
+   * @param col_names Vector of column names to read as datetime.
    * @return this for chaining.
    */
-  csv_reader_options_builder& infer_date_names(std::vector<std::string> col_names)
+  csv_reader_options_builder& parse_dates(std::vector<std::string> col_names)
   {
-    options._infer_date_names = std::move(col_names);
+    options._parse_dates_names = std::move(col_names);
     return *this;
   }
 
   /**
    * @brief Sets indexes of columns to read as datetime.
    *
-   * @param col_names Vector of column indices to infer as datetime.
+   * @param col_names Vector of column indices to read as datetime.
    * @return this for chaining.
    */
-  csv_reader_options_builder& infer_date_indexes(std::vector<int> col_ind)
+  csv_reader_options_builder& parse_dates(std::vector<int> col_ind)
   {
-    options._infer_date_indexes = std::move(col_ind);
+    options._parse_dates_indexes = std::move(col_ind);
+    return *this;
+  }
+
+  /**
+   * @brief Sets names of columns to parse as hexadecimal.
+   *
+   * @param col_names Vector of column names to parse as hexadecimal
+   * @return this for chaining.
+   */
+  csv_reader_options_builder& parse_hex(std::vector<std::string> col_names)
+  {
+    options._parse_hex_names = std::move(col_names);
+    return *this;
+  }
+
+  /**
+   * @brief Sets indexes of columns to parse as hexadecimal.
+   *
+   * @param col_names Vector of column indices to parse as hexadecimal
+   * @return this for chaining.
+   */
+  csv_reader_options_builder& parse_hex(std::vector<int> col_ind)
+  {
+    options._parse_hex_indexes = std::move(col_ind);
+    return *this;
+  }
+
+  /**
+   * @brief Sets per-column types.
+   *
+   * @param types Column name -> data type map specifying the columns' target data types
+   * @return this for chaining.
+   */
+  csv_reader_options_builder& dtypes(std::map<std::string, data_type> types)
+  {
+    options._dtypes = std::move(types);
     return *this;
   }
 

cpp/include/cudf/io/json.hpp

@@ -23,7 +23,9 @@
 
 #include <rmm/mr/device/per_device_resource.hpp>
 
+#include <map>
 #include <string>
+#include <variant>
 #include <vector>
 
 namespace cudf {
@@ -66,7 +68,8 @@ class json_reader_options {
   source_info _source;
 
   // Data types of the column; empty to infer dtypes
-  std::vector<std::string> _dtypes;
+  std::variant<std::vector<std::string>, std::vector<data_type>, std::map<std::string, data_type>>
+    _dtypes;
   // Specify the compression format of the source or infer from file extension
   compression_type _compression = compression_type::AUTO;
 
@@ -114,7 +117,13 @@ class json_reader_options {
   /**
    * @brief Returns data types of the columns.
    */
-  std::vector<std::string> const& get_dtypes() const { return _dtypes; }
+  std::variant<std::vector<std::string>,
+               std::vector<data_type>,
+               std::map<std::string, data_type>> const&
+  get_dtypes() const
+  {
+    return _dtypes;
+  }
 
   /**
    * @brief Returns compression format of the source.
@@ -146,14 +155,28 @@ class json_reader_options {
    *
    * @param types Vector dtypes in string format.
    */
-  void dtypes(std::vector<std::string> types) { _dtypes = std::move(types); }
+  void set_dtypes(std::vector<std::string> types) { _dtypes = std::move(types); }
+
+  /**
+   * @brief Set data types for columns to be read.
+   *
+   * @param types Vector dtypes in string format.
+   */
+
+  void set_dtypes(std::vector<data_type> types) { _dtypes = std::move(types); }
+  /**
+   * @brief Set data types for columns to be read.
+   *
+   * @param types Vector dtypes in string format.
+   */
+  void set_dtypes(std::map<std::string, data_type> types) { _dtypes = std::move(types); }
 
   /**
    * @brief Set the compression type.
    *
    * @param comp_type The compression type used.
    */
-  void compression(compression_type comp_type) { _compression = comp_type; }
+  void set_compression(compression_type comp_type) { _compression = comp_type; }
 
   /**
    * @brief Set number of bytes to skip from source start.
@@ -205,15 +228,39 @@ class json_reader_options_builder {
   /**
    * @brief Set data types for columns to be read.
    *
-   * @param types Vector dtypes in string format.
-   * @return this for chaining.
+   * @param types Vector of dtypes in string format
+   * @return this for chaining
    */
   json_reader_options_builder& dtypes(std::vector<std::string> types)
   {
     options._dtypes = std::move(types);
     return *this;
   }
 
+  /**
+   * @brief Set data types for columns to be read.
+   *
+   * @param types Vector of dtypes
+   * @return this for chaining
+   */
+  json_reader_options_builder& dtypes(std::vector<data_type> types)
+  {
+    options._dtypes = std::move(types);
+    return *this;
+  }
+
+  /**
+   * @brief Set data types for columns to be read.
+   *
+   * @param types Column name -> dtype map.
+   * @return this for chaining
+   */
+  json_reader_options_builder& dtypes(std::map<std::string, data_type> types)
+  {
+    options._dtypes = std::move(types);
+    return *this;
+  }
+
   /**
    * @brief Set the compression type.
    *

cpp/src/io/csv/reader_impl.cu

@@ -280,6 +280,41 @@ reader::impl::select_data_and_row_offsets(rmm::cuda_stream_view stream)
   return {rmm::device_uvector<char>{0, stream}, selected_rows_offsets{stream}};
 }
 
+std::vector<data_type> reader::impl::select_data_types(
+  std::map<std::string, data_type> const& col_type_map)
+{
+  std::vector<data_type> selected_dtypes;
+
+  for (int col = 0; col < num_actual_cols_; col++) {
+    if (column_flags_[col] & column_parse::enabled) {
+      auto const col_type_it = col_type_map.find(col_names_[col]);
+      CUDF_EXPECTS(col_type_it != col_type_map.end(),
+                   "Must specify data types for all active columns");
+      selected_dtypes.emplace_back(col_type_it->second);
+    }
+  }
+  return selected_dtypes;
+}
+
+std::vector<data_type> reader::impl::select_data_types(std::vector<data_type> const& dtypes)
+{
+  std::vector<data_type> selected_dtypes;
+
+  if (dtypes.size() == 1) {
+    // If it's a single dtype, assign that dtype to all active columns
+    selected_dtypes.resize(num_active_cols_, dtypes.front());
+  } else {
+    // If it's a list, assign dtypes to active columns in the given order
+    CUDF_EXPECTS(static_cast<int>(dtypes.size()) >= num_actual_cols_,
+                 "Must specify data types for all columns");
+
+    for (int col = 0; col < num_actual_cols_; col++) {
+      if (column_flags_[col] & column_parse::enabled) { selected_dtypes.emplace_back(dtypes[col]); }
+    }
+  }
+  return selected_dtypes;
+}
+
 table_with_metadata reader::impl::read(rmm::cuda_stream_view stream)
 {
   auto const data_row_offsets = select_data_and_row_offsets(stream);
@@ -355,20 +390,34 @@ table_with_metadata reader::impl::read(rmm::cuda_stream_view stream)
     }
   }
 
-  // User can specify which columns should be inferred as datetime
-  if (!opts_.get_infer_date_indexes().empty() || !opts_.get_infer_date_names().empty()) {
-    for (const auto index : opts_.get_infer_date_indexes()) {
+  // User can specify which columns should be read as datetime
+  if (!opts_.get_parse_dates_indexes().empty() || !opts_.get_parse_dates_names().empty()) {
+    for (const auto index : opts_.get_parse_dates_indexes()) {
       column_flags_[index] |= column_parse::as_datetime;
     }
 
-    for (const auto& name : opts_.get_infer_date_names()) {
+    for (const auto& name : opts_.get_parse_dates_names()) {
       auto it = std::find(col_names_.begin(), col_names_.end(), name);
       if (it != col_names_.end()) {
         column_flags_[it - col_names_.begin()] |= column_parse::as_datetime;
       }
     }
   }
 
+  // User can specify which columns should be inferred as datetime
+  if (!opts_.get_parse_hex_indexes().empty() || !opts_.get_parse_hex_names().empty()) {
+    for (const auto index : opts_.get_parse_hex_indexes()) {
+      column_flags_[index] |= column_parse::as_hexadecimal;
+    }
+
+    for (const auto& name : opts_.get_parse_hex_names()) {
+      auto it = std::find(col_names_.begin(), col_names_.end(), name);
+      if (it != col_names_.end()) {
+        column_flags_[it - col_names_.begin()] |= column_parse::as_hexadecimal;
+      }
+    }
+  }
+
   // Return empty table rather than exception if nothing to load
   if (num_active_cols_ == 0) { return {std::make_unique<table>(), {}}; }
 
@@ -382,11 +431,14 @@ table_with_metadata reader::impl::read(rmm::cuda_stream_view stream)
   if (has_to_infer_column_types) {
     column_types = infer_column_types(data, row_offsets, stream);
   } else {
-    column_types =
-      std::visit(VisitorOverload{
-                   [&](const std::vector<data_type>& data_types) { return data_types; },
-                   [&](const std::vector<string>& dtypes) { return parse_column_types(dtypes); }},
-                 opts_.get_dtypes());
+    column_types = std::visit(
+      VisitorOverload{
+        [&](const std::vector<data_type>& data_types) { return select_data_types(data_types); },
+        [&](const std::map<std::string, data_type>& data_types) {
+          return select_data_types(data_types);
+        },
+        [&](const std::vector<string>& dtypes) { return parse_column_types(dtypes); }},
+      opts_.get_dtypes());
   }
 
   out_columns.reserve(column_types.size());