diff --git a/cpp/src/io/json/column_tree_construction.cu b/cpp/src/io/json/column_tree_construction.cu index ae2406ae7f9..d9ff6353ad5 100644 --- a/cpp/src/io/json/column_tree_construction.cu +++ b/cpp/src/io/json/column_tree_construction.cu @@ -327,13 +327,15 @@ std::tuple reduce_to_column_tree( // create list's children max_row_offsets array // gather the max_row_offsets from children row offset array. { - auto max_row_offsets_it = thrust::make_transform_iterator(thrust::make_counting_iterator(0), + auto max_row_offsets_it = thrust::make_transform_iterator( + thrust::make_counting_iterator(0), cuda::proclaim_return_type( - [colidx = colidx.begin(), - max_row_offsets = max_row_offsets.begin()] __device__ (size_t i) { - if(colidx[i] == -1) return -1; - else return max_row_offsets[colidx[i]]; - })); + [colidx = colidx.begin(), max_row_offsets = max_row_offsets.begin()] __device__(size_t i) { + if (colidx[i] == -1) + return -1; + else + return max_row_offsets[colidx[i]]; + })); rmm::device_uvector max_children_max_row_offsets(num_columns, stream); size_t temp_storage_bytes = 0; cub::DeviceSegmentedReduce::Max(nullptr, @@ -361,29 +363,29 @@ std::tuple reduce_to_column_tree( rowidx.begin() + 1, colidx.begin()); - rmm::device_uvector list_ancestors(num_columns, stream); - thrust::for_each_n( - rmm::exec_policy_nosync(stream), - thrust::make_counting_iterator(0), - num_columns, - [rowidx = rowidx.begin(), - colidx = colidx.begin(), - column_categories = column_categories.begin(), - dev_num_levels_ptr, - list_ancestors = list_ancestors.begin()] __device__(NodeIndexT node) { - auto num_levels = *dev_num_levels_ptr; - list_ancestors[node] = colidx[rowidx[node]]; - for (int level = 0; level <= num_levels && list_ancestors[node] && column_categories[list_ancestors[node]] != NC_LIST; level++) { - list_ancestors[node] = colidx[rowidx[list_ancestors[node]]]; - } - }); + thrust::for_each_n(rmm::exec_policy_nosync(stream), + thrust::make_counting_iterator(0), + num_columns, + [rowidx = rowidx.begin(), + colidx = colidx.begin(), + column_categories = column_categories.begin(), + dev_num_levels_ptr, + list_ancestors = list_ancestors.begin()] __device__(NodeIndexT node) { + auto num_levels = *dev_num_levels_ptr; + list_ancestors[node] = colidx[rowidx[node]]; + for (int level = 0; level <= num_levels && list_ancestors[node] && + column_categories[list_ancestors[node]] != NC_LIST; + level++) { + list_ancestors[node] = colidx[rowidx[list_ancestors[node]]]; + } + }); thrust::gather(rmm::exec_policy_nosync(stream), - list_ancestors.begin(), - list_ancestors.end(), - max_children_max_row_offsets.begin(), - max_row_offsets.begin()); + list_ancestors.begin(), + list_ancestors.end(), + max_children_max_row_offsets.begin(), + max_row_offsets.begin()); } return std::tuple{ @@ -536,7 +538,6 @@ reduce_to_column_tree(tree_meta_t& tree, return parent_col_id != parent_node_sentinel and column_categories[parent_col_id] == node_t::NC_LIST; }); - } // copy lists' max_row_offsets to children. diff --git a/cpp/tests/io/json/json_tree_csr.cu b/cpp/tests/io/json/json_tree_csr.cu index 27ed0e00bbb..1d4b742ccc2 100644 --- a/cpp/tests/io/json/json_tree_csr.cu +++ b/cpp/tests/io/json/json_tree_csr.cu @@ -91,39 +91,27 @@ bool check_equality(cuio_json::tree_meta_t& d_a, stream.synchronize(); auto num_nodes = a.parent_node_ids.size(); - if (b.rowidx.size() != num_nodes + 1) { - return false; - } + if (b.rowidx.size() != num_nodes + 1) { return false; } for (auto pos = b.rowidx[0]; pos < b.rowidx[1]; pos++) { auto v = b.colidx[pos]; - if (a.parent_node_ids[b.column_ids[v]] != b.column_ids[0]) { - return false; - } + if (a.parent_node_ids[b.column_ids[v]] != b.column_ids[0]) { return false; } } for (size_t u = 1; u < num_nodes; u++) { auto v = b.colidx[b.rowidx[u]]; - if (a.parent_node_ids[b.column_ids[u]] != b.column_ids[v]) { - return false; - } + if (a.parent_node_ids[b.column_ids[u]] != b.column_ids[v]) { return false; } for (auto pos = b.rowidx[u] + 1; pos < b.rowidx[u + 1]; pos++) { v = b.colidx[pos]; - if (a.parent_node_ids[b.column_ids[v]] != b.column_ids[u]) { - return false; - } + if (a.parent_node_ids[b.column_ids[v]] != b.column_ids[u]) { return false; } } } for (size_t u = 0; u < num_nodes; u++) { - if (a.node_categories[b.column_ids[u]] != b.categories[u]) { - return false; - } + if (a.node_categories[b.column_ids[u]] != b.categories[u]) { return false; } } for (size_t u = 0; u < num_nodes; u++) { - if (a_max_row_offsets[b.column_ids[u]] != b_max_row_offsets[u]) { - return false; - } + if (a_max_row_offsets[b.column_ids[u]] != b_max_row_offsets[u]) { return false; } } return true; } @@ -189,7 +177,8 @@ void run_test(std::string const& input) row_array_parent_col_id, stream); - std::printf("\n========================================================================================\n"); + std::printf( + "\n========================================================================================\n"); auto [d_column_tree_csr, d_column_tree_properties] = cudf::io::json::experimental::detail::reduce_to_column_tree( gpu_tree, gpu_col_id, gpu_row_offsets, false, row_array_parent_col_id, stream);