Prepend source and row index columns in Hybrid scan reader

cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp

@@ -860,12 +860,13 @@ void hybrid_scan_reader_impl::reset_internal_state()
   _options.decimal_width  = type_id::EMPTY;
   _options.num_rows       = std::nullopt;
   _options.row_group_indices.clear();
-  _options.use_jit_filter       = false;
-  _options.case_sensitive_names = true;
-  _num_sources                  = 0;
-  _input_pass_read_limit        = 0;
-  _output_chunk_read_limit      = 0;
-  _strings_to_categorical       = false;
+  _options.use_jit_filter              = false;
+  _options.case_sensitive_names        = true;
+  _options.prepend_source_index_column = false;
+  _num_sources                         = 0;
+  _input_pass_read_limit               = 0;
+  _output_chunk_read_limit             = 0;
+  _strings_to_categorical              = false;
   _reader_column_schema.reset();
   _expr_conv = named_to_reference_converter{};
   _mr        = cudf::get_current_device_resource_ref();
@@ -877,10 +878,11 @@ void hybrid_scan_reader_impl::initialize_column_selection_options(
   // Strings may be returned as either string or categorical columns
   _strings_to_categorical = options.is_enabled_convert_strings_to_categories();
 
-  _options.timestamp_type       = cudf::data_type{options.get_timestamp_type().id()};
-  _options.decimal_width        = options.get_decimal_width();
-  _options.use_jit_filter       = options.is_enabled_use_jit_filter();
-  _options.case_sensitive_names = options.is_enabled_case_sensitive_names();
+  _options.timestamp_type              = cudf::data_type{options.get_timestamp_type().id()};
+  _options.decimal_width               = options.get_decimal_width();
+  _options.use_jit_filter              = options.is_enabled_use_jit_filter();
+  _options.case_sensitive_names        = options.is_enabled_case_sensitive_names();
+  _options.prepend_source_index_column = options.is_enabled_prepend_source_index_column();
 
   _use_pandas_metadata = options.is_enabled_use_pandas_metadata();
 }
