From c2ce174463cabda5e8389b008c4205606e318c03 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Thu, 9 Jul 2026 21:07:25 +0000 Subject: [PATCH] Supply stream and mr to column synthesizers --- cpp/src/io/parquet/reader_impl.cpp | 8 ++-- cpp/src/io/parquet/reader_impl.hpp | 11 +++++- cpp/src/io/parquet/reader_impl_preprocess.cu | 39 +++++++++++--------- 3 files changed, 36 insertions(+), 22 deletions(-) diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index 1155628159bf..1c686bcc7067 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -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}); diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 81ae58ab05da..666413f95f03 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -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 synthesize_source_index_column( - std::span num_rows_per_source); + std::span num_rows_per_source, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); /** * @brief Synthesize file-local row index column @@ -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 * @return Synthesized row index column */ - [[nodiscard]] std::unique_ptr synthesize_row_index_column(row_range const& read_info); + [[nodiscard]] std::unique_ptr 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. diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index bc0540e5a919..7634bb90cdcc 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -1152,7 +1152,9 @@ struct map_global_to_local_row_index { } // namespace -std::unique_ptr reader_impl::synthesize_row_index_column(row_range const& read_info) +std::unique_ptr 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; @@ -1161,16 +1163,16 @@ std::unique_ptr reader_impl::synthesize_row_index_column(row_range const } // Allocate column data vector - auto col_data = rmm::device_uvector(read_info.num_rows, _stream, _mr); + auto col_data = rmm::device_uvector(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(row_groups.size(), _stream); + cudf::detail::make_empty_pinned_vector(row_groups.size(), stream); auto host_rg_local_offsets = - cudf::detail::make_empty_pinned_vector(row_groups.size(), _stream); + cudf::detail::make_empty_pinned_vector(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); @@ -1178,9 +1180,9 @@ std::unique_ptr reader_impl::synthesize_row_index_column(row_range const // 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( @@ -1189,15 +1191,17 @@ std::unique_ptr 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(std::move(col_data), rmm::device_buffer{}, 0); + return std::make_unique(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0); } std::unique_ptr reader_impl::synthesize_source_index_column( - std::span num_rows_per_source) + std::span num_rows_per_source, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { using column_type = cudf::size_type; @@ -1211,12 +1215,13 @@ std::unique_ptr reader_impl::synthesize_source_index_column( // Single source if (num_sources == 1) { - auto const scalar = cudf::numeric_scalar(0, true, _stream, _mr); - return cudf::make_column_from_scalar(scalar, num_rows, _stream, _mr); + auto const scalar = + cudf::numeric_scalar(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(num_rows, _stream, _mr); + auto col_data = rmm::device_uvector(num_rows, stream, mr); // Label each output row with its source index via segment boundaries. { @@ -1228,13 +1233,13 @@ std::unique_ptr 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(std::move(col_data), rmm::device_buffer{}, 0); + return std::make_unique(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0); } } // namespace cudf::io::parquet::detail