diff --git a/cpp/src/json/json_path.cu b/cpp/src/json/json_path.cu index 146b54c0d87..2be5798098d 100644 --- a/cpp/src/json/json_path.cu +++ b/cpp/src/json/json_path.cu @@ -19,10 +19,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -903,7 +905,8 @@ template __launch_bounds__(block_size) CUDF_KERNEL void get_json_object_kernel(column_device_view col, path_operator const* const commands, - size_type* output_offsets, + size_type* d_sizes, + cudf::detail::input_offsetalator output_offsets, thrust::optional out_buf, thrust::optional out_validity, thrust::optional out_valid_count, @@ -934,7 +937,7 @@ __launch_bounds__(block_size) CUDF_KERNEL // filled in only during the precompute step. during the compute step, the offsets // are fed back in so we do -not- want to write them out - if (!out_buf.has_value()) { output_offsets[tid] = static_cast(output_size); } + if (!out_buf.has_value()) { d_sizes[tid] = output_size; } // validity filled in only during the output step if (out_validity.has_value()) { @@ -971,11 +974,6 @@ std::unique_ptr get_json_object(cudf::strings_column_view const& c if (col.is_empty()) return make_empty_column(type_id::STRING); - // allocate output offsets buffer. - auto offsets = cudf::make_fixed_width_column( - data_type{type_id::INT32}, col.size() + 1, mask_state::UNALLOCATED, stream, mr); - cudf::mutable_column_view offsets_view(*offsets); - // if the query is empty, return a string column containing all nulls if (!std::get<0>(preprocess).has_value()) { return std::make_unique( @@ -986,6 +984,11 @@ std::unique_ptr get_json_object(cudf::strings_column_view const& c col.size()); // null count } + // compute output sizes + auto sizes = + rmm::device_uvector(col.size(), stream, rmm::mr::get_current_device_resource()); + auto d_offsets = cudf::detail::offsetalator_factory::make_input_iterator(col.offsets()); + constexpr int block_size = 512; cudf::detail::grid_1d const grid{col.size(), block_size}; auto cdv = column_device_view::create(col.parent(), stream); @@ -994,20 +997,17 @@ std::unique_ptr get_json_object(cudf::strings_column_view const& c <<>>( *cdv, std::get<0>(preprocess).value().data(), - offsets_view.head(), + sizes.data(), + d_offsets, thrust::nullopt, thrust::nullopt, thrust::nullopt, options); // convert sizes to offsets - thrust::exclusive_scan(rmm::exec_policy(stream), - offsets_view.head(), - offsets_view.head() + col.size() + 1, - offsets_view.head(), - 0); - size_type const output_size = - cudf::detail::get_value(offsets_view, col.size(), stream); + auto [offsets, output_size] = + cudf::strings::detail::make_offsets_child_column(sizes.begin(), sizes.end(), stream, mr); + d_offsets = cudf::detail::offsetalator_factory::make_input_iterator(offsets->view()); // allocate output string column rmm::device_uvector chars(output_size, stream, mr); @@ -1024,7 +1024,8 @@ std::unique_ptr get_json_object(cudf::strings_column_view const& c <<>>( *cdv, std::get<0>(preprocess).value().data(), - offsets_view.head(), + sizes.data(), + d_offsets, chars.data(), static_cast(validity.data()), d_valid_count.data(),