Support casting of Map type to string in JSON reader

cpp/CMakeLists.txt

@@ -382,6 +382,7 @@ add_library(
   src/io/json/read_json.cu
   src/io/json/legacy/json_gpu.cu
   src/io/json/legacy/reader_impl.cu
+  src/io/json/parser_features.cpp
   src/io/json/write_json.cu
   src/io/orc/aggregate_orc_metadata.cpp
   src/io/orc/dict_enc.cu

cpp/include/cudf/io/json.hpp

@@ -323,6 +323,7 @@ class json_reader_options {
 
   /**
    * @brief Set whether to parse mixed types as a string column.
+   * Also enables forcing to read a struct as string column using schema.
    *
    * @param val Boolean value to enable/disable parsing mixed types as a string column
    */
@@ -473,6 +474,7 @@ class json_reader_options_builder {
 
   /**
    * @brief Set whether to parse mixed types as a string column.
+   * Also enables forcing to read a struct as string column using schema.
    *
    * @param val Boolean value to enable/disable parsing mixed types as a string column
    * @return this for chaining

cpp/src/io/json/json_column.cu

@@ -392,6 +392,41 @@ std::vector<std::string> copy_strings_to_host(device_span<SymbolT const> input,
   return to_host(d_column_names->view());
 }
 
+rmm::device_uvector<uint8_t> is_all_nulls_str_column(device_span<SymbolT const> input,
+                                                     tree_meta_t const& d_column_tree,
+                                                     tree_meta_t const& tree,
+                                                     device_span<NodeIndexT> col_ids,
+                                                     cudf::io::json_reader_options const& options,
+                                                     rmm::cuda_stream_view stream)
+{
+  auto const num_nodes = col_ids.size();
+  auto const num_cols  = d_column_tree.node_categories.size();
+  rmm::device_uvector<uint8_t> is_all_nulls(num_cols, stream);
+  thrust::fill(rmm::exec_policy(stream), is_all_nulls.begin(), is_all_nulls.end(), true);
+
+  auto parse_opt = parsing_options(options, stream);
+  thrust::for_each_n(
+    rmm::exec_policy(stream),
+    thrust::counting_iterator<size_type>(0),
+    num_nodes,
+    [options           = parse_opt.view(),
+     data              = input.data(),
+     column_categories = d_column_tree.node_categories.begin(),
+     col_ids           = col_ids.begin(),
+     range_begin       = tree.node_range_begin.begin(),
+     range_end         = tree.node_range_end.begin(),
+     is_all_nulls      = is_all_nulls.begin()] __device__(size_type i) {
+      auto const node_category = column_categories[col_ids[i]];
+      if (node_category == NC_STR or node_category == NC_VAL) {
+        auto const is_null_literal = serialized_trie_contains(
+          options.trie_na,
+          {data + range_begin[i], static_cast<size_t>(range_end[i] - range_begin[i])});
+        if (!is_null_literal) is_all_nulls[col_ids[i]] = false;
+      }
+    });
+  return is_all_nulls;
+}
+
 /**
  * @brief Holds member data pointers of `d_json_column`
  *
@@ -429,11 +464,12 @@ void make_device_json_column(device_span<SymbolT const> input,
                              bool is_array_of_arrays,
                              bool is_enabled_lines,
                              bool is_enabled_mixed_types_as_string,
+                             cudf::io::json_reader_options const& options,
                              rmm::cuda_stream_view stream,
                              rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  auto num_nodes = col_ids.size();
+  auto const num_nodes = col_ids.size();
   rmm::device_uvector<NodeIndexT> sorted_col_ids(col_ids.size(), stream);  // make a copy
   thrust::copy(rmm::exec_policy(stream), col_ids.begin(), col_ids.end(), sorted_col_ids.begin());
 
@@ -444,15 +480,16 @@ void make_device_json_column(device_span<SymbolT const> input,
     rmm::exec_policy(stream), sorted_col_ids.begin(), sorted_col_ids.end(), node_ids.begin());
 
   NodeIndexT const row_array_parent_col_id = [&]() {
-    if (!is_array_of_arrays) return parent_node_sentinel;
-    auto const list_node_index = is_enabled_lines ? 0 : 1;
-    NodeIndexT value;
-    CUDF_CUDA_TRY(cudaMemcpyAsync(&value,
-                                  col_ids.data() + list_node_index,
-                                  sizeof(NodeIndexT),
-                                  cudaMemcpyDefault,
-                                  stream.value()));
-    stream.synchronize();
+    NodeIndexT value = parent_node_sentinel;
+    if (!col_ids.empty()) {
+      auto const list_node_index = is_enabled_lines ? 0 : 1;
+      CUDF_CUDA_TRY(cudaMemcpyAsync(&value,
+                                    col_ids.data() + list_node_index,
+                                    sizeof(NodeIndexT),
+                                    cudaMemcpyDefault,
+                                    stream.value()));
+      stream.synchronize();
+    }
     return value;
   }();
 
@@ -540,6 +577,12 @@ void make_device_json_column(device_span<SymbolT const> input,
     col.column_order.clear();
   };
 
+  path_from_tree tree_path{column_categories,
+                           column_parent_ids,
+                           column_names,
+                           is_array_of_arrays,
+                           row_array_parent_col_id};
+
   // 2. generate nested columns tree and its device_memory
   // reorder unique_col_ids w.r.t. column_range_begin for order of column to be in field order.
   auto h_range_col_id_it =
@@ -548,6 +591,12 @@ void make_device_json_column(device_span<SymbolT const> input,
     return thrust::get<0>(a) < thrust::get<0>(b);
   });
 
+  std::vector<uint8_t> is_str_column_all_nulls{};
+  if (is_enabled_mixed_types_as_string) {
+    is_str_column_all_nulls = cudf::detail::make_std_vector_async(
+      is_all_nulls_str_column(input, d_column_tree, tree, col_ids, options, stream), stream);
+  }
+
   // use hash map because we may skip field name's col_ids
   std::unordered_map<NodeIndexT, std::reference_wrapper<device_json_column>> columns;
   // map{parent_col_id, child_col_name}> = child_col_id, used for null value column tracking
@@ -584,6 +633,7 @@ void make_device_json_column(device_span<SymbolT const> input,
       ignore_vals[this_col_id]          = 1;
       continue;
     }
+
     // If the child is already found,
     // replace if this column is a nested column and the existing was a value column
     // ignore this column if this column is a value column and the existing was a nested column
@@ -592,29 +642,36 @@ void make_device_json_column(device_span<SymbolT const> input,
     auto& parent_col = it->second.get();
     bool replaced    = false;
     if (mapped_columns.count({parent_col_id, name}) > 0) {
+      auto const old_col_id = mapped_columns[{parent_col_id, name}];
       // If mixed type as string is enabled, make both of them strings and merge them.
       // All child columns will be ignored when parsing.
       if (is_enabled_mixed_types_as_string) {
-        // VAL/STR or STRUCT or LIST
-        auto old_col_id = mapped_columns[{parent_col_id, name}];
-
-        is_mixed_type_column[this_col_id] = 1;
-        is_mixed_type_column[old_col_id]  = 1;
-        // if old col type (not cat) is list or struct, replace with string.
-        auto& col = columns.at(old_col_id).get();
-        if (col.type == json_col_t::ListColumn or col.type == json_col_t::StructColumn) {
-          reinitialize_as_string(old_col_id, col);
-          // all its children (which are already inserted) are ignored later.
+        bool is_mixed_type = true;
+        // If new or old is STR and they are all not null, make it mixed type, else ignore.
+        if (column_categories[this_col_id] == NC_VAL || column_categories[this_col_id] == NC_STR) {
+          if (is_str_column_all_nulls[this_col_id]) is_mixed_type = false;
+        }
+        if (column_categories[old_col_id] == NC_VAL || column_categories[old_col_id] == NC_STR) {
+          if (is_str_column_all_nulls[old_col_id]) is_mixed_type = false;
+        }
+        if (is_mixed_type) {
+          is_mixed_type_column[this_col_id] = 1;
+          is_mixed_type_column[old_col_id]  = 1;
+          // if old col type (not cat) is list or struct, replace with string.
+          auto& col = columns.at(old_col_id).get();
+          if (col.type == json_col_t::ListColumn or col.type == json_col_t::StructColumn) {
+            reinitialize_as_string(old_col_id, col);
+            // all its children (which are already inserted) are ignored later.
+          }
+          columns.try_emplace(this_col_id, columns.at(old_col_id));
+          continue;
         }
-        columns.try_emplace(this_col_id, columns.at(old_col_id));
-        continue;
       }
 
       if (column_categories[this_col_id] == NC_VAL || column_categories[this_col_id] == NC_STR) {
         ignore_vals[this_col_id] = 1;
         continue;
       }
-      auto old_col_id = mapped_columns[{parent_col_id, name}];
       if (column_categories[old_col_id] == NC_VAL || column_categories[old_col_id] == NC_STR) {
         // remap
         ignore_vals[old_col_id] = 1;
@@ -632,6 +689,17 @@ void make_device_json_column(device_span<SymbolT const> input,
                      "A mix of lists and structs within the same column is not supported");
       }
     }
+    if (is_enabled_mixed_types_as_string) {
+      // get path of this column, check if it is a struct forced as string, and enforce it
+      auto nt                          = tree_path.get_path(this_col_id);
+      std::optional<data_type> user_dt = get_path_data_type(nt, options);
+      if (column_categories[this_col_id] == NC_STRUCT and user_dt.has_value() and
+          user_dt.value().id() == type_id::STRING) {
+        is_mixed_type_column[this_col_id] = 1;
+        column_categories[this_col_id]    = NC_STR;
+      }
+    }
+
     CUDF_EXPECTS(parent_col.child_columns.count(name) == 0, "duplicate column name: " + name);
     // move into parent
     device_json_column col(stream, mr);
@@ -795,15 +863,6 @@ void make_device_json_column(device_span<SymbolT const> input,
   }
 }
 
-/**
- * @brief Retrieves the parse_options to be used for type inference and type casting
- *
- * @param options The reader options to influence the relevant type inference and type casting
- * options
- */
-cudf::io::parse_options parsing_options(cudf::io::json_reader_options const& options,
-                                        rmm::cuda_stream_view stream);
-
 std::pair<std::unique_ptr<column>, std::vector<column_name_info>> device_json_column_to_cudf_column(
   device_json_column& json_col,
   device_span<SymbolT const> d_input,
@@ -1023,6 +1082,7 @@ table_with_metadata device_parse_nested_json(device_span<SymbolT const> d_input,
                           is_array_of_arrays,
                           options.is_enabled_lines(),
                           options.is_enabled_mixed_types_as_string(),
+                          options,
                           stream,
                           mr);
 

cpp/src/io/json/nested_json.hpp

@@ -25,6 +25,10 @@
 #include <map>
 #include <vector>
 
+// Forward declaration of parse_options from parsing_utils.cuh
+namespace cudf::io {
+struct parse_options;
+}
 namespace cudf::io::json {
 
 /**
@@ -284,6 +288,15 @@ reduce_to_column_tree(tree_meta_t& tree,
                       device_span<size_type> row_offsets,
                       rmm::cuda_stream_view stream);
 
+/**
+ * @brief Retrieves the parse_options to be used for type inference and type casting
+ *
+ * @param options The reader options to influence the relevant type inference and type casting
+ * options
+ */
+cudf::io::parse_options parsing_options(cudf::io::json_reader_options const& options,
+                                        rmm::cuda_stream_view stream);
+
 /** @copydoc host_parse_nested_json
  * All processing is done in device memory.
  *
@@ -293,6 +306,32 @@ table_with_metadata device_parse_nested_json(device_span<SymbolT const> input,
                                              rmm::cuda_stream_view stream,
                                              rmm::mr::device_memory_resource* mr);
 
+/**
+ * @brief Get the path data type of a column by path if present in input schema
+ *
+ * @param path path of the column
+ * @param options json reader options which holds schema
+ * @return data type of the column if present
+ */
+std::optional<data_type> get_path_data_type(
+  host_span<std::pair<std::string, cudf::io::json::NodeT>> path,
+  cudf::io::json_reader_options const& options);
+
+/**
+ * @brief Helper class to get path of a column by column id from reduced column tree
+ *
+ */
+struct path_from_tree {
+  host_span<NodeT const> column_categories;
+  host_span<NodeIndexT const> column_parent_ids;
+  host_span<std::string const> column_names;
+  bool is_array_of_arrays;
+  NodeIndexT const row_array_parent_col_id;
+
+  using path_rep = std::pair<std::string, cudf::io::json::NodeT>;
+  std::vector<path_rep> get_path(NodeIndexT this_col_id);
+};
+
 /**
  * @brief Parses the given JSON string and generates table from the given input.
  *

cpp/src/io/json/nested_json_gpu.cu

@@ -2042,7 +2042,8 @@ void make_json_column(json_column& root_column,
  * options
  * @param stream The CUDA stream to which kernels are dispatched
  */
-auto parsing_options(cudf::io::json_reader_options const& options, rmm::cuda_stream_view stream)
+cudf::io::parse_options parsing_options(cudf::io::json_reader_options const& options,
+                                        rmm::cuda_stream_view stream)
 {
   auto parse_opts = cudf::io::parse_options{',', '\n', '\"', '.'};