diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 26b7ee68aad7..3aeda3cf03ad 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -805,6 +805,7 @@ add_library( src/io/parquet/reader_impl_preprocess.cu src/io/parquet/reader_impl_preprocess_utils.cu src/io/parquet/stats_filter_helpers.cpp + src/io/parquet/synthetic_column_helpers.cu src/io/parquet/writer_impl.cu src/io/parquet/writer_impl_helpers.cpp src/io/parquet/decode_fixed.cu diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 9732c2dddf09..a3090c1393fe 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -8,6 +8,7 @@ #include "cudf/io/text/byte_range_info.hpp" #include "hybrid_scan_helpers.hpp" #include "io/parquet/reader_impl_chunking_utils.cuh" +#include "io/parquet/synthetic_column_helpers.hpp" #include #include @@ -1379,15 +1380,16 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( // Prepend the source and row index columns to filter columns only if (read_columns_mode == read_columns_mode::FILTER_COLUMNS) { if (_options.prepend_row_index_column) { - out_columns.emplace(out_columns.begin(), - synthesize_row_index_column(read_info, _stream, _mr)); + out_columns.emplace( + out_columns.begin(), + synthesize_row_index_column(_file_itm_data.row_groups, 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, _stream, _mr)); + out_columns.emplace(out_columns.begin(), + parquet::detail::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.cpp b/cpp/src/io/parquet/reader_impl.cpp index f9f69d270b31..b5cf2005f0cf 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -8,6 +8,7 @@ #include "column_path_helpers.hpp" #include "error.hpp" #include "runtime/context.hpp" +#include "synthetic_column_helpers.hpp" #include #include @@ -989,8 +990,9 @@ 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, _stream, _mr)); + out_columns.emplace( + out_columns.begin(), + synthesize_row_index_column(_file_itm_data.row_groups, read_info, _stream, _mr)); out_metadata.schema_info.emplace(out_metadata.schema_info.begin(), column_name_info{.name = "row_index", .is_nullable = false}); } diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index e811971d0583..6fbd4789c6dd 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -494,34 +494,6 @@ class reader_impl { [[nodiscard]] std::vector calculate_output_num_rows_per_source(size_t chunk_start_row, size_t chunk_num_rows); - /** - * @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, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - - /** - * @brief Synthesize file-local row index column - * - * For each output row, the column contains the row's index within its parquet source file, - * accounting for row group selection and row bounds. - * - * @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, cuda::stream_ref 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_helpers.cpp b/cpp/src/io/parquet/reader_impl_helpers.cpp index 5b51c2fd871e..f229a6120721 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.cpp +++ b/cpp/src/io/parquet/reader_impl_helpers.cpp @@ -13,6 +13,7 @@ #include "ipc/Schema_generated.h" #include "parquet_common.hpp" #include "row_group_stats_helpers.hpp" +#include "synthetic_column_helpers.hpp" #include #include diff --git a/cpp/src/io/parquet/reader_impl_helpers.hpp b/cpp/src/io/parquet/reader_impl_helpers.hpp index 6e9eb5e18266..50279594164d 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.hpp +++ b/cpp/src/io/parquet/reader_impl_helpers.hpp @@ -114,30 +114,6 @@ struct row_group_size_info { */ [[nodiscard]] std::size_t derive_pass_read_limit(std::size_t chunk_read_limit); -/** - * @brief Synthesizes a source-index column. - * - * @param num_rows_per_source Number of rows per 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, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - -/** - * @brief Synthesizes row-group indices from a sorted source-index column. - * - * @param source_indices Source-index column containing one row per 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-group index column - */ -[[nodiscard]] std::unique_ptr synthesize_row_group_index_column( - column_view const& source_indices, cuda::stream_ref stream, rmm::device_async_resource_ref mr); - /** * @brief Find the offset of the column chunk with the given schema index in the specified row group * diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index b1880aa71dcc..6637d54bc0cd 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -1164,158 +1164,4 @@ cudf::detail::host_vector reader_impl::calculate_page_string_offsets() return cudf::detail::make_pinned_vector(d_col_sizes, _stream); } -namespace { - -/** - * @brief Maps each global row index to its corresponding file-local row index - */ -struct map_global_to_local_row_index { - std::size_t const* global_row_offsets; ///< Global row offsets for each row group - std::size_t const* local_row_offsets; ///< Source-local start row for each row group - std::size_t num_row_groups; - - __device__ std::size_t operator()(std::size_t row_idx) const noexcept - { - auto const row_group_idx = - cuda::std::distance( - global_row_offsets, - thrust::upper_bound( - thrust::seq, global_row_offsets, global_row_offsets + num_row_groups, row_idx)) - - 1; // Subtract 1 to get the index of the selected row group - return row_idx - global_row_offsets[row_group_idx] + local_row_offsets[row_group_idx]; - } -}; - -} // namespace - -std::unique_ptr reader_impl::synthesize_row_index_column(row_range const& read_info, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - using column_type = size_t; - - if (read_info.num_rows == 0) { - return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); - } - - // Allocate column data vector - 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); - auto host_rg_local_offsets = - 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); - } - - // 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()); - auto const rg_local_offsets = cudf::detail::make_device_uvector_async( - 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( - cuda::counting_iterator(read_info.skip_rows), - col_data.begin(), - read_info.num_rows, - map_global_to_local_row_index{ - rg_global_offsets.data(), rg_local_offsets.data(), rg_global_offsets.size()}, - stream.get())); - stream.sync(); - } - - return std::make_unique(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0); -} - -std::unique_ptr synthesize_source_index_column( - std::span num_rows_per_source, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - using column_type = cudf::size_type; - - auto const num_sources = num_rows_per_source.size(); - auto const num_rows = - std::accumulate(num_rows_per_source.begin(), num_rows_per_source.end(), std::size_t{0}); - - if (num_rows == 0) { - return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); - } - - // Single source - if (num_sources == 1) { - 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); - - // Label each output row with its source index via segment boundaries. - { - // Host per-source row offsets, including the final total row count. - auto host_row_offsets = - cudf::detail::make_empty_pinned_vector(num_sources + 1, stream); - host_row_offsets.resize(num_sources + 1); - host_row_offsets.front() = cudf::size_type{0}; - 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()); - cudf::detail::label_segments( - row_offsets.begin(), row_offsets.end(), col_data.begin(), col_data.end(), stream); - stream.sync(); - } - - return std::make_unique(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0); -} - -std::unique_ptr synthesize_row_group_index_column(column_view const& source_indices, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - using column_type = cudf::size_type; - - CUDF_EXPECTS(source_indices.type().id() == type_id::INT32, - "Source index column must have INT32 type", - std::invalid_argument); - CUDF_EXPECTS(source_indices.null_count() == 0, - "Source index column must not contain null values", - std::invalid_argument); - - if (source_indices.is_empty()) { - return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); - } - - auto const output_type = data_type{cudf::type_to_id()}; - auto output = cudf::make_fixed_width_column( - output_type, source_indices.size(), mask_state::UNALLOCATED, stream, mr); - auto output_view = output->mutable_view(); - thrust::exclusive_scan_by_key( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - source_indices.begin(), - source_indices.end(), - cuda::make_constant_iterator(column_type{1}), - output_view.begin(), - column_type{0}); - return output; -} - -std::unique_ptr reader_impl::synthesize_source_index_column( - std::span num_rows_per_source, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - return ::cudf::io::parquet::detail::synthesize_source_index_column( - num_rows_per_source, stream, mr); -} - } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/synthetic_column_helpers.cu b/cpp/src/io/parquet/synthetic_column_helpers.cu new file mode 100644 index 000000000000..c656bac8dcee --- /dev/null +++ b/cpp/src/io/parquet/synthetic_column_helpers.cu @@ -0,0 +1,176 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "reader_impl_helpers.hpp" +#include "synthetic_column_helpers.hpp" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +namespace cudf::io::parquet::detail { + +namespace { + +/** + * @brief Maps each global row index to its corresponding file-local row index + */ +struct map_global_to_local_row_index { + std::size_t const* global_row_offsets; ///< Global row offsets for each row group + std::size_t const* local_row_offsets; ///< Source-local start row for each row group + std::size_t num_row_groups; + + __device__ std::size_t operator()(std::size_t row_idx) const noexcept + { + auto const row_group_idx = + cuda::std::distance( + global_row_offsets, + thrust::upper_bound( + thrust::seq, global_row_offsets, global_row_offsets + num_row_groups, row_idx)) - + 1; // Subtract 1 to get the index of the selected row group + return row_idx - global_row_offsets[row_group_idx] + local_row_offsets[row_group_idx]; + } +}; + +} // namespace + +std::unique_ptr synthesize_row_index_column( + std::span row_groups, + row_range const& read_info, + cuda::stream_ref stream, + cudf::memory_resources mr) +{ + using column_type = size_t; + + if (read_info.num_rows == 0) { + return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); + } + + // Allocate column data vector + auto col_data = rmm::device_uvector(read_info.num_rows, stream, mr.get_output_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 host_rg_global_offsets = + 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); + 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, mr.get_temporary_mr()); + auto const rg_local_offsets = + cudf::detail::make_device_uvector_async(host_rg_local_offsets, stream, mr.get_temporary_mr()); + + // For each output row, binary search its row group and compute the (file-local) row index + CUDF_CUDA_TRY(cub::DeviceTransform::Transform( + cuda::counting_iterator(read_info.skip_rows), + col_data.begin(), + read_info.num_rows, + map_global_to_local_row_index{ + rg_global_offsets.data(), rg_local_offsets.data(), rg_global_offsets.size()}, + stream.get())); + stream.sync(); + } + + return std::make_unique( + std::move(col_data), rmm::device_buffer{0, stream, mr.get_output_mr()}, 0); +} + +std::unique_ptr synthesize_source_index_column( + std::span num_rows_per_source, + cuda::stream_ref stream, + cudf::memory_resources mr) +{ + using column_type = cudf::size_type; + + auto const num_sources = num_rows_per_source.size(); + auto const num_rows = + std::accumulate(num_rows_per_source.begin(), num_rows_per_source.end(), std::size_t{0}); + + if (num_rows == 0) { + return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); + } + + // Single source + if (num_sources == 1) { + auto const scalar = cudf::numeric_scalar(0, true, stream, mr.get_temporary_mr()); + return cudf::make_column_from_scalar(scalar, num_rows, stream, mr.get_output_mr()); + } + + // Allocate column data vector + auto col_data = rmm::device_uvector(num_rows, stream, mr.get_output_mr()); + + // Label each output row with its source index via segment boundaries. + { + // Host per-source row offsets, including the final total row count. + auto host_row_offsets = + cudf::detail::make_empty_pinned_vector(num_sources + 1, stream); + host_row_offsets.resize(num_sources + 1); + host_row_offsets.front() = cudf::size_type{0}; + 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, mr.get_temporary_mr()); + cudf::detail::label_segments( + row_offsets.begin(), row_offsets.end(), col_data.begin(), col_data.end(), stream); + stream.sync(); + } + + return std::make_unique( + std::move(col_data), rmm::device_buffer{0, stream, mr.get_output_mr()}, 0); +} + +std::unique_ptr synthesize_row_group_index_column( + cudf::column_view const& source_indices, cuda::stream_ref stream, cudf::memory_resources mr) +{ + using column_type = cudf::size_type; + + CUDF_EXPECTS(source_indices.type().id() == type_id::INT32, + "Source index column must have INT32 type", + std::invalid_argument); + CUDF_EXPECTS(source_indices.null_count() == 0, + "Source index column must not contain null values", + std::invalid_argument); + + if (source_indices.is_empty()) { + return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); + } + + auto const output_type = data_type{cudf::type_to_id()}; + auto output = cudf::make_fixed_width_column( + output_type, source_indices.size(), mask_state::UNALLOCATED, stream, mr.get_output_mr()); + auto output_view = output->mutable_view(); + thrust::exclusive_scan_by_key(rmm::exec_policy_nosync(stream, mr.get_temporary_mr()), + source_indices.begin(), + source_indices.end(), + cuda::make_constant_iterator(column_type{1}), + output_view.begin(), + column_type{0}); + return output; +} + +} // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/synthetic_column_helpers.hpp b/cpp/src/io/parquet/synthetic_column_helpers.hpp new file mode 100644 index 000000000000..ad437a9a8708 --- /dev/null +++ b/cpp/src/io/parquet/synthetic_column_helpers.hpp @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "reader_impl_chunking.hpp" +#include "reader_impl_helpers.hpp" + +#include +#include +#include + +#include + +#include +#include +#include + +namespace cudf::io::parquet::detail { + +/** + * @brief Synthesizes a source-index column. + * + * @param num_rows_per_source Number of rows per source + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Memory resources used for output and temporary device allocations + * @return Synthesized source-index column + */ +[[nodiscard]] std::unique_ptr synthesize_source_index_column( + std::span num_rows_per_source, + cuda::stream_ref stream, + cudf::memory_resources mr); + +/** + * @brief Synthesizes row-group indices from a sorted source-index column. + * + * @param source_indices Source-index column containing one row per row group + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Memory resources used for output and temporary device allocations + * @return Synthesized row-group index column + */ +[[nodiscard]] std::unique_ptr synthesize_row_group_index_column( + cudf::column_view const& source_indices, cuda::stream_ref stream, cudf::memory_resources mr); + +/** + * @brief Synthesizes a row-index column for the selected row groups. + * + * @param row_groups Selected row groups that map global rows to source-local rows + * @param read_info Row range of the output chunk relative to the first selected row group + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Memory resources used for output and temporary device allocations + * @return Synthesized row-index column + */ +[[nodiscard]] std::unique_ptr synthesize_row_index_column( + std::span row_groups, + row_range const& read_info, + cuda::stream_ref stream, + cudf::memory_resources mr); + +} // namespace cudf::io::parquet::detail