Prepend row index column in Parquet reader

cpp/include/cudf/io/parquet.hpp

@@ -109,6 +109,8 @@ class parquet_reader_options {
   bool _case_sensitive_names = true;
   // Whether to prepend a source file index column to the output
   bool _prepend_source_index_column = false;
+  // Whether to prepend a file-local row index column to the output
+  bool _prepend_row_index_column = false;
 
   std::optional<std::vector<reader_column_schema>> _reader_column_schema;
 
@@ -311,6 +313,20 @@ class parquet_reader_options {
     return _prepend_source_index_column;
   }
 
+  /**
+   * @brief Returns whether to prepend a file-local row index column to the output.
+   *
+   * The row index column contains, for each output row, the row's index within its parquet
+   * source file. If the source index column is also enabled, the column order is: source index,
+   * row index, data columns.
+   *
+   * @return `true` if a row index column should be prepended
+   */
+  [[nodiscard]] bool is_enabled_prepend_row_index_column() const
+  {
+    return _prepend_row_index_column;
+  }
+
   /**
    * @brief Set a new source location
    *
@@ -562,6 +578,13 @@ class parquet_reader_options {
    * @param val Boolean indicating whether to prepend the source file index column.
    */
   void enable_prepend_source_index_column(bool val) { _prepend_source_index_column = val; }
+
+  /**
+   * @brief Sets whether to prepend a file-local row index column to the output.
+   *
+   * @param val Boolean indicating whether to prepend the row index column.
+   */
+  void enable_prepend_row_index_column(bool val) { _prepend_row_index_column = val; }
 };
 
 /**
@@ -835,6 +858,18 @@ class parquet_reader_options_builder {
     return *this;
   }
 
+  /**
+   * @brief Sets whether to prepend a file-local row index column to the output.
+   *
+   * @param val Boolean indicating whether to prepend a row index column
+   * @return this for chaining
+   */
+  parquet_reader_options_builder& prepend_row_index_column(bool val)
+  {
+    options._prepend_row_index_column = val;
+    return *this;
+  }
+
   /**
    * @brief move parquet_reader_options member once it's built.
    */

cpp/src/io/parquet/reader_impl.cpp

@@ -513,7 +513,8 @@ reader_impl::reader_impl(std::size_t chunk_read_limit,
              options.get_row_groups(),
              options.is_enabled_use_jit_filter(),
              options.is_enabled_case_sensitive_names(),
-             options.is_enabled_prepend_source_index_column()},
+             options.is_enabled_prepend_source_index_column(),
+             options.is_enabled_prepend_row_index_column()},
     _sources{std::move(sources)},
     _output_chunk_read_limit{chunk_read_limit},
     _input_pass_read_limit{pass_read_limit}
@@ -876,6 +877,13 @@ table_with_metadata reader_impl::finalize_output(read_mode mode,
     _output_metadata = std::make_unique<table_metadata>(out_metadata);
   }
 
+  // Row-range of the current output chunk relative to the current row group selection.
+  auto const read_info =
+    (_file_itm_data._current_input_pass < _file_itm_data.num_passes())
+      ? _pass_itm_data->subpass
+          ->output_chunk_read_info[_pass_itm_data->subpass->current_output_chunk]
+      : row_range{0, 0};
+
   // advance output chunk/subpass/pass info for non-empty tables if and only if we are in bounds
   if (_file_itm_data._current_input_pass < _file_itm_data.num_passes()) {
     auto& pass    = *_pass_itm_data;
@@ -886,15 +894,25 @@ table_with_metadata reader_impl::finalize_output(read_mode mode,
   // increment the output chunk count
   _file_itm_data._output_chunk_count++;
 
-  // Prepend the source index column if requested
-  if (_options.prepend_source_index_column) {
-    prepend_source_index_column(out_metadata.num_rows_per_source, out_columns);
-    out_metadata.schema_info.emplace(out_metadata.schema_info.begin(),
-                                     column_name_info{.name = "src_idx", .is_nullable = false});
+  // 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_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_metadata.schema_info.emplace(
+        out_metadata.schema_info.begin(),
+        column_name_info{.name = "source_index", .is_nullable = false});
+    }
   }
 
   // Offset column references in `_expr_conv` by the number of prepended columns
