Skip to content

Supply stream and mr to column synthesizers #23209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 1 commit into from
Jul 10, 2026
Merged

Supply stream and mr to column synthesizers #23209

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions cpp/src/io/parquet/reader_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -897,13 +897,15 @@ table_with_metadata reader_impl::finalize_output(read_mode mode,
// Prepend the source and row index columns if requested
{
if (_options.prepend_row_index_column) {
out_columns.emplace(out_columns.begin(), synthesize_row_index_column(read_info));
out_columns.emplace(out_columns.begin(),
synthesize_row_index_column(read_info, _stream, _mr));
out_metadata.schema_info.emplace(out_metadata.schema_info.begin(),
column_name_info{.name = "row_index", .is_nullable = false});
}
if (_options.prepend_source_index_column) {
out_columns.emplace(out_columns.begin(),
synthesize_source_index_column(out_metadata.num_rows_per_source));
out_columns.emplace(
out_columns.begin(),
synthesize_source_index_column(out_metadata.num_rows_per_source, _stream, _mr));
out_metadata.schema_info.emplace(
out_metadata.schema_info.begin(),
column_name_info{.name = "source_index", .is_nullable = false});
Expand Down
11 changes: 9 additions & 2 deletions cpp/src/io/parquet/reader_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -424,10 +424,14 @@ class reader_impl {
* @brief Synthesize source index column
*
* @param num_rows_per_source Number of rows per parquet source
* @param stream CUDA stream used for device memory operations and kernel launches
* @param mr Device memory resource to use for device memory allocation
* @return Synthesized source index column
*/
[[nodiscard]] std::unique_ptr<column> synthesize_source_index_column(
std::span<std::size_t const> num_rows_per_source);
std::span<std::size_t const> num_rows_per_source,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr);

/**
* @brief Synthesize file-local row index column
Expand All @@ -437,9 +441,12 @@ class reader_impl {
*
* @param read_info Row range of the output chunk relative to the first row of the first
* selected row group
* @param stream CUDA stream used for device memory operations and kernel launches
* @param mr Device memory resource to use for device memory allocation
Comment thread
mhaseeb123 marked this conversation as resolved.
* @return Synthesized row index column
*/
[[nodiscard]] std::unique_ptr<column> synthesize_row_index_column(row_range const& read_info);
[[nodiscard]] std::unique_ptr<column> synthesize_row_index_column(
row_range const& read_info, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr);

/**
* @brief Computes the names of columns to be read from the file, if specified.
Expand Down
39 changes: 22 additions & 17 deletions cpp/src/io/parquet/reader_impl_preprocess.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1152,7 +1152,9 @@ struct map_global_to_local_row_index {

} // namespace

std::unique_ptr<column> reader_impl::synthesize_row_index_column(row_range const& read_info)
std::unique_ptr<column> reader_impl::synthesize_row_index_column(row_range const& read_info,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
using column_type = size_t;

Expand All @@ -1161,26 +1163,26 @@ std::unique_ptr<column> reader_impl::synthesize_row_index_column(row_range const
}

// Allocate column data vector
auto col_data = rmm::device_uvector<column_type>(read_info.num_rows, _stream, _mr);
auto col_data = rmm::device_uvector<column_type>(read_info.num_rows, stream, mr);

// Map global row indices in the current row-range to corresponding source-local row indices
{
// Collect global and file-local start rows for each selected row group
auto const& row_groups = _file_itm_data.row_groups;
auto host_rg_global_offsets =
cudf::detail::make_empty_pinned_vector<std::size_t>(row_groups.size(), _stream);
cudf::detail::make_empty_pinned_vector<std::size_t>(row_groups.size(), stream);
auto host_rg_local_offsets =
cudf::detail::make_empty_pinned_vector<size_t>(row_groups.size(), _stream);
cudf::detail::make_empty_pinned_vector<size_t>(row_groups.size(), stream);
for (auto const& rg : row_groups) {
host_rg_global_offsets.push_back(rg.start_row);
host_rg_local_offsets.push_back(rg.source_start_row);
}

// Copy to device
auto const rg_global_offsets = cudf::detail::make_device_uvector_async(
host_rg_global_offsets, _stream, cudf::get_current_device_resource_ref());
host_rg_global_offsets, stream, cudf::get_current_device_resource_ref());
auto const rg_local_offsets = cudf::detail::make_device_uvector_async(
host_rg_local_offsets, _stream, cudf::get_current_device_resource_ref());
host_rg_local_offsets, stream, cudf::get_current_device_resource_ref());

// For each output row, binary search its row group and compute the (file-local) row index
CUDF_CUDA_TRY(cub::DeviceTransform::Transform(
Expand All @@ -1189,15 +1191,17 @@ std::unique_ptr<column> reader_impl::synthesize_row_index_column(row_range const
read_info.num_rows,
map_global_to_local_row_index{
rg_global_offsets.data(), rg_local_offsets.data(), rg_global_offsets.size()},
_stream.value()));
_stream.synchronize();
stream.value()));
stream.synchronize();
}

return std::make_unique<cudf::column>(std::move(col_data), rmm::device_buffer{}, 0);
return std::make_unique<cudf::column>(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0);
Comment thread
mhaseeb123 marked this conversation as resolved.
}

std::unique_ptr<column> reader_impl::synthesize_source_index_column(
std::span<std::size_t const> num_rows_per_source)
std::span<std::size_t const> num_rows_per_source,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
using column_type = cudf::size_type;

Expand All @@ -1211,12 +1215,13 @@ std::unique_ptr<column> reader_impl::synthesize_source_index_column(

// Single source
if (num_sources == 1) {
auto const scalar = cudf::numeric_scalar<column_type>(0, true, _stream, _mr);
return cudf::make_column_from_scalar(scalar, num_rows, _stream, _mr);
auto const scalar =
cudf::numeric_scalar<column_type>(0, true, stream, cudf::get_current_device_resource_ref());
return cudf::make_column_from_scalar(scalar, num_rows, stream, mr);
}

// Allocate column data vector
auto col_data = rmm::device_uvector<column_type>(num_rows, _stream, _mr);
auto col_data = rmm::device_uvector<column_type>(num_rows, stream, mr);

// Label each output row with its source index via segment boundaries.
{
Expand All @@ -1228,13 +1233,13 @@ std::unique_ptr<column> reader_impl::synthesize_source_index_column(
std::inclusive_scan(
num_rows_per_source.begin(), num_rows_per_source.end(), host_row_offsets.begin() + 1);
auto const row_offsets = cudf::detail::make_device_uvector_async(
host_row_offsets, _stream, cudf::get_current_device_resource_ref());
host_row_offsets, stream, cudf::get_current_device_resource_ref());
cudf::detail::label_segments(
row_offsets.begin(), row_offsets.end(), col_data.begin(), col_data.end(), _stream);
_stream.synchronize();
row_offsets.begin(), row_offsets.end(), col_data.begin(), col_data.end(), stream);
stream.synchronize();
}

return std::make_unique<cudf::column>(std::move(col_data), rmm::device_buffer{}, 0);
return std::make_unique<cudf::column>(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0);
}

} // namespace cudf::io::parquet::detail
Loading