@@ -969,12 +971,9 @@ table_with_metadata hybrid_scan_reader_impl::read_chunk_internal(
 
   // no work to do (this can happen on the first pass if we have no rows to read)
   if (!has_more_work()) {
-    // Check if number of rows per source should be included in output metadata.
-    if (include_output_num_rows_per_source()) {
-      // Empty dataframe case: Simply initialize to a list of zeros
-      out_metadata.num_rows_per_source =
-        std::vector<std::size_t>(_file_itm_data.num_rows_per_source.size(), 0);
-    }
+    // Empty dataframe case: Simply initialize to a list of zeros
+    out_metadata.num_rows_per_source =
+      std::vector<std::size_t>(_file_itm_data.num_rows_per_source.size(), 0);
 
     // Finalize output
     return finalize_output(read_columns_mode, out_metadata, out_columns, row_mask);
@@ -1036,18 +1035,13 @@ table_with_metadata hybrid_scan_reader_impl::read_chunk_internal(
   out_columns =
     cudf::structs::detail::enforce_null_consistency(std::move(out_columns), _stream, _mr);
 
-  // Check if number of rows per source should be included in output metadata.
-  if (include_output_num_rows_per_source()) {
-    // For chunked reading, compute the output number of rows per source
-    if (mode == read_mode::CHUNKED_READ) {
-      out_metadata.num_rows_per_source =
-        calculate_output_num_rows_per_source(read_info.skip_rows, read_info.num_rows);
-    }
-    // Simply move the number of rows per file if reading all at once
-    else {
-      // Move is okay here as we are reading in one go.
-      out_metadata.num_rows_per_source = std::move(_file_itm_data.num_rows_per_source);
-    }
+  // Compute the output number of rows per source
+  if (mode == read_mode::CHUNKED_READ) {
+    out_metadata.num_rows_per_source =
+      calculate_output_num_rows_per_source(read_info.skip_rows, read_info.num_rows);
+  } else {
+    // Move is okay here as we are reading in one go.
+    out_metadata.num_rows_per_source = std::move(_file_itm_data.num_rows_per_source);
   }
 
   // Add empty columns if needed. Filter output columns based on filter.
@@ -1090,6 +1084,18 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output(
 
   apply_decimal_width_cast(out_columns);
 
+  // Prepend the source index column to filter columns or all columns
+  if (_options.prepend_source_index_column and
+      read_columns_mode == read_columns_mode::FILTER_COLUMNS) {
+    prepend_source_index_column(out_metadata.num_rows_per_source, out_columns);
+  }
+
+  // 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 final_filter =
+    parquet::detail::offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols);
+  auto const final_filter_expr = final_filter.get_converted_expr();
+
   // Create a table from the output columns.
   auto read_table = std::make_unique<table>(std::move(out_columns));
 
@@ -1107,15 +1113,17 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output(
   _row_mask_offset += read_table->num_rows();
   _output_chunk_produced = true;
 
+  // Clear the number of rows per source as it is not valid after filtering
+  out_metadata.num_rows_per_source.clear();
+
   // For filter columns, apply the filter expression and update the input row mask
   if constexpr (std::is_same_v<RowMaskView, cudf::mutable_column_view>) {
     CUDF_EXPECTS(read_columns_mode == read_columns_mode::FILTER_COLUMNS, "Invalid read mode");
 
-    auto final_row_mask =
-      cudf::detail::compute_column(*read_table,
-                                   _expr_conv.get_converted_expr().value().get(),
-                                   _stream,
-                                   cudf::get_current_device_resource_ref());
+    auto final_row_mask = cudf::detail::compute_column(*read_table,
+                                                       final_filter_expr.value().get(),
+                                                       _stream,
+                                                       cudf::get_current_device_resource_ref());
     CUDF_EXPECTS(final_row_mask->view().type().id() == type_id::BOOL8,
                  "Predicate filter should return a boolean");
 

cpp/src/io/parquet/reader_impl.hpp

@@ -395,16 +395,6 @@ class reader_impl {
     return _file_itm_data._output_chunk_count == 0;
   }
 
-  /**
-   * @brief Check if number of rows per source should be included in output metadata.
-   *
-   * @return True if AST filter is not present
-   */
-  [[nodiscard]] bool include_output_num_rows_per_source() const
-  {
-    return not _expr_conv.get_converted_expr().has_value();
-  }
-
   [[nodiscard]] cudf::detail::hostdevice_span<bool> subpass_page_mask_span() const
   {
     return _subpass_page_mask ? *_subpass_page_mask : cudf::detail::hostdevice_span<bool>{};

cpp/tests/io/experimental/hybrid_scan_test.cpp

@@ -878,6 +878,148 @@ TEST_F(HybridScanTest, StructChildFilterColumn)
     std::invalid_argument);
 }
 
+TEST_F(HybridScanTest, PrependSourceIndexColumn)
+{
+  using T = uint32_t;
+  using cudf::io::parquet::experimental::use_data_page_mask;
+
+  // Single source parquet file with columns col0 (filter), col1 and col2 (payload)
+  auto constexpr num_concat            = 1;
+  auto [written_table, parquet_buffer] = create_parquet_with_stats<T, num_concat>();
+
+  // Filtering AST - col0 < 100
+  auto literal_value     = cudf::numeric_scalar<T>(100);
+  auto literal           = cudf::ast::literal(literal_value);
+  auto col_ref_0         = cudf::ast::column_name_reference("col0");
+  auto filter_expression = cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref_0, literal);
+
+  auto const stream = cudf::get_default_stream();
+  auto const mr     = cudf::get_current_device_resource_ref();
+
+  // Expected table
+  auto const expected = [&]() {
+    auto const options =
+      cudf::io::parquet_reader_options::builder(
+        cudf::io::source_info(cudf::host_span<char>(parquet_buffer.data(), parquet_buffer.size())))
+        .filter(filter_expression)
+        .prepend_source_index_column(true)
+        .build();
+    return cudf::io::read_parquet(options, stream, mr).tbl;
+  }();
+
+  // Hybrid scan reader options
+  auto const options = cudf::io::parquet_reader_options::builder()
+                         .filter(filter_expression)
+                         .prepend_source_index_column(true)
+                         .build();
+
+  auto datasource          = cudf::io::datasource::create(cudf::host_span<std::byte const>(
+    reinterpret_cast<std::byte const*>(parquet_buffer.data()), parquet_buffer.size()));
+  auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*datasource);
+
+  auto reader =
+    std::make_unique<cudf::io::parquet::experimental::hybrid_scan_reader>(*footer_buffer, options);
+
+  auto const all_row_groups = reader->all_row_groups(options);
+
+  auto const current_row_groups = cudf::host_span<cudf::size_type const>(all_row_groups);
+  auto const mid                = current_row_groups.size() / 2;
+  auto const row_group_passes =
+    std::array{current_row_groups.subspan(0, mid),
+               current_row_groups.subspan(mid, current_row_groups.size() - mid)};
+  auto const table_views = [](auto const& tables) {
+    auto views = std::vector<cudf::table_view>{};
+    std::transform(tables.begin(), tables.end(), std::back_inserter(views), [](auto const& t) {
+      return t->view();
+    });
+    return views;
+  };
+
+  // Two-step chunked read
+  {
+    auto filter_tables  = std::vector<std::unique_ptr<cudf::table>>{};
+    auto payload_tables = std::vector<std::unique_ptr<cudf::table>>{};
+
+    auto const read_pass = [&](auto const row_groups) {
+      auto row_mask      = reader->build_all_true_row_mask(row_groups, stream, mr);
+      auto row_mask_view = row_mask->mutable_view();
+
+      auto const filter_byte_ranges = reader->filter_column_chunks_byte_ranges(row_groups, options);
+      auto [filter_bufs, filter_data, filter_tasks] =
+        cudf::io::parquet::fetch_byte_ranges_to_device_async(
+          *datasource, filter_byte_ranges, stream, mr);
+      filter_tasks.get();
+      reader->setup_chunking_for_filter_columns(1024,
+                                                10240,
+                                                row_groups,
+                                                row_mask_view,
+                                                use_data_page_mask::NO,
+                                                filter_data,
+                                                options,
+                                                stream,
+                                                mr);
+      while (reader->has_next_table_chunk()) {
+        filter_tables.push_back(reader->materialize_filter_columns_chunk(row_mask_view).tbl);
+      }
+
+      auto const payload_byte_ranges =
+        reader->payload_column_chunks_byte_ranges(row_groups, options);
+      auto [payload_bufs, payload_data, payload_tasks] =
+        cudf::io::parquet::fetch_byte_ranges_to_device_async(
+          *datasource, payload_byte_ranges, stream, mr);
+      payload_tasks.get();
+
+      reader->setup_chunking_for_payload_columns(1024,
+                                                 10240,
+                                                 row_groups,
+                                                 row_mask_view,
+                                                 use_data_page_mask::NO,
+                                                 payload_data,
+                                                 options,
+                                                 stream,
+                                                 mr);
+      while (reader->has_next_table_chunk()) {
+        payload_tables.push_back(reader->materialize_payload_columns_chunk(row_mask_view).tbl);
+      }
+    };
+
+    for (auto const& pass : row_group_passes) {
+      read_pass(pass);
+    }
+
+    auto const filter_table  = cudf::concatenate(table_views(filter_tables), stream, mr);
+    auto const payload_table = cudf::concatenate(table_views(payload_tables), stream, mr);
+
+    CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->select({0, 1}), filter_table->view());
+    CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->select({2, 3}), payload_table->view());
+  }
+
+  // Single-step chunked read
+  {
+    auto all_tables = std::vector<std::unique_ptr<cudf::table>>{};
+
+    auto const read_pass = [&](cudf::host_span<cudf::size_type const> row_groups) {
+      auto const chunk_ranges = reader->all_column_chunks_byte_ranges(row_groups, options);
+      auto [bufs, col_data, tasks] =
+        cudf::io::parquet::fetch_byte_ranges_to_device_async(*datasource, chunk_ranges, stream, mr);
+      tasks.get();
+
+      reader->setup_chunking_for_all_columns(
+        1024, 10240, row_groups, col_data, options, stream, mr);
+      while (reader->has_next_table_chunk()) {
+        all_tables.push_back(reader->materialize_all_columns_chunk().tbl);
+      }
+    };
+
+    for (auto const& pass : row_group_passes) {
+      read_pass(pass);
+    }
+
+    auto const all_columns_table = cudf::concatenate(table_views(all_tables), stream, mr);
+    CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), all_columns_table->view());
+  }
+}
+
 TEST_F(HybridScanTest, ChunkedReadRowMaskPerPass)
 {
   using T = uint32_t;