-  auto const num_prepended_cols = static_cast<size_type>(_options.prepend_source_index_column);
+  auto const num_prepended_cols = static_cast<size_type>(_options.prepend_source_index_column) +
+                                  static_cast<size_type>(_options.prepend_row_index_column);
   auto const final_filter =
     offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols);
   auto const final_filter_expr = final_filter.get_converted_expr();

cpp/src/io/parquet/reader_impl.hpp

@@ -421,13 +421,25 @@ class reader_impl {
                                                                          size_t chunk_num_rows);
 
   /**
-   * @brief Construct and prepend the source index column to the output columns
+   * @brief Synthesize source index column
    *
    * @param num_rows_per_source Number of rows per parquet source
-   * @param out_columns Current output columns
+   * @return Synthesized source index column
    */
-  void prepend_source_index_column(std::span<std::size_t const> num_rows_per_source,
-                                   std::vector<std::unique_ptr<column>>& out_columns);
+  [[nodiscard]] std::unique_ptr<column> synthesize_source_index_column(
+    std::span<std::size_t const> num_rows_per_source);
+
+  /**
+   * @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
+   * @return Synthesized row index column
+   */
+  [[nodiscard]] std::unique_ptr<column> synthesize_row_index_column(row_range const& read_info);
 
   /**
    * @brief Computes the names of columns to be read from the file, if specified.
@@ -469,6 +481,8 @@ class reader_impl {
     bool case_sensitive_names = true;
     // Whether to prepend the source file index column to the output
     bool prepend_source_index_column = false;
+    // Whether to prepend the file-local row index column to the output
+    bool prepend_row_index_column = false;
   } _options;
 
   // name to reference converter to extract AST output filter

cpp/src/io/parquet/reader_impl_helpers.cpp

@@ -558,6 +558,22 @@ std::vector<size_type> aggregate_reader_metadata::get_num_row_groups_per_file()
   return per_file_num_row_groups;
 }
 
+std::vector<size_t> aggregate_reader_metadata::compute_source_row_group_offsets(
+  size_type src_idx) const
+{
+  CUDF_EXPECTS(src_idx >= 0 && std::cmp_less(src_idx, per_file_metadata.size()),
+               "invalid source index");
+  auto const& row_groups = per_file_metadata[src_idx].row_groups;
+  std::vector<size_t> source_row_offsets(row_groups.size());
+  std::transform_exclusive_scan(row_groups.cbegin(),
+                                row_groups.cend(),
+                                source_row_offsets.begin(),
+                                size_t{0},
+                                std::plus<size_t>{},
+                                [](auto const& rg) { return rg.num_rows; });
+  return source_row_offsets;
+}
+
 // Copies info from the column and offset indexes into the passed in row_group_info.
 void aggregate_reader_metadata::column_info_for_row_group(row_group_info& rg_info,
                                                           size_t chunk_start_row) const
@@ -1604,6 +1620,10 @@ aggregate_reader_metadata::select_row_groups(
     [&](auto const& src_idx) {
       auto const& file_metadata = per_file_metadata[src_idx];
 
+      // File-local row group row offsets
+      auto const source_row_offsets =
+        compute_source_row_group_offsets(static_cast<size_type>(src_idx));
+
       // For each row group in this data source
       std::for_each(
         current_row_group_indices[src_idx].begin(),
@@ -1645,6 +1665,7 @@ aggregate_reader_metadata::select_row_groups(
           selection.emplace_back(
             row_group_info{.index               = rg_idx,
                            .start_row           = row_group_start_row,
+                           .source_start_row    = source_row_offsets[rg_idx],
                            .unadjusted_num_rows = num_rows,
                            .source_index        = static_cast<cudf::size_type>(src_idx),
                            .compressed_size     = compressed_size,

cpp/src/io/parquet/reader_impl_helpers.hpp

@@ -67,6 +67,7 @@ struct row_group_info {
   size_type index;  // row group index within a file. aggregate_reader_metadata::get_row_group() is
                     // called with index and source_index
   size_t start_row;
+  size_t source_start_row;     // file-local start row of this row group within its source file
   size_t unadjusted_num_rows;  // number of unadjusted rows in the row group
   size_type source_index;      // file index.
   size_t compressed_size;      // compressed size of the row group
@@ -452,6 +453,14 @@ class aggregate_reader_metadata {
    */
   [[nodiscard]] std::vector<size_type> get_num_row_groups_per_file() const;
 
