Skip to content

Commit

Permalink
formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
@shrshi
shrshi committed Sep 17, 2024
1 parent 82c9ebe commit 8dd6877
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 46 deletions.
55 changes: 28 additions & 27 deletions cpp/src/io/json/column_tree_construction.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -327,13 +327,15 @@ std::tuple<csr, column_tree_properties> 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<row_offset_t>(
[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<row_offset_t> max_children_max_row_offsets(num_columns, stream);
size_t temp_storage_bytes = 0;
cub::DeviceSegmentedReduce::Max(nullptr,
Expand Down Expand Up @@ -361,29 +363,29 @@ std::tuple<csr, column_tree_properties> reduce_to_column_tree(
rowidx.begin() + 1,
colidx.begin());


rmm::device_uvector<NodeIndexT> 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{
Expand Down Expand Up @@ -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.
Expand Down
27 changes: 8 additions & 19 deletions cpp/tests/io/json/json_tree_csr.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit 8dd6877

Please sign in to comment.