Refactor Parquet synthetic column helpers

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

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 <cudf/copying.hpp>
 #include <cudf/detail/stream_compaction.hpp>
@@ -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});

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 <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/stream_compaction.hpp>
@@ -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});
     }

cpp/src/io/parquet/reader_impl.hpp

@@ -494,34 +494,6 @@ class reader_impl {
   [[nodiscard]] std::vector<size_t> 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<column> synthesize_source_index_column(
-    std::span<std::size_t const> 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<column> 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.
    *

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 <cudf/column/column.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>

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<column> synthesize_source_index_column(
-  std::span<std::size_t const> 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<column> 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
  *

cpp/src/io/parquet/reader_impl_preprocess.cu

@@ -1164,158 +1164,4 @@ cudf::detail::host_vector<size_t> 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<column> 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<column_type>()});
-  }
-
-  // Allocate column data vector
-  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);
-    auto host_rg_local_offsets =
-      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());
-    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<std::size_t>(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<cudf::column>(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0);
-}
-
-std::unique_ptr<column> synthesize_source_index_column(
-  std::span<std::size_t const> 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<column_type>()});
-  }
-
-  // Single source
-  if (num_sources == 1) {
-    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);
-
-  // 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<cudf::size_type>(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<cudf::column>(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0);
-}
-
-std::unique_ptr<column> 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<column_type>()});
-  }
-
-  auto const output_type = data_type{cudf::type_to_id<column_type>()};
-  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<column_type>(),
-    source_indices.end<column_type>(),
-    cuda::make_constant_iterator(column_type{1}),
-    output_view.begin<column_type>(),
-    column_type{0});
-  return output;
-}
-
-std::unique_ptr<column> reader_impl::synthesize_source_index_column(
-  std::span<std::size_t const> 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