+  /**
+   * @brief Computes file-local row group row offsets for the specified source
+   *
+   * @param src_idx The source (per_file_metadata) index
+   * @return Vector of file-local row group row offsets
+   */
+  [[nodiscard]] std::vector<size_t> compute_source_row_group_offsets(size_type src_idx) const;
+
   /**
    * @brief Checks if a schema index from 0th source is mapped to the specified file index
    *

cpp/src/io/parquet/reader_impl_preprocess.cu

@@ -22,7 +22,10 @@
 
 #include <rmm/exec_policy.hpp>
 
+#include <cub/device/device_transform.cuh>
 #include <cuda/iterator>
+#include <thrust/binary_search.h>
+#include <thrust/execution_policy.h>
 #include <thrust/fill.h>
 #include <thrust/iterator/transform_iterator.h>
 #include <thrust/scan.h>
@@ -1125,56 +1128,113 @@ cudf::detail::host_vector<size_t> reader_impl::calculate_page_string_offsets()
   return cudf::detail::make_pinned_vector(d_col_sizes, _stream);
 }
 
-void reader_impl::prepend_source_index_column(std::span<std::size_t const> num_rows_per_source,
-                                              std::vector<std::unique_ptr<column>>& out_columns)
+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)
 {
-  if (not _options.prepend_source_index_column) { return; }
+  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.value()));
+    _stream.synchronize();
+  }
 
-  using column_type    = cudf::size_type;
-  auto constexpr dtype = cudf::data_type{cudf::type_to_id<column_type>()};
+  return std::make_unique<cudf::column>(std::move(col_data), rmm::device_buffer{}, 0);
+}
+
+std::unique_ptr<column> reader_impl::synthesize_source_index_column(
+  std::span<std::size_t const> num_rows_per_source)
+{
+  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});
 
-  auto const prepend_column = [&](auto&& column) {
-    out_columns.emplace(out_columns.begin(), std::move(column));
-  };
-
-  // Empty column
   if (num_rows == 0) {
-    prepend_column(cudf::make_empty_column(dtype));
-    return;
+    return cudf::make_empty_column(cudf::data_type{cudf::type_to_id<column_type>()});
   }
 
   // Single source
-  auto const num_sources = num_rows_per_source.size();
   if (num_sources == 1) {
     auto const scalar = cudf::numeric_scalar<column_type>(0, true, _stream, _mr);
-    prepend_column(cudf::make_column_from_scalar(scalar, num_rows, _stream, _mr));
-    return;
+    return cudf::make_column_from_scalar(scalar, num_rows, _stream, _mr);
   }
 
-  // Multiple sources
+  // Allocate column data vector
   auto col_data = rmm::device_uvector<column_type>(num_rows, _stream, _mr);
-  {
-    auto temp_mr = cudf::get_current_device_resource_ref();
 
+  // 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, temp_mr);
-
+    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.synchronize();
   }
 
-  prepend_column(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);
 }
 
 }  // namespace cudf::io::parquet::detail