fixing max row offsets

cpp/CMakeLists.txt

@@ -393,7 +393,7 @@ add_library(
   src/io/functions.cpp
   src/io/json/byte_range_info.cu
   src/io/json/json_column.cu
-  src/io/json/json_column_csr.cu
+  src/io/json/column_tree_construction.cu
   src/io/json/json_normalization.cu
   src/io/json/json_tree.cu
   src/io/json/nested_json_gpu.cu

cpp/src/io/json/json_column_csr.cu → cpp/src/io/json/column_tree_construction.cu

@@ -52,7 +52,22 @@
 #include <thrust/transform_scan.h>
 #include <thrust/unique.h>
 
-namespace cudf::io::json::experimental::detail {
+namespace cudf::io::json {
+
+template <typename T>
+void print(device_span<T const> d_vec, std::string name, rmm::cuda_stream_view stream)
+{
+  stream.synchronize();
+  auto h_vec = cudf::detail::make_std_vector_async(d_vec, stream);
+  stream.synchronize();
+  std::cout << name << " = ";
+  for (auto e : h_vec) {
+    std::cout << e << " ";
+  }
+  std::cout << std::endl;
+}
+
+namespace experimental::detail {
 
 using row_offset_t = size_type;
 
@@ -77,17 +92,6 @@ struct parent_nodeids_to_colids {
   }
 };
 
-template <typename T>
-void print(device_span<T const> d_vec, std::string name, rmm::cuda_stream_view stream)
-{
-  auto h_vec = cudf::detail::make_std_vector_async(d_vec, stream);
-  std::cout << name << " = ";
-  for (auto e : h_vec) {
-    std::cout << e << " ";
-  }
-  std::cout << std::endl;
-}
-
 /**
  * @brief Reduces node tree representation to column tree CSR representation.
  *
@@ -168,10 +172,9 @@ std::tuple<csr, column_tree_properties> reduce_to_column_tree(
     rmm::exec_policy_nosync(stream),
     level_ordered_col_ids.begin(),
     level_ordered_col_ids.end(),
-    thrust::make_zip_iterator(thrust::make_tuple(ordered_row_offsets, ordered_node_categories)),
+    thrust::make_zip_iterator(ordered_row_offsets, ordered_node_categories),
     thrust::make_discard_iterator(),
-    thrust::make_zip_iterator(
-      thrust::make_tuple(max_row_offsets.begin(), column_categories.begin())),
+    thrust::make_zip_iterator(max_row_offsets.begin(), column_categories.begin()),
     thrust::equal_to<NodeIndexT>(),
     [] __device__(auto a, auto b) {
       auto row_offset_a = thrust::get<0>(a);
@@ -195,8 +198,8 @@ std::tuple<csr, column_tree_properties> reduce_to_column_tree(
       } else
         ctg = NC_ERR;
 
-      thrust::maximum<size_type> row_offset_op;
-      return thrust::make_tuple(row_offset_op(row_offset_a, row_offset_b), ctg);
+      thrust::maximum<row_offset_t> row_offset_op;
+      return thrust::make_pair(row_offset_op(row_offset_a, row_offset_b), ctg);
     });
 
   // 4. construct parent_col_ids using permutation iterator
@@ -277,8 +280,7 @@ std::tuple<csr, column_tree_properties> reduce_to_column_tree(
 
   // Mixed types in List children go to different columns,
   // so all immediate children of list column should have same max_row_offsets.
-  //   create list's children max_row_offsets array. (initialize to zero)
-  //   atomicMax on  children max_row_offsets array.
+  //   create list's children max_row_offsets array
   //   gather the max_row_offsets from children row offset array.
   {
     auto max_row_offsets_it =
@@ -315,15 +317,14 @@ std::tuple<csr, column_tree_properties> reduce_to_column_tree(
        list_ancestors = list_ancestors.begin()] __device__(NodeIndexT node) {
         auto num_levels      = *dev_num_levels_ptr;
         list_ancestors[node] = node;
-        for (int level = 0; level < num_levels; level++) {
+        for (int level = 0; level <= num_levels; level++) {
           if (list_ancestors[node] > 0)
             list_ancestors[node] = colidx[rowidx[list_ancestors[node]]];
-          else
-            list_ancestors[node] = -1;
-          if (list_ancestors[node] == -1 || column_categories[list_ancestors[node]] == NC_LIST)
+          if (list_ancestors[node] == 0 || column_categories[list_ancestors[node]] == NC_LIST)
             break;
         }
       });
+
     thrust::gather_if(rmm::exec_policy_nosync(stream),
                       list_ancestors.begin(),
                       list_ancestors.end(),
@@ -339,4 +340,211 @@ std::tuple<csr, column_tree_properties> reduce_to_column_tree(
       std::move(column_categories), std::move(max_row_offsets), std::move(mapped_col_ids)}};
 }
 
-}  // namespace cudf::io::json::experimental::detail
+}  // namespace experimental::detail
+
+namespace detail {
+/**
+ * @brief Reduces node tree representation to column tree representation.
+ *
+ * @param tree Node tree representation of JSON string
+ * @param original_col_ids Column ids of nodes
+ * @param sorted_col_ids Sorted column ids of nodes
+ * @param ordered_node_ids Node ids of nodes sorted by column ids
+ * @param row_offsets Row offsets of nodes
+ * @param is_array_of_arrays Whether the tree is an array of arrays
+ * @param row_array_parent_col_id Column id of row array, if is_array_of_arrays is true
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ * @return A tuple of column tree representation of JSON string, column ids of columns, and
+ * max row offsets of columns
+ */
+std::tuple<tree_meta_t, rmm::device_uvector<NodeIndexT>, rmm::device_uvector<size_type>>
+reduce_to_column_tree(tree_meta_t& tree,
+                      device_span<NodeIndexT> original_col_ids,
+                      device_span<NodeIndexT> sorted_col_ids,
+                      device_span<NodeIndexT> ordered_node_ids,
+                      device_span<size_type> row_offsets,
+                      bool is_array_of_arrays,
+                      NodeIndexT const row_array_parent_col_id,
+                      rmm::cuda_stream_view stream)
+{
+  CUDF_FUNC_RANGE();
+  // 1. column count for allocation
+  auto const num_columns =
+    thrust::unique_count(rmm::exec_policy(stream), sorted_col_ids.begin(), sorted_col_ids.end());
+
+  // 2. reduce_by_key {col_id}, {row_offset}, max.
+  rmm::device_uvector<NodeIndexT> unique_col_ids(num_columns, stream);
+  rmm::device_uvector<size_type> max_row_offsets(num_columns, stream);
+  auto ordered_row_offsets =
+    thrust::make_permutation_iterator(row_offsets.begin(), ordered_node_ids.begin());
+  thrust::reduce_by_key(rmm::exec_policy(stream),
+                        sorted_col_ids.begin(),
+                        sorted_col_ids.end(),
+                        ordered_row_offsets,
+                        unique_col_ids.begin(),
+                        max_row_offsets.begin(),
+                        thrust::equal_to<size_type>(),
+                        thrust::maximum<size_type>());
+
+  // 3. reduce_by_key {col_id}, {node_categories} - custom opp (*+v=*, v+v=v, *+#=E)
+  rmm::device_uvector<NodeT> column_categories(num_columns, stream);
+  thrust::reduce_by_key(
+    rmm::exec_policy(stream),
+    sorted_col_ids.begin(),
+    sorted_col_ids.end(),
+    thrust::make_permutation_iterator(tree.node_categories.begin(), ordered_node_ids.begin()),
+    unique_col_ids.begin(),
+    column_categories.begin(),
+    thrust::equal_to<size_type>(),
+    [] __device__(NodeT type_a, NodeT type_b) -> NodeT {
+      auto is_a_leaf = (type_a == NC_VAL || type_a == NC_STR);
+      auto is_b_leaf = (type_b == NC_VAL || type_b == NC_STR);
+      // (v+v=v, *+*=*,  *+v=*, *+#=E, NESTED+VAL=NESTED)
+      // *+*=*, v+v=v
+      if (type_a == type_b) {
+        return type_a;
+      } else if (is_a_leaf) {
+        // *+v=*, N+V=N
+        // STRUCT/LIST + STR/VAL = STRUCT/LIST, STR/VAL + FN = ERR, STR/VAL + STR = STR
+        return type_b == NC_FN ? NC_ERR : (is_b_leaf ? NC_STR : type_b);
+      } else if (is_b_leaf) {
+        return type_a == NC_FN ? NC_ERR : (is_a_leaf ? NC_STR : type_a);
+      }
+      // *+#=E
+      return NC_ERR;
+    });
+
+  // 4. unique_copy parent_node_ids, ranges
+  rmm::device_uvector<TreeDepthT> column_levels(0, stream);  // not required
+  rmm::device_uvector<NodeIndexT> parent_col_ids(num_columns, stream);
+  rmm::device_uvector<SymbolOffsetT> col_range_begin(num_columns, stream);  // Field names
+  rmm::device_uvector<SymbolOffsetT> col_range_end(num_columns, stream);
+  rmm::device_uvector<size_type> unique_node_ids(num_columns, stream);
+  thrust::unique_by_key_copy(rmm::exec_policy(stream),
+                             sorted_col_ids.begin(),
+                             sorted_col_ids.end(),
+                             ordered_node_ids.begin(),
+                             thrust::make_discard_iterator(),
+                             unique_node_ids.begin());
+  thrust::copy_n(
+    rmm::exec_policy(stream),
+    thrust::make_zip_iterator(
+      thrust::make_permutation_iterator(tree.parent_node_ids.begin(), unique_node_ids.begin()),
+      thrust::make_permutation_iterator(tree.node_range_begin.begin(), unique_node_ids.begin()),
+      thrust::make_permutation_iterator(tree.node_range_end.begin(), unique_node_ids.begin())),
+    unique_node_ids.size(),
+    thrust::make_zip_iterator(
+      parent_col_ids.begin(), col_range_begin.begin(), col_range_end.begin()));
+
+  // convert parent_node_ids to parent_col_ids
+  thrust::transform(
+    rmm::exec_policy(stream),
+    parent_col_ids.begin(),
+    parent_col_ids.end(),
+    parent_col_ids.begin(),
+    [col_ids = original_col_ids.begin()] __device__(auto parent_node_id) -> size_type {
+      return parent_node_id == parent_node_sentinel ? parent_node_sentinel
+                                                    : col_ids[parent_node_id];
+    });
+
+  // condition is true if parent is not a list, or sentinel/root
+  // Special case to return true if parent is a list and is_array_of_arrays is true
+  auto is_non_list_parent = [column_categories = column_categories.begin(),
+                             is_array_of_arrays,
+                             row_array_parent_col_id] __device__(auto parent_col_id) -> bool {
+    return !(parent_col_id == parent_node_sentinel ||
+             column_categories[parent_col_id] == NC_LIST &&
+               (!is_array_of_arrays || parent_col_id != row_array_parent_col_id));
+    return (parent_col_id != parent_node_sentinel) && 
+             (column_categories[parent_col_id] != NC_LIST) || 
+               (is_array_of_arrays == true && parent_col_id == row_array_parent_col_id);
+  };
+
+  // Mixed types in List children go to different columns,
+  // so all immediate children of list column should have same max_row_offsets.
+  //   create list's children max_row_offsets array. (initialize to zero)
+  //   atomicMax on  children max_row_offsets array.
+  //   gather the max_row_offsets from children row offset array.
+  {
+    rmm::device_uvector<NodeIndexT> list_parents_children_max_row_offsets(num_columns, stream);
+    thrust::fill(rmm::exec_policy(stream),
+                 list_parents_children_max_row_offsets.begin(),
+                 list_parents_children_max_row_offsets.end(),
+                 0);
+    thrust::for_each(rmm::exec_policy(stream),
+                     unique_col_ids.begin(),
+                     unique_col_ids.end(),
+                     [column_categories = column_categories.begin(),
+                      parent_col_ids    = parent_col_ids.begin(),
+                      max_row_offsets   = max_row_offsets.begin(),
+                      list_parents_children_max_row_offsets =
+                        list_parents_children_max_row_offsets.begin()] __device__(auto col_id) {
+                       auto parent_col_id = parent_col_ids[col_id];
+                       if (parent_col_id != parent_node_sentinel and
+                           column_categories[parent_col_id] == node_t::NC_LIST) {
+                         cuda::atomic_ref<NodeIndexT, cuda::thread_scope_device> ref{
+                           *(list_parents_children_max_row_offsets + parent_col_id)};
+                         ref.fetch_max(max_row_offsets[col_id], cuda::std::memory_order_relaxed);
+                       }
+                     });
+    thrust::gather_if(
+      rmm::exec_policy(stream),
+      parent_col_ids.begin(),
+      parent_col_ids.end(),
+      parent_col_ids.begin(),
+      list_parents_children_max_row_offsets.begin(),
+      max_row_offsets.begin(),
+      [column_categories = column_categories.begin()] __device__(size_type parent_col_id) {
+        return parent_col_id != parent_node_sentinel and
+               column_categories[parent_col_id] == node_t::NC_LIST;
+      });
+  }
+
+  // copy lists' max_row_offsets to children.
+  // all structs should have same size.
+  thrust::transform_if(
+    rmm::exec_policy(stream),
+    unique_col_ids.begin(),
+    unique_col_ids.end(),
+    max_row_offsets.begin(),
+    [column_categories = column_categories.begin(),
+     is_non_list_parent,
+     parent_col_ids  = parent_col_ids.begin(),
+     max_row_offsets = max_row_offsets.begin()] __device__(size_type col_id) {
+      auto parent_col_id = parent_col_ids[col_id];
+      // condition is true if parent is not a list, or sentinel/root
+      while (is_non_list_parent(parent_col_id)) {
+        col_id        = parent_col_id;
+        parent_col_id = parent_col_ids[parent_col_id];
+      }
+      return max_row_offsets[col_id];
+    },
+    [column_categories = column_categories.begin(),
+     is_non_list_parent,
+     parent_col_ids = parent_col_ids.begin()] __device__(size_type col_id) {
+      auto parent_col_id = parent_col_ids[col_id];
+      // condition is true if parent is not a list, or sentinel/root
+      return is_non_list_parent(parent_col_id);
+    });
+
+  // For Struct and List (to avoid copying entire strings when mixed type as string is enabled)
+  thrust::transform_if(
+    rmm::exec_policy(stream),
+    col_range_begin.begin(),
+    col_range_begin.end(),
+    column_categories.begin(),
+    col_range_end.begin(),
+    [] __device__(auto i) { return i + 1; },
+    [] __device__(NodeT type) { return type == NC_STRUCT || type == NC_LIST; });
+
+  return std::tuple{tree_meta_t{std::move(column_categories),
+                                std::move(parent_col_ids),
+                                std::move(column_levels),
+                                std::move(col_range_begin),
+                                std::move(col_range_end)},
+                    std::move(unique_col_ids),
+                    std::move(max_row_offsets)};
+}
+
+} // namespace detail
+} // namespace cudf::io::json